aigym 0.0.0

Reinforcement learning environments for fine-tuning language models for reasoning tasks.

🤖 AI Gym

Reinforcement learning environments for AI fine-tuning

aigym is a library that provides a suite of reinforcement learning (RL) environments primarily for the purpose of fine-tuning pre-trained models - namely language models - for various reasoning tasks.

Built on top of the gymnasium API, the objective of this project is to expose a light-weight and extensible environments to fine-tune language models with techniques like PPO and GRPO.

It is designed to complement training frameworks like trl, transformers, pytorch, and pytorch lightning.

See the project roadmap here

Development Installation

Install uv:

pip install uv

Create a virtual environment:

uv venv

Activate the virtual environment:

source .venv/bin/activate

Install the package:

uv sync --extra ollama

Install ollama to run a local model: https://ollama.com/download

Usage

Run an ollama-based agent on the Wikipedia maze RL environment:

python examples/ollama_agent.py