`PPOAdapter` reference

PPOAdapter is the canonical leaf that turns a network’s forward output into a PPONetworkOutput. It is a regular StatefulModule that lives inside a Sequential (or any other container) like any other layer. Its forward output’s output field is the PPONetworkOutput.

Two-port router

PPOAdapter has two ports — action and value — and both receive the same upstream input. The action port emits a tree of sampler dicts (each {"action", "log_likelihood"}); the value port emits whatever value-shape your critic produces. The adapter walks the action port’s output to assemble PPONetworkOutput.actions and loglikelihoods, takes the value port’s output as value_estimates (trailing singleton axis squeezed), and packages everything.

PPOAdapter(action: StatefulModule, value: StatefulModule)

Minimal one-actor / one-critic

Each port owns its full chain. Both run on the same upstream tensor.

network = Sequential([
    Normalizer(obs_size),
    PPOAdapter(
        action=Sequential([actor_mlp, NormalTanhSampler(rngs, entropy_weight=1e-2)]),
        value=critic_mlp,
    ),
])

The factory make_mlp_actor_critic() builds exactly this shape.

Multi-head with per-key dispatch

When the trunk emits a dict of named action-params and named values, the adapter ports use Filter (to extract / rename) and Map (to dispatch per key):

network = Sequential([
    Normalizer(obs_size),
    graph,                                   # emits {"action_a": ..., "action_b": ...,
                                             #         "value_a": ...,  "value_b": ...}
    PPOAdapter(
        action=Sequential([
            Filter({"a": "action_a", "b": "action_b"}),   # rename
            Map({"a": sampler_a, "b": sampler_b}),         # per-key dispatch
        ]),
        value=Filter({"a": "value_a", "b": "value_b"}),
    ),
])

out.actions is then {"a": ..., "b": ...}, matching the env’s per-key action spec.

Shared trunk

To share computation between the action and value pathways, prepend the trunk to the outer Sequential and let both ports consume its output:

Sequential([
    Normalizer(obs_size),
    shared_trunk,                            # heavy computation
    PPOAdapter(
        action=Sequential([actor_head, sampler]),
        value=critic_head,
    ),
])

Action port output shape

The action port’s output should be a tree of sampler dicts. A sampler dict is a leaf-position {"action": ..., "log_likelihood": ...} (the standard shape that ActionSampler emits). The adapter walks the tree with jax.tree.map(..., is_leaf=...) so:

A bare sampler dict → out.actions is a bare array, out.loglikelihoods is a bare array.
Map({k: sampler}) → {k: sampler_dict} → out.actions = {k: array}, out.loglikelihoods = {k: array}.

Value port output shape

The value port’s output is used as-is for PPONetworkOutput.value_estimates, with each leaf squeezed if it has a trailing length-1 axis. So a critic emitting (B, 1) lands as a (B,) value estimate; a dict of critics emitting {k: (B, 1)} lands as {k: (B,)}.

Carry state

The adapter’s carry state is:

{"action": <action port's carry>,
 "value":  <value port's carry>}

initialize_state and reset_state route per port.

`rollout_extras` and `update_statistics`

The adapter routes both ports’ rollout_extras analogously to state: rollout_extras["action"] to the action port, rollout_extras["value"] to the value port. The same routing is used by update_statistics().

`train()` / `eval()`

nnx.Module.eval() and nnx.Module.train() set the deterministic attribute recursively across the module tree. The action sampler reads it to decide between sampling and returning the mean. Conventional usage:

Default: network.train() — samples stochastically (this is what the training loop does between rollouts).
For deterministic eval / video / deployment: network.eval() before the inference call; the PPO training loop already does this bookend around eval rollouts and restores train() afterwards.

During loss replay, the stored raw action from rollout_extras is used and the sampler’s fresh draw is discarded, so train() / eval() only affect calls that pass rollout_extras=None.

PPOAdapter reference