Modrax

A compact, modular reinforcement learning library built with JAX and Flax NNX.

Modrax is based around three main object.

• Env provides standardized batched environment interaction for different environment suites.
• Network is an NNX model that contains the train/evaluation policy
• Alg is initialized with an Env and a Network, then defines a training step that updates the network given interactions with the Env.

These three are brought together by training.py that contains a training loop which handles logging, checkpointing and evaluation.

Modrax requires Python 3.11 or newer and uses uv.

git clone https://github.com/marcojira/modrax.git
cd modrax

# Install Modrax with the Gymnax environment integration
uv sync --python 3.11 --extra gymnax

# Check the available configuration options
uv run examples/pqn_minatar.py --help

Environment integrations are optional. Use --extra all to install every available integration. The base dependencies currently install CUDA-enabled JAX.

See the examples folder for more.

[WIP 🚧] More documentation coming soon.