Custom networks 2: graph networks

The containers in Custom networks 1: composing with containers are enough for any stack of layers — even with branching, the topology stays a tree. They start to hurt when the topology has:

  • multiple connections feeding the same node (the node needs to integrate incoming signals, not just receive one);

  • recurrent loops between nodes;

  • per-connection delays.

These are common in modular controllers (one population per body part, communicating across joints) and in GNN-style message-passing networks. PopulationGraph is a StatefulModule that gives you a declarative API for that class of network.

What PopulationGraph does

A graph is a set of named populations (nodes), each with a single size, connected by directed connections (edges) with default linear Dense transforms. At each step:

  1. For each population in topological order:

  2. The population sums all incoming connection outputs. If it was registered with add_input(), the corresponding obs[key] value is added too. This is the population’s integrated input.

  3. The population’s activation function (if any) is applied once to the integrated input, giving the post-activation result. Connections are linear; activations live on populations — so a node fed by two connections is not double-non-linear.

  4. The post-activation output is written into the population’s shared delay buffer (sized to the longest outgoing delay from that population), so future steps can read delayed copies.

Populations registered with add_output() have their post-activation value emitted in the graph’s forward output dict (under output_to if set, else under the population’s own name).

Delays are first-class. Each connect() takes a delay argument. delay=0 reads the source’s freshly-computed output in the same step (the topological order guarantees that the source has already run); delay=k reads the output from k steps ago. Before the buffer has filled, delayed reads return zeros.

A delay=0 cycle is a hard error: it would force a node to read its own output from the current step, which is undefined. Break the cycle with a delay=1 self-loop or a delay=1 back-edge.

A worked example

The smallest non-trivial graph: an input population, a recurrent hidden population, and two output read-outs. The hidden population has a self-loop with delay=1 so it acts like a one-step recurrent unit.

from flax import nnx
from nnx_ppo.networks.graph import PopulationGraph
from nnx_ppo.networks.containers import Sequential
from nnx_ppo.networks.normalizer import Normalizer
from nnx_ppo.networks.adapter import PPOAdapter
from nnx_ppo.networks.sampling_layers import NormalTanhSampler

rngs = nnx.Rngs(0)

OBS, H, A = 8, 32, 2

g = PopulationGraph(rngs)

# Input population: reads obs["x"]. Same size as obs["x"].
g.add_input("in_x", size=OBS, input_from="x")

# Recurrent hidden population: sum-integrates input + own delayed
# output, passes the sum through swish.
g.add_population("hidden", size=H, activation=nnx.swish)
g.connect("in_x", "hidden")                  # default Dense(OBS, H)
g.connect("hidden", "hidden", delay=1)       # delayed self-loop

# Output populations: action-params head and value head, each
# exposed in the graph's forward output under its own name.
g.add_output("action_params", size=2 * A)
g.connect("hidden", "action_params")         # default Dense(H, 2A)

g.add_output("value", size=1)
g.connect("hidden", "value")                 # default Dense(H, 1)

g.finalize()

A few things to notice:

  • add_input() registers a population that reads obs["x"] directly. Its size must match the obs entry it consumes; downstream connections handle any size change (in this case Dense(OBS, H)).

  • g.connect("in_x", "hidden") does not need an explicit transform=. The default is a linear Dense(in_x.size, hidden.size).

  • g.connect("hidden", "hidden", delay=1) is the recurrence. Without the delay=1 this would be a delay-0 self-loop and finalize() would reject it.

  • add_output() registers a population whose post-activation activation appears in the graph’s forward output dict. With no explicit output_to, the key in the output dict is the population’s own name.

  • Both output populations have no activation — they’re linear read-outs.

The graph’s forward output is {"action_params": (B, 2A), "value": (B, 1)}, ready to feed a PPOAdapter.

Wrap with normalization and the adapter

The graph itself is a StatefulModule, so you can put it behind a Normalizer using Sequential, and then wrap the whole thing in a PPOAdapter:

from nnx_ppo.networks.utils import Filter, Map

sampler = NormalTanhSampler(rngs, entropy_weight=1e-2)
nets = Sequential([
    Normalizer({"x": OBS}),
    g,                                  # emits {"action_params": ..., "value": ...}
    PPOAdapter(
        # Each port sees the graph's full dict and picks what it needs.
        # `Map` dispatches per key and drops keys it doesn't consume.
        action=Map({"action_params": sampler}),
        value=Filter({"value": "value"}),
    ),
])

The resulting PPONetworkOutput.actions is {"action_params": (B, A)} — a single-key dict because the graph names every output. For the per-body-module case in the next section, you would have one sampler per population name in the action port’s Map.

You can now feed this network to train_ppo() exactly as in Quickstart, provided your env’s observation is a dict containing key "x".

Per-population delay buffers

A node’s outgoing connections may have different delays. The graph allocates a single circular buffer per population, sized to the longest outgoing delay from that population. All delayed connections sharing that source read from the same buffer at their own offset — so the cost of a delayed connection is the connection’s own parameters, plus a constant-size buffer per source population. (No duplicate storage per connection.)

What you’d build next

The toy graph above has one input population and one hidden population. Real modular networks typically have one population per body part, with reciprocal connections to a central “root” population:

g.connect("hand_L", "arm_L", delay=1, reciprocal=True)
g.connect("arm_L",  "root",  delay=1, reciprocal=True)
g.connect("torso",  "root",  delay=1, reciprocal=True)
...

reciprocal=True is sugar for two connect() calls (one in each direction, each with its own default Dense and the same delay). Activation lives on each population, connections are linear, and the entire topology is declared up-front. Once finalize() is called you can JIT the whole network and step it like any other StatefulModule.

If you need a layer that the existing modules and containers do not cover — say, your own custom integration rule, or a recurrent cell the library does not ship — see Custom networks 3: writing your own StatefulModule.

For details on per-connection delays, Delay as a standalone layer, and how Normalizer interacts with the surrounding network, see Delay and Normalizer placement.