Buffers¶

Off-policy agents store transitions in a flashbax replay buffer. With memory=, they store per-environment trajectories and replay sequences R2D2-style. The world models keep their own replay (DreamerV3, TD-MPC2).

ajax.buffers.utils ¶

ReplayBatch ¶

Bases: NamedTuple

One replay sample: its rows' fields, the optional ones None where the buffer does not store them (:func:init_buffer).

get_buffer ¶

get_buffer(
    buffer_size, batch_size, n_envs=1, sequence_length=None
)

Replay buffer factory: a trajectory buffer of buffer_size // n_envs rows per env, whose add takes the single-step (n_envs, ...) transitions the collector writes.

sequence_length=None (default) samples rows alone, uniformly over every stored one: a row holds its next observation, so it needs no successor (flashbax's flat buffer samples a row with its successor, which leaves each env's newest row out). When set (recurrent agents), sample yields contiguous per-env sequences of that length, shaped (batch_size, sequence_length, ...).

init_buffer ¶

init_buffer(
    buffer,
    env_args,
    add_train_frac=False,
    action_dim_override=None,
    expert_state_aug_dim=0,
    include_expert_fields=False,
    actor_carry_dim=0,
    include_raw_next_obs=False,
    include_raw_obs=False,
    include_is_expert=False,
    include_return_to_go=False,
    include_next_action=False,
)

Initialize the flashbax buffer state with correctly shaped dummy transitions.

Observations are flat (build_env_from_id flattens gymnax ones).

whether to append the train_time_fraction scalar as the

last observation dimension.

expert_state_aug_dim: extra trailing obs dims reserved for the flattened expert internal state when running with augment_obs_with_expert_state=True. 0 disables. include_raw_obs, include_raw_next_obs: also keep the environment's own observation, raw_obs, and next observation (the final one where an episode ended), raw_next_obs, for an expert to read. include_is_expert: also keep whether an expert acted, is_expert. include_return_to_go: also keep a dataset row's return to its episode's end, return_to_go (a collected row's is -inf: unknown). include_next_action: also keep a dataset row's next action, its episode's next row's, next_action (NaN where unknown).

A field is stored only for a reader: a row is about 40% raw observation, which a run without an expert never reads.

normalized_rewards ¶

normalized_rewards(reward, stats)

The env's own reward divided by a reward normaliser's statistics, which it leaves as they are; while they are empty (no step yet), the reward as it is.

get_batch_from_buffer ¶

get_batch_from_buffer(
    buffer, buffer_state, key, reward_stats=None
)

A replay batch. The replay keeps the env's own rewards; with reward_stats, a reward normaliser's current statistics (:func:ajax.environments.interaction.reward_statistics), they come divided by them, as SB3's replay buffer does under VecNormalize.

get_sequence_batch_from_buffer ¶

get_sequence_batch_from_buffer(buffer, buffer_state, key)

Sample contiguous sequences from a trajectory buffer (see get_buffer(sequence_length=...)) and return them time-major.

Returns a dict of arrays shaped (sequence_length, batch_size, ...), ready for the recurrent (scan-based) network path.

ajax.buffers.dataset ¶

Transition datasets: logged or recorded transitions, kept on disk and added to a replay.

A :class:TransitionDataset holds whole episodes as flat rows in time order, each episode contiguous (the replay agents read single transitions; TD-MPC2's offline datasets, :mod:ajax.agents.TDMPC2.dataset, hold [N, L, ...] episodes for its sequence sampler):

================ ========== =============================================== field shape content ================ ========== =============================================== obs [N, O] f32 the env's own observation (not normalised) action [N, A] f32 the action taken, in agent units ([-1, 1]) reward [N] f32 the env's own reward terminated [N] bool a true terminal (no bootstrap) truncated [N] bool a time limit (bootstrap on next_obs) next_obs [N, O] f32 the next observation, the final one at an end is_expert [N] bool whether the expert chose action a_expert [N, A] f32 the expert's action at obs (NaN: unknown) return_to_go [N] f32 sum_k gamma^k r_{t+k} to the episode's end episode [N] i32 the episode's index ================ ========== ===============================================

return_to_go and episode follow from the others. return_to_go is Cal-QL's Monte-Carlo return (Nakamoto et al. 2023), cut at a time limit.

:func:record rolls an expert out, :func:export reads a run's replay, :func:fill_replay adds a dataset to one. :func:save and :func:load follow TD-MPC2's file conventions: an uncompressed .npz of plain arrays and a JSON meta string, no pickle.

TransitionDataset `dataclass` ¶

TransitionDataset(
    obs,
    action,
    reward,
    terminated,
    truncated,
    next_obs,
    is_expert=None,
    a_expert=None,
    gamma=0.99,
    env="env",
)

Whole episodes of transitions (the module docstring's schema).

is_expert defaults to False and a_expert to NaN (logged data); gamma discounts return_to_go; env names the environment.

returns_to_go ¶

returns_to_go(reward, ends, gamma)

sum_k gamma^k reward[t + k] to the end of t's episode.

record ¶

record(
    expert_policy,
    env_args,
    key,
    n_steps,
    gamma=0.99,
    env="env",
)

The whole episodes of n_steps steps per env of expert_policy (Ajax's expert collector, so its state is kept for a stateful one): the actions in agent units, as a replay keeps them. Each env's last episode, cut by the end of the recording, is left out.

export ¶

export(state, gamma=0.99, env='env', seed_index=None)

The whole episodes in a replay agent's replay (state as its train returned it; seed_index picks a seed), oldest first, env after env. A ring that wrapped has lost its oldest episode's start: that episode is left out. The observations are the replay's raw_obs and raw_next_obs where it keeps them (else, after a termination, next_obs is the reset observation it stores).

next_actions ¶

next_actions(data)

Each row's next action, its episode's next row's; NaN on an episode's last row (a terminal one needs none; a time limit's is unknown).

fill_replay ¶

fill_replay(buffer, buffer_state, data)

buffer_state with data's rows added, one row per env per add as collection adds its own: each env's replay row gets a contiguous run of data, so a sequence replay reads its episodes whole. A replay already holding rows drops its oldest when full, as collection does.

Raises ValueError when the replay keeps fields data does not hold or of another width, when data's rows do not divide among its envs (flashbax advances every env's row together), or when they outnumber its capacity.

save ¶

save(data, path)

Write data to path, through a temporary file renamed, so an interrupted save leaves no partial file there.

load ¶

load(path)

The dataset :func:save wrote to path; raises ValueError on another format or version. return_to_go and episode are recomputed from the rows.

ajax.agents.recurrent ¶

Shared machinery for recurrent off-policy updates (R2D2-style).

Every off-policy agent with a Q-critic (SAC, ASAC, REDQ, TD3) trains its memory the same way:

Sample contiguous per-env sequences of length burn_in + sequence_length + 1 from a trajectory buffer.
Warm up ("burn in") the actor / online-critic / target-critic carries from zero over the first burn_in steps with the CURRENT params, under stop_gradient — the carries are constants inside the losses.
Train on the next sequence_length steps with BPTT; the final step is not read (each row stores the observation its target bootstraps on, interaction.bootstrap_obs).

sample_and_burnin_sequences implements steps 1-3 once for all agents; the per-agent losses then consume the returned :class:RecurrentCarries through :func:actor_dist and :func:q_values, which run a network on a feedforward batch or, given carries, on a replayed sequence.

A critic's memory reads each row's observation and the action before it, the one taken; the action a loss asks about joins after the memory, at the critic's head, so a taken or a sampled action is a query on the same carry and never feeds it (Critic.query_dim; the recurrent off-policy critic of Ni et al. 2022, cited from memory, unverified).

RecurrentReplayConfig ¶

Bases: BaseAgentConfig

An off-policy agent's sequence-replay settings, read only when it has memory: replayed windows of burn_in + sequence_length + 1 steps.

RecurrentCarries ¶

Per-update recurrent context for replayed sequences (R2D2-style).

All carries are burned in from zero on the sequence prefix with the CURRENT params and treated as constants inside the losses (they are stop-gradiented); only the training segment backpropagates through time. The carries hold for the training segment's first row; its rows (obs, the actions taken, resets their episode starts) are the history a bootstrap runs over (:func:_bootstrap_stream).

previous_actions ¶

previous_actions(actions, resets)

Each step's previous action in its episode, (T, B, A): zero at an episode start (resets) and at the sequence's first step.

stream_positions ¶

stream_positions(cut)

Each row's position in its window's bootstrap stream, from the rows cut (S, B) where a time limit alone ended the episode: each such row's final observation follows it (:func:_bootstrap_stream).

sample_and_burnin_sequences ¶

sample_and_burnin_sequences(
    agent_state,
    buffer,
    sample_key,
    burn_in,
    burn_target_actor=False,
    stored_state=False,
)

Sample sequences and burn in all carries; see the module docstring.

stored_state=True enables R2D2's stored-state strategy for the ACTOR carries: instead of burning them in from zero, the carries the actor actually had at collection time are read back from the buffer (actor_carry field, written by collect_experience with store_hidden=True). Kapturowski et al. 2019 show stored state beats zero-init + burn-in at mitigating recurrent-state staleness. The critic carries are still burned in from zero: the critic never runs at collection time, so there is nothing stored for it.

Returns a time-major :class:Transition over the training segment (leaves shaped (sequence_length, batch, ...)) and the matching stop-gradiented :class:RecurrentCarries.

sample_replay ¶

sample_replay(
    agent_state,
    buffer,
    key,
    recurrent=False,
    burn_in=0,
    stored_state=False,
    burn_target_actor=False,
)

One replay batch: transitions, or burned-in sequences when recurrent.

Returns (batch, carries). A feedforward batch keeps its pre-normalisation observations in batch.raw_obs and has no carries (None); sequences are time-major, with the carries of :func:sample_and_burnin_sequences and no raw observations.

actor_dist ¶

actor_dist(
    actor_state,
    params,
    obs,
    carries=None,
    *,
    bootstrap=False,
    target_actor=False,
)

The actor's distribution over obs.

Without carries, one feedforward batch. With them, a replayed sequence started from its burned-in carry (the target actor's when target_actor, TD3's bootstrap); when bootstrap, obs are the rows' next observations, each read after its row (:func:_bootstrap_stream).

actor_dist_train ¶

actor_dist_train(actor_state, params, obs)

:func:actor_dist in training, feedforward: the actor's distribution over obs and its running statistics after this pass (:meth:~ajax.state.LoadedTrainState.apply_train).

q_values ¶

q_values(
    critic_state,
    params,
    obs,
    actions,
    carries=None,
    *,
    bootstrap=False,
    quantiles=False,
)

Every critic's Q(obs, actions), the ensemble axis first (quantiles: a distributional critic's quantiles, last axis).

With carries, on a replayed sequence from the online critic's burned-in carry, or the target critic's when bootstrap (obs then being the rows' next observations, each read after its row, :func:_bootstrap_stream): the memory reads the actions taken before the rows, actions are queried at the head. Every action reaches the critic in its units (:func:in_critic_units; :func:predict_value maps a feedforward query's).

stored_actor_carry_dim ¶

stored_actor_carry_dim(memory, stored_state)

Width of the actor carry a stored-state replay buffer keeps per step (R2D2 stored state), 0 when nothing is stored.

make_replay_buffer ¶

make_replay_buffer(
    agent_config, n_envs, memory, buffer_size, batch_size
)

The replay buffer of an off-policy agent configured by agent_config (a :class:RecurrentReplayConfig with learning_starts): transitions, or per-env trajectories sampled as windows when memory is set.

unsupported_recurrent_options ¶

unsupported_recurrent_options(agent_name, **options)

Raise if any expert-guidance style option is combined with memory.