aigym 0.0.1

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

Installation

pip install aigym

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 --group dev

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

Usage

The examples directory contains examples on how to use the aigym environments. Run an ollama-based agent on the Wikipedia maze environment:

python examples/ollama_agent.py