Networks

Core interfaces and layers for building PPO networks.

Types and Interfaces

class nnx_ppo.networks.types.PPONetworkOutput(actions, loglikelihoods, value_estimates)[source]

Bases: JaxDataclass

PPO-specific forward output. Produced by PPOAdapter.

Lives inside the output field of a StatefulModuleOutput. regularization_loss and metrics flow through the enclosing StatefulModuleOutput — they are not duplicated here.

actions: Any

loglikelihoods: PyTree[jaxtyping.Float[Array, '*time batch']]

value_estimates: Any

__init__(actions, loglikelihoods, value_estimates)

class nnx_ppo.networks.types.StatefulModuleOutput(next_state: jaxtyping.PyTree, output: Any, regularization_loss: jaxtyping.Float[Array, '*batch'], metrics: dict[Union[str, int], Any], rollout_extras: Any = None)[source]

Bases: JaxDataclass

next_state: PyTree

output: Any

regularization_loss: Float[Array, '*batch']

metrics: dict[str | int, Any]

rollout_extras: Any = None

__init__(next_state, output, regularization_loss, metrics, rollout_extras=None)

class nnx_ppo.networks.types.StatefulModule(*args, **kwargs)[source]

Bases: ABC, Module

Abstract base class for network modules and layers, specifying the interface with the RL algorithm. Each module treated as stateful, with non-stateful modules having empty states.

There are two types of module state. First, there is the state of the nnx.Module which follows the common NNX patterns. This state stores trainable parameters (of type nnx.Param) as well as any RNGStreams and other variables. This state is _not_ reset when the RL environment is reset, and may not be written from inside __call__ if those writes affect the forward output.

Stats-bearing modules (e.g. Normalizer) accumulate state by overriding update_statistics(), which is called once per training step after the loss / gradient update.

Second, there is an explicit carry state that is intended for stateful network layers, e.g. the hidden activations of RNNs. This state _is_ reset when the RL environment is reset.

__call__ takes rollout_extras as a third positional argument (default None). It is a pytree shaped like the module’s own contribution to the network’s ROLLOUT → LOSS_REPLAY communication channel. In ROLLOUT the module emits its contribution as part of the returned StatefulModuleOutput; in LOSS_REPLAY the same value is fed back in via this argument. Modules that don’t need replay information leave it None. The phase a module is in is derivable from this argument: None means ROLLOUT or INFERENCE (sample fresh / emit); a non-None value means LOSS_REPLAY (consume).

abstractmethod __call__(module_state, obs, rollout_extras=None)[source]

Parameters:

module_state (PyTree) – The current state of the module.
obs (PyTree) – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: PyTree

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

Containers

class nnx_ppo.networks.containers.Sequential(*args, **kwargs)[source]

Bases: StatefulModule

__init__(layers)[source]

__call__(network_state, obs, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs (Any) – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: list[PyTree]

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

class nnx_ppo.networks.containers.Concat(*args, **kwargs)[source]

Bases: StatefulModule

Per-key dispatch + concat: dict input, single-tensor output.

Each named sub-module sees the upstream’s same-named entry as input; the per-component outputs are concatenated along the last axis. Accepts either keyword arguments or a positional dict Concat({...}) (when keys are not valid Python identifiers).

__init__(modules=None, /, **kwargs)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: dict[str, PyTree]

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

class nnx_ppo.networks.containers.Parallel(*args, **kwargs)[source]

Bases: StatefulModule

Runs several sub-modules on the same input and returns their outputs as a dict keyed by sub-module name.

Typical use: assemble a trunk that produces both action-distribution parameters and value estimates from shared upstream features:

trunk = Sequential([
    shared_encoder,
    Parallel(action_params=actor_head, value=critic_head),
])
# trunk(state, x).output is {"action_params": ..., "value": ...}

__init__(modules=None, /, **kwargs)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: dict[str, PyTree]

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

class nnx_ppo.networks.containers.Splitter(*args, **kwargs)[source]

Bases: StatefulModule

Splits a single input tensor into a dict of named slices along the last axis.

Used at the end of a stack to turn a flat tensor head into a structured dict output that an adapter can route to samplers / value specs:

Sequential([
    trunk,
    Dense(hidden, 2 * action_size + 1, rngs),
    Splitter(action_params=2 * action_size, value=1),
])

With a single keyword (Splitter(action_params=N)) the layer simply relabels the input as a dict, taking the first N features.

The slices are taken in keyword-argument insertion order. The sum of declared sizes must not exceed the input’s last-axis size; any excess input features are silently ignored, matching plain slicing semantics.

__init__(**sizes)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

Utilities

Small utility StatefulModule s.

These are stateless layers that handle pytree projection / reshaping / scalar arithmetic. Container utilities route their children’s rollout_extras the same way they route state.

Available:

Flattener — flatten a pytree into one tensor (depth 0), or preserve the top preserve_levels levels of dict/list/tuple structure and flatten below.
Filter — declarative pytree extraction/projection. Takes a dict spec keyed by output name; each entry is a string (top-level key), a tuple of strings/ints (nested path), or a callable applied to the full input.
Scale — multiply by a fixed scalar.
Merge — run several named sub-modules on the same input, each producing a dict, and merge them into one flat dict. The natural complement to Parallel when downstream consumers (e.g. PPOAdapter) want one flat dict of named heads.
Map — per-key dispatch: dict input, dict output. Each named sub-module sees the upstream’s same-named entry.

class nnx_ppo.networks.utils.Flattener(*args, **kwargs)[source]

Bases: StatefulModule

Flatten a pytree into a tensor (or a dict-of-tensors).

With the default preserve_levels=0 every leaf is reshaped to (B, -1) and concatenated along the last axis, producing one flat tensor.

With preserve_levels=N, the top N levels of dict / list / tuple structure are preserved and only the sub-trees below them are flattened. So Flattener(preserve_levels=1) applied to {"a": {"p": (B, 4), "t": (B, 8)}, "b": (B, 6)} returns {"a": (B, 12), "b": (B, 6)}.

Idempotent on already-flat inputs at the appropriate depth: passing {"a": (B, 12), "b": (B, 6)} through Flattener(preserve_levels=1) is a no-op (each value is reshape-and-concat of a single leaf).

__init__(preserve_levels=0)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

class nnx_ppo.networks.utils.Filter(*args, **kwargs)[source]

Bases: StatefulModule

Declarative pytree extraction / projection.

spec is a dict {output_key: extraction} where each extraction is one of:

a string k — take x[k];
a tuple of strings/ints (k1, k2, ...) — nested path, equivalent to x[k1][k2]...;
a callable fn — applied to the full input x; fn(x) becomes the value for output_key.

The result is a dict with the same keys as spec. Anything in the input not named (directly or via a callable) is dropped.

__init__(spec)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

class nnx_ppo.networks.utils.Scale(*args, **kwargs)[source]

Bases: StatefulModule

Multiply the input by a fixed scalar factor.

__init__(factor)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

class nnx_ppo.networks.utils.Merge(*args, **kwargs)[source]

Bases: StatefulModule

Run named sub-modules on the same input; merge their dict outputs.

Each sub-module must return a dict. The outputs are merged into one flat dict; duplicate keys across components are a hard error.

Accepts either keyword arguments (when names are valid Python identifiers) or a positional dict Merge({...}) (when they are not).

Carry state is a dict {name: component_state} — same shape as Parallel.

__init__(modules=None, /, **kwargs)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: dict[str, PyTree]

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

class nnx_ppo.networks.utils.Map(*args, **kwargs)[source]

Bases: StatefulModule

Per-key dispatch: dict input, dict output.

Each named sub-module sees the upstream’s same-named entry as input and produces the same-named entry of the output. Distinct from:

Parallel — same input fed to every component, dict output.
Concat — per-key dispatch, concatenated output (no dict).

The input dict must contain at least every key in modules; extra keys are dropped.

Accepts either keyword arguments or a positional dict Map({...}) (necessary when keys are not valid Python identifiers).

Carry state is a dict {name: component_state}.

__init__(modules=None, /, **kwargs)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: dict[str, PyTree]

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

Action Sampling

Action samplers.

Action samplers turn pre-distribution parameters (e.g. concatenated mean and log_std) into a sampled action plus its log-likelihood. They are ordinary StatefulModule s: they live inside the network just like any other layer, and the user composes them with the other containers (typically as the last layer of the action= port of a PPOAdapter).

Sampler behaviour is driven by rollout_extras:

rollout_extras is None (ROLLOUT and INFERENCE): sample fresh. In the returned StatefulModuleOutput.rollout_extras field, emit the sampled raw action so a later LOSS_REPLAY pass can reproduce it.
rollout_extras is not None (LOSS_REPLAY): use the stored raw action to compute the log-likelihood under the current policy. The RNG stream still advances so any downstream stochastic layers stay in lockstep with the rollout.

The per-instance deterministic flag is orthogonal: when set to True (typically by network.eval()), the sampler returns the mean instead of sampling. It applies regardless of whether rollout_extras is provided.

Forward output is a small dict {"action", "log_likelihood"}. The enclosing PPOAdapter lifts each field into the matching dict on PPONetworkOutput. Mean / std live in the sampler’s metrics for logging.

class nnx_ppo.networks.sampling_layers.ActionSampler(*args, **kwargs)[source]

Bases: StatefulModule, ABC

deterministic: bool = False

abstractmethod __call__(state, mean_and_std, rollout_extras=None)[source]

Apply the sampler.

Parameters:

state (tuple[()]) – Empty tuple (stateless sampler).
mean_and_std (Float[Array, 'batch mean_std_dim']) – Concatenated mean and std, shape [batch, 2 * action_dim].
rollout_extras (Float[Array, 'batch action_dim'] | None) – None to sample fresh (ROLLOUT / INFERENCE); the stored action to reuse (LOSS_REPLAY).

class nnx_ppo.networks.sampling_layers.NormalTanhSampler(*args, **kwargs)[source]

Bases: ActionSampler

Normal distribution followed by tanh.

__init__(rng, entropy_weight, min_std=0.001, std_scale=1.0)[source]

__call__(state, mean_and_std, rollout_extras=None)[source]

Apply the sampler.

Parameters:

state (tuple[()]) – Empty tuple (stateless sampler).
mean_and_std (Float[Array, 'batch mean_std_dim']) – Concatenated mean and std, shape [batch, 2 * action_dim].
rollout_extras (Float[Array, 'batch action_dim'] | None) – None to sample fresh (ROLLOUT / INFERENCE); the stored action to reuse (LOSS_REPLAY).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: tuple[()]

PPO Adapter

PPOAdapter: two-port router from network output to PPONetworkOutput.

PPOAdapter is a tiny StatefulModule (~50 LoC) that takes two sub-modules — an action port and a value port — runs both on its upstream input, and packages their outputs into a PPONetworkOutput. It is the canonical way to make a Sequential trunk into a PPO-shaped network:

pipeline = Sequential([
    Normalizer(obs_shape),
    trunk,                            # produces {action_params, value}
    PPOAdapter(
        action=Sequential([
            Filter({"action_params": "action_params"}),
            NormalTanhSampler(rngs, entropy_weight=1e-2),
        ]),
        value=Filter({"value": "value"}),
    ),
])

The action port’s forward output must be a tree of sampler dicts. A sampler dict is the small {"action", "log_likelihood"} payload each ActionSampler returns. For a single-sampler action port the output IS a sampler dict; for a per-key sampler bank (Map({pop: sampler})) the output is {pop: sampler_dict}. The adapter extracts fields uniformly via jax.tree.map with an is_leaf recogniser.

The value port’s forward output is taken as-is for PPONetworkOutput.value_estimates. Trailing singleton axes are squeezed for ergonomic shape ([B, 1] → [B]).

class nnx_ppo.networks.adapter.PPOAdapter(*args, **kwargs)[source]

Bases: StatefulModule

Two-port router producing PPONetworkOutput.

Parameters:

action – The action port. Its forward output is a tree of sampler dicts {"action", "log_likelihood"}.
value – The value port. Its forward output is used directly as value_estimates (trailing singleton axes are squeezed).

__init__(action, value)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: dict[str, PyTree]

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

Feedforward Layers

Feedforward network layers.

class nnx_ppo.networks.feedforward.Dense(*args, **kwargs)[source]

Bases: StatefulModule

Single dense layer with optional activation, implementing StatefulModule.

This is a thin wrapper around nnx.Linear that conforms to the StatefulModule interface, allowing it to be used in Sequential containers.

__init__(in_features, out_features, rngs, activation=None, **linear_kwargs)[source]

Initialize the dense layer.

Parameters:

in_features (int) – Number of isnput features.
out_features (int) – Number of output features.
rngs (Rngs) – NNX random number generators.
activation (Callable | None) – Optional activation function applied after the linear transform.
**linear_kwargs – Additional arguments passed to nnx.Linear (e.g., kernel_init).

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

Recurrent Layers

Recurrent network modules for PPO.

class nnx_ppo.networks.recurrent.LSTM(*args, **kwargs)[source]

Bases: StatefulModule

LSTM layer that conforms to the StatefulModule interface.

Wraps flax.nnx.LSTMCell to provide proper state management for RL rollouts. The hidden state is reset when the environment resets.

Example usage:: lstm = LSTM(in_features=64, hidden_features=128, rngs=nnx.Rngs(0)) state = lstm.initialize_state(batch_size=32) output = lstm(state, x) # x has shape (32, 64) # output.output has shape (32, 128) # output.next_state is (h, c) tuple for next timestep

__init__(in_features, hidden_features, rngs, *, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function tanh>>, kernel_init=None, recurrent_kernel_init=None, bias_init=None, use_optimized=True, trainable_initial_state=False)[source]

Initialize the LSTM layer.

Parameters:

in_features (int) – Number of input features.
hidden_features (int) – Number of hidden units (output size).
rngs (Rngs) – NNX random number generators.
gate_fn (Callable) – Activation function for gates (default: sigmoid).
activation_fn (Callable) – Activation function for cell state (default: tanh).
kernel_init (Callable | None) – Initializer for input-to-hidden weights.
recurrent_kernel_init (Callable | None) – Initializer for hidden-to-hidden weights.
bias_init (Callable | None) – Initializer for biases.
use_optimized (bool) – If True, use OptimizedLSTMCell which is faster for hidden_features <= 2048.
trainable_initial_state (bool) – If True, the initial hidden and cell states are learnable parameters. Otherwise, they are zeros.

__call__(state, x, rollout_extras=None)[source]

Process input through the LSTM.

Parameters:

state (tuple[Float[Array, '*batch hidden'], Float[Array, '*batch hidden']]) – LSTM carry tuple (hidden_state, cell_state), each with shape (batch_size, hidden_features).
x (Float[Array, 'batch {self.in_features}']) – Input array with shape (batch_size, in_features).

Returns:

next_state: Updated (h, c) carry tuple
output: LSTM output with shape (batch_size, hidden_features)
regularization_loss: Zero (LSTM has no regularization)
metrics: Empty dict

Return type:

StatefulModuleOutput with

initialize_state(batch_size)[source]

Initialize the LSTM hidden state.

Parameters:: batch_size (int) – Number of parallel environments/sequences.
Returns:: Tuple of (hidden_state, cell_state), each with shape (batch_size, hidden_features). If trainable_initial_state=True, these are learned parameters; otherwise zeros.
Return type:: tuple[Float[Array, ‘*batch hidden’], Float[Array, ‘*batch hidden’]]

reset_state(prev_state)[source]

Reset LSTM state (called when environment resets).

Parameters:: prev_state (tuple[Float[Array, '*batch hidden'], Float[Array, '*batch hidden']]) – Previous carry state (used to preserve shape).
Returns:: Initial carry with same shape as prev_state. If trainable_initial_state=True, returns learned initial state broadcast to match prev_state shape; otherwise returns zeros.
Return type:: tuple[Float[Array, ‘*batch hidden’], Float[Array, ‘*batch hidden’]]

Delays

k-step delay layer.

Generalises the per-network DelayedObsNetwork wrapper (previously in vnl-experiments) into a composable StatefulModule that can be inserted anywhere a layer fits: as an element of a Sequential, as the transform of a graph connection, or as a wrapper around any module producing an array / pytree of arrays.

class nnx_ppo.networks.delay.Delay(*args, **kwargs)[source]

Bases: StatefulModule

k-step delay.

Output at time t is the input from time t - k_steps. Before the buffer fills (t < k_steps), output is initial_value (default zero).

Carry state is a dict:

{"buffer": <pytree mirroring input, leaves [B, k_steps, *leaf]>,
 "idx":    <[B] int32 circular write pointer>}

reset_state zeros both buffer and idx, which is the correct behaviour at episode boundaries.

Example

sample_obs = jax.jit(env.reset)(jax.random.key(0)).obs net = Sequential([Delay(sample_obs, k_steps=5), inner_network])

__init__(sample_input, k_steps, initial_value=0.0)[source]

Initialise the delay buffer’s pytree spec from a sample input.

Parameters:

sample_input (Any) – A single unbatched example of the input PyTree. Used only to capture the leaf shapes, dtypes, and tree structure for buffer allocation. The values themselves are not retained.
k_steps (int) – Delay length in steps. Must be >= 1.
initial_value (float) – Value used to fill the buffer before it has been written k_steps times (and on reset_state).

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: dict

reset_state(prev_state)[source]

Population Graphs

Population-graph network container.

Public surface: PopulationGraph. Population and Connection are internal node/edge specs constructed by the graph’s build methods (add_population / connect).

class nnx_ppo.networks.graph.PopulationGraph(*args, **kwargs)[source]

Bases: StatefulModule

Declarative population/connection graph as a StatefulModule.

Construction is two-phase:

add_population / add_input / add_output / connect to describe the graph.
finalize() to validate (cycle detection, shape inference, buffer sizing) and freeze the graph for forward passes.

After finalize(), the graph behaves as any other StatefulModule: it can be a layer in Sequential, an inner module of an adapter, etc.

__init__(rngs)[source]

add_population(name, size, *, activation=None)[source]

Register an internal population.

Parameters:

name (str) – Unique identifier.
size (int) – Size of the population’s activation vector. All incoming connections must produce vectors of this size after their transform.
activation (Callable | None) – Optional transfer function applied elementwise once after the sum integration of incoming connections.

add_input(name, size, *, input_from, activation=None)[source]

Register an input population that reads from obs[input_from].

The corresponding obs[input_from] value must have shape [B, size] at call time. The obs value is added to the population’s integrated input — incoming connections (if any) are summed in alongside it before the activation.

add_output(name, size, *, output_to=None, activation=None)[source]

Register an output population whose activation is exposed.

The population’s post-activation activation appears in the graph’s forward output dict under output_to (or under name if output_to is None).

connect(src, dst, *, transform=None, delay=0, reciprocal=False)[source]

Add a directed connection from src to dst.

Parameters:

src (str) – Source population name. Must already be registered.
dst (str) – Destination population name. Must already be registered.
transform (StatefulModule | None) – A StatefulModule mapping [B, src.size] -> [B, dst.size]. Defaults to a linear Dense of the appropriate shape.
delay (int) – Integer step delay. 0 reads the source’s output from the current step; k >= 1 reads from k steps ago. Before the buffer fills, delayed reads return zeros.
reciprocal (bool) – If True, additionally add the reverse connection dst -> src with the same delay and an independent default Dense transform. transform must be None when reciprocal is True (use two explicit connect() calls if you need custom transforms in both directions).

finalize()[source]

Validate and freeze the graph.

Performs:

delay-0 cycle detection (any cycle is a hard error);
topological ordering of populations (delay-0 edges only);
per-population max_outgoing_delay computation;
promotion of populations / connections to NNX containers so their parameters are tracked.

__call__(state, obs, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs (Any) – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: PyTree

reset_state(prev_state)[source]

class nnx_ppo.networks.graph.Population(*args, **kwargs)[source]

Bases: Module

Internal node spec for PopulationGraph.

__init__(name, size, activation, input_from, output_to)[source]

class nnx_ppo.networks.graph.Connection(*args, **kwargs)[source]

Bases: Module

Internal edge spec for PopulationGraph.

__init__(src, dst, transform, delay)[source]

Population-graph container.

PopulationGraph is a StatefulModule that owns a set of named Population nodes and typed Connection edges between them.

Each population has a single size. Connections between populations default to a linear Dense transform sized from source to destination. Each population sum-integrates its incoming connections (and optionally the corresponding obs entry, for input populations), then applies its activation (transfer function) once. Connections carry an integer delay; delay=0 reads the source’s freshly computed output in the same step (the topological order guarantees the source has already run), delay=k reads from k steps ago via a per-population shared circular buffer.

The build API is three methods:

add_population(name, size, *, activation=None)() — internal pop.
add_input(name, size, *, input_from=key, activation=None)() — reads obs[key] as the population’s integrated input.
add_output(name, size, *, output_to=None, activation=None)() — the population’s post-activation activation is exposed under output[output_to or name] in the graph’s forward output.

Plus connect(src, dst, *, transform=None, delay=0, reciprocal=False)() to wire populations together, and finalize() to validate.

Use PopulationGraph when you have a graph with multiple connections per node, recurrent loops, or per-connection delays. For straight encoder/decoder stacks Sequential is simpler.

class nnx_ppo.networks.graph.graph.PopulationGraph(*args, **kwargs)[source]

Bases: StatefulModule

Declarative population/connection graph as a StatefulModule.

Construction is two-phase:

add_population / add_input / add_output / connect to describe the graph.
finalize() to validate (cycle detection, shape inference, buffer sizing) and freeze the graph for forward passes.

After finalize(), the graph behaves as any other StatefulModule: it can be a layer in Sequential, an inner module of an adapter, etc.

__init__(rngs)[source]

add_population(name, size, *, activation=None)[source]

Register an internal population.

Parameters:

name (str) – Unique identifier.
size (int) – Size of the population’s activation vector. All incoming connections must produce vectors of this size after their transform.
activation (Callable | None) – Optional transfer function applied elementwise once after the sum integration of incoming connections.

add_input(name, size, *, input_from, activation=None)[source]

Register an input population that reads from obs[input_from].

The corresponding obs[input_from] value must have shape [B, size] at call time. The obs value is added to the population’s integrated input — incoming connections (if any) are summed in alongside it before the activation.

add_output(name, size, *, output_to=None, activation=None)[source]

Register an output population whose activation is exposed.

The population’s post-activation activation appears in the graph’s forward output dict under output_to (or under name if output_to is None).

connect(src, dst, *, transform=None, delay=0, reciprocal=False)[source]

Add a directed connection from src to dst.

Parameters:

src (str) – Source population name. Must already be registered.
dst (str) – Destination population name. Must already be registered.
transform (StatefulModule | None) – A StatefulModule mapping [B, src.size] -> [B, dst.size]. Defaults to a linear Dense of the appropriate shape.
delay (int) – Integer step delay. 0 reads the source’s output from the current step; k >= 1 reads from k steps ago. Before the buffer fills, delayed reads return zeros.
reciprocal (bool) – If True, additionally add the reverse connection dst -> src with the same delay and an independent default Dense transform. transform must be None when reciprocal is True (use two explicit connect() calls if you need custom transforms in both directions).

finalize()[source]

Validate and freeze the graph.

Performs:

delay-0 cycle detection (any cycle is a hard error);
topological ordering of populations (delay-0 edges only);
per-population max_outgoing_delay computation;
promotion of populations / connections to NNX containers so their parameters are tracked.

__call__(state, obs, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs (Any) – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: PyTree

reset_state(prev_state)[source]

Variational Layers

class nnx_ppo.networks.variational.VariationalBottleneck(*args, **kwargs)[source]

Bases: StatefulModule

Variational bottleneck with KL regularization.

Takes input of size 2*latent_size (mean and log_std concatenated), samples from the latent distribution using the reparameterization trick, and adds KL divergence against a standard normal prior to the regularization loss.

__init__(latent_size, rng, kl_weight=1.0, min_std=1e-06)[source]

Initialize the variational bottleneck.

Parameters:

latent_size (int) – Size of the latent space. Input is expected to be 2*latent_size.
rngs – NNX random number generators.
kl_weight (float) – Weight for the KL divergence regularization term.
min_std (float) – Minimum standard deviation to prevent numerical instability.

__call__(key, x, rollout_extras=None)[source]

Sample from the variational distribution.

Parameters:

key (Key[Array, 'batch']) – rng key.
x (Float[Array, 'batch {2*self.latent_size}']) – Input array of shape (…, 2*latent_size) containing concatenated mean and log_std.

Returns:

next_state: RNG keys
output: Sampled latent vector of shape (…, latent_size)
regularization_loss: KL divergence weighted by kl_weight
metrics: Dictionary with mu, sigma, and kl_divergence

Return type:

StatefulModuleOutput with

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: Key[Array, ‘batch’]

reset_state(prev_state)[source]

class nnx_ppo.networks.variational.AR1VariationalBottleneck(*args, **kwargs)[source]

Bases: StatefulModule

Variational bottleneck with KL and auto-regressive regularization.

Takes input of size 2*latent_size (mean and log_std concatenated), samples from the latent distribution using the reparameterization trick, adds KL divergence against a standard normal prior to the regularization loss, and adds a first-order autoregressive loss. The autoregressive loss encourages smoother trajectories in the latent space.

__init__(latent_size, rng, kl_weight=1.0, min_std=1e-06, ar1_weight=1.0, backprop_through_time=True)[source]

Initialize the variational bottleneck.

Parameters:

latent_size (int) – Size of the latent space. Input is expected to be 2*latent_size.
rngs – NNX random number generators.
kl_weight (float) – Weight for the KL divergence regularization term.
min_std (float) – Minimum standard deviation to prevent numerical instability.
ar1_weight (float) – Weight for the autoregressive loss
backprop_through_time (bool) – Conceptually, the AR1 loss can be minimized in two ways: either by making the current latent vector (z) closer to the previous latent vector (prev_z), or by making prev_z closer to z. The latter requires the gradient to flow back in time, which is perfectly fine in most instances. However, this param can be set to False to turn that off.

__call__(state, x, rollout_extras=None)[source]

Sample from the variational distribution.

Parameters:

state (dict[str, Any]) – Dict with ‘keys’ (PRNGKeyArray) and ‘last_z’ (Float[batch, latent]).
x (Float[Array, 'batch {2*self.latent_size}']) – Input array of shape (batch, 2*latent) containing concatenated mean and log_std.

Returns:

next_state: RNG keys and prev_z
output: Sampled latent vector of shape (batch, latent)
regularization_loss: sum of KL loss and AR1 loss
metrics: Dictionary with mu, sigma, kl_divergence, and squared diff

Return type:

StatefulModuleOutput with

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: dict[str, Any]

reset_state(prev_state)[source]

Observation Normalization

class nnx_ppo.networks.normalizer.NormalizerStatistics(*args, **kwargs)[source]: Bases: Variable

class nnx_ppo.networks.normalizer.Normalizer(*args, **kwargs)[source]

Bases: StatefulModule

Online (Welford) input normalizer.

Forward pass standardises x using the running mean and standard deviation derived from M2 and counter. The running statistics are read-only in __call__ — never written from the forward path.

Stats are updated once per training step via update_statistics(), which receives the rollout’s history of normalizer inputs (one [T, B, *feat] slice per leaf) and folds it in with a single batched Welford merge. The values it sees are exactly the activations the Normalizer received during ROLLOUT (the forward pass emits them as rollout_extras), so placing the Normalizer anywhere — behind a Delay, inside a graph population, after an encoder — works automatically.

__init__(shape)[source]

__call__(state, x, rollout_extras=None)[source]

Parameters:

module_state – The current state of the module.
obs – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of normalizer inputs into running stats.

rollout_extras is a pytree matching the structure the Normalizer emitted on each step, with an additional leading time dimension — leaves have shape [T, B, *feat]. We flatten T*B and apply one batched Welford merge.

Factory Functions

nnx_ppo.networks.factories.make_mlp_layers(sizes, rngs, activation=<jax._src.custom_derivatives.custom_jvp object>, activation_last_layer=True, **linear_kwargs)[source]

Create a list of Dense layers for an MLP.

Use this when embedding MLP layers within another Sequential:

Sequential([: Flattener(), *make_mlp_layers([64, 32, 16], rngs), SomeOtherModule(),

])

Parameters:

sizes (list[int]) – List of layer sizes including input and output.
rngs (Rngs) – NNX random number generators.
activation (Callable) – Activation function to use between layers.
activation_last_layer (bool) – Whether to apply activation after the last layer.
**linear_kwargs – Additional arguments passed to nnx.Linear.

Returns:

A list of Dense layers.

Return type:

list[Dense]

nnx_ppo.networks.factories.make_mlp(sizes, rngs, activation=<jax._src.custom_derivatives.custom_jvp object>, activation_last_layer=True, **linear_kwargs)[source]

Create an MLP as a Sequential of Dense layers.

Parameters:

sizes (list[int]) – List of layer sizes including input and output.
rngs (Rngs) – NNX random number generators.
activation (Callable) – Activation function to use between layers.
activation_last_layer (bool) – Whether to apply activation after the last layer.
**linear_kwargs – Additional arguments passed to nnx.Linear.

Returns:

A Sequential container of Dense layers.

Return type:

Sequential

nnx_ppo.networks.factories.make_mlp_actor_critic(obs_size, action_size, actor_hidden_sizes, critic_hidden_sizes, rngs, activation=<jax._src.custom_derivatives.custom_jvp object>, normalize_obs=True, initializer_scale=1.0, entropy_weight=0.01, min_std=0.1, std_scale=1.0)[source]

Build a standard one-actor / one-critic PPO network.

Returns a Sequential whose forward output is a PPONetworkOutput. Pass it straight to train_ppo().

The constructed pipeline is:

Sequential([
    Normalizer(obs_size)?,        # if normalize_obs
    PPOAdapter(
        action=Sequential([actor, NormalTanhSampler(...)]),
        value=critic,
    ),
])

Both adapter ports receive the same upstream input (the normalised obs), so there is no shared trunk; the actor and critic each run independently. Insert a shared trunk by prepending it to the Sequential and pointing the ports at it.

Activation Recording

Activation recording for network modules.

Recording is an eval/analysis-only feature for inspecting the per-unit activations of a network. The design keeps all recording logic out of the modules themselves: a leaf module’s forward output already is its activation, and the containers already propagate each child’s metrics up the tree into a nested, named dict. We exploit both.

with_recording() returns a separate copy of a network in which every leaf StatefulModule is wrapped in a Recorder. A Recorder runs its wrapped module unchanged and injects that module’s output into the returned StatefulModuleOutput.metrics under the reserved key ACTIVATION_KEY. The activation therefore rides the existing metrics channel to the top-level call for free, where extract_activations() pulls it back out, keyed by each module’s structural path in the container tree (e.g. action/0, action/1, value/...).

Because the original network is never modified, recording is “off” by absence: there is no if record: branch anywhere, and the unwrapped network behaves identically jitted, non-jitted, and in tests.

The graph network (PopulationGraph) is a special case: its interesting units are internal populations (plain dataclasses, not sub-modules), and it drops its children’s metrics. Instead of being leaf-wrapped, it has its own record_activations flag that with_recording() enables; when set it emits all population activations into its metrics under ACTIVATION_KEY.

Per-timestep activations can only leave an nnx.scan via its ys return (stacking, not reduction). eval_rollout reduces metrics to scalars and is therefore not usable for per-unit recording. Use a plain Python eval loop (calling extract_activations() on out.metrics each step) or the convenience record_activations_rollout().

class nnx_ppo.networks.recording.Recorder(*args, **kwargs)[source]

Bases: StatefulModule

Wraps a single leaf module and records its forward output.

Delegates every part of the StatefulModule interface to the wrapped module verbatim; the only addition is the wrapped module’s output placed into the returned metrics under ACTIVATION_KEY. The forward output / next_state / rollout_extras pass through unchanged, so a recording network’s forward pass is bit-identical to the original.

__init__(wrapped)[source]

__call__(module_state, obs, rollout_extras=None)[source]

Parameters:

module_state (PyTree) – The current state of the module.
obs (Any) – Observation(s) to process. Leading dimension is batch.
rollout_extras (Any) – Replay snapshots for this module (and, via container routing, its descendants). None in ROLLOUT / INFERENCE; the stored value from Transition.rollout_extras in LOSS_REPLAY.

Returns StatefulModuleOutput with next_state, output, regularization_loss, metrics, and rollout_extras (the snapshot to be stored on the transition for later replay).

initialize_state(batch_size)[source]

Create a new state for this module.

Parameters:: batch_size (int) – Batch size (leading dimension on any returned arrays).
Returns:: A ModuleState with batch_size as the leading dimension. Default is an empty tuple (stateless module).
Return type:: PyTree

reset_state(prev_state)[source]

update_statistics(rollout_extras)[source]

Fold the rollout’s worth of replay snapshots into any running statistics this module owns. Called once per training step after the loss / gradient update.

Containers override this to route rollout_extras per child the same way they route state. Stats-bearing leaves (e.g. Normalizer) consume their [T, B, *feat] history slice and update their NNX variables in place. Default is a no-op.

nnx_ppo.networks.recording.with_recording(net)[source]

Return a recording copy of net; the original is left untouched.

Every leaf module is wrapped in a Recorder; every PopulationGraph has its record_activations flag enabled. The rebuild uses nnx.recursive_map (bottom-up), so leaves are wrapped exactly once and parameters are shared with the original (read-only at eval).

If net already contains Recorder layers it is already recordable; a warning is issued and it is returned unchanged (wrapping it again would double-wrap the inner leaves).

nnx_ppo.networks.recording.extract_activations(metrics)[source]

Pull the recorded activations out of a network’s out.metrics tree.

Walks the nested metrics dict and returns a tree of the same container shape in which each recorded module is replaced by its activation (the value stored under ACTIVATION_KEY). Real scalar metrics and branches without any activations are dropped. Returns None if there are no activations (e.g. the network was not made recordable).