Modrax

An easy-to-use, performant RL library based on Jax and NNX.

Quick start

git clone https://github.com/marcojira/modrax.git
cd modrax
git checkout dev
uv venv --python 3.10
source .venv/bin/activate
uv pip install -e .

# Launch example script
python examples/pqn/minatar.py