Ajax

Reinforcement-learning agents in JAX, compiled end to end with the environment,
including average-reward, expert-guided and world-model agents.

A whole training run, environment and every seed included, one compiled program.

Ajax is a library of reinforcement-learning agents written in JAX. Compared with Stable-Baselines3 and its JAX version SBX, it differs in four ways:

• End to end on the accelerator: the environment and the training loop compile into one JAX program, as in PureJaxRL, Rejax, Stoix and Brax's training agents, with no Python step between them, on a CPU or a GPU.
• Average-reward agents: ASAC and APO maximise the reward per step, with no discount, for tasks that never end.
• Expert-guided agents: learn from a controller you already have, or from recorded transitions, by cloning, imitation, residual policies, hand-over, prior data in the replay or offline training (methods below).
• World models: DreamerV3 and TD-MPC2, following the code behind their papers' results.

Why Ajax shows each with an example and compares Ajax with SB3, SBX, PureJaxRL, Brax, Stoix and Rejax, including when to use them instead.

It has sixteen agents, from SAC and PPO to DreamerV3 and TD-MPC2, each implementing a published paper, with the places where it departs from that paper written down. Several seeds train side by side in one program, on a CPU or a GPU. Research features (learning from an expert, exploration rules, diagnostics) are extensions, passed as extensions=(...) rather than added as new arguments to the agents. An agent refuses, when built, an extension it cannot run; the compatibility table shows which agent runs which.

Installation

Ajax is on PyPI as ajax-rl, for Python 3.12 or newer, and imports as ajax:

pip install ajax-rl
pip install "ajax-rl[playground]"  # with Brax and MuJoCo Playground

JAX comes for the CPU; on a Linux machine with an NVIDIA GPU, add CUDA with ajax-rl[cuda]. Brax and MuJoCo Playground environments need an extra: ajax-rl[playground], which brings Brax, or ajax-rl[brax] alone (Installation). pip takes the newest versions that pyproject.toml allows; main installs with pip install "ajax-rl @ git+https://github.com/YannBerthelot/Ajax.git".

From source, uv installs the versions in uv.lock, which CI tests, with the development tools and the playground extra (uv sync --extra cuda adds CUDA):

git clone https://github.com/YannBerthelot/Ajax.git
cd Ajax
uv sync   # creates .venv with the locked dependencies

Installation covers the extras (cuda, compare, reproduction) and the dependency groups.

Quickstart

from ajax import PPO
from ajax.logging.wandb_logging import LoggingConfig

agent = PPO("CartPole-v1")  # 4 environments x 2048-step rollouts: one update per 8,192 steps
# Evaluate after every update; use_wandb=False keeps the results in memory only.
logging = LoggingConfig(config={}, log_frequency=8_192, use_wandb=False)
state, evaluations = agent.train(seed=[1, 2, 3], n_timesteps=40_960, logging_config=logging)

print(evaluations["timestep"][0])                # when each evaluation ran
print(evaluations["Eval/episodic mean reward"])  # one row per seed, one column per evaluation

The three seeds train together, and the run takes about fifteen seconds on a laptop CPU, compilation included. evaluations maps each logged name to an array of shape (seeds, evaluations); "Eval/episodic mean reward" is the mean return of 10 test episodes (num_episode_test). With use_wandb=True (the default) or use_tensorboard=True, the same records also go to Weights & Biases or TensorBoard. To train further, pass the state back: agent.train(seed=[1, 2, 3], n_timesteps=more, initial_state=state).

Features are arguments of the same agents:

from ajax import PPO, SAC
from ajax.extensions.instrumentation import ConditioningMetrics

# A recurrent actor and critic, for when one observation is not enough.
agent = SAC("Pendulum-v1", memory={"kind": "lstm", "hidden_size": 64})

# A research feature as an extension: the critic's health at each evaluation.
agent = PPO("CartPole-v1", expose_recent_rollout=True, extensions=(ConditioningMetrics(),))

# Brax and MuJoCo Playground environments by id (their extras), here 8 copies in parallel.
agent = PPO("ant", n_envs=8)

The quickstart guide goes on from here.

Agents

Memory says whether the agent accepts memory=; DreamerV3's world model is recurrent by design. Every agent is importable from ajax (from ajax import TD3). The algorithms overview helps choose one.

Learning from an expert or from data

Each method is an extension or an agent option. The last column lists the agents that accept it, as the compatibility table, generated from the code, does.

Expert guidance explains the expert's conventions and how the pieces combine.

What else is in the box

• Extensions: twenty-four of them, for expert guidance, exploration, critic-target changes, actor and critic pre-training, prior data, offline RL, adversarial imitation, ensemble diversity and critic diagnostics, passed as extensions=(...); the compatibility table says which agents run each one.
• Expert guidance: the expert reads the environment's raw observations and answers in the environment's action units; an expert with a state, such as a PID's integral, gets it back at every step.
• Recurrent memory: a GRU, LSTM, transformer or Mamba block for PPO, SAC, TD3, REDQ, ASAC and APG, with R2D2-style sequence replay (burn-in, stored state) for the replay agents.
• Differentiable simulation: APG back-propagates the return through gymnax dynamics, across system classes (distributions over physical parameters), with Busetto et al.'s in-context controller and its staged curriculum.
• World models: DreamerV3 and TD-MPC2 (single- and multi-task), following the paper-era official code, with specifications, deviations and parity tests against that code.
• Networks and the PID policy: layers and initialisers named by strings, a CNN encoder for image observations (PPO, DQN, PQN, UDRL), and a PID controller as the actor of SAC, TD3, REDQ, ASAC, AVG, PPO or APO, with learned fixed gains that deploy as a plain PID, or gains scheduled on the observation. Where an agent starts from its reference code's initialisation and optimiser settings, its algorithm page says so and cites the code.
• Environments: gymnax, Brax and MuJoCo Playground by id, n_envs copies in parallel, terminations kept apart from time limits (a truncated episode still bootstraps), episode_length and action_repeat on every agent, and wrappers for normalisation, observation noise, fixed initial states, model-reference tracking and PID error terms. Every continuous agent acts in [-1, 1], mapped onto the environment's action bounds.
• Logging: train returns every evaluation as arrays, with training diagnostics (gradient norms, PPO's approximate KL, explained variance, ...) and each extension's own metrics, and a background worker sends them to Weights & Biases or TensorBoard.
• Comparing agents: ajax.compare runs AdaStop (Mathieu et al., 2024), which trains more seeds only while the answer is unclear, and ajax.compare.estimate gives one agent's metrics confidence intervals that stay valid however many seeds are added.
• Reproduction: Ajax agents trained beside a reference implementation (Stable-Baselines3, Brax, the authors' code) on the same JAX environment, run on demand from benchmarks/reference, under a written protocol; the results so far cover PPO and PQN.
• Checkpoints and resuming: pass a state back to train further, or save it and resume in another process; the tests check that a split run matches the run left whole.
• Tests: unit tests check one component against an exact contract with no learning, written once for every agent or extension that shares the contract. Probes are tiny generic environments with known answers: each agent trains on them briefly, and each probe checks one capability of the whole agent. A catalogue of planted bugs lists the probes that should catch each one; CI checks that every bug still applies to the code, and planting them to run the probes is done by hand. Conformance checks run one update of an agent and of its original implementation from the same state and compare them (PQN, TD3, TD3+BC, REDQ and AVG so far).
• Performance: seeds vectorised with jax.vmap, a persistent compile cache, one per machine and JAX version, and a per-agent benchmark script, run by hand, that flags slowdowns over 10% against a recorded baseline.

Documentation

The site, with search and navigation, is at yannberthelot.github.io/Ajax. The links above point at the Markdown in this repository, which also reads on GitHub.

Related projects

• TargetGym: JAX environments for reaching and holding a setpoint, with tuned PID and MPC baselines. A scenario of Ajax's benchmarks/perf_bench.py trains on one of its aircraft tasks (uv sync --group plane).
• ProbingEnvironments: small environments with known answers, used by Ajax's probing tests.
• gymnax, Brax and MuJoCo Playground: the environment libraries Ajax trains on.
• flashbax: the replay buffers of SAC, ASAC, REDQ, TD3, DQN, IQL and DAgger (UDRL and the world models keep their own).

Contributing

Bug reports, new agents and new extensions are welcome. After installing, run uv run pre-commit install once, so ruff and mypy check every commit. CONTRIBUTING.md explains where the code lives, how to add an agent or an extension, and how the tests check them. Every pull request runs pre-commit, the tests, the conformance checks and a strict build of the documentation site in CI; the slow tier, every test that trains, also runs daily on main, on a pull request labelled run-slow-tier, and before each release. make ci runs the checks locally (about an hour of CPU).

Citation

@misc{ajax2025,
  title        = {Ajax: Reinforcement Learning Agents in Jax},
  author       = {Yann Berthelot},
  year         = {2025},
  url          = {https://github.com/YannBerthelot/Ajax},
}

License

MIT, see LICENSE.