Pytree utilities

The nnx_ppo.networks.utils module collects small stateless StatefulModule s for pytree projection, reshaping, and per-key dispatch. They complement the four canonical containers in nnx_ppo.networks.containers.

Flattener

Pytree → flat tensor (or dict-of-flat-tensors).

from nnx_ppo.networks.utils import Flattener

# depth-0: every leaf is reshaped to (B, -1) and concatenated.
Flattener()
# depth-1: preserves the top dict, flattens each value.
Flattener(preserve_levels=1)

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.

Filter

Declarative pytree extraction / projection. Output is a dict whose keys are exactly the keys of the spec and whose values are pulled from the input per spec entry.

from nnx_ppo.networks.utils import Filter

Filter({
    "arm_L":  ("arm_L", "proprioception"),          # nested path
    "arm_R":  ("arm_R", "proprioception"),
    "head":   "head",                               # top-level key
    "joystick": ("root", "future_target", "pos"),
    "zero":   lambda obs: jp.zeros_like(obs["x"]),  # callable
})

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

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.

Anything in the input not named is dropped. Use Filter to ablate obs streams, give the critic privileged info, or reshape structured observations before a downstream consumer.

Map

Per-key dispatch: dict input, dict output, with a different sub-module per key.

from nnx_ppo.networks.utils import Map

Map({pop: NormalTanhSampler(rngs, ...) for pop in POPULATIONS})

Map({k: f for k in keys}) applies each f to the upstream’s same-named entry. Extra keys in the upstream are dropped; missing keys raise KeyError.

Distinct from Parallel (same input fed to every component) and Concat (per-key dispatch with concatenated output).

Carry state is a dict {name: component_state}.

Merge

Run named sub-modules on the same input; each must return a dict; the dicts are merged into one flat dict.

from nnx_ppo.networks.utils import Merge

inner = Merge(
    motors=graph,                       # emits {motor_arm_L, ...}
    critic=detached_critic_stack,       # emits {value_arm_L, ...}
)

Duplicate keys across components are a hard error. The natural complement to Parallel when downstream consumers want one flat dict of named heads but the heads come from independent sub-networks.

Scale

Multiply input by a fixed scalar factor.

from nnx_ppo.networks.utils import Scale

Sequential([motor_head_dense, Scale(0.1)])  # damp motor output

Use this on the output of a head when you want to scale action means or value estimates without baking the factor into the Dense initialiser — keeps the factor visible to introspection / logging.

Construction forms

Map and Merge accept either keyword arguments or a positional dict, e.g.:

Map({pop: sampler 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.