Networks and modules¶

The actor and critic networks, Squash, the tanh wrapper over any policy distribution, the pluggable memory block, the PID policy and the blocks the world models share. See Networks and Memory.

ajax.networks.networks ¶

CNNSpec ¶

Bases: NamedTuple

CNN encoder architecture: conv stack + projection width.

CNNSpec() is :class:CNNEncoder's default architecture. Each kernel_sizes / strides entry may be an int (square) or an (h, w) tuple. A NamedTuple -> hashable, so it is safe as a flax module field and as a :class:NetworkConfig field. padding is flax's nn.Conv padding; layer_norm puts a LayerNorm after every convolution and the projection, before its ReLU; kernel_init names every layer's initializer as :func:~ajax.networks.utils.dense reads it, "torch_linear" drawing kernel and bias uniform in +-1/sqrt(fan_in) (torch's and Torch7's default); None is flax's lecun_normal with zero biases.

CNNEncoder ¶

Bases: Module

Conv encoder for image observations packed flat as (H*W*C + extra_dim,).

The flat layout (rather than native NHWC) keeps Ajax's existing collector and command-augmentation paths working unchanged: extra trailing scalar dims (e.g. UDRL's (d_r, d_h) command) are concatenated to the embedding AFTER the convolutions, so the conv stack only sees the image.

The image is assumed to be stored in NHWC order: a Flatten wrapper in user code must transpose a stack-first array (T, H, W) into (H, W, T) before flattening so reshape(H, W, C) recovers the right layout.

Actor ¶

Bases: Module

The stochastic policy of SAC, REDQ, ASAC, AVG, PPO, APO and UDRL: a Normal head for continuous actions (:class:Squashed bounds it), a Categorical one for discrete. Expert guidance acts at the loss level, not in the architecture, which keeps the policy unconstrained.

Squashed ¶

Bases: Module

actor with its policy bounded by tanh (:func:squash): any actor, stochastic or deterministic, feedforward or recurrent. The two share one scope, so the parameters are actor's alone.

Critic ¶

Bases: Module

call ¶

__call__(x, hidden_state=None, done=None)

The value: the head's output, or the mean of its quantiles.

quantiles ¶

quantiles(x, hidden_state=None, done=None)

The head's n_quantiles outputs (and the carry, with memory).

apply_encoder ¶

apply_encoder(x)

The encoder's features alone, without the value head (read by :class:ajax.extensions.instrumentation.ConditioningMetrics).

MultiCritic ¶

Bases: Module

Ensemble of num critics, vmapped over a leading parameter axis.

quantiles ¶

quantiles(x, hidden_state=None, done=None)

Every member's quantiles, the ensemble axis first.

apply_encoder_ensemble ¶

apply_encoder_ensemble(x)

Per-ensemble-member encoder features, shape (num, ..., d) (:class:ajax.extensions.instrumentation.ConditioningMetrics).

reference_cnn ¶

reference_cnn(image_shape)

Mnih et al.'s network when the image fits it, 36 pixels a side or more (its three VALID convolutions leave nothing of less), else MinAtar's, made for MinAtar's 10 x 10 images.

check_log_std_head ¶

check_log_std_head(
    state_independent, log_std_init, log_std_kernel_init
)

Refuse a log-std head whose start is set twice or not at all: the learned parameter needs log_std_init and has no kernel; the Dense head takes log_std_init as its bias unless its kernel init is "torch_linear", whose layer draws the bias (log_std_init None).

action_value_input ¶

action_value_input(obs, action, previous_action=None)

A Q-critic's input: (obs, action), or, given the previous action, the recurrent critic's (obs, previous_action, action): its memory reads the observation and the previous action, the current action joins at its head (Critic.query_dim).

get_initialized_actor_critic ¶

get_initialized_actor_critic(
    key,
    env_config,
    actor_optimizer_config,
    critic_optimizer_config,
    network_config,
    continuous=False,
    action_value=False,
    squash=False,
    num_critics=4,
    max_timesteps=None,
    extra_obs_dim=0,
    action_dim_override=None,
    actor=None,
)

Create actor and critic networks, their initializers and log-std head as network_config sets them; a prebuilt actor (TD3's deterministic one) replaces the :class:Actor built here. squash wraps a continuous actor in :class:Squashed (a discrete actor keeps its Categorical).

extra_obs_dim: dimensions the agent appends to obs at run time (UDRL's command; SAC's expert action and expert state when it augments with them), so the network is initialised on the input width it will be fed.

A :class:PIDPolicy as network_config.actor_architecture makes the actor a :class:PIDActor (see :mod:ajax.modules.pid_actor).

get_initialized_critic ¶

get_initialized_critic(
    key,
    env_config,
    critic_optimizer_config,
    network_config,
    num_critics=2,
    max_timesteps=None,
    extra_obs_dim=0,
)

Initialize a standalone critic network (no actor). Same architecture as the online critic so predict_value can be called on its params directly. Always feedforward: its one reader (MCPretrain's frozen expert critic) only calls predict_value.

init_network_carry ¶

init_network_carry(network, memory, key, batch_size)

Fresh carry for network. Ensembles (modules exposing num) get one carry per member, stacked on a leading (num,) axis.

init_network_state ¶

init_network_state(
    init_x, network, key, tx, n_envs=1, memory=None
)

Initialise network on init_x into a :class:LoadedTrainState.

With a memory config the network is recurrent: it is initialised on a single-step time-major sequence and the state carries a fresh carry batch-sized to n_envs.

in_critic_units ¶

in_critic_units(critic_state, actions)

actions as the critic reads them: critic_state.action_units' affine map of the agent's [-1, 1] actions, or the actions themselves.

predict_value_sequence ¶

predict_value_sequence(
    critic_state,
    critic_params,
    x,
    resets,
    initial_hidden,
    *,
    quantiles=False,
)

Run a recurrent critic over a time-major sequence (quantiles: a distributional critic's quantiles, last axis, in place of the values).

Parameters:

Name	Type	Description	Default
`x`	`Array`	(T, B, features) critic input (obs, or concat(obs, action)).	required
`resets`	`Array`	(T, B) episode-start flags aligned with `x` (resets[t] means x[t] is the first observation of a new episode).	required
`initial_hidden`	`Carry`	carry valid for x[0]; leaves are (num, B, hidden) for critic ensembles.	required

Returns:

Type	Description
`(values, final_hidden)`	values (num, T, B, 1); final_hidden is the
`Carry`	carry after consuming the whole sequence.

phi_star_critic_state ¶

phi_star_critic_state(critic_state)

The critic state the frozen expert critic phi's parameters (expert_critic_params) are applied with: the online critic's network without the agent's observation statistics, as MC pretraining and PhiRefresh fit phi on the observations as stored.

predict_value ¶

predict_value(
    critic_state, critic_params, x, *, quantiles=False
)

The critic's values at x: obs, or concat(obs, action) with the agent's [-1, 1] action, which a critic with action_units reads in its units (TD3's, in the env's). quantiles: a distributional critic's quantiles (last axis) in place of their mean, the value.

predict_value_train ¶

predict_value_train(critic_state, critic_params, x)

:func:predict_value in training: the critic's values at x and its running statistics after this pass (:meth:LoadedTrainState.apply_train).

state_value ¶

state_value(agent_state, obs, action)

V(obs) of the agent's critic, its ensemble's mean (action unused): an evaluation's value estimate (:class:ajax.evaluate.ValueEstimate).

action_value ¶

action_value(agent_state, obs, action)

Q(obs, action) of the agent's critic, its ensemble's mean: an evaluation's value estimate (:class:ajax.evaluate.ValueEstimate).

ajax.networks.distributions ¶

Policy distributions shared by the agents.

:class:Squash bounds any continuous policy to (-1, 1) with tanh (:func:squash also takes a :class:Deterministic one). The functions below it read any policy, squashed or not, so callers need not check which they hold.

Deterministic ¶

Tiny distrax-compatible wrapper for deterministic policies (TD3's).

Exposes .mean() so existing code (the policy loss, the target action computation, the ImitationLoss extension) keeps working without branches. A pytree of its action, as distrax distributions are of their parameters (a replayed bootstrap gathers its steps, recurrent.actor_dist).

log_prob ¶

log_prob(x)

For BC pretraining: a delta distribution's log-density is +∞ at x=action, -∞ elsewhere. We use the negative-MSE proxy so that maximizing log_prob is equivalent to MSE-minimizing the actor toward the target action. Shape: (..., action_dim) -> (..., action_dim) (per-dim, summed by caller).

GuardedTanh ¶

GuardedTanh(eps)

Bases: Tanh

tanh whose log-det-Jacobian is log(1 - tanh(x)^2 + eps), as stable-baselines3's squashed Gaussian (eps 1e-6) and REDQ (core.py:232) compute it; distrax's own is exact.

Squash ¶

Squash(distribution, log_det_eps=0.0)

Bases: Transformed

tanh of a policy distribution, its latent: any distrax distribution (a Block spans its event dimensions, if it has any). log_det_eps > 0 takes :class:GuardedTanh's log-det instead.

Exact (with log_det_eps 0): sampling and log-probabilities. log_prob(action) inverts tanh, finite only while |action| < 1, which float32 loses once the latent sample passes about 9; :meth:log_prob_from_raw reads the latent sample itself and stays finite.

Not exact, as neither has a closed form: :meth:mean is tanh of the latent's mean, the deterministic action, not the mean of the squashed action; :meth:entropy is sampled.

log_prob_from_raw ¶

log_prob_from_raw(raw)

The log-probability of tanh(raw), from the latent sample raw (no arctanh, so finite at saturation).

sample_with_raw ¶

sample_with_raw(seed)

(action, log_prob, raw): raw the latent sample, the action tanh(raw).

entropy ¶

entropy(seed=None)

The squashed action's entropy, estimated as brax does: the latent's entropy plus the tanh log-det-Jacobian at one latent sample drawn with seed. Unbiased, and its gradient reaches the mean; the latent's entropy alone grows without bound in its scale.

squash ¶

squash(pi, log_det_eps=0.0)

pi bounded by tanh: :class:Squash of it, or, for a :class:Deterministic one, tanh of its action (a point mass has no density to correct).

latent ¶

latent(pi)

The distribution pi squashes, or pi itself.

sample_with_raw ¶

sample_with_raw(pi, seed)

(action, log_prob, raw): raw is the latent sample, the action itself when pi is not squashed.

log_prob_from_raw ¶

log_prob_from_raw(pi, raw)

The log-probability of the action whose latent sample is raw.

entropy ¶

entropy(pi, seed=None)

pi's entropy: sampled with seed when squashed, else exact.

ajax.networks.utils ¶

BatchRenorm ¶

Bases: Module

Batch Renormalization (Ioffe, NeurIPS 2017) as CrossQ's authors run it (adityab/CrossQ@a318a26, sbx/sac/policies.py:37-186): batch norm for its first warmup_steps training passes, then the batch's statistics corrected towards the running ones by r and d, clipped to r_max and d_max and held constant.

It trains, normalising by the batch and updating its running statistics, exactly when its batch_stats collection is mutable (apply(..., mutable=["batch_stats"]), as the authors pass train=True exactly then); otherwise it normalises by the running statistics. An init traces the training branch without updating them.

TorchLinear ¶

Bases: Module

A dense layer that starts as torch's nn.Linear does: its kernel and its bias each uniform in ±1 / sqrt(fan_in).

torch 1.5.1 (torch/nn/modules/linear.py:79-84, at its v1.5.1 tag) draws the weight with kaiming_uniform_(a=sqrt(5)), whose bound sqrt(3) * sqrt(2 / (1 + a**2)) / sqrt(fan_in) is 1 / sqrt(fan_in) (torch/nn/init.py:61-69, :322-327), and the bias with the same bound. A layer, not an initializer: a flax bias initializer never sees the fan-in. Its parameters are a Dense layer's, at the same path.

clip_by_member_norm ¶

clip_by_member_norm(max_norm)

optax.clip_by_global_norm applied to each member of an ensemble on its own: every leaf's leading axis indexes a member, as a :class:~ajax.networks.networks.MultiCritic's vmapped parameters do. Each member is its own network in the references, clipped alone (ASAC: BobMak/asac @4b68862, darer/Models.py:299-301).

get_adam_tx ¶

get_adam_tx(
    learning_rate=0.001,
    max_grad_norm=None,
    eps=1e-05,
    beta_1=0.9,
    beta_2=0.999,
    weight_decay=0.0,
    per_member=False,
    optimizer="adam",
)

Return the optimizer optimizer names, Adam by default, clipping the gradients' global norm to max_grad_norm when one is given; with per_member, each ensemble member's norm on its own (:func:clip_by_member_norm).

Unclipped by default. Global-norm clipping is a PPO-family convention (PPO passes max_grad_norm=0.5; APO clips both its networks' gradients at once, train_APO.update_networks); the off-policy agents and the supervised ones run plain Adam, as in their reference implementations.

Parameters:

Name	Type	Description	Default
`weight_decay`	`float`	Decoupled weight decay, multiplied by the learning rate (`OptimizerConfig.weight_decay`), between the update rule and the learning rate as `optax.adamw` places it; `0` leaves it out, so `"adam"` is `optax.adam`.	`0.0`
`optimizer`	`str`	The update rule (`OptimizerConfig.optimizer`): `"adam"`, `"radam"` or `"centered_rmsprop"`.	`'adam'`

parse_activation ¶

parse_activation(activation)

Parse an activation name, or pass an activation callable through.

"silu" is x * sigmoid(x) (DreamerV3's activation, Table 4 of arXiv:2301.04104v2). "mish" is x * tanh(softplus(x)) (TD-MPC2's activation, Sec. 3.1 of arXiv:2310.16828v2; torch nn.Mish).

parse_function_string ¶

parse_function_string(s, context=None)

Split "name" / "name(expr)" into (name, value of expr).

A string that does not start with an identifier (a letter or _), such as "0.5" or "-1.0", is evaluated as a whole and returned as (None, value). One that does but is neither form, such as "orthogonal(2.0" or "np.sqrt(2)", raises ValueError rather than being read as a bare name.

trunc_normal_fan_in ¶

trunc_normal_fan_in(scale=1.0)

Fan-in truncated normal: DreamerV3's weight initializer.

DreamerV3 (arXiv:2301.04104v2; danijar/dreamerv3@2411f7d, dreamerv3/nets.py:842-848, :870, fans :873-884; dreamerv3_spec 1.5) initialises every Linear, BlockLinear and Conv kernel as::

W = scale * 1.1368 * sqrt(1 / fan_in) * TruncatedNormal(-2, 2)

with a unit-sigma normal truncated at +-2 and scale the layer's outscale. This is variance_scaling(scale**2, "fan_in", "truncated_normal"): flax divides by the exact standard deviation of the truncated unit normal, 0.87962566, where DreamerV3 multiplies by the rounded 1.1368. Drawn with the same key the two kernels differ by the constant factor 1.13684723 / 1.1368, a residual relative difference of 4.16e-5 (tested; deviation D23 in docs/world_models/deviations.md). Both give Std[W] = scale / sqrt(fan_in), and flax's fan-in matches the reference for every kernel rank >= 2: shape[-2] * prod(shape[:-2]), so a BlockLinear kernel (g, I/g, U/g) gets the full input width I.

scale = 0 gives an all-zero kernel (DreamerV3's reward and value heads).

get_initializer ¶

get_initializer(name)

Factory for the initializer called name.

Looks in Ajax's own registry first, then in flax.linen.initializers.

parse_initialization ¶

parse_initialization(initialization, default=None)

Parse an initializer string such as "orthogonal(1.0)", or pass a callable initializer through; None gives default.

"name" calls the factory name (see :func:get_initializer) with no argument and "name(expr)" with the value of expr. A string that does not start with a name, such as "0.5" or "-1.0", gives a constant initializer.

parse_layer ¶

parse_layer(layer, kernel_init=None, bias_init=None)

Parse one architecture token: a width such as "256" gives a Dense layer (:func:dense, kernel orthogonal(1.0) by default), "layernorm" a new LayerNorm with its own scale and bias, any other name or a callable an activation (:func:parse_activation).

("256", "layernorm", "relu") per hidden layer is RLPD's critic MLP (Ball et al., ICML 2023; ikostrikov/rlpd, rlpd/networks/mlp.py). "batchrenorm" gives a :class:BatchRenorm, whose running statistics only an agent training them in its losses may have (CrossQ).

dense ¶

dense(
    features,
    kernel_init,
    bias_init,
    default_kernel_init,
    name=None,
)

A Dense layer with kernel_init (default_kernel_init for None) and bias_init (zeros for None), as :func:parse_initialization reads them; kernel_init="torch_linear" builds a :class:TorchLinear, which draws its bias too.

parse_architecture ¶

parse_architecture(
    architecture, kernel_init=None, bias_init=None
)

Parse a list of string/module architecture into a list of jax modules

ajax.networks.memory ¶

Pluggable memory (recurrent) modules for actor/critic networks.

Every memory architecture unifies behind one interface, operating on time-major sequences (T, B, features) with per-step reset flags (T, B):

init_carry(config, rng, batch_size) -> carry returns the recurrent state as a pytree (always a tuple with one entry per layer).
__call__(carry, x, resets) -> (new_carry, y) runs the cell over the whole sequence, isolating episodes wherever resets[t] is set, and returns the final carry plus the per-step outputs (T, B, hidden).

Step-wise inference (collection, evaluation) is the same code path with T=1, so training and acting can never diverge. The step-vs-sequence equivalence tests in tests/networks/test_memory.py are the acceptance harness for every kind.

Carries are zero at initialization for every supported kind; callers rely on this to build fresh carries with zeros_carry_like (e.g. at eval time or for replayed sequences).

Kinds and their compute/memory profile:

"gru" / "lstm": nn.scan over the classic cells — sequential in T, carry is O(hidden) per env.
"transformer": sliding-window causal self-attention. Training runs FULL PARALLEL attention over the chunk (no scan) with an episode-segment mask; acting keeps a rolling cache of the last window - 1 embeddings per layer, so per-step memory is O(window * hidden) and constant in episode length.
"mamba": selective state-space (S6-style) block. Training runs a parallel associative_scan over T; the carry is the SSM state plus a small causal-conv tail — RNN-sized memory with transformer-like training parallelism.

The transformer and mamba outputs end with a LayerNorm, as their references' pre-norm stacks do; the GRU and LSTM outputs are the cells' own.

MemoryConfig `dataclass` ¶

MemoryConfig(
    kind="gru",
    hidden_size=64,
    num_layers=1,
    gradient_checkpoint=False,
    window=32,
    num_heads=4,
    d_state=16,
    d_conv=4,
    expand=2,
    final_norm=True,
    init=None,
)

Single-hyperparameter surface for adding memory to any network.

Frozen (hashable) so it can be a Flax module field and a jit-static argument. kind selects the architecture; everything else has a sane default. Kind-specific knobs are ignored by the other kinds.

ends_normalised `property` ¶

ends_normalised

Whether the cell's output leaves a final LayerNorm.

MemoryCell ¶

Bases: Module

Reset-aware memory block; dispatches to the core selected by config.kind. See the module docstring for the interface contract.

parse_memory_config ¶

parse_memory_config(memory)

Coerce user input (None | MemoryConfig | dict) into a MemoryConfig.

with_default_init ¶

with_default_init(memory, inits)

memory with init set to inits[memory.kind] when unset: an agent's reference initialisation, which an explicit init overrides. A kind inits does not name keeps flax's.

init_carry ¶

init_carry(config, rng, batch_size)

Fresh (all-zero) carry for config: a tuple with one entry per layer. Shapes are derived from the config alone — no params needed.

zeros_carry_like ¶

zeros_carry_like(carry, batch_size, batch_axis=0)

Fresh carry with the batch dimension resized to batch_size.

Valid because every supported cell initializes its carry to zeros. batch_axis is 0 for a single network's carry and 1 for an ensemble's stacked carry (leading (num_critics,) axis).

flatten_carry ¶

flatten_carry(carry)

Pack a carry pytree into one (batch, D) float array.

Used to write per-step actor carries into the replay buffer (R2D2-style stored-state replay): flashbax schemas hold flat arrays, not pytrees. Every leaf keeps its leading batch axis; the trailing dims are flattened and concatenated in jax.tree leaf order. Reversible via :func:unflatten_carry given any same-structure template. Bool leaves (e.g. the transformer cache's validity flags) are cast to float and recovered by dtype from the template.

unflatten_carry ¶

unflatten_carry(flat, template)

Inverse of :func:flatten_carry.

template supplies the pytree structure, per-leaf trailing shapes and dtypes (e.g. from :func:zeros_carry_like with the right batch size); flat is (batch, D) with the same leaf order.

flat_carry_dim ¶

flat_carry_dim(config)

Width D of :func:flatten_carry's output for config (per env).

tree_where_reset ¶

tree_where_reset(resets, fresh, tree)

tree with fresh in every batch entry where resets is set: a recurrent carry, or a stateful expert's state, restarting with its episode. resets has shape (B,) (or (B, 1)); every leaf leads with the batch axis.

gate_block_orthogonal ¶

gate_block_orthogonal(n_gates)

One gate's kernel as torch's nn.init.orthogonal_ leaves it on a GRU's or LSTM's stacked (n_gates * H, fan_in) weight: a (fan_in, H) block of a random orthogonal (fan_in, n_gates * H) matrix, of std 1/sqrt(max(n_gates * H, fan_in)). Ni et al. start every GRU and LSTM weight so, with zero biases (pomdp-baselines @e7c19c3 policies/models/recurrent_actor.py:70-74, recurrent_critic.py:64-68; Memory-RL @c9b8b39 policies/seq_models/rnn_vanilla.py:23-31), where flax's per-gate kernels have std 1/sqrt(fan_in) (lecun_normal) and 1/sqrt(H) (orthogonal). Each gate draws its own matrix, so the gates are independent, where torch's are jointly orthogonal.

ajax.modules.pid_actor ¶

PID policies: an actor that is a PID controller with learned gains.

The actor acts on the error terms (e, I, D) that :class:ajax.environments.pid_terms.PIDTermsWrapper puts first in the observation; its pre-squash action, per action dimension a, is::

u_a = sum_j Kp[a, j] e_j + Ki[a, j] I_j + Kd[a, j] D_j

Passing a :class:PIDPolicy as actor_architecture selects it in any continuous-action agent (SAC, TD3, REDQ, ASAC, AVG, PPO, APO). Two kinds of gains:

gains="fixed": the gains are the actor's parameters, a linear map with no bias and no hidden layer, so the trained policy deploys as a plain (saturated) PID with the gains :func:pid_gains reads off it, in normalised units (below) (Lawrence et al. 2022, "Deep reinforcement learning with shallow controllers: an experimental application to PID tuning", cited from memory, unverified).
gains="observation": an MLP predicts the gains from the observation (gain scheduling); with a constant network it is the fixed policy.

Each agent's own output rules apply to u as to any actor's mean: SAC, REDQ, ASAC and AVG squash it with tanh, TD3 too (deterministic); APO takes tanh(u) as its mean (its tanh_mean) and samples around it unsquashed, PPO around u (unless squash=True). Like every action, u is in agent units, [-1, 1], mapped onto the env's box [low, high] (:func:ajax.environments.utils.agent_action_to_env): the deployed controller is c + h * sat(u), with centre c = (high + low) / 2, half-range h = (high - low) / 2 and sat the agent's tanh or clip to [-1, 1]. The gains are therefore normalised: they map the error terms to agent units, and a PID in the env's units has gains h * K (small-signal, tanh near 0) and offset c; zero gains command the box's centre (no control on a symmetric box). The stochastic agents explore with a learned, state-independent log-std, clipped to the agent's log_std_bounds (PPO's and APO's default: no clip). Gains start at zero. Observation normalisation (env- or agent-side) leaves the terms raw, as the wrapper declares them, so the gains act on the errors in the env's units; the gain network of gains="observation" reads the rest normalised.

PIDPolicy `dataclass` ¶

PIDPolicy(n_errors=1, gains='fixed', architecture=())

A PID actor, passed as an agent's actor_architecture.

errors the :class:PIDTermsWrapper tracks (its terms are

the first 3 * n_errors observation entries).

gains: "fixed" (learned constants) or "observation" (an MLP of architecture predicts them from the observation). architecture: the gain network's hidden layers ("observation" only; empty: gains affine in the observation).

PIDLaw ¶

Bases: Module

The PID output u for an observation, shape (..., action_dim).

gains ¶

gains(obs)

Kp, Ki, Kd stacked: (..., 3, action_dim, n_errors).

PIDActor ¶

Bases: Module

Stochastic PID actor: Normal around the PID output (squashed by :class:~ajax.networks.networks.Squashed).

check_pid_actor ¶

check_pid_actor(continuous, memory)

Reject what a PID actor cannot be: a discrete policy or a recurrent one (its memory is the wrapper's integral and difference).

pid_gains ¶

pid_gains(actor_params)

{"kp", "ki", "kd"} of a fixed-gain PID actor, each (..., action_dim, n_errors) (leading axes, e.g. seeds, kept).

They are normalised gains, from the errors to agent units [-1, 1]: on an env box [low, high] the env-unit gains are (high - low) / 2 times these (the module docstring), equal on a [-1, 1] box.

Observation-conditioned gains depend on the observation: read them with PIDActor(policy, action_dim).apply(params, obs, method="gains") (a squashed actor's parameters are the same).

ajax.modules.pid_head ¶

PID output head: a stateful, learnable PID layer on a network's output.

Busetto et al. 2024 (arXiv:2411.06482, §4.1) close their in-context controller with a PID layer: the transformer emits a signal z_k and the applied input is

u_k = K_P z_k + K_I sum_{j<=k} z_j + K_D (z_k - z_{k-1})

with three learnable gains. The integral term promotes zero steady-state error on step references, and the authors report faster training. This is NOT :mod:ajax.modules.pid_actor (a PID policy acting on the observation's tracking-error terms): here the PID acts on the network's own output, so it needs a carry (integral and previous signal) like a memory cell.

Interface mirrors :class:ajax.networks.memory.MemoryCell: time-major (T, B, n) inputs with (T, B) reset flags, initialize_carry, and __call__(carry, z, resets) -> (new_carry, u); step-wise use is the T=1 case of the same code path.

PIDOutputHead ¶

Bases: Module

Learnable PID on a time-major signal z; see the module docstring.

Attributes:

Name	Type	Description
`n_outputs`	`int`	signal width `n` (one gain triple per channel; the paper's single-input case is `n_outputs=1`).
`use_p`	`bool`	enable the proportional term; a disabled term has no parameter. With only `use_p` and `kp_init=1` the head is the identity at initialisation. `use_i` and `use_d` likewise enable the integral and derivative terms.
`anti_windup`	`float \| None`	pre-squash output magnitude beyond which the integral stops accumulating in the saturating direction (conditional integration, the textbook anti-windup). `None` disables it. With a downstream `tanh` bound, a value of about 3 (where `tanh` is within 0.5 % of its rail) is a sensible choice: it stops the integral from winding up while the applied input is pinned, so the loop does not overshoot when it comes off the rail.

call ¶

__call__(carry, z, resets)

z: (T, B, n), resets: (T, B) -> (new_carry, u (T, B, n)).

A reset at step t clears the integral and the previous signal before z_t is consumed, so u_t is computed as at the start of an episode.

init_pid_carry ¶

init_pid_carry(batch_size, n_outputs)

Zero carry: no accumulated integral, previous signal z_{-1} = 0.

ajax.networks.blocks ¶

Dense layers and normalised MLP shared by DreamerV3 and TD-MPC2.

Both world-model agents build every network from the same hidden layer, Dense -> Norm -> Activation with the bias kept, and differ only in its configuration:

DreamerV3 (arXiv:2301.04104v2 Table 4, "RMSNorm + SiLU"; danijar/dreamerv3@2411f7d:dreamerv3/nets.py:367-403 MLP, :594-635 Linear, :724-758 Norm, :816-884 Initializer): RMSNorm with epsilon 1e-4 and a learnable scale only, SiLU, kernels from the fan-in truncated normal, output layers scaled by outscale (0 for the reward and value heads, 0.01 for the policy).
TD-MPC2 (arXiv:2310.16828v2 Sec. 3.1 and App. A, "LayerNorm + Mish", 1 % dropout after the first linear layer of each Q-function; nicklashansen/tdmpc2@5f6fade:tdmpc2/common/layers.py:85-122 NormedLinear and mlp, common/init.py:4-17 weight_init): LayerNorm with epsilon 1e-5 and a learnable scale and shift, Mish, dropout between the first Linear and its norm, kernels from N(0, 0.02^2), final weights of the reward and Q heads zeroed.

flax's defaults match neither paper (epsilon 1e-6, and LayerNorm's fast variance E[x^2] - E[x]^2, off by ~0.2 on mean-1000 unit-variance inputs where torch's two-pass variance is within 1e-4), so :class:NormedMLP sets them explicitly. Statistics are computed in float32: DreamerV3 at 2411f7d computed the RMS in its bf16 compute dtype, and Ajax runs these agents in float32 throughout (deviation D4 in docs/world_models/deviations.md). Output layers and heads are composed by the agents with :func:linear.

NormedMLP ¶

Bases: Module

Stack of layers hidden layers Dense -> [Dropout] -> Norm -> Act.

Every layer has width units. Dropout (inverted, rate dropout) is applied only after the Dense of the first layer and only when deterministic is False, using the 'dropout' RNG stream (TD-MPC2 layers.py:120, dropout*(i==0)). There is no output layer: compose one with :func:linear. Use :meth:dreamerv3 and :meth:tdmpc2 for the papers' configurations.

Attributes:

Name	Type	Description
`act`	`Union[str, ActivationFunction]`	activation or activation name (see `parse_activation`).
`norm`	`Literal['layer', 'rms']`	`"rms"` (scale only) or `"layer"` (scale and shift, exact two-pass variance).
`norm_eps`	`float`	epsilon added to the mean square or variance.
`kernel_init`	`Union[str, InitializationFunction]`	initializer or initializer name of the Dense kernels.
`dropout`	`float`	dropout rate of the first layer; 0 disables it.

dreamerv3 `classmethod` ¶

dreamerv3(layers, units, **module_kwargs)

DreamerV3's MLP trunk: RMSNorm(eps 1e-4) and SiLU, fan-in init.

nets.py:390-391 applies layers x Linear(units, act='silu', norm='rms') with winit: normal and fan in (configs.yaml:113-129, nets.py:600-601); RMSNorm's default epsilon is 1e-4 (nets.py:728).

tdmpc2 `classmethod` ¶

tdmpc2(layers, units, dropout=0.0, **module_kwargs)

TD-MPC2's hidden NormedLinear layers: LayerNorm(eps 1e-5) and Mish.

layers.py:85-100 (Linear -> Dropout -> LayerNorm -> Mish, torch LayerNorm epsilon 1e-5) and init.py:7-9: torch's trunc_normal_(std=0.02) truncates at the absolute bounds +-2, i.e. 100 standard deviations, so it is N(0, 0.02^2), hence "normal(0.02)"; jax's truncated_normal(0.02) would cut at +-0.04 and shrink the standard deviation to 0.0176. TD-MPC2 sets dropout=0.01 for the Q-functions only.

call ¶

__call__(x, deterministic=None)

Apply the trunk to x [..., in]; returns [..., units].

deterministic must be given when dropout > 0: True disables dropout, False enables it and needs a 'dropout' RNG. Which passes use dropout is the agent's choice (paper-era TD-MPC2 enables it in every Q pass, docs/world_models/deviations.md §2).

linear ¶

linear(
    features,
    kernel_init,
    bias_init="zeros",
    outscale=1.0,
    name=None,
)

nn.Dense with initializers given by name and an output scale.

The papers' Linear layer: the bias is always present and zero-initialised, and never scaled (DreamerV3 nets.py:594-635, bias: True; TD-MPC2 nn.Linear with weight_init setting the bias to 0). outscale multiplies the kernel initializer; outscale = 0 gives an exactly zero kernel, as in DreamerV3's reward and value heads (configs.yaml:122, :129) and TD-MPC2's reward and Q heads (common/world_model.py:30, zero_ on their last weights). Parameters are float32 kernel [in, features] and bias [features]; leading input dimensions are batch dimensions. Call it where a module would be constructed (inside setup or an nn.compact method).

Parameters:

Name	Type	Description	Default
`kernel_init`	`Union[str, InitializationFunction]`	initializer or initializer name (see `parse_initialization`), e.g. `"trunc_normal_fan_in"` (DreamerV3) or `"normal(0.02)"` (TD-MPC2).	required
`bias_init`	`Union[str, InitializationFunction]`	initializer or name for the bias; both papers use zeros.	`'zeros'`
`name`	`Optional[str]`	the module name; flax's automatic `Dense_<i>` when None.	`None`

ajax.distributional ¶

Symlog transforms and two-hot discrete regression (DreamerV3, TD-MPC2).

Both world-model agents predict rewards and values as a categorical distribution over fixed bins and train it with a cross-entropy against a two-hot target: the scalar target is spread over its two neighbouring bins by linear interpolation. The papers share the idea but not the details, and the details matter for fidelity:

DreamerV3 (Hafner et al., arXiv:2301.04104v2, Eqs. 9-12 and p.18; paper-era code danijar/dreamerv3@2411f7d) places 255 bins at symexp(linspace(-20, 20, 255)) (dreamerv3/nets.py:462-466), interpolates the raw target between them and predicts the expectation of the raw bins (dreamerv3/jaxutils.py:210-274).
TD-MPC2 (Hansen et al., arXiv:2310.16828v2, Sec. 3.1 and Table 8; paper-era code nicklashansen/tdmpc2@5f6fade) places 101 bins at linspace(-10, 10, 101) in symlog space, interpolates clip(symlog(y), -10, 10) between them and predicts symexp(E_p[bins]) (tdmpc2/common/math.py:5-9, 66-95).

In raw space both bin sets are symexp(linspace(-limit, limit, n)); they differ in the bin count, the limit and the space in which the target is interpolated and the expectation taken. :class:TwoHot captures exactly these three parameters. For the same target the two interpolation spaces give different weights even on identical bins (y = 1 on DreamerV3's bins: upper weight 0.3827 in raw space, 0.4015 in symlog space), and symexp(E[b]) != E[symexp(b)], so neither can be simplified into the other.

Both bin sets are built by mirroring their negative half, so the middle bin is exactly 0 and the bins are exactly antisymmetric. The expectation is DreamerV3's symmetric sum over mirror bins, written so that uniform probabilities give exactly 0 in eager and in jitted code alike: a zero-initialised head predicts 0 (:meth:TwoHot.decode).

Fidelity, tested against literal ports of both references (tests/world_models/reference_impls.py), with the float32-level departures registered as deviations D22 (DreamerV3) and T24 (TD-MPC2) in docs/world_models/deviations.md:

The bins are the float32 roundings of their float64 values, computed once with numpy, so they are the same constants in eager and jitted code and on every backend. The references evaluate the same expressions in float32 on the device: DreamerV3's jnp bins differ from these by up to 1.3e-6 relative (mostly float32 linspace rounding near 0), TD-MPC2's torch.linspace bins by up to 4.8e-7 (its middle bin is -1.5e-7, not 0).
Given the same bins, :meth:TwoHot.encode and :meth:TwoHot.loss are bit-identical to DreamerV3's code run op by op, and :meth:TwoHot.decode is within 3 eps E_p|b| of it (eps the float32 machine epsilon), eagerly and under jax.jit.
TD-MPC2 encodes with a floor formula on its uniform grid. The generic interpolation used here is equal in exact arithmetic, and in float32 the more accurate of the two (within 2e-6 of the exact weights, against up to 8.6e-6 for the floor formula, whose positions near 100 have an ulp of 7.6e-6); they agree to 1.1e-5. TD-MPC2's naive expectation is matched within eps (1 + 6 E_p|b|) in symlog space (4.8e-6 at most; the eps term is its exp(|x|) - 1 rounding near 0), and its F.log_softmax (max subtracted first) within 3 eps (max|logits| + loss).

TwoHot `dataclass` ¶

TwoHot(num_bins, limit, transform)

Two-hot discrete regression over symlog-spaced bins.

Frozen and hashable, so it can be a jax.jit static argument or a field of a frozen module config. Use the paper constructors :meth:dreamerv3 and :meth:tdmpc2 rather than the raw fields.

Parameters:

Name	Type	Description	Default
`num_bins`	`int`	number of bins `n`; odd and at least 3, so that the middle bin is 0 and the symmetric summation pairs every bin.	required
`limit`	`float`	the bins span `[-limit, limit]` in symlog space, i.e. `[-symexp(limit), symexp(limit)]` in raw space. Only symmetric ranges are representable.	required
`transform`	`Literal['identity', 'symlog']`	the space in which targets are interpolated and the expectation is taken. `"identity"` (DreamerV3): raw space, the bins are `symexp(linspace(-limit, limit, n))`. `"symlog"` (TD-MPC2): symlog space, the bins are `linspace(-limit, limit, n)`, the target is `symlog(y)` and the prediction is `symexp` of the expectation.	required

All methods compute in float32 and return float32.

dreamerv3 `classmethod` ¶

dreamerv3(num_bins=255)

DreamerV3's symexp_twohot: raw-space bins and interpolation.

Bins symexp(linspace(-20, 20, num_bins)) (range about ±4.85e8), raw targets, raw-space expectation (2411f7d:dreamerv3/nets.py:462-466, jaxutils.py:210-274; bins: 255 in configs.yaml).

tdmpc2 `classmethod` ¶

tdmpc2(num_bins=101, limit=10.0)

TD-MPC2's two_hot: symlog-space bins and interpolation.

Bins linspace(-limit, limit, num_bins) in symlog space (step 0.2 by default, raw range about ±22025), target clip(symlog(y), -limit, limit), prediction symexp(E_p[bins]) (5f6fade:tdmpc2/common/math.py:66-95). TD-MPC2's config has separate vmin: -10 and vmax: +10 (config.yaml:48-50); the range must be symmetric here, which every TD-MPC2 config is: limit = vmax = -vmin (deviation T24).

bins ¶

bins()

The [num_bins] bin locations in the interpolation space.

Built as DreamerV3 builds its bins (nets.py:464-466): the negative half linspace(-limit, 0, (n + 1) // 2) (mapped by symexp for the identity transform) concatenated with its negated mirror, so bins[n // 2] == 0 and bins == -bins[::-1] exactly. The half is computed in float64 with numpy and rounded to float32 once, so the bins are trace-time constants, identical in eager and jitted code (deviations D22, T24; module docstring).

encode ¶

encode(y)

Two-hot weights [..., num_bins] for scalar targets y [...].

DreamerV3's algorithm (2411f7d:dreamerv3/jaxutils.py:257-271; dreamerv3_spec Algorithm D) applied to the transformed target: linear interpolation between the neighbouring bins b_below <= y < b_above; a target outside the bin range puts all its mass on the edge bin; a target equal to a bin puts all its mass on that bin. On TD-MPC2's uniform symlog grid this equals, in exact arithmetic, its floor formula k = floor((y - vmin) / step), t[k] = 1 - o, t[k + 1] = o after clip(y, vmin, vmax) (5f6fade:tdmpc2/common/math.py:72-77): a clipped target lands exactly on the edge bin. See the module docstring for the float32 agreement (deviation T24).

decode ¶

decode(logits)

Predicted scalars [...] from logits [..., num_bins].

inverse_transform(E_p[bins]) with p = softmax(logits) and the expectation summed symmetrically over mirror bins, as DreamerV3 does (2411f7d:dreamerv3/jaxutils.py:226-243; paper p.18) so that a zero-initialised head predicts exactly 0. DreamerV3 sums p_m b_m + sum_i (p_i b_i + p_j b_j) over the mirror pairs (i, j). Under jax.jit XLA contracts each pair into a fused multiply-add, which skips the rounding of one of the two products, so the pairs no longer cancel: on DreamerV3's bins the reference code then gives a constant at zero logits that depends on the compiled program (-0.17 to +0.16 on CPU) and, once the logits move, noise of std about 0.085 on every prediction, enough to slow the actor's early entropy decline (deviation D22). With b_m = 0 and b_i = -b_j exactly, the same sum is sum_j (p_j - p_i) b_j; the difference is exactly 0 for mirror-equal p before any multiplication, so the result is exactly 0 however the operations are fused. It stays noise-free only while the outer logits are bitwise mirror-symmetric, which depends on the compiled program: the outer bins (up to 4.85e8 on DreamerV3's bins) hold about 1/255 of the mass early on, so last-ulp differences between mirrored logits move the prediction by 1e-2 to 1e-1. A naive jnp.sum of p * bins gives 2.0 on DreamerV3's bins (CPU).

loss ¶

loss(logits, y)

Soft cross-entropy -sum(w * log_softmax(logits), -1), shape [...].

The log-probabilities are logits - logsumexp(logits) as in DreamerV3 (2411f7d:dreamerv3/jaxutils.py:272-274); TD-MPC2's F.log_softmax (5f6fade:tdmpc2/common/math.py:5-9) subtracts the max first, the same up to float32 rounding (deviation T24). The weights w = encode(y) are stop-gradiented. In both references the callers pass targets that carry no gradient (DreamerV3 agent.py:257 replay data, :349-350 and :371-373 log_prob(sg(...)); TD-MPC2 tdmpc2.py:201-202 and :231-233, TD targets under torch.no_grad, and replay rewards at :257); Ajax makes that explicit. There is no reduction over the batch; callers apply each paper's own weighting and averaging.

symlog ¶

symlog(x)

sign(x) * log(1 + |x|), elementwise and dtype-preserving.

DreamerV3 Eq. 9 (2411f7d:dreamerv3/jaxutils.py:76-77); TD-MPC2 5f6fade:tdmpc2/common/math.py:48-54 writes log(1 + |x|), equal to rounding. As in both references the derivative at exactly 0 is 0 (sign'(0) = sign(0) = 0); the limit from either side is 1.

symexp ¶

symexp(x)

sign(x) * (exp(|x|) - 1), the inverse of :func:symlog.

DreamerV3 2411f7d:dreamerv3/jaxutils.py:80-81; TD-MPC2 5f6fade:tdmpc2/common/math.py:57-63 (exp(|x|) - 1).

ajax.normalizers ¶

Percentile-range scale normalisers (DreamerV3 retnorm, TD-MPC2 RunningScale).

Both world-model agents divide a quantity by a running estimate of the spread between its 5th and 95th percentiles, floored at 1, so that the actor's learning signal is invariant to the reward scale without amplifying small returns. The papers describe the same idea ("an EMA of the 5-95 percentile range, floored at 1"), but the two implementations differ, and each normaliser below reproduces its own reference:

:class:ReturnNormalizer, DreamerV3's return normaliser (arXiv:2301.04104v2 p.6 and Eq. 7, Table 4; Moments(impl='perc', rate=0.01, limit=1.0) at danijar/dreamerv3@2411f7d:dreamerv3/jaxutils.py:301-395, configs.yaml:141): separate EMAs of the 5th and 95th percentiles, both starting at 0, no bias correction; the floor is applied when the scale is read, max(1, hi - lo). Bit-identical to the reference run op by op.
:class:RunningScale, TD-MPC2's Q normaliser (arXiv:2310.16828v2 Sec. 3.1, Table 8; nicklashansen/tdmpc2@5f6fade:tdmpc2/common/scale.py:4-45): one EMA starting at 1 of the range floored before averaging, S <- lerp(S, max(1, p95 - p5), tau). Equal to the reference up to float32 rounding of the EMA (deviation T25 in docs/world_models/deviations.md).

Clamping before or after the EMA, and starting at 0 or 1, give different scales on the same data (ranges alternating 0 and 3: 1.5075 vs 2.0050). Both EMAs are optax.incremental_update.

In both references the statistics are updated before the scale is read in the same training step, so the current batch already counts: DreamerV3 agent.py:329 (retnorm(ret, update)), TD-MPC2 tdmpc2.py:188-189 (update then divide). Percentiles are jnp.percentile with linear interpolation over all elements of the stop-gradiented float32 input; TD-MPC2's sort-and-interpolate _percentile (scale.py:23-35) is the same formula, and it always receives a (B, 1) batch, so "all elements" is "over the batch". Neither normaliser passes gradients.

ReturnNormalizer ¶

Bases: Replaceable

DreamerV3 return normaliser state; see the module docstring.

lo and hi are the EMAs of the 5th and 95th percentiles of the returns. Use as norm = norm.update(returns) followed by norm.scale() in the same step (update, then read). DreamerV3 divides advantages by the scale and does not subtract the offset lo.

Attributes:

Name	Type	Description
`lo,`	`hi`	float32 scalars, both 0 at :meth:`create`.
`rate`	`float`	EMA rate, 0.01 (decay 0.99).
`limit`	`float`	floor of the scale, 1.0.

create `classmethod` ¶

create(rate=0.01, limit=1.0)

Initial state: lo = hi = 0 (jaxutils.py:319-321).

lo and hi are two arrays, not one aliased twice, so that a state holding the normaliser can be donated to a jitted function.

update ¶

update(x)

Fold the percentiles of all elements of x into the EMAs.

lo <- (1 - rate) lo + rate P5(x), hi likewise with P95 (jaxutils.py:356-359). No bias correction. optax.incremental_update adds the same two products in the other order, which is exact: floating-point addition commutes.

scale ¶

scale()

max(limit, hi - lo), stop-gradiented (jaxutils.py:382-386).

RunningScale ¶

Bases: Replaceable

TD-MPC2 running Q scale state; see the module docstring.

Use as scale = scale.update(q) then q / scale.scale() in the same step (update before divide), with q the (B, 1) batch of t = 0 policy-loss Q values (tdmpc2.py:188-189).

Attributes:

Name	Type	Description
`value`	`Array`	the float32 scale `S`, 1 at :meth:`create`; always >= 1.
`rate`	`float`	EMA rate `tau`, 0.01 (TD-MPC2 reuses the target-network `tau`).

create `classmethod` ¶

create(rate=0.01)

Initial state: S = 1 (scale.py:9).

update ¶

update(x)

S <- (1 - rate) S + rate max(1, P95(x) - P5(x)) (scale.py:37-40).

The range is floored at 1 before it enters the average, so S never drops below 1. torch's lerp_ computes S + rate (v - S), the same average up to float32 rounding (deviation T25).

scale ¶

scale()

The current scale S, stop-gradiented.

Networks and modules¶

ajax.networks.networks ¶

CNNSpec ¶

CNNEncoder ¶

Actor ¶

Squashed ¶

Critic ¶

__call__ ¶

quantiles ¶

apply_encoder ¶

MultiCritic ¶

quantiles ¶

apply_encoder_ensemble ¶

reference_cnn ¶

check_log_std_head ¶

action_value_input ¶

get_initialized_actor_critic ¶

get_initialized_critic ¶

init_network_carry ¶

init_network_state ¶

in_critic_units ¶

predict_value_sequence ¶

phi_star_critic_state ¶

predict_value ¶

predict_value_train ¶

state_value ¶

action_value ¶

ajax.networks.distributions ¶

Deterministic ¶

log_prob ¶

GuardedTanh ¶

Squash ¶

log_prob_from_raw ¶

sample_with_raw ¶

entropy ¶

squash ¶

latent ¶

sample_with_raw ¶

log_prob_from_raw ¶

entropy ¶

ajax.networks.utils ¶

BatchRenorm ¶

TorchLinear ¶

clip_by_member_norm ¶

get_adam_tx ¶

parse_activation ¶

parse_function_string ¶

trunc_normal_fan_in ¶

get_initializer ¶

parse_initialization ¶

parse_layer ¶

dense ¶

parse_architecture ¶

ajax.networks.memory ¶

MemoryConfig dataclass ¶

ends_normalised property ¶

MemoryCell ¶

parse_memory_config ¶

with_default_init ¶

init_carry ¶

zeros_carry_like ¶

flatten_carry ¶

unflatten_carry ¶

flat_carry_dim ¶

tree_where_reset ¶

gate_block_orthogonal ¶

ajax.modules.pid_actor ¶

PIDPolicy dataclass ¶

PIDLaw ¶

gains ¶

PIDActor ¶

check_pid_actor ¶

pid_gains ¶

ajax.modules.pid_head ¶

PIDOutputHead ¶

__call__ ¶

init_pid_carry ¶

ajax.networks.blocks ¶

NormedMLP ¶

dreamerv3 classmethod ¶

call ¶

MemoryConfig `dataclass` ¶

ends_normalised `property` ¶

PIDPolicy `dataclass` ¶

call ¶

dreamerv3 `classmethod` ¶

tdmpc2 `classmethod` ¶

call ¶

TwoHot `dataclass` ¶

dreamerv3 `classmethod` ¶

tdmpc2 `classmethod` ¶

create `classmethod` ¶

create `classmethod` ¶