Ajax

Fast, modular reinforcement learning in JAX.
The environment, the agent and every seed run together in one program, on a CPU or a GPU, and agents are built from parts that snap together.

Install · Quickstart · Agents · Expert and data · Reproduction · Documentation

Ajax has nineteen agents, each from a published paper. Compared with Stable-Baselines3 and SBX, it offers:

• Fast and parallel: the environment steps beside the agent on a CPU or a GPU, with no Python in the loop, so many seeds and environment copies train at once on one device. The benchmarks run on a CPU; GPU speed is not yet measured.
• Modular: GRU, LSTM, transformer and Mamba memories, CNN encoders, sequence replay with burn-in, a PID policy and thirty-four extensions, combined in the agent's constructor. The compatibility table shows which parts go together.
• Multi-paradigm: discounted and average-reward agents (ASAC, RVI-SAC, APO), model-free and model-based, online and offline, all behind one interface.
• Expert-guided: cloning, imitation, prior data and offline training.
• World models, ready to use: DreamerV3 and TD-MPC2 are among the hardest agents to get right, and their papers run them on over a hundred tasks with one set of hyperparameters. Ajax checks them against the official code, update by update.

Why Ajax compares it with SB3, SBX, PureJaxRL, Brax, Stoix, Rejax, TorchRL and Tianshou, including when to use them instead.

Install

pip install ajax-rl                 # imports as ajax; Python 3.12 or newer
pip install "ajax-rl[playground]"   # adds Brax and MuJoCo Playground environments
pip install "ajax-rl[cuda]"         # adds CUDA, on Linux with an NVIDIA GPU

ajax-rl[brax] adds Brax alone; Installation covers the other extras and uv sync from a clone.

Quickstart

from ajax import PPO
from ajax.logging import LoggingConfig

agent = PPO("CartPole-v1")  # 4 environments x 2048-step rollouts: one update per 8,192 steps
# Evaluate after every update; backends=() keeps the results in memory only.
logging = LoggingConfig(log_frequency=8_192, backends=())
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, in about fifteen seconds on a laptop CPU, compilation included. Parts combine in the constructor:

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", extensions=(ConditioningMetrics(),))

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

Agents

• Off-policy actor-critic: SAC, TD3, REDQ, CrossQ (batch-renormalised, no target networks), TQC (truncated quantile critics), AVG (no replay), ASAC and RVISAC (RVI-SAC; average reward), IQL (offline)
• On-policy: PPO, APO (average reward)
• Value-based: DQN, PQN
• World models: DreamerV3, TDMPC2, TDMPC2MultiTask (multi-task, offline)
• Supervised: DAgger (imitation of an expert), UDRL (command-conditioned)
• Differentiable simulation: APG

DQN and PQN take discrete actions; PPO, APO, DAgger, UDRL and DreamerV3 take both; the others continuous ones. SAC, TD3, REDQ, ASAC and DQN also train offline (train_offline). Each name opens the agent's page, with its paper; the compatibility table lists every extension per agent.

Learning from an expert or from data

• Before training: behaviour cloning (CloningPretrain), expert transitions in the replay (ExpertPrefill), prior data as in RLPD (PriorData).
• During training: TD3+BC's actor term, kickstarting and Q-filtered cloning (ImitationLoss), IBRL (IBRL), residual RL (ResidualPolicy), JSRL hand-over (JSRLCurriculum), GAIL (AdversarialReward), and the DAgger agent.
• Instead of online data: agent.train_offline(dataset, n_updates), then online fine-tuning if wanted: IQL, CQL and Cal-QL (ConservativeQ, CalQL), ReBRAC (TargetBC), TD3+BC; TD-MPC2's multi-task agent (TDMPC2MultiTask) on pooled datasets.

Choosing the pieces gives each method's paper and the agents that take it.

Beside reference implementations

Ajax's PPO beside Stable-Baselines3's and its PQN beside purejaxql's, on the same JAX environment, 96 seeds each; both comparisons find the returns similar within the registered margin. Results has two more comparisons (one inconclusive) and the speed of each side; the protocol says how they are run.

Also in the box

The documentation site has the user guide, one page per agent, the API reference and the developer guide.

Related projects

gymnax, Brax and MuJoCo Playground provide the environments, flashbax most replay buffers. TargetGym has setpoint-tracking tasks with PID and MPC baselines, and ProbingEnvironments the small environments with known answers that Ajax's probes train on.

Contributing, citing, license

Bug reports, new agents and new extensions are welcome; see CONTRIBUTING.md. Ajax is MIT-licensed (LICENSE).

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