The rollout_extras channel

PPO runs the network twice per training iteration: once during rollout to drive the environment and collect data, then once per minibatch during loss replay to compute gradients on the recorded observations. A handful of modules (action samplers, normalizers) need to communicate between these two passes — the sampler needs to recompute the log-likelihood of the actually-taken action under updated weights; the normalizer needs the rollout’s activations to fold into its running statistics.

rollout_extras is the channel that carries this communication. It is a pytree shaped exactly like the network’s state tree. Each container routes children’s extras the same way it routes their state.

The three phases

Rollout

Driving the environment to collect data. unroll_env() is the entry point. Callers pass rollout_extras=None. Each module emits its replay snapshot into the returned StatefulModuleOutput.rollout_extras. The rollout scan stacks these over T into Transition.rollout_extras.

Loss replay

Re-running the rollout to compute the PPO loss and its gradient. Threaded into the scan body of ppo_loss(). The per-step rollout_extras slice from the stored Transition is fed back in. Modules that need it (notably action samplers) consume it to reproduce the actually-taken action’s log-likelihood under the current (updated) policy.

Inference

Anything outside of data collection or training: eval_rollout(), debugging, ad-hoc forward passes. Callers pass nothing. Modules still emit extras but the caller drops them on the floor.

A module that needs to distinguish “fresh sample” from “use stored value” reads if rollout_extras is None. There is no Context enum.

How containers route rollout_extras

Every container in nnx_ppo.networks.containers (Sequential, Parallel, Concat, Splitter) and nnx_ppo.networks.utils (Flattener, Filter, Scale, Merge, Map) plus Delay, PopulationGraph, and PPOAdapter accepts rollout_extras as the third positional argument. Each container slices its incoming rollout_extras per child the same way it slices state, calls each child, and reassembles the emitted extras into a tree mirroring state.

Leaf modules either:

• ignore rollout_extras and emit None (most layers — Dense, LSTM, Filter, Flattener, Scale, Splitter, Delay, VariationalBottleneck),

• or use it (Normalizer, ActionSampler).

Sampler replay rule

A module that produces a sample whose log-likelihood will be evaluated under updated weights — i.e. an action sampler in PPO — must store the sample in rollout_extras. RNG-in-state is wrong because the same RNG under updated weights would produce a different sample, but PPO needs the log-likelihood of the actually-taken action under the new policy.

RNG-in-state is only valid for modules whose sample is consumed as a forward activation (reparameterised gradient), like VariationalBottleneck. See Randomness for the per-env-RNG-in-carry-state pattern.

The write rule

A forward pass through any StatefulModule must be fully reproducible — gradients computed during loss replay are gradients for the activations that drove the environment during rollout. Concretely:

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

Stats-bearing modules accumulate state by overriding update_statistics(), called once per training step after the gradient update with the full rollout’s stacked rollout_extras history.

Per-module behaviour table

Module	rollout_extras passed in	rollout_extras emitted
Normalizer	ignored — normalise with live mean / M2 / counter	the input x (every call) — consumed by update_statistics()
ActionSampler (e.g. NormalTanhSampler)	if not None: use as raw_action and compute the log-likelihood under current policy	the freshly-sampled raw_action (every call)
VariationalBottleneck / AR1VariationalBottleneck	ignored — RNG is in carry state; reparameterised gradient	None
everything else (Dense, LSTM, containers, Delay, PopulationGraph)	threaded to children	None at leaves; assembled tree at containers