Skip to content

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.

Compared with Stable-Baselines3 and its JAX version SBX, Ajax 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. Choosing the pieces names each method and its paper: TD3+BC, Q-filtered cloning, kickstarting, residual RL, JSRL, IBRL, RLPD's prior data, IQL, CQL, Cal-QL, ReBRAC, GAIL, DAgger.
  • 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.

Ajax has eighteen agents, from SAC and PPO to DreamerV3 and TD-MPC2. Each implements a published paper, and its page lists where it departs from it. 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; Compatibility shows which agent runs which.

git clone https://github.com/YannBerthelot/Ajax.git
cd Ajax
uv sync

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["Eval/episodic mean reward"])  # one row per seed, one column per evaluation

Quickstart → Why Ajax → Choose an algorithm →

The agents

Kind Agents
Off-policy actor-critic SAC, ASAC (average reward), REDQ, CrossQ, TQC, TD3
Offline, then online IQL
Incremental, no replay AVG
Value-based, discrete actions DQN, PQN
On-policy PPO, APO (average reward)
Supervised, command-conditioned UDRL
Imitation of an expert, no reward DAgger
Differentiable simulation APG
World models DreamerV3, TD-MPC2 and its multi-task variant

What else is in the box

  • Extensions: expert guidance, exploration, critic-target changes, actor and critic pre-training, offline RL, adversarial imitation and diagnostics, each usable with the agents that run the training phases it touches.
  • Training offline: a replay agent trains on a transition dataset before it acts, then fine-tunes online.
  • Expert guidance: learn from an existing controller, which reads the environment's raw observations and answers in its action units.
  • Memory: GRU, LSTM, transformer or Mamba blocks, with R2D2-style sequence replay.
  • PID policy: a PID controller as the actor, whose learned gains deploy as a plain PID.
  • Differentiable simulation: APG through gymnax dynamics, across distributions of physical parameters.
  • World models: DreamerV3 and TD-MPC2, following the paper-era official code.
  • Environments: gymnax, Brax and MuJoCo Playground by id, with wrappers for normalisation, observation noise, reference tracking and more; every continuous agent's [-1, 1] actions are mapped onto the environment's bounds.
  • Logging and checkpoints: evaluations as arrays, W&B or TensorBoard, resuming in another process.
  • Performance: every seed in one program, a persistent compile cache, one per machine and JAX version.
  • Comparing agents: AdaStop, which trains more seeds only while the answer is unclear, and confidence intervals for one agent's metrics.
  • Reproduction: Ajax agents trained beside reference implementations on the same JAX environment, under a written protocol.

Documentation

Ajax is MIT-licensed. TargetGym, its sibling project, provides JAX environments for holding a setpoint.