Training loop¶
The agents share one loop, TrainLoop: it collects experience, waits for
learning_starts, folds the extensions, evaluates, logs and resumes the same
way for every agent. An agent supplies only its init and update.
Architecture explains how the pieces fit.
ajax.agents.loop ¶
The training loop the agents share.
An agent's make_train builds a :class:TrainLoop and hands it what
makes the agent its algorithm: how to initialise its state and how to
update it. The loop owns everything else, the same for every agent::
fresh run: state = init(key), then every extension's initial state
and one-shot pretraining
resumed: state = the checkpointed state
iterate: collect one step per env (off-policy) or a rollout (on-policy)
if timestep >= learning_starts (on-policy: always):
state, aux = update(state, experience)
fold the extensions' post_update
else:
aux = the update's metrics, filled with NaN
evaluate and log every log_frequency steps
An off-policy agent that acts, trains and evaluates its own way (the world
models: a planner or a filtered latent acts through the row collector, a
train ratio schedules the updates) supplies those steps too, and the loop
runs them in the same order (:meth:TrainLoop.off_policy)::
iterate: state, experience = collect(state, tick)
repeat n_updates(tick, state) times:
state, _ = update(state, experience)
fold the extensions' post_update
evaluate and log every evaluation.every ticks
A replay agent given a dataset trains offline instead
(:meth:TrainLoop.offline)::
fresh run: the replay and the normalisers start from the dataset
iterate: state, aux = update(state, None) # no environment step
fold the extensions' post_update
evaluate and log every log_frequency updates
The result is :func:ajax.perf_utils.build_resumable_train's jitted
train(key, index, initial_state, resume_from_state), which
:meth:ajax.agents.base.ActorCritic.train vmaps over seeds. It returns
(state, None) without a logging config, the evaluation not even traced
(tracing alone rebuilds the evaluation env). With one it evaluates every
log_frequency steps, whether or not a backend (W&B, TensorBoard) records
the logs, and returns (state, evaluations): every logged key's values at
the evaluations, in order, padded to the most a call makes with -1 (integer
leaves, the timestep) and NaN. They are written into rows the scan carries,
not stacked per iteration, which would cost seeds x iterations memory.
Evaluation
dataclass
¶
An agent's own evaluation, every every ticks (None: never).
metrics(agent_state, aux) are its training metrics (timestep
included), updated with its policy's evaluate(agent_state, key)
when it has one and the extensions' eval_metrics;
after_log(agent_state, logged) runs on every tick after the gate.
TrainLoop
dataclass
¶
TrainLoop(
env_args,
total_timesteps,
num_episode_test,
stack,
log,
log_fn,
log_frequency,
start_timestep=0,
end_timestep=None,
expert_policy=None,
offline_data=None,
)
One call's environment, budget, extensions and logging (compared and hashed by identity: the environment's parameters are arrays).
A call trains from start_timestep (a resumed run's timestep, else 0)
to end_timestep (None: the horizon); total_timesteps is the
run's horizon: every schedule (the extensions' total_steps, the
training fraction, the logging's end) runs over the whole run, so a run
split into calls follows the schedules of one call; the call's own
iterations come from its :attr:budget.
n_evaluations
property
¶
The most evaluations one call ending at its budget makes
(:meth:evaluations_until).
create
classmethod
¶
create(
env_args,
total_timesteps,
num_episode_test,
run_ids=None,
logging_config=None,
extensions=(),
start_timestep=0,
end_timestep=None,
expert_policy=None,
offline_data=None,
)
The loop of one make_train call; starts the logging worker
when logging_config names a backend.
warn_if_partial ¶
Warn when the budget is not a whole number of iterations of
per_iteration env steps (unit names them): the iterations
stop at a multiple of it, not at the call's end, so the run trains
another number of steps than asked and an evaluation due at that end
never happens (its row stays NaN).
n_rollouts ¶
The iterations of :meth:on_policy for n_steps rollouts over
timesteps (None: the budget, warning when it is not whole
rollouts): in whole rollouts, plus one (as PPO always ran).
evaluations_until ¶
The most evaluations one call reaching timestep makes:
ajax.log's gate counts each one as a log point, so the k-th of
the run comes at timestep k * log_frequency or later, and none
comes past the horizon.
collect_kwargs ¶
:func:collect_experience's arguments for this run. With the
agent's expert, its action pipeline also stores the expert's actions
from the expert's tracked state (:func:expert_collect_kwargs).
maybe_update ¶
update and fold post_update once learning_starts steps
are collected; before, the metrics have the update's structure,
filled with NaN (which the logger drops).
repeat_update ¶
n times update then post_update; the updates keep what
they log on the state. An n derived from unbatched values (the
tick, static hyperparameters) keeps the loop one while under the
seed vmap; one read from the per-seed state runs it to the
largest count, each seed's state frozen after its own.
evaluate_and_log ¶
Evaluate and log every log_frequency steps
(:func:ajax.log.evaluate_and_log), the extensions' eval metrics
included; returns the state and the metrics (None without
logging).
evaluate_every ¶
An agent's own evaluation and logs on the ticks t with (t +
1) % evaluation.every == 0 (:func:ajax.log.maybe_eval_and_log,
the extensions' eval metrics included); returns the state and the
metrics (None without logging).
continued ¶
A resumed state, its training fraction measured against the run's horizon (a fresh state is built with it).
train ¶
train(
init,
iteration,
num_updates,
last_rollout=None,
evaluations=None,
update=None,
buffer=None,
)
The resumable train function: init (fresh runs only), then
num_updates scan iterations iteration(agent_state, index,
tick), recording the evaluations when logging (see the module
docstring): :attr:n_evaluations rows, or evaluations.
last_rollout, (length, collect_kwargs), pre-allocates
agent_state.last_rollout for an agent exposing its rollouts.
The extensions' pretraining gets the env, update and the replay
buffer (:class:ajax.extensions.base.ExtensionContext).
off_policy ¶
off_policy(
init,
update,
learning_starts=0,
*,
uniform_steps=None,
recurrent=False,
expose_rollout=False,
after_collect=None,
collect_kwargs=None,
eval_kwargs=None,
collect=None,
n_updates=None,
evaluation=None,
num_ticks=None,
)
Train on one environment step per env per iteration (a tick).
Each tick collects a step (uniform actions before uniform_steps,
by default learning_starts; 0: the policy's from the first step,
as TQC's authors act), applies after_collect(agent_state,
transition) when given, then from learning_starts runs
update(agent_state, transition), then evaluates the actor
(:meth:evaluate_and_log). expose_rollout keeps the step on
agent_state.last_rollout as a T = 1 rollout. The budget is
:attr:budget // n_envs ticks.
An agent acting, training and evaluating its own way (the module
docstring) supplies those steps: collect replaces the collection
and the arguments configuring it; n_updates(tick, agent_state)
updates (of the absolute, unbatched tick: a train ratio, a burst;
or of the collected state: the env steps the tick took) replace the
learning_starts gate (:meth:repeat_update); evaluation
replaces the actor's evaluation (:meth:evaluate_every);
num_ticks the budget.
With :attr:offline_data it trains offline instead (:meth:offline).
offline ¶
Train on :attr:offline_data alone, with no environment step.
A fresh run's replay and normalisers start from the dataset's rows
(:func:seed_replay). Each tick runs update(agent_state, None)
on that replay, folds post_update and counts the update in the
collector's timestep: the schedules run over the updates, and the
evaluations (:meth:evaluate_and_log) come every log_frequency
updates and log the updates done as their timestep. The budget is
:attr:budget updates.
on_policy ¶
on_policy(
init,
update,
n_steps,
*,
recurrent=False,
expose_rollout=False,
collect_kwargs=None,
eval_kwargs=None,
)
Train on an n_steps rollout per env per iteration.
Each iteration collects the rollout from the state start
(:func:collect_experience with collect_kwargs), then runs
update(agent_state, rollout, start) and folds post_update;
expose_rollout keeps the rollout on agent_state.last_rollout.
The budget is :meth:n_rollouts; the last one may reach past the
call's end, and evaluates there when the horizon is further.
record ¶
Write metrics into row count of rows when they are an
evaluation (their timestep is not the -1 sentinel), else nothing.
Returns the rows and the count of evaluations written.
fresh_state ¶
A fresh run's state: init(init_key) on split(key), then the
extensions' initial states and one-shot pretraining on
extension_key (:meth:ExtensionStack.fold_init; context: the
pretraining's static context fields).
with_obs_norm ¶
An agent normalising its own observations (the collector's
obs_norm_info) with those statistics mirrored on its networks, which
apply them; collection refreshes the mirror at every step, and setting
it at init keeps the scan carry's structure fixed.
seed_replay ¶
agent_state with data's rows in its replay, where they stay
(:func:ajax.buffers.dataset.fill_replay), and in the normalisers it
has, each row counted once as collection counts its own: its own
observation statistics, mirrored on its networks, and the env's reward
statistics, of the reward or, with the normaliser's gamma, of the
discounted return of the episode so far (its running return left as it
is).
fold_post_update ¶
Fold the extensions' post_update at step on a fresh key split
from agent_state.rng (no key is drawn without extensions).
when ¶
fn(state) where pred holds, else state unchanged: a policy
delay, a target refresh, a warm-up.
fn returns the new state, or the new state and its metrics. Skipped,
the metrics have fn's structure, read from its trace (with the names
the extensions log), filled with fill: by default NaN, -1 in integer
leaves (:func:ajax.log.unevaluated), which the logger drops. Under the
seed vmap a pred read from per-seed state runs fn every time
and selects; one from unbatched counters runs it only where it holds.
extension_key ¶
The extensions' key of a step on key, which the step's own draws
(on key and the keys it splits off it) never get.
gradient_step ¶
One optimiser step down loss_fn(params) -> (loss, aux).
Returns the stepped train state, the loss's aux and the global norm
of the gradients, before any clipping the optimiser does (a caller that
does not log it leaves it to dead-code elimination).
fold_on_target ¶
fold_on_target(
stack,
agent_state,
batch,
target_q,
key,
total_timesteps,
*,
rewards,
dones,
gamma,
reward_scale,
target_extras=None,
)
target_q with the extensions' on_target folded in, held
constant, and their metrics: :func:critic_step's first half, which an
agent whose target comes out of its own critic pass (CrossQ) calls
inside its loss.
batch is a :class:Transition. rewards, dones (where the
bootstrap is cut), gamma (None for an average-reward agent) and
reward_scale are what the target used; target_extras join the
extensions' batch (TD3's smoothed next actions). The fold draws on a key
derived from key, the step's (:func:extension_key).
fold_critic_loss ¶
The extensions' critic_loss terms at the critic's params (the
differentiated ones) on batch against the held target_q, and
their metrics: :func:critic_step's second half, on a key derived from
key, the step's.
critic_step ¶
critic_step(
agent_state,
batch,
target_q,
value_loss,
extension_stack,
key,
total_timesteps,
**target_args,
)
A replay agent's critic step on batch (a :class:Transition).
Folds the extensions' on_target into target_q, held constant
(:func:fold_on_target, which target_args reach: what the target
used), then steps the critic down value_loss(params, target_q) ->
(loss, aux) plus their critic_loss terms
(:func:fold_critic_loss). The folds draw on keys derived from key,
the step's, on which the agent may draw too. Returns the stepped critic
state, aux, the gradients' norm (:func:gradient_step) and the two
folds' metrics.
actor_step ¶
A replay agent's actor step.
Steps the actor down policy_loss(params) -> (loss, aux, pi_mean)
(pi_mean: the policy's mean action at the rows, None for a discrete
policy) plus the extensions' actor_loss terms, folded on a key
derived from key, the step's (:func:extension_key). Their batch is
batch, the rows' entries the agent passes (observations,
raw_observations, actions, a_expert, critic_state), with
the differentiated actor_params, the actor_state, pi_mean and
agent_loss.
Returns the stepped actor state, aux, the gradients' norm
(:func:gradient_step) and the fold's metrics.
ajax.perf_utils ¶
Shared JAX-pure performance helpers for Ajax agents.
Patterns the perf audits found across agents (PERFORMANCE_LOG.md),
written once so a fix lands in one place instead of in every agent's
train_*.py.
build_resumable_train ¶
build_resumable_train(
*,
init_fn,
make_scan_fn,
num_updates,
resume_transform=None,
carry_out=None,
)
Build the canonical init-or-resume + lax.scan inner train fn.
Every Ajax agent shares this skeleton, written once here so it cannot
drift between agents: build a fresh agent_state with init_fn
(its one-shot initialisation included, so a continued run does not
re-run it) or, with resume_from_state=True, reuse the
initial_state handed in from a checkpoint; then jax.lax.scan
the per-iteration body num_updates times, its input the iteration
index. Each agent's make_train keeps its own setup (networks,
optimisers, action pipeline, logging) and hands the pieces over.
The returned train function has the signature
train(key, index=None, initial_state=None, resume_from_state=False,
iteration_offset=0, shared=None)
and is jitted with resume_from_state static (the init-fresh vs
load-from-checkpoint branch is dead-code-eliminated) and
initial_state donated: XLA reuses the incoming agent-state buffers
on the resume path, saving one agent state of peak memory (donating
None, the init path, is a no-op).
The scan input is the iteration index iteration_offset +
arange(num_updates): a resumed run passes the number of iterations
already done so the body sees absolute iteration indices and every
schedule computed from them (seed phase, update bursts, static resets,
logging cadence) continues instead of restarting. A passed offset is a
traced scalar (one compilation serves every offset) and must be
unbatched -- pass it through jax.vmap with in_axes=None -- so
the index, and every predicate derived from it, stays unbatched under
the seed vmap. The default (the Python int 0, not passed) adds nothing
to the program.
shared is an optional pytree the body reads but never carries -- an
offline dataset, say -- handed to make_scan_fn as the keyword
argument shared (only when it is not None, so builders that do
not take it are unchanged). Pass it through the seed vmap with
in_axes=None: it then enters the compiled program once, as an
argument shared by every seed. Closing over it instead would bake it into
the program as a constant (HLO bloat, a recompilation per dataset), and
a batched argument would copy it per seed.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
init_fn
|
Callable[[Any, Any], Any]
|
|
required |
make_scan_fn
|
Callable[..., Callable[..., Tuple[Any, Any]]]
|
a builder
|
required |
num_updates
|
int
|
number of scan iterations (static int). |
required |
resume_transform
|
Callable[[Any, Any], Any] | None
|
optional one-shot |
None
|
carry_out
|
Callable[[Any, Any], Any] | None
|
optional |
None
|
final_aux_scan ¶
Run lax.scan and expose only the final-step aux from the body.
Replaces the common idiom
carry, ys = jax.lax.scan(body, init, xs, length=length)
last_aux = jax.tree.map(lambda x: x[-1], ys)
with a carry-only variant that never materializes the leading scan
axis on device for the aux. Critical under vmap-over-seeds because
the discarded ys dimension would otherwise grow as
[seeds, length, *aux_shape] purely to be thrown away.
The body maps (carry, x) -> (new_carry, aux), aux a pytree of
arrays (any structure).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
xs
|
Any
|
optional pytree of per-step inputs; if None, the body
receives |
None
|
Returns:
| Type | Description |
|---|---|
|
|
|
|
return value on the final iteration only. |
ajax.schedule ¶
Schedules: hyperparameters that change over training.
linear ¶
From start at timestep 0 to end at steps (at least one),
then end.
warmup_cosine_schedule ¶
warmup_cosine_schedule(
peak_value,
warmup_steps,
total_steps,
end_value_fraction=0.1,
init_value_fraction=0.0,
)
Linear warmup to peak_value then cosine decay to end_value.
An optax step schedule (step -> learning_rate) usable as
OptimizerConfig.learning_rate. The warmup-cosine profile is the
one the in-context controller reference code trains with (Busetto et
al. 2024, GPT-style: 5k warmup steps, decay to peak / 10).
warmup_steps=0 gives pure cosine decay from peak_value.
Agent state and configuration¶
ajax.state.BaseAgentState ¶
Bases: Replaceable
next_keys ¶
(state, *keys): n fresh keys and the state whose rng
has moved past them (*keys, rng = split(rng, n + 1)).
ajax.state.BaseAgentConfig ¶
Bases: Replaceable
ajax.state.NetworkConfig ¶
Bases: Replaceable
ajax.state.OptimizerConfig ¶
Bases: Replaceable