Batching contract

nnx-ppo’s batching contract is uniform but stratified by interface: StatefulModule __call__ runs over a flat batch with no time dimension; the time scan is done outside, by the training and rollout machinery. update_statistics() is the one place where a time dimension does surface. Environments are written for a single env and are vmapped to n_envs automatically.

Networks: [B, *feat] per leaf

Inside StatefulModule.__call__(), every input leaf has shape (B, *feature_dims). There is no leading time axis — the scan over T rollout / replay steps happens outside the module, in unroll_env() (during data collection) and in the loss-replay scan inside ppo_loss() (during the gradient phase).

Concretely, for an env with a flat (O,) observation and a training config with n_envs=512:

• obs inside a network __call__ has shape (512, O).

• state["something"] has its leading axis equal to 512 too.

• The return arrays in StatefulModuleOutput.output, next_state, rollout_extras, etc. all carry the same leading axis.

For modules that need to aggregate across the batch (e.g. a running-mean welford step would do jp.mean(x, axis=0)), do that in update_statistics(), not in __call__(). The forward path should be pure per-sample — the algorithm and JAX transforms assume it.

For dict-structured observations the same rule holds per leaf: an obs of {"proprio": (B, P), "goal": (B, G)} has 2-D leaves with B as the leading axis.

update_statistics: [T, B, *feat]

StatefulModule.update_statistics() is the only entry point that sees a time dimension. It is called once per training step after the gradient update, with the rollout’s stacked rollout_extras history. Leaves of that history have shape (T, B, *feat), where T is the rollout length and B is the number of envs.

A stats-bearing module typically reshapes the leading two axes together and processes the result as a single batch:

def update_statistics(self, rollout_extras):
    # rollout_extras leaves: (T, B, *feat)
    flat = jax.tree.map(
        lambda v: v.reshape((-1,) + v.shape[2:]), rollout_extras
    )
    # Now treat `flat` as a flat (T*B, *feat) batch.
    ...

If your module’s emission is multi-dimensional or structured, the same rule applies leaf-wise: drop the leading T into B and operate on the combined axis.

Containers route update_statistics per child the same way they route __call__’s state and rollout_extras — the time axis flows through transparently and only the leaves see it. See Containers for the custom-container contract.

Environments: one env, vmapped automatically

Environments follow the MuJoCo Playground convention: env.reset(key) and env.step(state, action) are written for a single environment. The training loop (and any user-side eval rollout helpers) vmaps the env across n_envs automatically:

• unroll_env() wraps each env.step in jax.vmap(), and a fresh batch of reset keys via jax.random.split().

• eval_rollout() does the same for evaluation rollouts.

• new_training_state() initialises the n_envs env states by nnx.vmap(env.reset) over independent RNG keys.

You never need to write a “vectorised” version of your env — the single-env step / reset are enough. If your env’s internal data has fields without a batch dimension (e.g. shared MuJoCo model handles), those are passed through unchanged; see tree_where() for the per-env selection on done-flag that the rollout uses to keep state shapes consistent across the vmap.

Summary

Interface	Leading axes on each leaf	Who scans / vmaps
StatefulModule.__call__()	(B, *feat)	The library’s scan runs T steps outside __call__
StatefulModule.update_statistics()	(T, B, *feat)	The rollout scan stacked T per-step emissions
env.reset(key) / env.step(state, action)	unbatched single-env state	The training loop vmaps to n_envs automatically