Custom networks 3: writing your own StatefulModule

The library ships a small but growing set of layers (Dense, LSTM, Normalizer, Delay, VariationalBottleneck, …). When you need something else, write your own StatefulModule. This tutorial walks through the four things you need to know to do that.

What we’ll build

A MovingAverage layer that maintains, for each batch element, a running average of its last k inputs and emits that average. We will also give it a running global mean across all observed inputs — not because the layer needs it, but to show how a layer can update NNX-tracked statistics safely.

The three things to learn

  1. Carry state vs. NNX state.

  2. The ``rollout_extras`` argument and the ``update_statistics`` hook.

  3. NNX containers for variable-size children.

Carry state vs. NNX state

Every StatefulModule has two kinds of state:

  • NNX-tracked state lives on self as attributes typed nnx.Param, nnx.Variable, or sub-modules. NNX traces these through nnx.jit and nnx.grad automatically. Params get gradients; non-Param variables (like running statistics) do not. NNX state is not reset on environment reset.

  • Carry state is the explicit state argument to __call__ and the next_state field of the returned StatefulModuleOutput. It carries per-batch quantities — LSTM hidden states, delay buffers, RNG carries for stochastic layers. The training loop does reset carry state when the environment resets via reset_state().

Rule of thumb: if it’s “per-batch transient state that should be zeroed when an episode ends”, it goes in carry state. If it’s “per-network state that survives episodes” (params, RNG streams, running stats), it goes in NNX state.

For the moving average:

  • The ring buffer of last-k inputs is carry state (each env’s window resets when that env’s episode resets).

  • The running global mean / count are NNX state (a single per-network running statistic, not per-batch).

A first cut: forward pass and carry state

import jax.numpy as jp
from flax import nnx

from nnx_ppo.networks.types import (
    StatefulModule,
    StatefulModuleOutput,
)


class MovingAverage(StatefulModule):
    """Per-batch running mean over the last `k` inputs."""

    def __init__(self, feature_size: int, k: int):
        self.feature_size = feature_size
        self.k = k

    def __call__(self, state, x, rollout_extras=None):
        # state["buffer"] : (B, k, feature_size)
        # state["idx"]    : (B,)  circular write pointer
        idx = state["idx"]
        batch = jp.arange(idx.shape[0])

        new_buffer = state["buffer"].at[batch, idx].set(x)
        new_idx = (idx + 1) % self.k

        mean = jp.mean(new_buffer, axis=1)
        return StatefulModuleOutput(
            next_state={"buffer": new_buffer, "idx": new_idx},
            output=mean,
            regularization_loss=jp.array(0.0),
            metrics={},
            rollout_extras=None,
        )

    def initialize_state(self, batch_size: int):
        return {
            "buffer": jp.zeros((batch_size, self.k, self.feature_size)),
            "idx": jp.zeros(batch_size, jp.int32),
        }

    def reset_state(self, prev_state):
        return {
            "buffer": jp.zeros_like(prev_state["buffer"]),
            "idx": jp.zeros_like(prev_state["idx"]),
        }

Some things to notice:

  • __call__ accepts and returns explicit carry state. It does not mutate self. The caller is responsible for keeping the returned next_state alive for the next step.

  • __init__ only stores config (feature_size, k). For real parameters you would assign nnx.Param or sub-modules here.

  • reset_state() is what the training loop calls on per-environment episode boundaries (via unroll_env()). It must preserve shapes, since it runs under vmap.

The rollout_extras argument

Every StatefulModule __call__ takes a third positional argument, rollout_extras. It is a pytree shaped like state, threaded through every container. Modules that need to communicate information from the rollout pass back to a later loss-replay pass emit a value in the returned StatefulModuleOutput.rollout_extras and consume one via the argument. The three phases:

  • Rollout — caller passes rollout_extras=None. Each module emits its replay snapshot (action samplers store the sampled raw action; normalizers store the input they just normalised). The rollout scan stacks these over T into Transition.rollout_extras.

  • Loss replay — caller passes the stored slice back in. Modules that need replay info (notably action samplers) consume it to reproduce the actually-taken action’s log-likelihood under the updated policy.

  • Inference — caller passes nothing. Modules still emit but the caller drops the emission on the floor.

A module that needs to switch behaviour between “fresh sample” and “use stored value” checks if rollout_extras is None.

For the full per-module behaviour table see The rollout_extras channel.

The library is built on one rule about what __call__ is allowed to do to NNX variables:

No writes to NNX variables that affect the forward output in ``__call__``, ever.

Stats-bearing modules accumulate state by overriding update_statistics(), which is called once per training step after the loss / gradient update, with the full rollout’s [T, B, ...] slice of rollout_extras. This keeps every __call__ pure and the rollout / loss-replay numerically identical.

Adding a running mean via update_statistics

Extend MovingAverage to also track a global running mean. The forward pass stays pure; the update is folded in by overriding update_statistics():

class GlobalMeanStats(nnx.Variable):
    pass


class MovingAverageWithStats(MovingAverage):
    def __init__(self, feature_size, k):
        super().__init__(feature_size, k)
        self.mean = GlobalMeanStats(jp.zeros(feature_size))
        self.count = GlobalMeanStats(jp.zeros((), jp.int32))

    def __call__(self, state, x, rollout_extras=None):
        # Forward stays pure; always emit `x` so update_statistics
        # can see it. Eval-only callers drop the emission on the floor.
        out = super().__call__(state, x)
        return out.replace(rollout_extras=x)

    def update_statistics(self, rollout_extras):
        # rollout_extras: [T, B, feature_size]
        flat = rollout_extras.reshape((-1,) + rollout_extras.shape[2:])
        batch_count = flat.shape[0]
        new_count = self.count[...] + batch_count
        delta = jp.mean(flat, axis=0) - self.mean[...]
        self.mean[...] = self.mean[...] + delta * (batch_count / new_count)
        self.count[...] = new_count

A few things to call out:

  • nnx.Variable (as opposed to nnx.Param) is the right base for running statistics: NNX tracks it but it does not receive gradients.

  • The update_statistics() override is the only place writes happen to self.mean / self.count. Forward stays pure.

  • The training loop calls network.update_statistics(rollout.rollout_extras) once per step, after the gradient update. Containers route per child. Normalizer uses this same pattern in production.

Containers for variable-size children

If your module holds a fixed set of sub-modules, you can just assign them as attributes and NNX will trace them:

class TwoLayer(StatefulModule):
    def __init__(self, rngs):
        self.layer_a = Dense(8, 16, rngs)
        self.layer_b = Dense(16, 4, rngs)

If you hold a list or dict of sub-modules — particularly one whose size depends on a constructor argument — wrap it in nnx.List or nnx.Dict. A plain Python list / dict of modules will not be traced by NNX and your parameters will silently disappear from the trainable set:

class StackOfLayers(StatefulModule):
    def __init__(self, sizes, rngs):
        self.layers = nnx.List([
            Dense(d_in, d_out, rngs)
            for d_in, d_out in zip(sizes[:-1], sizes[1:])
        ])

Sequential is itself written exactly this way — it stores its layers in an nnx.List.

Where rollout_extras comes from

You will almost never pass rollout_extras by hand. The library wires it in at three places:

  • unroll_env() (data collection) calls the network with rollout_extras=None and stores each module’s emission into Transition.rollout_extras;

  • eval_rollout() (evaluation) also passes rollout_extras=None and discards the emissions;

  • ppo_loss() (the gradient phase) feeds each per-step slice of the stored extras back into the network.

If you want deterministic actions in evaluation, that is a separate knob on the action sampler (nets.eval() flips it; see PPOAdapter reference).

Recap

To write a StatefulModule:

  1. Decide what is carry state (per-batch, reset on episode boundary) and what is NNX state (per-network, persistent).

  2. Implement __init__, __call__(state, x, rollout_extras=None), initialize_state(batch_size), and — if your carry state needs episode resets — reset_state(prev_state).

  3. Read rollout_extras only if your behaviour depends on it (consume the stored value when non-None; sample fresh when None). Respect the “no writes to NNX variables that affect forward output in __call__, ever” rule; defer accumulation to update_statistics().

  4. Wrap variable-size lists/dicts of sub-modules in nnx.List / nnx.Dict.

Your new module is now a drop-in StatefulModule — composable with Sequential, embeddable as a population’s compute or a connection’s transform in a PopulationGraph, inspectable via nnx.iter_modules(), traceable through nnx.jit / nnx.grad.