Containers

The nnx_ppo.networks.containers module provides the four core composition primitives: Sequential, Parallel, Concat, and Splitter. Every container is itself a StatefulModule, so they nest freely.

Each container threads three things through its children:

carry state — sliced per-child the same way state is structured (list for Sequential, dict for Parallel/Concat, leaf for Splitter);
rollout_extras — same routing as state;
regularization loss & metrics — accumulated from all children and surfaced on the container’s own StatefulModuleOutput.

Sequential

Chain layers. Sequential([a, b, c]) runs a → b → c and routes each layer’s carry state through a list keyed by position.

from nnx_ppo.networks.containers import Sequential

stack = Sequential([encoder, body, head])
state = stack.initialize_state(batch_size)
out = stack(state, obs)
# out.next_state == [encoder.next_state, body.next_state, head.next_state]

Sequential is the workhorse: most “chains of layers” are Sequentials. Use it as the action / value port of a PPOAdapter, as the inner body of a Parallel branch, etc.

Parallel

Run several named sub-modules on the same input. Returns a dict keyed by sub-module name.

from nnx_ppo.networks.containers import Parallel

# Same input fed to every component.
branches = Parallel(action_params=actor, value=critic)
out = branches((), obs)
# out.output == {"action_params": <actor_out>, "value": <critic_out>}

Typical use: assemble per-key heads on top of a shared trunk, e.g. per-body-module value heads (Parallel({k: Dense(H, 1) for k in keys})) fed by the same trunk features.

Concat

Per-key dispatch + concat: dict input, single-tensor output. Each named sub-module sees the upstream’s same-named entry and the per-component outputs are concatenated along the last axis.

from nnx_ppo.networks.containers import Concat

# Per-stream encoders → single concatenated feature vector.
encoders = Concat(proprio=proprio_encoder, goal=goal_encoder)
out = encoders(state, {"proprio": (B, P), "goal": (B, G)})
# out.output : (B, H_proprio + H_goal)

Use at the start of a stack when the observation is structured.

Splitter

Inverse of Concat. Takes a flat tensor and splits it into named slices along the last axis.

from nnx_ppo.networks.containers import Splitter

s = Splitter(action_params=2 * action_size, value=1)
out = s((), x)  # x : (B, 2A + 1)
# out.output == {"action_params": (B, 2A), "value": (B, 1)}

Use at the end of a stack to turn a flat head into a dict output an adapter / downstream can route to samplers and value specs. With a single keyword (Splitter(action_params=N)) it 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.

Construction forms

Concat and Parallel accept either keyword arguments or a positional dict, e.g.:

Parallel({pop: head for pop in POPULATIONS})

Use the positional form when keys aren’t valid Python identifiers or when you’re building the container from a comprehension.

Writing your own container

The four shipped containers are not load-bearing on the rest of the library — they are ordinary StatefulModule s that thread plumbing through their children. If none of them captures your topology you can write your own. The contract has four parts.

1. Be a valid Flax NNX module

A container holds child modules and their parameters. To make NNX trace those parameters, every sub-module must be either a direct attribute or sit inside an NNX-tracked container. For a fixed set of children, plain attribute assignment is enough:

class TwoChild(StatefulModule):
    def __init__(self, a, b):
        self.a = a
        self.b = b

For a variable-size list or dict of children, wrap them in flax.nnx.List or flax.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:

from flax import nnx

class MyMap(StatefulModule):
    def __init__(self, modules: dict[str, StatefulModule]):
        self.components = nnx.Dict(modules)   # NNX-tracked

This is what every shipped container does internally (Sequential uses nnx.List; Concat / Parallel use nnx.Dict).

2. Forward state and `rollout_extras` per child

In __call__, slice the incoming state and rollout_extras per child the same way, call each child, and reassemble both into output containers shaped exactly like the input. A typical loop:

def __call__(self, state, x, rollout_extras=None):
    new_state, new_extras = {}, {}
    ...
    for name, child in self.components.items():
        child_extras = None if rollout_extras is None else rollout_extras[name]
        out = child(state[name], <child input>, child_extras)
        new_state[name] = out.next_state
        new_extras[name] = out.rollout_extras
        ...

The crucial invariant is that the shape of returned next_state mirrors the shape of the incoming state, and the same for rollout_extras. If a caller passes rollout_extras=None, you pass None to every child (each child knows how to fall back). If a caller passes a non-None value, the structure must match what your container emitted on a previous step.

You also need to implement initialize_state() and (if your children carry resettable state) reset_state() with the same shape — for the same reasons.

3. Aggregate `regularization_loss` and `metrics`

Each child’s StatefulModuleOutput.regularization_loss should flow up to your container’s output. In most cases this is just a sum:

regularization_loss = jp.array(0.0)
for ...:
    out = child(...)
    regularization_loss += out.regularization_loss

If your container has a reason to scale or otherwise transform a child’s regularization (e.g. a hypothetical container that runs children alternately and only counts the active one’s loss), do that here — but the default is a plain sum.

metrics are typically collected into a per-child dict {child_name: out.metrics} so downstream logging can disambiguate. Sequential uses positional integer keys; the named-dispatch containers use string keys matching their children.

4. Route `update_statistics` per child

Stats accumulation runs outside __call__. Override update_statistics() to route the rollout’s stacked rollout_extras per child analogously to how __call__ routes the per-step extras:

def update_statistics(self, rollout_extras):
    for name, child in self.components.items():
        child.update_statistics(rollout_extras[name])

Without this override the default implementation is a no-op for the container, which means any stats-bearing module nested inside it will never see its data and silently fail to update. This is the most common mistake when writing a new container — easy to forget because every other rule fails loudly while this one fails silently.

Worked example

The shipped Parallel is the smallest container that exercises all four rules; reading [its source](../_modules/nnx_ppo/networks/containers.html) is the fastest way to see the pattern end-to-end (~50 LoC). Sequential is similar but lists rather than dicts; the per-key dispatch containers Map and Merge are also good references.