Skip to content

Contributing to Ajax

This guide is for changing Ajax: where the code lives, how an agent trains, how to add an agent or an extension, and how the tests check them. To use Ajax, read the README. CLAUDE.md holds the checks a change must pass and the refactoring rules; this guide does not repeat them.

Development setup

Install as the README says, then run uv run pre-commit install once, so ruff and mypy check every commit. Prefix commands with uv run, or activate .venv.

Where things live

Architecture gives the overview. These are the pieces every agent shares:

src/ajax/
├── agents/base.py       ActorCritic: environment and network setup, extensions=, train()
├── agents/loop.py       TrainLoop, the shared loop; critic_step, actor_step, gradient_step
├── agents/recurrent.py  replay buffers and R2D2-style sequence replay for memory=
├── agents/cloning.py    the behaviour-cloning loss and fit (CloningPretrain's)
├── extensions/base.py   Extension, ExtensionStack, ExtensionContext, the phases
├── environments/interaction.py  collect_experience: act, step and store
├── perf_utils.py        build_resumable_train: initialise or resume, then one jitted scan
├── log.py, evaluate.py  when to evaluate and log; the evaluation episodes
├── state.py             BaseAgentState, BaseAgentConfig, the configuration dataclasses
└── checkpoint.py        save_checkpoint, restore_into

Agent anatomy. Each folder src/ajax/agents/<AGENT>/ holds <AGENT>.py, the public class (the paper's hyperparameters and extensions=; it names its make_train and fills train_kwargs, which get_make_train() binds); train_<AGENT>.py, the algorithm and a make_train that hands its initialisation and update to TrainLoop; and state.py, <AGENT>State and <AGENT>Config. SAC, PPO and TDMPC2 also have a core.py: the maths their descendants import instead of copying (ASAC, REDQ and AVG from SAC; APO from PPO; TDMPC2MultiTask from it).

The shared training loop

The agents share one loop, TrainLoop, so collecting experience, waiting for learning_starts, applying the extensions, evaluating, logging and resuming work the same way everywhere and are fixed in one place. An agent supplies only its algorithm:

  • init(key) -> agent_state: a fresh state. One-shot pretraining (behaviour cloning) is an extension's, on a key the loop splits off.
  • update(...) -> (agent_state, aux): one update. aux holds the metrics as flax dataclasses nested exactly one level, as TD3's AuxiliaryLogs(policy=..., value=...), logged as value/critic_loss; the logger expects that shape.

and picks one of two iterations:

  • loop.off_policy(init, update, learning_starts, ...) (SAC, TD3, REDQ, ASAC, AVG, DQN): each iteration collects one step per environment, then, from learning_starts, runs update(agent_state, transition) and folds the extensions' post_update. Before learning_starts the actions are uniform and the metrics have the update's structure, filled with NaN, which the logger drops.
  • loop.on_policy(init, update, n_steps, ...) (PPO, APO, PQN): each iteration collects an n_steps rollout per environment from the state start, then runs update(agent_state, rollout, start) and folds post_update. A recurrent agent replays the rollout from start's memory.

The loop does the rest: init and the extensions' initial state and pre-training (ExtensionStack.fold_init) on a fresh run, the given state on a resumed one, then the iterations in build_resumable_train's jitted scan, evaluating and logging every log_frequency steps (with eval_metrics).

Continuous agents act in [-1, 1]; never scale to the env's bounds yourself. Every Ajax env step maps the action onto them (agent_action_to_env), and ActorCritic reads an expert, written in env units, in the same [-1, 1] (ExpertInAgentUnits): replay, losses and extensions see one unit.

The world models (DreamerV3, TDMPC2) also give off_policy their own collect(agent_state, tick), n_updates(tick, agent_state) (such as a train ratio) and Evaluation. UDRL runs its own iteration on TrainLoop.train; APG and the offline TDMPC2MultiTask call build_resumable_train themselves, with fresh_state and the post_update fold.

ActorCritic.train runs the seeds at once (jax.vmap). With a LoggingConfig it returns (state, evaluations), each logged key's values per seed and evaluation, kept whether or not a backend (use_wandb, use_tensorboard, the only case that starts the logging worker) records them. Without one it does not evaluate and returns (state, None); APG returns every update's metrics instead. initial_state=state continues a run.

Extensions

An extension is a research feature (an expert to learn from, an exploration rule, an extra loss term, a measurement) that plugs into any agent without changing it, so an agent's surface stays its algorithm's hyperparameters plus extensions=() (see CLAUDE.md, "Extensions are self-contained"). The extensions guide lists the ones Ajax ships.

An extension is a frozen dataclass subclassing Extension. It holds its own settings and overrides only the phases it touches; the other phases keep their defaults, which change nothing. An agent folds its ExtensionStack through a phase in list order, mostly with a stack.fold_<phase>(...) helper (fold_init, fold_on_target, fold_critic_loss, fold_actor_loss, fold_post_update, fold_eval_metrics) that builds the ExtensionContext (step, rng, total_steps) and threads the extensions' state. On an empty stack a helper changes nothing (a loss term is 0.0), so an agent without extensions traces nothing extra and never guards a fold.

Phase What it does Folded today by
init_state(agent_state, rng) returns the extension's state ((): none) every agent, fresh runs only
pretrain(agent_state, ext_state, ctx) one-shot step before training every agent, fresh runs only
on_target(agent_state, ext_state, batch, target, ctx) transforms the TD or value target SAC, REDQ, ASAC, AVG, TD3, DQN, PQN, PPO, APO
critic_loss(agent_state, ext_state, batch, ctx) adds a term to the critic loss SAC, REDQ, ASAC, AVG, TD3, DQN, PQN, PPO, APO
actor_loss(agent_state, ext_state, batch, ctx) adds a term to the actor loss SAC, REDQ, ASAC, AVG, TD3, PPO, APO, APG, UDRL
action(agent_state, ext_state, obs, rng, ctx) overrides the collection action (None defers) SAC with an expert_policy, for extensions with an action_slot
eval_action(agent_state, ext_state, obs, rng, ctx) overrides the evaluation action no agent yet
post_update(agent_state, ext_state, ctx) runs after each update every agent
eval_metrics(agent_state, ext_state, rng, ctx) adds metrics to each evaluation every agent

tests/probing/test_extensions.py checks the phases on the agents with test-only extensions of known effect, and that an extension implementing a phase the agent does not fold is rejected.

Every agent lists the phases it folds in supported_extension_phases: the shared loop's LOOP_PHASES (pretrain, post_update, eval_metrics, in ajax.agents.loop) plus those its update folds; the base class's default is empty. An extension implementing any other phase is rejected when the agent is built instead of being silently ignored.

What changes during training lives in agent_state.ext_state (one entry per extension), never on self: extensions are static arguments of the compiled program, so they stay hashable and unchanged. An extension that needs what the agent builds (environment, network config, buffer, discount) overrides bind_to_agent(**agent_context) to return a new instance holding it, as MCPretrain does; SAC, TD3, REDQ, APO and DreamerV3 call it. The batch a phase receives is a dictionary: the keys every agent puts in it are listed in ajax/extensions/base.py; read the agent's fold call before relying on any other.

Adding an agent

Boilerplate goes in the shared code; the algorithm goes in one file, so a reader can follow train_<AGENT>.py like the paper's pseudocode. Lineage rule: an agent descending from another (REDQ from SAC) imports the parent's maths from its core.py. Use critic_step (a replay critic step), actor_step (an actor step) and gradient_step (an optimiser step) from agents/loop.py rather than writing your own.

These templates for an agent FOO follow the replay agents TD3 and REDQ; an on-policy agent follows APO instead (no buffer, loop.on_policy).

1. state.py. BaseAgentState already holds the random key, the actor, critic and collector states, the update count and the extensions' state. A replay agent's config extends RecurrentReplayConfig, a BaseAgentConfig for sequence replay; kw_only lets fields without defaults follow others.

@partial(struct.dataclass, kw_only=True)  # Partial from jax.tree_util
class FOOState(BaseAgentState):
    """FOO carries only actor + critic + collector."""


@partial(struct.dataclass, kw_only=True)
class FOOConfig(RecurrentReplayConfig):
    gamma: float
    learning_starts: int = 100
    reward_scale: float = 1.0

2. train_FOO.py. The paper's maths, then make_train. The critic step goes through critic_step, which folds on_target and critic_loss; its keywords are what the target used (dones: where it cut the bootstrap, as TD3's batch.terminated). The actor step goes through actor_step, which folds actor_loss; its keywords are the rows the extensions read:

def update_value_functions(agent_state, batch, agent_config, extension_stack, total_timesteps):
    agent_state, key = agent_state.next_keys()  # key, rng = split(rng)
    target_q = compute_foo_td_target(agent_state, key, batch, agent_config)  # the paper's target

    def value_loss(params, target_q):  # -> (loss, aux)
        return value_loss_function(params, agent_state.critic_state, batch.obs, batch.action, target_q)

    critic_state, aux, grad_norm, metrics = critic_step(
        agent_state, batch, target_q, value_loss, extension_stack, key, total_timesteps,
        rewards=batch.reward, dones=batch.terminated,
        gamma=agent_config.gamma, reward_scale=agent_config.reward_scale,
    )
    return agent_state.replace(critic_state=critic_state), aux, metrics


def update_policy(agent_state, batch, extension_stack, total_timesteps):
    agent_state, key = agent_state.next_keys()

    def policy_loss(params):  # -> (loss, aux, pi_mean)
        return foo_policy_loss(params, agent_state, batch.obs)  # the paper's actor loss

    actor_state, aux, grad_norm, metrics = actor_step(
        agent_state, policy_loss, extension_stack, key, total_timesteps,
        observations=batch.obs, actions=batch.action,
    )
    return agent_state.replace(actor_state=actor_state), aux, metrics


def make_train(
    env_args, actor_optimizer_args, critic_optimizer_args, network_args, buffer,
    agent_config, total_timesteps, num_episode_test,
    run_ids=None, logging_config=None, start_timestep=0, extensions=(),
):
    loop = TrainLoop.create(
        env_args, total_timesteps, num_episode_test, run_ids, logging_config, extensions,
        start_timestep=start_timestep,
    )

    def init(key):
        return init_FOO(key, env_args, actor_optimizer_args, critic_optimizer_args, network_args, buffer)

    def update(agent_state, _transition):
        # The step just collected is in the buffer: FOO samples it from there.
        return update_agent(agent_state, buffer, agent_config, loop.stack, total_timesteps)

    return loop.off_policy(
        init, update, agent_config.learning_starts,
        collect_kwargs={"buffer": buffer},
    )

update_agent samples a batch (sample_replay), takes the critic step, then the actor step, and returns the state with an AuxiliaryLogs of both steps' metrics (the extensions' under extensions, merge_metrics of the two steps'). ActorCritic.train calls make_train with the keyword names above (start_timestep when it resumes a run), so keep them. A step takes its key from the state with agent_state.next_keys(), which splits rng and writes the new one back; critic_step and actor_step derive the extensions' keys from it (Randomness), so an agent's own random stream is the same with or without extensions.

3. FOO.py.

class FOO(ActorCritic):
    """FOO (Author et al., year) for continuous action spaces; defaults as in the paper."""

    name: str = "FOO"
    # The phases the loop folds, plus those FOO's update folds.
    supported_extension_phases: frozenset = LOOP_PHASES | {
        "on_target", "critic_loss", "actor_loss"
    }
    make_train = staticmethod(make_train)  # ActorCritic.get_make_train binds it

    def __init__(
        self, env_id: str | EnvType, n_envs: int = 1, gamma: float = 0.99,
        buffer_size: int = int(1e6), batch_size: int = 100, learning_starts: int = int(1e4),
        extensions: Sequence[Extension] = (),
    ) -> None:
        self.config = run_config(locals(), self.name)  # what train() logs with the run
        super().__init__(env_id=env_id, n_envs=n_envs, extensions=extensions)
        self.agent_config = FOOConfig(gamma=gamma, learning_starts=learning_starts)
        self.buffer = make_replay_buffer(
            self.agent_config, n_envs, self.network_args.memory, buffer_size, batch_size
        )
        # make_train's own keywords, beyond those every agent's takes.
        self.train_kwargs = {"buffer": self.buffer}

Do not add a keyword for a feature (use_critic_blend=True, ...): it belongs on an extension. Set supports_memory = True only once the agent trains with memory=; until then ActorCritic rejects it.

4. Export and document. Import it in src/ajax/__init__.py, add it to __all__, and add a row to the README's agent table.

5. Tests and benchmark. tests/agents/FOO/: the maths on hand-computed values. tests/agents/test_agent_contract.py: a row, so FOO builds, refuses what it cannot act in and folds extensions on tiny runs. tests/probing/: presets in agents.py, then the cases that fit (test_framework.py fails while an exported agent has none). tests/agents/test_agent_pins.py: a tiny run's fingerprint, so a later restructuring can show it changed nothing. benchmarks/agent_bench.py: an entry in AGENTS and a baseline (see Performance).

Adding an extension

Add an extension, never a flag or a keyword on an agent.

  1. Pick the phases from the table and check that the agents you target fold them. If one does not, adding the fold to that agent is a change of its own, and the agent still must not know about your extension.
  2. Write the dataclass in src/ajax/extensions/<topic>.py, its settings as fields. This one counts updates and logs the count; it runs as is:
from dataclasses import dataclass

import jax.numpy as jnp

from ajax import PPO
from ajax.extensions.base import Extension
from ajax.logging.wandb_logging import LoggingConfig


@dataclass(frozen=True)
class UpdateCounter(Extension):
    """Counts the updates and logs the count at each evaluation."""

    name: str = "update-counter"

    def init_state(self, agent_state, rng):
        return jnp.asarray(0)

    def post_update(self, agent_state, ext_state, ctx):
        return agent_state, ext_state + 1

    def eval_metrics(self, agent_state, ext_state, rng, ctx):
        # A float: the steps that do not evaluate log NaN in its place.
        return {"Counter/updates": ext_state.astype(jnp.float32)}


agent = PPO("CartPole-v1", n_steps=128, extensions=(UpdateCounter(),))
logging = LoggingConfig(config={}, log_frequency=1024, use_wandb=False)
state, evaluations = agent.train(seed=[0, 1], n_timesteps=4096, logging_config=logging)
print(state.ext_state[0])  # the updates each seed made: [9 9]
print(evaluations["Counter/updates"])  # per seed, at each evaluation: [2. 4. 6. 8.]
  1. Test it in tests/extensions/: the feature fires and does what its paper says (a metric appears, a term moves the loss the right way).

Callable hooks. A few Optional[Callable] keywords predate extensions and stay: TD3's action_pipeline (AjaxExperiments' TD3 variants), DQN's td_target_fn and td_loss_fn and PQN's td_loss_fn (Double DQN and Huber, which their agents' tests run). tests/modules/test_hook_composition.py pins them. Add none.

Testing

Two kinds of tests check the code, each answering its own question (Testing is the full guide):

  • Unit tests (tests/<area>/, tests/agents/<AGENT>/) check one component against an exact contract, with no learning. Where agents or extensions share a contract, the test is generic (tests/agents/test_*_contract.py, tests/extensions/test_phase_contract.py, ...): it runs every one of them, and a new one fails it until it joins.
  • Probing tests (tests/probing) train an agent briefly on a tiny generic environment whose answer is known, and check only what learning shows: the right value, the right action, the right count.

A convention that differs from a reference implementation (an ordering, a default, a constant) is for code review, not for a test. Pins record what a tiny run produces today, so a restructuring can show it changed nothing: tests/agents/test_agent_pins.py (tiny runs' fingerprints), tests/extensions/test_sac_extensions_equivalence.py and tests/test_downstream_api.py (the names downstream projects import). Slow tests (@pytest.mark.slow) are the long trainings.

The tests come in two tiers. The fast tier is the tests not marked slow: every pull request must pass it. A fast test may run a tiny training, of at most 128 gradient steps: tests/conftest.py counts each test's gradient steps (the larger of the agent's update count and each train state's step) and fails a fast test past the line (AJAX_FAST_MAX_UPDATES); a test that trains belongs in tests/probing, or is marked slow. It also times each test and lists those past 30 s (AJAX_FAST_MAX_SECONDS) without failing them: compiles, not trainings, set those times. The slow tier is everything that trains: tests/probing, which tests/conftest.py marks slow but for the harness's own checks, and the tests marked slow elsewhere; a workflow of its own runs it (below).

Every directory of tests/ that holds Python files is a package (an empty __init__.py; tests/test_repo_hygiene.py checks it), so tests import shared helpers by their full name: from tests.agents.extension_stubs import CounterExt. The benchmark scripts the tests import are on the path through pythonpath in pyproject.toml.

uv run pytest --collect-only -q   # everything imports and collects
uv run pytest tests/agents/TD3    # the tests of what you changed
uv run pytest -m "not slow" -n auto   # the fast tier, a worker per core
make ci       # the pull-request gate: ci-precommit, ci-test, ci-conformance, docs
make ci-all   # the gate and the slow tier (ci-slow)

The make targets hide the GPU (CPU_ENV) so tests never compete with a running experiment; elsewhere set JAX_PLATFORMS=cpu while a GPU is busy.

CI (.github/workflows/ci.yml) runs the fast tier on CPU for every pull request and push to main: Pre-commit; Tests (three shards: the tests not marked slow, the probing harness's checks among them, under coverage, a pytest-xdist worker per vCPU, at XLA's optimisation level 0 since these programs are tiny; the last shard runs whatever the first two do not name, so a new test file always runs); Coverage (the shards combined, at least 90%); Conformance (each agent's training step against its original implementation); Docs (the documentation site, mkdocs build --strict). Lint and Test, the one check branch protection requires, passes only when all of them passed. They run on GitHub's arm64 runners, which all come on one CPU model, so a job finds main's compile cache (JAX keys it by the CPU).

The Slow tier workflow (slow.yml) runs everything that trains (pytest -m slow) on x86, the platform most users train on, in six shards by recorded duration (--shard I/N, tests/.durations, make durations): daily on main when main changed, by hand, and on a pull request each time the label run-slow-tier is added. A pull request that changes how an agent trains gets the label before it merges. A red run on main opens an issue, and bisects the merges since the last green run when that is cheap.

As the gate runs no probe, every planted fault (tests/probing/faults/<family>.jsonl) names in its fast field the fast tests that catch it, measured by planting the fault in a scratch copy and running the fast tests that execute its lines (uv run python -m tests.fault_catches SCRATCH [FAMILY ...]); tests/probing/test_framework.py checks that each names live tests outside tests/probing. A new fault comes with its fast catch.

The probing tests

A unit test checks a contract; it cannot show that an agent learns what it should. A probe trains an agent on a tiny problem whose answer is known exactly (a value, a best action, a count) and compares what it learned with the right answer and named wrong ones (a discount applied twice), so a failure names the mistake.

envs.py builds the tiny problems, agents.py holds each agent's presets (pinned by DIGEST), runs.py trains every seed in one compiled program, readouts.py reads values and actions off a trained agent and oracles.py computes the right and wrong answers. In verdict.py, a Query is a reading's right and named wrong answers and a Case adds the budget and tolerances; a check trains seeds 0-7: 7 or more within tolerance pass, 4 or fewer fail, otherwise seeds 8-15 run and 13 of 16 must pass. faults.py and faults/<module>.jsonl plant known bugs in a scratch copy of the source, each naming the tests that must catch it. The test_*.py modules group the cases by subject (value chains, episode ends, bookkeeping, extensions, memory, world models, resuming a split run); test_framework.py tests the harness. A planted fault quotes the source it edits, so any change under src/ajax can break one: such a change runs tests/probing/test_framework.py locally before the push (text and collection only, no training).

A new case's budget and tolerances come from calibration on seeds no check uses (uv run python -m tests.probing.verdict ladder|choose|certify ...; verdict.py describes the procedure); then its module's ANSWER_DIGEST is re-pinned. A known bug stays visible as a strict xfail naming the defect, never as a looser tolerance. The directory is capped at 7,000 lines.

Performance

CLAUDE.md ("Performance no-regression guardrail") sets the rule; here is the check. benchmarks/agent_bench.py times a fixed run per agent; --compare fails when one is slower than --tol (10% by default) against a baseline. Capturing requires --out, so a quick check never overwrites the baseline.

export JAX_PLATFORMS=cpu
uv run python benchmarks/agent_bench.py --only TD3 --out benchmarks/my_change.jsonl
uv run python benchmarks/agent_bench.py --compare benchmarks/agent_baseline.jsonl benchmarks/my_change.jsonl

Dependencies and releases

Ajax is a library: pyproject.toml says which versions work, and uv.lock fixes the ones CI and development use. Users' installers never read the lock.

  • Floors only. Each lower bound is the oldest version Ajax works with. Raise it when Ajax needs something newer or an older version is known to break, with a comment saying why. Add an upper bound only after a known break, with its reason, and remove it once Ajax is fixed.
  • The ranges are tested. The Compatibility workflow (compat.yml) runs the tests not marked slow every Monday, and on pull requests that change the dependencies, twice without the lock: at the newest versions allowed, on the newest Python, and with every direct dependency at its floor (uv sync --resolution lowest-direct), on the oldest.
  • The lock is refreshed monthly by Dependabot (dependabot.yml), the JAX stack in one pull request.
  • Python: 3.12 and newer, the versions JAX supports. A version goes when JAX drops it.
  • Releases. Versions follow semantic versioning; before 1.0, a breaking change (a raised floor, a removed keyword) bumps the minor version. A release is an entry in CHANGELOG.md and a vX.Y.Z tag on main, which downstream projects can pin. Public names change by expand then contract (CLAUDE.md, rule 10). Tags come only from the Release workflow (release.yml, by hand, with the version): it tags main's head once the slow tier has passed there, running it first if needed, then publishes that commit to PyPI as ajax-rl.

Style

Pre-commit (.pre-commit-config.yaml) runs on every commit and in CI: ruff lint with rules I F E W B C RUF ARG (ARG, unused arguments, is off under tests/), ruff-format (settings in pyproject.toml), and mypy on everything but benchmarks/, run in the project's environment (uv run mypy) so it sees jax and flax.