ctde_v0.nets

nets

Grounded CTDE v0 networks (Equinox) — LPAC backbone + GNN-KB, multi-level goal head, decentralized λ̂₂ head, and a centralized critic.

This is the TeamBlue agent (agent_architecture.md), NOT a flat actor-critic. The learning stack does NOT emit raw moves: L3 picks a goal off the belief; a fixed L1 controller (controller.py) turns it into the env move.

Pipeline (decentralized, runs per agent at execution):

obs_i (C,H,W) └─[1] CNN local-perception (depth conv, same-pad, ReLU) ─ GAP ──▶ f_i (W,) (GAP, not flatten: latent dim independent of H/W -> scale-invariant; conv weight-sharing is translation-equivariant in perception.) └─[2] GNN message-passing KB: fuse in-range NEIGHBOURS' features over the comm graph (adjacency from positions at comm_r), mp_rounds rounds, a configurable NORMALIZED aggregator (mean | max | multihead) ──▶ z_i (the "KB" + "comms/aggregation agent-count-invariant" modules). └─ off z_i: (a) GOAL head -> logits over K candidate relative waypoints (L3 intent) (b) λ̂₂ head -> the decentralized local-Fiedler estimate (one scalar) (c) value head -> per-agent baseline (diagnostic / IPPO fallback)

The goal policy is what PPO optimizes; the controller is fixed. The centralized Critic (CTDE, training only) reads central_obs (Cg,H,W).

The GNN aggregator is the heart of scale-invariance: it never raw-sums neighbours (which would scale with team size); mean / max / softmax-attention are all agent-count-invariant. Adjacency is derived from positions at comm_r so the KB fuses exactly the in-range team — the comm graph the formalism's bridge exposes.

MPLayer

MPLayer(width, agg, heads, *, key, message_content='learned')

Bases: Module

One message-passing round over the comm graph.

Each node updates from (its own feature) + (aggregated neighbour messages). The aggregator is configurable and ALWAYS normalized / size-invariant:

• "mean" — degree-normalized average of neighbour messages.
• "max" — elementwise max over neighbours (default; magnitude-invariant).
• "multihead" — softmax attention over neighbours (heads heads), the attention weights sum to 1 so the readout is agent-count-invariant.

message_content (the I2 "message design" dial) selects WHAT each agent puts in its comm message BEYOND the learned msg feature transform — an EXTRA per-edge channel appended to the aggregated summary before the update (the receiver fuses it alongside the learned messages). It is ALWAYS aggregated with a degree-normalized MEAN (count-invariant regardless of agg) so adding it never breaks size-transfer:

• "learned" (default) — NOTHING extra; the message is msg(feats) exactly, the update reads [self || agg] (2W) and the layer is BYTE-IDENTICAL to v0.
• "edge_distance" — append the (comm_r-normalized) sender→receiver Chebyshev distance per edge, summarized to the receiver as its [mean, min] neighbour distance (2 channels). The receiver thus knows HOW FAR each neighbour is; the normalization by comm_r keeps it in [0,1] = scale-invariant.
• "index" — append a fixed sinusoidal embedding of the SENDER's normalized index (_index_signal; _IDX_DIM channels), mean-pooled over neighbours, so the receiver can tell its neighbours apart. Fixed (not a learned per-N table) and a function of i/N -> agent-count-invariant.

For the two non-default modes the update Linear widens to 2W + extra and a dist (N,N) matrix (normalized sender→receiver distance, diagonal 0) is threaded in by the caller; the learned path ignores dist and keeps the (2W -> W) update.

A boolean adj (N,N, self-loops removed by the caller for neighbour msgs) selects who is in range. __call__(feats, adj_off, dist=None) -> (N, width).

Source code in experiments/ctde_v0/nets.py

def __init__(self, width: int, agg: str, heads: int, *, key,
             message_content: str = "learned"):
    km, ku, kq, kk = jax.random.split(key, 4)
    extra = _EXTRA_DIM[str(message_content)]
    self.msg = eqx.nn.Linear(width, width, key=km)
    # update input = [self (W) || agg (W) || message-content extra (E)]; E==0 for
    # the default 'learned' -> Linear(2W, W) exactly as v0 (byte-identical param tree).
    self.upd = eqx.nn.Linear(2 * width + extra, width, key=ku)
    self.q = eqx.nn.Linear(width, width, key=kq)
    self.k = eqx.nn.Linear(width, width, key=kk)
    self.agg = agg
    self.heads = int(heads)
    self.width = int(width)
    self.message_content = str(message_content)

Backbone

Backbone(in_ch, width, depth, mp_rounds, agg, heads, norm, dropout, *, key, message_content='learned', position_ground=False)

Bases: Module

LPAC backbone: per-agent CNN -> GAP -> feature, then mp_rounds of GNN message passing over the comm graph -> per-agent belief z_i (N, width).

__call__(obs, adj_off, *, dist=None, key) with obs (N,C,H,W) and adj_off (N,N) bool (in-range neighbours, diagonal cleared). dist (N,N) is the NORMALIZED sender→receiver distance (in [0,1], diagonal 0) the non-default message_content modes append to each message; the default learned mode ignores it entirely (so the forward is byte-identical to v0 whether or not a dist is supplied). Optional LayerNorm + dropout on the belief. Returns z (N, width).

message_content (the I2 message-design dial) is threaded into every MPLayer; see :class:MPLayer for the modes (learned | edge_distance | index).

position_ground (the position-grounding dial): when True, ADD a projection of each agent's boundary-relative NORMALIZED position (_norm_position, recovered from the own_pos one-hot centroid, in [-1,1]) to its post-GAP feature BEFORE the message passing — hence before every head — grounding the otherwise position-blind policy. The posground Linear (2 -> W) is ALWAYS built (its key is fold_in-derived so the conv / MP keys are byte-UNCHANGED) but only USED when the flag is on; with it off the branch is never traced and the backbone forward is byte-identical to the pre-grounding version. Projecting-and-ADDING (rather than concatenating the 2 dims into the first MP Linear, which would widen it and perturb its init) keeps the WHOLE MP / head param surface identical whether grounding is on or off — the established compass idiom.

Source code in experiments/ctde_v0/nets.py

def __init__(self, in_ch: int, width: int, depth: int, mp_rounds: int,
             agg: str, heads: int, norm: str, dropout: float, *, key,
             message_content: str = "learned", position_ground: bool = False):
    kc, kmp = jax.random.split(key)
    # position-grounding projection key: fold_in (NOT a widened split) so kc / kmp — and
    # therefore the conv + every MP layer — stay byte-identical to the pre-grounding
    # backbone (position_ground='off' is bit-for-bit the pre-grounding network).
    kpos = jax.random.fold_in(key, 0x9051)
    self.conv = _conv_stack(in_ch, width, depth, kc)
    mp_keys = jax.random.split(kmp, max(mp_rounds, 1))
    self.mp = [MPLayer(width, agg, heads, key=mp_keys[i],
                       message_content=message_content)
               for i in range(mp_rounds)]
    self.ln = eqx.nn.LayerNorm(width) if norm == "layer" else None
    self.drop = eqx.nn.Dropout(dropout) if dropout > 0 else None
    # 2-dim normalized position -> belief width; always built (stable param surface),
    # only USED when position_ground (added to the post-GAP feature before MP).
    self.posground = eqx.nn.Linear(2, width, key=kpos)
    self.width = int(width)
    self.message_content = str(message_content)
    self.position_ground = bool(position_ground)

GCRNCell

GCRNCell(width, agg, heads, *, key)

Bases: Module

Graph-Convolutional Recurrent belief cell — the recurrence == 'gcrn' path.

Carries a per-NODE recurrent hidden state h ACROSS the episode THROUGH the comm graph, a recurrent MESSAGE-PASSING belief:

h_next = GRU( MP(z + h_prev), h_prev )

where MP is one message-passing round over the in-range comm graph (a reused :class:MPLayer) and GRU is a per-node :class:eqx.nn.GRUCell (W -> W). This is DISTINCT from the per-agent recurrence == 'recurrent' GRU (h_next = GRU(z, h), which never fuses over the graph): here node i's memory is updated from a graph-fused summary of its NEIGHBOURS' beliefs AND hiddens, so memory / coverage history PROPAGATES across the team, not just along each agent's own trajectory.

SIZE-INVARIANT by construction — the MPLayer aggregator is normalized (never a raw-sum) and the GRU is applied per node, so a cell trained @16²/4 transfers @32²/10, matching the LPAC backbone. The internal MP uses the default message_content='learned' (self-contained; it ignores any dist supplied) so the GCRN path is independent of the backbone's I2 message-design dial. Always BUILT (stable param surface, fold_in key) but only USED when recurrence == 'gcrn'. Pure JAX (vmap/scan/jit-safe).

__call__(z (N,W), h_prev (N,W), adj_off (N,N) bool, dist=None) -> h_next (N,W).

Source code in experiments/ctde_v0/nets.py

def __init__(self, width: int, agg: str, heads: int, *, key):
    km, kg = jax.random.split(key)
    self.mp = MPLayer(width, agg, heads, key=km, message_content="learned")
    self.gru = eqx.nn.GRUCell(width, width, key=kg)
    self.width = int(width)

FrontierAttn

FrontierAttn(width, *, d=32, F_feat=2, sharp=4.0, alpha_init=1.0, key)

Bases: Module

Frontier-biased goal-sector attention — the explorer's "disperse" tool.

Queries the belief z and keys the K per-sector frontier features, producing one additive logit per compass sector that PULLS the goal policy toward the sector with the most informative unexplored ground. The combined goal logits are

goal_logits = goal_head(z) + alpha · frontier_logits

so PPO still samples a goal from a DISTRIBUTION (the attention biases, never argmaxes — the policy keeps training). alpha is a learned scalar gate (softplus, ≥0) so the network can dial the frontier pull up or down per the reward signal, starting near a configured value.

Construction (F_feat = per-sector feature dim, d = attention dim): * q Linear(W -> d) query from the belief z * k Linear(F_feat -> d) key from each sector's frontier feature

score_k = q(z)·k(feat_k)/√d gives a belief-conditioned attention weight per sector (softmax over K). The contributed logit MULTIPLIES that learned weight by the sector's own frontier FRACTION (feat[k,0] — a non-negative, frontier- peaked scalar): frontier_logit_k = K · attn_k · frontier_frac_k (the K restores unit scale since Σ attn = 1). This is high where the belief asks "explore here" AND sector k is frontier-rich, and — crucially — it is frontier-POSITIVE BY CONSTRUCTION: even at random init (attn ≈ uniform) the largest additive logit lands on the most-uncovered sector, an INDUCTIVE BIAS the reward then sharpens (via q/k and alpha) rather than having to discover from scratch. SIZE-INVARIANT: K is fixed, the frontier fraction is a normalized fraction, and the dot-product readout is independent of H, W and team size. Pure JAX (vmap/jit-safe).

Source code in experiments/ctde_v0/nets.py

def __init__(self, width: int, *, d: int = 32, F_feat: int = 2,
             sharp: float = 4.0, alpha_init: float = 1.0, key):
    kq, kk = jax.random.split(key, 2)
    self.q = eqx.nn.Linear(width, d, key=kq)
    self.k = eqx.nn.Linear(F_feat, d, key=kk)
    # invert softplus so alpha starts ≈ alpha_init: softplus(x)=alpha_init.
    a0 = float(max(alpha_init, 1e-4))
    self.log_alpha = jnp.asarray(jnp.log(jnp.expm1(a0)), dtype=jnp.float32)
    self.d = int(d)
    self.sharp = float(sharp)
    self.F_feat = int(F_feat)

sector_logits

sector_logits(z_i, feats_i)

(K,) additive goal logits for ONE agent: belief z_i (W,) attending over the per-sector frontier features feats_i (K, F_feat). The learned attention weights the sector's own frontier fraction (feats_i[:,0]), so the readout is non-negative and peaks at the most-uncovered sector.

Source code in experiments/ctde_v0/nets.py

def sector_logits(self, z_i: jax.Array, feats_i: jax.Array) -> jax.Array:
    """(K,) additive goal logits for ONE agent: belief ``z_i`` (W,) attending
    over the per-sector frontier features ``feats_i`` (K, F_feat). The learned
    attention weights the sector's own frontier fraction (``feats_i[:,0]``), so
    the readout is non-negative and peaks at the most-uncovered sector."""
    K = feats_i.shape[0]
    qz = self.q(z_i)                                          # (d,) query
    kf = jax.vmap(self.k)(feats_i)                            # (K,d) sector keys
    scores = (kf @ qz) / jnp.sqrt(float(self.d))             # (K,) attention scores
    attn = jax.nn.softmax(scores, axis=0)                    # (K,) sector weights, Σ=1
    frac = feats_i[:, 0]                                      # (K,) frontier fraction >=0
    return float(K) * attn * frac                            # (K,) frontier logits

GoalResidual

GoalResidual(width, K, *, hidden=32, alpha_init=1.0, key)

Bases: Module

Per-agent identity-conditioned goal-logit residual — the B-dico diversity tool.

residual_i = alpha · ( g(z_i, id_i) − mean_j g(z_j, id_j) )      (mean-zero over agents)

where id_i = _index_signal is the agent-count-invariant i/N code, g is a tiny MLP, and alpha = softplus(log_alpha) >= 0 is a learned gate (starts ≈ alpha_init). Centering across agents keeps the TEAM-MEAN goal policy unchanged — the residual only SPREADS agents apart (controlled diversity), so it can raise behavioural diversity (SND) without biasing the average behaviour. SIZE-INVARIANT: the identity is a function of i/N and the readout is per-agent, so a model trained @16²/4 transfers @32²/10. Built ALWAYS (stable param surface) and only USED when diversity_residual == 'on'. Pure JAX.

Source code in experiments/ctde_v0/nets.py

def __init__(self, width: int, K: int, *, hidden: int = 32,
             alpha_init: float = 1.0, key):
    k1, k2 = jax.random.split(key, 2)
    self.l1 = eqx.nn.Linear(width + _IDX_DIM, hidden, key=k1)
    self.l2 = eqx.nn.Linear(hidden, K, key=k2)
    a0 = float(max(alpha_init, 1e-4))
    self.log_alpha = jnp.asarray(jnp.log(jnp.expm1(a0)), dtype=jnp.float32)
    self.K = int(K)
    self.hidden = int(hidden)

Compass

Compass(width, K, *, sharp=4.0, explore_decay=0.5, beta_init=1.0, key)

Bases: Module

The compass directional-feature module — an explicit, scale-free navigation signal added to the per-agent belief z before the heads.

It reads each agent's own obs to build two soft K-sector DIRECTION distributions (compass_features: GATHER = toward in-range teammates, EXPLORE = toward the nearest uncovered cell), flattens them to 2K scalars, and PROJECTS+GATES them into the belief width to ADD to z:

z' = z + beta · proj( [gather(K) , explore(K)] )

The contribution is gated by a learned scalar beta = softplus(log_beta) >= 0 (so the network dials the navigation pull per the reward, starting near a configured value). Projecting-and-adding (rather than concatenating + widening the heads) keeps the head input width — and therefore the WHOLE param surface of the goal / role / λ̂₂ / value heads — IDENTICAL whether the compass is on or off; the module is ALWAYS built (mirrors FrontierAttn / the role head) and only its USE is gated, so with compass == 'off' the belief z is byte-identical to the pre-compass actor. SIZE-INVARIANT: K is fixed and the features are normalized directions, so a model trained @16²/4 transfers up the scale ladder. Pure JAX (vmap/jit-safe).

Source code in experiments/ctde_v0/nets.py

def __init__(self, width: int, K: int, *, sharp: float = 4.0,
             explore_decay: float = 0.5, beta_init: float = 1.0, key):
    self.proj = eqx.nn.Linear(2 * K, width, key=key)
    # invert softplus so beta starts ≈ beta_init: softplus(x)=beta_init.
    b0 = float(max(beta_init, 1e-4))
    self.log_beta = jnp.asarray(jnp.log(jnp.expm1(b0)), dtype=jnp.float32)
    self.K = int(K)
    self.width = int(width)
    self.sharp = float(sharp)
    self.explore_decay = float(explore_decay)

Actor

Actor(in_ch, K, *, backbone_cfg, dropout, key, n_roles=2, explorer_tool='goal_head', compass='off', recurrence='feedforward', diversity_residual='off', selector='off', flock='scripted', comm_r=5.0, flock_sharp=2.0, hold_floor=1.0)

Bases: Module

Decentralized per-agent actor: LPAC backbone -> belief z_i -> four heads (+ the frontier-attention explorer tool + the compass directional feature).

• goal_head (W -> K) L3 goal-pointer logits over candidate waypoints.
• role_head (W -> R) L3 role-picker logits over {explorer, relay} (R = n_roles; the Increment-1 labor-division head off the belief).
• frontier_attn (the L4 "disperse" skill) biases the goal logits toward the most frontier-rich compass sector (FrontierAttn).
• compass (the directional feature) ADDS a scale-free gather/explore navigation term to the belief z BEFORE the heads (Compass).
• aux_head (W -> 1) decentralized local-Fiedler λ̂₂ estimate (raw).
• value_head (W -> 1) per-agent baseline (diagnostic / IPPO fallback).

__call__(obs, adj_off, *, key) -> (goal_logits (N,K), role_logits (N,R), value (N,), lambda2_hat (N,), z (N,W)). The role head, frontier_attn AND compass are ALWAYS built (cheap, stable param surface) so the parameter tree is invariant to the role_picker / explorer_tool / compass knobs; each is only used when its knob is on. With compass == 'off' (default) the compass term is never added, so the belief z — and therefore EVERY head's output — is byte-identical to the pre-compass actor; with 'on' the belief becomes z + compass(obs, z) before all heads (giving them an explicit directional cue). With explorer_tool == 'goal_head' (default) the frontier term is never added, so goal_logits is byte-identical to the pre-tool behaviour; with 'frontier_attn' the goal logits become goal_head(z) + frontier_attn(obs,z). The role head is likewise sampled only when role_picker == 'expl_relay'. The (post-compass) belief z is returned so the trainer can compute the degree regularizer.

Recurrence (the recurrence axis): a per-agent gru (eqx.nn.GRUCell, W -> W) is ALWAYS built (stable param surface) but only USED when recurrence == 'recurrent'. In that mode each step folds the (post-compass) belief z into a carried hidden state h — h_next = GRUCell(z, h) per agent (vmap over N) — and EVERY head (goal / role / λ̂₂ / value, incl. the frontier/compass tools' belief input) reads h_next INSTEAD of z, so the agent remembers its own trajectory / coverage history across the episode. The incoming hidden h is threaded by the caller (the rollout scan carry; reset to zeros at each episode start, and recomputed under the current params along the trajectory in the PPO loss). With recurrence == 'feedforward' (default) the GRU is never traced, the heads read z exactly as before, and h_next is the zero passthrough — so the actor forward is BYTE-IDENTICAL to the pre-recurrence actor (the gru params just sit unused). With recurrence == 'gcrn' the carried hidden is instead updated by a recurrent MESSAGE-PASSING cell (:class:GCRNCell, h_next = GRU(MP(z + h_prev), h_prev)) so node i's memory is fused over the comm graph — the always-built gcrn cell (a stable, fold_in-keyed param surface) is only USED in this mode; the hidden threads through the SAME rollout/loss plumbing as the per-agent GRU (the PPO loss BPTTs it per episode, like 'recurrent').

Selector (the selector axis): a hierarchical mode-picker over a 3-skill library {0=disperse, 1=flock, 2=hold}. The selector_head (W -> 3) and a learned flock_head (FlockHead, W -> K) are ALWAYS built (cheap, stable param surface, fold_in-derived keys) but ONLY used by :meth:skill_forward — never by :meth:__call__. When selector == 'on' the PPO trainer calls :meth:skill_forward to (1) sample a skill m off the belief (a categorical PPO action), (2) take skill m's (N,K) goal-offset logits, (3) sample the offset (the second PPO action, replacing the goal-head sample), routed through the same fixed L1 controller. The skills: disperse = goal_head + frontier_attn (the validated explorer); flock = scripted_flock_logits or the learned flock_head per the flock flavor; hold = a STAY scorer with a soft-degree reconnect fallback (:meth:_hold_logits). With selector == 'off' (default) skill_forward is never called and the two extra heads sit unused, so the actor — and every byte of :meth:__call__ — is identical to the pre-selector network. The selector SUPERSEDES the role picker (assume role_picker off when selector on).

Init note: the compass / gru / goal_residual / selector_head / flock_head keys are all derived via jax.random.fold_in (NOT by widening the split) so the backbone / goal / role / frontier / aux / value keys are byte-IDENTICAL to the pre-recurrence actor — an actor built with compass='off' / recurrence='feedforward' / selector='off' is bit-for-bit the same network as before these modules existed.

Source code in experiments/ctde_v0/nets.py

def __init__(self, in_ch: int, K: int, *, backbone_cfg, dropout: float, key,
             n_roles: int = 2, explorer_tool: str = "goal_head",
             compass: str = "off", recurrence: str = "feedforward",
             diversity_residual: str = "off", selector: str = "off",
             flock: str = "scripted", comm_r: float = 5.0,
             flock_sharp: float = 2.0, hold_floor: float = 1.0):
    kb, kg, kr, kf, ka, kv = jax.random.split(key, 6)
    # Derive the compass / gru / goal-residual / selector / flock keys by folding fixed
    # constants into the ORIGINAL key, so the six keys above are unchanged -> compass='off'
    # / recurrence='feedforward' / diversity_residual='off' / selector='off' is byte-
    # identical to the pre-module actor (split(key,7+) would have perturbed all six).
    kcomp = jax.random.fold_in(key, 0xC0)
    kgru = jax.random.fold_in(key, 0x60)
    kres = jax.random.fold_in(key, 0xD1C0)
    ksel = jax.random.fold_in(key, 0x5E1)        # selector head key (fold_in, not split)
    kflk = jax.random.fold_in(key, 0xF10C)       # learned flock-head key (fold_in)
    kgcrn = jax.random.fold_in(key, 0x6C64)      # GCRN recurrent-MP cell key (fold_in)
    self.backbone = Backbone(
        in_ch, backbone_cfg.width, backbone_cfg.depth, backbone_cfg.mp_rounds,
        backbone_cfg.agg, backbone_cfg.heads, backbone_cfg.norm, dropout, key=kb,
        message_content=getattr(backbone_cfg, "message_content", "learned"),
        position_ground=getattr(backbone_cfg, "position_ground", False),
    )
    W = backbone_cfg.width
    self.goal_head = eqx.nn.Linear(W, K, key=kg)
    self.role_head = eqx.nn.Linear(W, n_roles, key=kr)
    self.frontier_attn = FrontierAttn(W, key=kf)
    self.compass = Compass(W, K, key=kcomp)
    # B-dico per-agent diversity residual; always built (stable param surface), only
    # USED when diversity_residual == 'on' (mean-zero -> team-mean policy unchanged).
    self.goal_residual = GoalResidual(W, K, key=kres)
    # per-agent recurrent cell over the belief width (W -> W); always built so the
    # param tree is invariant to the recurrence knob, only USED when recurrent.
    self.gru = eqx.nn.GRUCell(W, W, key=kgru)
    # GCRN recurrent MESSAGE-PASSING belief cell (h_next = GRU(MP(z + h_prev), h_prev));
    # always built (stable param surface, fold_in key), only USED when recurrence=='gcrn'.
    # Distinct from the per-agent gru above: it fuses the hidden over the comm graph.
    self.gcrn = GCRNCell(W, backbone_cfg.agg, backbone_cfg.heads, key=kgcrn)
    # SELECTOR head (the L3 mode-picker over the 3-skill library {disperse,flock,hold})
    # AND the learned FlockHead are ALWAYS built (cheap, stable param surface) — exactly
    # like goal_residual / gru / compass — so the parameter tree is invariant to the
    # ``selector`` / ``flock`` knobs; both are only USED via ``skill_forward`` when
    # selector == 'on'. Their keys are fold_in-derived (above), so the backbone / goal /
    # role / frontier / aux / value keys are UNCHANGED and a selector='off' actor is
    # bit-for-bit the pre-selector network.
    from .flock import FlockHead as _FlockHead   # lazy: breaks the nets<->flock cycle
    self.selector_head = eqx.nn.Linear(W, 3, key=ksel)
    self.flock_head = _FlockHead(W, K, key=kflk)
    self.aux_head = eqx.nn.Linear(W, 1, key=ka)
    self.value_head = eqx.nn.Linear(W, 1, key=kv)
    self.K = int(K)
    self.n_roles = int(n_roles)
    self.explorer_tool = str(explorer_tool)
    self.compass_on = (str(compass) == "on")
    self.diversity_on = (str(diversity_residual) == "on")
    self.recurrent = (str(recurrence) == "recurrent")
    self.gcrn_on = (str(recurrence) == "gcrn")
    self.selector_on = (str(selector) == "on")
    self.flock_flavor = str(flock)
    self.comm_r = float(comm_r)
    self.flock_sharp = float(flock_sharp)
    self.hold_floor = float(hold_floor)
    self.width = int(W)

init_hidden

init_hidden(n)

Zero per-agent hidden state (N, W) — the episode-start carry for the recurrent path (and the inert passthrough returned by the feedforward path).

Source code in experiments/ctde_v0/nets.py

def init_hidden(self, n: int) -> jax.Array:
    """Zero per-agent hidden state ``(N, W)`` — the episode-start carry for the
    recurrent path (and the inert passthrough returned by the feedforward path)."""
    return jnp.zeros((n, self.width), dtype=jnp.float32)

skill_forward

skill_forward(obs, adj_off, position, *, dist=None, h=None, key=None, inference=False)

The hierarchical SELECTOR forward (used only when selector == 'on').

Runs the SAME backbone / compass / recurrence path as :meth:__call__ (:meth:_belief_and_hidden), then off the per-agent feature emits (1) the SELECTOR head — a categorical over the 3-skill library {0=disperse, 1=flock, 2=hold} — and (2) each skill's (N, K) goal-offset logits, stacked to (3, N, K):

• skill 0 — disperse: the validated explorer, goal_head(z) + frontier_attn(obs, z, K) (the same disperse the explorer_tool='frontier_attn' path uses — frontier-positive by construction so it spreads the swarm).
• skill 1 — flock: the connectivity-repair skill — flock.scripted_flock_logits(position, K, comm_r, flock_sharp) when flock == 'scripted' (the weakest-link-repair primitive, no params) else the learned self.flock_head(z) (a tiny belief-conditioned head). The flavor is a STATIC field so only the selected branch is traced.
• skill 2 — hold: the STAY scorer with a reconnect fallback (:meth:_hold_logits) — hold the post unless about to isolate, then step toward the nearest neighbour.

The PPO trainer samples the skill m from the selector logits (one PPO action), gathers skill m's (N, K) offset-logits, masks + samples the offset (the SECOND PPO action, replacing the goal-head sample), and routes the offset-goal through the fixed L1 controller (role_idx=None — ALL skills route the same greedy way; no relay bypass).

Args mirror :meth:__call__ plus position (N,2) — the agent cells the scripted flock / hold skills read for their geometry (the belief carries no absolute coordinate, so the position is threaded explicitly, exactly like the controller's inputs).

Returns (skill_logits (N,3), offset_logits (3,N,K), feat (N,W), h_next (N,W)). SCALE-INVARIANT (every skill reads only the belief, normalized fractions, or unit bearings / unit compass directions); pure JAX (vmap/scan/jit-safe). Does NOT touch :meth:__call__.

Source code in experiments/ctde_v0/nets.py

def skill_forward(self, obs, adj_off, position, *, dist=None, h=None, key=None,
                  inference: bool = False):
    """The hierarchical SELECTOR forward (used only when ``selector == 'on'``).

    Runs the SAME backbone / compass / recurrence path as :meth:`__call__`
    (:meth:`_belief_and_hidden`), then off the per-agent feature emits (1) the SELECTOR
    head — a categorical over the 3-skill library {0=disperse, 1=flock, 2=hold} — and
    (2) each skill's ``(N, K)`` goal-offset logits, stacked to ``(3, N, K)``:

      * skill 0 — **disperse**: the validated explorer, ``goal_head(z) +
        frontier_attn(obs, z, K)`` (the same disperse the ``explorer_tool='frontier_attn'``
        path uses — frontier-positive by construction so it spreads the swarm).
      * skill 1 — **flock**: the connectivity-repair skill —
        ``flock.scripted_flock_logits(position, K, comm_r, flock_sharp)`` when
        ``flock == 'scripted'`` (the weakest-link-repair primitive, no params) else the
        learned ``self.flock_head(z)`` (a tiny belief-conditioned head). The flavor is a
        STATIC field so only the selected branch is traced.
      * skill 2 — **hold**: the STAY scorer with a reconnect fallback
        (:meth:`_hold_logits`) — hold the post unless about to isolate, then step toward
        the nearest neighbour.

    The PPO trainer samples the skill m from the selector logits (one PPO action), gathers
    skill m's ``(N, K)`` offset-logits, masks + samples the offset (the SECOND PPO action,
    replacing the goal-head sample), and routes the offset-goal through the fixed L1
    controller (``role_idx=None`` — ALL skills route the same greedy way; no relay bypass).

    Args mirror :meth:`__call__` plus ``position`` (N,2) — the agent cells the scripted
    flock / hold skills read for their geometry (the belief carries no absolute coordinate,
    so the position is threaded explicitly, exactly like the controller's inputs).

    Returns ``(skill_logits (N,3), offset_logits (3,N,K), feat (N,W), h_next (N,W))``.
    SCALE-INVARIANT (every skill reads only the belief, normalized fractions, or unit
    bearings / unit compass directions); pure JAX (vmap/scan/jit-safe). Does NOT touch
    :meth:`__call__`."""
    from .flock import scripted_flock_logits  # lazy: breaks the nets<->flock cycle
    feat, h_next = self._belief_and_hidden(obs, adj_off, dist, h, key, inference)
    skill_logits = jax.vmap(self.selector_head)(feat)              # (N,3) mode picker

    # skill 0 — disperse: the validated frontier-biased explorer (goal_head + frontier).
    disperse = jax.vmap(self.goal_head)(feat) + self.frontier_attn(obs, feat, self.K)  # (N,K)

    # skill 1 — flock: scripted weakest-link repair OR the learned belief head (STATIC).
    if self.flock_flavor == "learned":
        flock = self.flock_head(feat)                              # (N,K) learned head
    else:
        flock = scripted_flock_logits(position, self.K, self.comm_r, self.flock_sharp)  # (N,K)

    # skill 2 — hold: STAY (offset 0) unless about to isolate -> reconnect toward nbr.
    hold = self._hold_logits(position)                            # (N,K)

    offset_logits = jnp.stack([disperse, flock, hold], axis=0)    # (3,N,K)
    return skill_logits, offset_logits, feat, h_next

GroupedActor

GroupedActor(subs)

Bases: Module

B-fork: G independent sub-:class:Actors with SEPARATE parameters, applied to a FIXED contiguous partition of the team (G=2 → group 0 = first half, group 1 = the rest).

Each sub-actor runs over the WHOLE team (so its backbone still fuses the full comm graph), and agent i's outputs are taken from ITS group's sub-actor; the other sub-actors' outputs for agent i are DISCARDED — so each sub-actor's parameters receive gradient ONLY from its own group's agents. That is the fork: two groups with separate policies that specialize independently (the lit's CTDE-bootstrap → fork → specialize; SePS/Kaleidoscope selective sharing), while the CTDE critic stays single and shared (it is NOT wrapped here — that is what keeps training stable).

Drop-in for :class:Actor: identical __call__ signature + 6-tuple return + an init_hidden, so the PPO rollout/loss are UNCHANGED — only construction differs (built in ppo.init_state from scratch, or in ppo.init_state_from_checkpoint by REPLICATING a single shared bootstrap into G copies, copies>0 lightly perturbed to break symmetry so the groups diverge). The partition is recomputed from N at call time (a fraction of N, never a baked array), so a GroupedActor trained @16²/4 transfers @32²/10 exactly like a single Actor. Pure JAX (vmap/scan/jit-safe).

Source code in experiments/ctde_v0/nets.py

def __init__(self, subs):
    self.subs = list(subs)
    self.G = len(self.subs)
    self.width = int(self.subs[0].width)

init_hidden

init_hidden(n)

Zero per-agent hidden state (N, W) — delegated to a sub-actor (all share the same width), matching :meth:Actor.init_hidden for the rollout carry.

Source code in experiments/ctde_v0/nets.py

def init_hidden(self, n: int) -> jax.Array:
    """Zero per-agent hidden state ``(N, W)`` — delegated to a sub-actor (all share
    the same width), matching :meth:`Actor.init_hidden` for the rollout carry."""
    return self.subs[0].init_hidden(n)

Critic

Critic(in_ch, width, depth, norm, dropout, *, key)

Bases: Module

Centralized critic over the team central_obs (Cg,H,W) -> value ().

Source code in experiments/ctde_v0/nets.py

def __init__(self, in_ch: int, width: int, depth: int, norm: str,
             dropout: float, *, key):
    kc, kv = jax.random.split(key)
    self.conv = _conv_stack(in_ch, width, depth, kc)
    self.ln = eqx.nn.LayerNorm(width) if norm == "layer" else None
    self.drop = eqx.nn.Dropout(dropout) if dropout > 0 else None
    self.value_head = eqx.nn.Linear(width, 1, key=kv)

AttnCritic

AttnCritic(feat_dim, n_actions, hid, heads, *, key)

Bases: Module

MAAC-style centralized attention critic (CTDE, training only) → a PER-AGENT Q_i, where the scalar :class:Critic gives one team value shared by everyone.

Each agent encodes its own (belief feat_i, action a_i) into e_i; agent i then ATTENDS over every other agent's encoding (self masked out — full team graph, CTDE) to a teammate context x_i, and reads Q_i = head([e_i ‖ x_i]). Because e_i (and the query that forms x_i) carry agent i's own action, Q_i depends on a_i — which is what makes the COMA counterfactual baseline (vary a_i, hold the rest) a real per-agent advantage. Symmetric over agents (no identity), attention weights sum to 1 → agent-count-invariant, same as the rest of the stack.

__call__(feats (N,D), act (N,K)) -> q (N,). act is a per-agent action one-hot (hard, for a taken joint action) or soft distribution; D = belief width, K = action-space size (e.g. goals or moves).

Source code in experiments/ctde_v0/nets.py

def __init__(self, feat_dim: int, n_actions: int, hid: int, heads: int, *, key):
    ke, kq, kk, kv, k1, k2 = jax.random.split(key, 6)
    self.enc = eqx.nn.Linear(feat_dim + n_actions, hid, key=ke)
    self.q = eqx.nn.Linear(hid, hid, key=kq)
    self.k = eqx.nn.Linear(hid, hid, key=kk)
    self.v = eqx.nn.Linear(hid, hid, key=kv)
    self.head1 = eqx.nn.Linear(2 * hid, hid, key=k1)
    self.head2 = eqx.nn.Linear(hid, 1, key=k2)
    self.heads = int(heads)
    self.hid = int(hid)

DeepSetsCritic

DeepSetsCritic(in_ch, width, depth, norm, dropout, *, pool='mean', key)

Bases: Module

Permutation- & count-invariant centralized critic (CTDE, training only).

Where the conv :class:Critic reads the single 3-channel GLOBAL central_obs (Cg,H,W), this critic reads the per-agent obs stack (N,C,H,W) — the same (C,H,W) tensors the actor sees — and is invariant to the number of agents N by construction (Deep Sets / permutation-invariant pooling), matching the LPAC backbone's scale-invariance. Pipeline::

obs (N,C,H,W)
  └─ SHARED CNN encoder (``_conv_stack`` / ``_encode``, GAP), vmapped over agents ─▶ (N,width)
  └─ PERMUTATION-INVARIANT pool over the N agents (``pool``):
       "mean" -> plain mean over agents
       "attn" -> single-head self-attention (AttnCritic-style q/k) weights each agent
                 over the team, then mean over the resulting per-agent contexts
     ─▶ pooled (width,)
  └─ concat [pooled ‖ team scalars (coverage_fraction, mean λ̂₂)] ─▶ (width+2,)
  └─ 2-layer MLP ─▶ value ()

Drop-in for :class:Critic's RETURN (a scalar ()) and its key/inference kwargs; the difference is the FIRST positional input — the per-agent obs stack (N,C,H,W) plus a (2,) team-scalar vector — wired at train time only (ppo._make_critic / the rollout+loss call sites branch on cfg.critic_arch). The exec/decentralized path never touches it.

Source code in experiments/ctde_v0/nets.py

def __init__(self, in_ch: int, width: int, depth: int, norm: str,
             dropout: float, *, pool: str = "mean", key):
    kc, kq, kk, k1, k2 = jax.random.split(key, 5)
    self.conv = _conv_stack(in_ch, width, depth, kc)
    self.ln = eqx.nn.LayerNorm(width) if norm == "layer" else None
    self.drop = eqx.nn.Dropout(dropout) if dropout > 0 else None
    if pool == "attn":
        self.attn_q = eqx.nn.Linear(width, width, key=kq)
        self.attn_k = eqx.nn.Linear(width, width, key=kk)
    else:
        self.attn_q = None
        self.attn_k = None
    self.mlp1 = eqx.nn.Linear(width + 2, width, key=k1)          # [pooled ‖ 2 team scalars]
    self.mlp2 = eqx.nn.Linear(width, 1, key=k2)
    self.pool = str(pool)

sector_frontier_features

sector_frontier_features(obs_i, K, sharp=4.0)

(K, 2) float32 per-sector frontier features for ONE agent's obs (C,H,W).

SCALE-INVARIANT by construction — every quantity is a fraction or a unit direction; no absolute coordinate or grid-size-dependent magnitude survives:

frontier(cell) = 1 - known(cell) (ch _CH_KNOWN; uncovered) (cr, cc) = centroid of own_pos (ch _CH_OWN_POS; the agent) u(cell) = (row-cr, col-cc) / ||·|| (UNIT displacement, scale-free) m_k(cell) = softmax_k( sharp · u(cell)·dir_k ) over the K compass dirs feat[k,0] = Σ_cell m_k·frontier / Σ_cell m_k (frontier FRACTION in sector k) feat[k,1] = Σ_cell m_k·frontier / (H·W) (frontier DENSITY toward k)

feat[:,0] answers "of the cells lying toward compass-dir k, what fraction is unexplored?" and feat[:,1] "how much of my whole view's frontier sits toward k?" — both bounded in [0,1] regardless of H,W or team size. The agent's own cell (zero displacement) carries no direction; the soft sector membership lets it fall to the "here" sector (dir 0) so it never spuriously votes for a compass heading. Pure JAX (vmap/jit-safe).

Source code in experiments/ctde_v0/nets.py

def sector_frontier_features(obs_i: jax.Array, K: int, sharp: float = 4.0) -> jax.Array:
    """(K, 2) float32 per-sector frontier features for ONE agent's obs (C,H,W).

    SCALE-INVARIANT by construction — every quantity is a fraction or a unit
    direction; no absolute coordinate or grid-size-dependent magnitude survives:

      frontier(cell)  = 1 - known(cell)                 (ch ``_CH_KNOWN``; uncovered)
      (cr, cc)        = centroid of own_pos              (ch ``_CH_OWN_POS``; the agent)
      u(cell)         = (row-cr, col-cc) / ||·||         (UNIT displacement, scale-free)
      m_k(cell)       = softmax_k( sharp · u(cell)·dir_k )  over the K compass dirs
      feat[k,0]       = Σ_cell m_k·frontier / Σ_cell m_k     (frontier FRACTION in sector k)
      feat[k,1]       = Σ_cell m_k·frontier / (H·W)          (frontier DENSITY toward k)

    ``feat[:,0]`` answers "of the cells lying toward compass-dir k, what fraction is
    unexplored?" and ``feat[:,1]`` "how much of my whole view's frontier sits toward
    k?" — both bounded in [0,1] regardless of H,W or team size. The agent's own cell
    (zero displacement) carries no direction; the soft sector membership lets it fall
    to the "here" sector (dir 0) so it never spuriously votes for a compass heading.
    Pure JAX (vmap/jit-safe)."""
    C, H, W = obs_i.shape
    frontier = 1.0 - obs_i[_CH_KNOWN]                              # (H,W) 1=uncovered
    own = obs_i[_CH_OWN_POS]                                       # (H,W) one-hot

    # Recover the agent's (row, col) as the centroid of its own-position one-hot —
    # exact for a one-hot, and robust if it were ever smoothed. NO absolute coord is
    # exported; only per-cell displacement (a relative, translation-free geometry).
    rows = jnp.arange(H, dtype=jnp.float32)[:, None]              # (H,1)
    cols = jnp.arange(W, dtype=jnp.float32)[None, :]             # (1,W)
    mass = jnp.maximum(own.sum(), 1.0)
    cr = (own * rows).sum() / mass                                # () agent row
    cc = (own * cols).sum() / mass                                # () agent col

    dr = rows - cr                                                # (H,1) row displacement
    dc = cols - cc                                                # (1,W) col displacement
    dr = jnp.broadcast_to(dr, (H, W))
    dc = jnp.broadcast_to(dc, (H, W))
    dist = jnp.sqrt(dr ** 2 + dc ** 2)                            # (H,W) Euclidean radius
    inv = 1.0 / jnp.maximum(dist, 1e-6)
    ur = dr * inv                                                 # (H,W) unit row dir
    uc = dc * inv                                                 # (H,W) unit col dir

    dirs = _compass_unit_dirs(K)                                  # (K,2) unit compass dirs
    # cosine of each cell's direction with each sector direction -> (K,H,W).
    cos = dirs[:, 0][:, None, None] * ur[None] + dirs[:, 1][:, None, None] * uc[None]
    # the "here" sector (dir 0 -> cos==0 everywhere) should win only for cells AT the
    # agent (tiny radius); give it a closeness score so near-cell frontier lands there
    # rather than leaking into an arbitrary compass heading. Directional sectors keep
    # their cosine. score_k(cell): closeness for sector 0, cosine for sectors >=1.
    is_here = (jnp.abs(dirs[:, 0]) + jnp.abs(dirs[:, 1])) < 1e-6  # (K,) True for dir 0
    closeness = jnp.exp(-dist)[None]                              # (1,H,W) in (0,1], 1 at agent
    score = jnp.where(is_here[:, None, None], closeness, cos)     # (K,H,W)

    member = jax.nn.softmax(sharp * score, axis=0)               # (K,H,W) soft sector assign
    fmass = (member * frontier[None]).sum(axis=(1, 2))           # (K,) frontier mass / sector
    smass = member.sum(axis=(1, 2))                              # (K,) cell mass / sector
    frac = fmass / jnp.maximum(smass, 1e-6)                       # (K,) frontier FRACTION
    dens = fmass / float(H * W)                                   # (K,) frontier DENSITY
    return jnp.stack([frac, dens], axis=-1)                       # (K,2) per-sector feats

compass_features

compass_features(obs_i, K, sharp=4.0, explore_decay=0.5)

(2, K) float32 directional compass features for ONE agent's obs (C,H,W):

row 0 — GATHER direction: a soft K-sector one-hot pointing toward the centroid of the agent's IN-RANGE TEAMMATES (the neighbors channel one-hots); "which way is my team". row 1 — EXPLORE direction: a soft K-sector one-hot pointing toward the NEAREST UNCOVERED cell in view (frontier = 1 - known, distance-decayed so the nearest fresh ground dominates); "which way is fresh ground".

Both rows are normalized soft sector distributions over controller._COMPASS (Σ_k = 1, index 0 = "here"/no-direction) — DIRECTIONS ONLY. SCALE-INVARIANT by construction: every quantity is a cosine of unit displacements or a normalized softmax, so no absolute coordinate or grid-size magnitude survives and the SAME relative layout yields the same features at any H, W or team size. When the agent has no in-range teammate (gather) or no frontier in view (explore) that row falls to the "here" sector. Pure JAX (vmap/jit-safe).

Source code in experiments/ctde_v0/nets.py

def compass_features(obs_i: jax.Array, K: int, sharp: float = 4.0,
                     explore_decay: float = 0.5) -> jax.Array:
    """(2, K) float32 directional compass features for ONE agent's obs (C,H,W):

      row 0 — GATHER direction: a soft K-sector one-hot pointing toward the centroid
              of the agent's IN-RANGE TEAMMATES (the ``neighbors`` channel one-hots);
              "which way is my team".
      row 1 — EXPLORE direction: a soft K-sector one-hot pointing toward the NEAREST
              UNCOVERED cell in view (frontier = ``1 - known``, distance-decayed so the
              nearest fresh ground dominates); "which way is fresh ground".

    Both rows are normalized soft sector distributions over ``controller._COMPASS``
    (Σ_k = 1, index 0 = "here"/no-direction) — DIRECTIONS ONLY. SCALE-INVARIANT by
    construction: every quantity is a cosine of unit displacements or a normalized
    softmax, so no absolute coordinate or grid-size magnitude survives and the SAME
    relative layout yields the same features at any H, W or team size. When the agent
    has no in-range teammate (gather) or no frontier in view (explore) that row falls
    to the "here" sector. Pure JAX (vmap/jit-safe)."""
    own = obs_i[_CH_OWN_POS]                                      # (H,W) own one-hot
    neigh = obs_i[_CH_NEIGHBORS]                                  # (H,W) teammate one-hots
    frontier = 1.0 - obs_i[_CH_KNOWN]                            # (H,W) 1 = uncovered
    ur, uc, dist = _agent_unit_dirs(own)                         # (H,W) unit dirs + radius
    dirs = _compass_unit_dirs(K)                                  # (K,2) unit compass headings
    gather = _soft_sector_dir(neigh, ur, uc, dist, dirs, sharp, decay=0.0)        # (K,)
    explore = _soft_sector_dir(frontier, ur, uc, dist, dirs, sharp, decay=explore_decay)  # (K,)
    return jnp.stack([gather, explore], axis=0)                  # (2,K) directions

coma_counterfactual

coma_counterfactual(critic, feats, act_onehot, act_logits)

COMA per-agent counterfactual advantage = the contribution signal (#70).

For each agent i, holding every other agent's action fixed at the taken joint action, marginalize agent i's own action over its policy π_i::

A_i = Q_i(s, a) − Σ_{a'} π_i(a') · Q_i(s, a₋ᵢ, a')

feats (N,D), act_onehot (N,K) the taken actions, act_logits (N,K) the policy logits (-inf for masked-invalid actions is fine — softmax zeros them). Returns (adv (N,), q_taken (N,)): adv_i is how much agent i's chosen action beat its own average given the team — positive = it pulled its weight, ≈0 = interchangeable with its default, negative = it hurt. This is exactly the per-agent credit a shared team value cannot give, and the difference-reward / resilience contribution #70 measures.

Source code in experiments/ctde_v0/nets.py

def coma_counterfactual(critic: "AttnCritic", feats: jax.Array,
                        act_onehot: jax.Array, act_logits: jax.Array):
    """COMA per-agent counterfactual advantage = the contribution signal (#70).

    For each agent i, holding every *other* agent's action fixed at the taken joint action,
    marginalize agent i's own action over its policy ``π_i``::

        A_i = Q_i(s, a) − Σ_{a'} π_i(a') · Q_i(s, a₋ᵢ, a')

    ``feats (N,D)``, ``act_onehot (N,K)`` the taken actions, ``act_logits (N,K)`` the policy
    logits (``-inf`` for masked-invalid actions is fine — softmax zeros them). Returns
    ``(adv (N,), q_taken (N,))``: ``adv_i`` is how much agent i's chosen action beat its own
    average given the team — positive = it pulled its weight, ≈0 = interchangeable with its
    default, negative = it hurt. This is exactly the per-agent credit a shared team value
    cannot give, and the difference-reward / resilience contribution #70 measures."""
    n, K = act_onehot.shape
    q_taken = critic(feats, act_onehot)                              # (N,) Q at the joint action
    pi = jax.nn.softmax(act_logits, axis=-1)                         # (N,K)
    cand = jnp.eye(K, dtype=feats.dtype)                            # (K,K) candidate one-hots

    def q_i_over_actions(i):
        def one(a_prime):                                           # swap agent i -> a', recompute Q_i
            a2 = act_onehot.at[i].set(a_prime)
            return critic(feats, a2)[i]
        return jax.vmap(one)(cand)                                  # (K,) Q_i for each a'_i
    qi_all = jax.vmap(q_i_over_actions)(jnp.arange(n))              # (N,K)
    baseline = jnp.sum(pi * qi_all, axis=-1)                        # (N,) expected Q_i over own action
    return q_taken - baseline, q_taken