ctde_v0.ppo

ppo

MAPPO-CTDE trainer for the grounded v0 agent.

PPO optimizes the GOAL policy (the L3 goal head), NOT raw moves. Each step:

backbone(obs, kb_adj) -> belief z_i -> goal_logits (N,K) └─ [mechanism] action_mask: mask candidate goals whose greedy first move would drop true λ₂ below the floor (forbid-disconnect); soft_lambda: no mask, a λ-penalty enters the reward instead. └─ sample goal index g_i ~ masked-softmax(goal_logits) (PPO action) └─ goal cell = pos + stencil[g_i] └─ L1 greedy controller -> env move (only valid moves, STAY fallback) └─ env.step(move) -> reward terms, true λ₂, coverage

GAE runs on the centralized critic (team reward + team value = CTDE). Total loss:

total = PPO(goal) + vf_coefvalue + aux_betaaux(λ̂₂, λ₂_true) + degree_regVar_batch(mean-degree) - entropy_coefH(goal)

The aux head is supervised against the simulator's true λ₂ (mse | huber knob); the degree regularizer (SizeShiftReg-style) penalizes the across-batch variance of the per-node aggregated degree statistic to protect GNN size-transfer.

All JAX/Equinox: eqx.filter_jit rollouts + update; optax AdamW (decoupled weight_decay) + global-norm clip. CPU-friendly (JAX_PLATFORMS=cpu).

Shapes (T horizon, B rollouts, N agents, K candidate goals): obs (B,T,N,C,H,W) central (B,T,Cg,H,W) goal (B,T,N) goal index sampled (PPO action) goal_logp (B,T,N) masked log-prob of the sampled goal goal_mask (B,T,N,K) the safe-goal mask used at sample time (replayed) rew_agent (B,T,N) composed reward rew_team (B,T) mean-over-agents v_team (B,T) centralized critic value true_l2 (B,T) true Fiedler value (aux target, broadcast to agents) l2_hat (B,T,N) per-agent local λ̂₂ estimate (head output) degree (B,T,N) per-node comm degree (degree regularizer input)

DualState

Bases: Module

Functional state for the ADAPTIVE connectivity mechanisms — carried THROUGH jax.lax/filter_jit in :class:TrainState (never host-side mutation), so the dual variable survives the jitted update.

• lam — the dual variable λ ≥ 0 (the connectivity-penalty weight the rollout reads; updated each PPO iteration by dual ascent / PID).
• integral — PID integral term Σ v (pid_lagrangian only; 0 otherwise).
• prev_v — previous iteration's violation, for the PID derivative term.

Inert for action_mask / soft_lambda (λ never enters their reward, and the update is gated to the two adaptive mechanisms) — so I1 behaviour is unchanged.

collect

collect(env, actor, critic, cfg, stencil, key, dual_lambda, mvplanner=None)

Vmap _single_rollout over B seeds -> batched trajectory (leading B,T).

dual_lambda (scalar) is the current train-state dual variable, broadcast to every rollout (it weights the adaptive connectivity penalty; see :func:_single_rollout). mvplanner (default None) is the frozen learned L1 planner used only when action_head.controller == 'mvprop'.

Source code in experiments/ctde_v0/ppo.py

def collect(env, actor, critic, cfg: CTDEConfig, stencil, key, dual_lambda, mvplanner=None):
    """Vmap ``_single_rollout`` over B seeds -> batched trajectory (leading B,T).

    ``dual_lambda`` (scalar) is the current train-state dual variable, broadcast to
    every rollout (it weights the adaptive connectivity penalty; see
    :func:`_single_rollout`). ``mvplanner`` (default None) is the frozen learned L1
    planner used only when ``action_head.controller == 'mvprop'``."""
    keys = jax.random.split(key, cfg.rollouts_per_iter)
    return jax.vmap(
        lambda k: _single_rollout(env, actor, critic, cfg, stencil, k, dual_lambda, mvplanner)
    )(keys)

compute_agent_advantages

compute_agent_advantages(traj, cfg, reward_field='rew_agent')

(B,T,N) PER-AGENT advantage for the v3 per-agent credit schemes. Each agent's OWN reward traj[reward_field][:,:,i] runs through GAE with the SHARED team value v_team as the baseline (a state-only baseline -> UNBIASED for the per-agent policy gradient, it only reshapes variance) and the team v_last bootstrap. The critic keeps training on the team return (:func:compute_advantages); only the POLICY advantage becomes per-agent, so each agent is credited for ITS contribution instead of the team mean. vmap over B (episodes) then over N (agents).

reward_field selects the per-agent reward source (a mutually-exclusive credit axis): - "rew_agent" (default) — the full per-agent composed reward, used by the top-level cfg.credit == 'agent' axis (the balthar A2 experiment). - "rew_agent_credit" — the EXACT submodular difference-reward variant (coverage magnitude swapped new_coverage_i -> D_i, uniquely-provided new cells), used by cfg.loss.credit == 'difference'; each agent's policy gradient uses its own marginal coverage contribution while the connectivity/collision terms stay shared.

Source code in experiments/ctde_v0/ppo.py

def compute_agent_advantages(traj, cfg: CTDEConfig, reward_field: str = "rew_agent"):
    """(B,T,N) PER-AGENT advantage for the v3 per-agent credit schemes. Each agent's OWN
    reward ``traj[reward_field][:,:,i]`` runs through GAE with the SHARED team value
    ``v_team`` as the baseline (a state-only baseline -> UNBIASED for the per-agent policy
    gradient, it only reshapes variance) and the team ``v_last`` bootstrap. The critic keeps
    training on the team return (:func:`compute_advantages`); only the POLICY advantage
    becomes per-agent, so each agent is credited for ITS contribution instead of the team
    mean. vmap over B (episodes) then over N (agents).

    ``reward_field`` selects the per-agent reward source (a mutually-exclusive credit axis):
    - ``"rew_agent"`` (default) — the full per-agent composed reward, used by the top-level
      ``cfg.credit == 'agent'`` axis (the balthar A2 experiment).
    - ``"rew_agent_credit"`` — the EXACT submodular difference-reward variant (coverage
      magnitude swapped new_coverage_i -> D_i, uniquely-provided new cells), used by
      ``cfg.loss.credit == 'difference'``; each agent's policy gradient uses its own
      marginal coverage contribution while the connectivity/collision terms stay shared."""
    t = cfg.trainer

    def per_episode(rew_a, v, vl):                    # rew_a (T,N), v (T,), vl ()
        return jax.vmap(
            lambda r: _gae(r, v, vl, t.gamma, t.gae_lambda)[0], in_axes=1, out_axes=1
        )(rew_a)                                       # (T,N)

    return jax.vmap(per_episode)(
        traj[reward_field], traj["v_team"], traj["v_last"]
    )                                                  # (B,T,N)

loss_fn

loss_fn(actor, critic, batch, cfg, key)

Total loss = PPO(goal) [+ PPO(role)] + vfvalue + betaaux + degreeReg - ent*(goal entropy [+ role entropy]). The role terms are added ONLY when role_picker == 'expl_relay' (off -> identical to v0).

Recurrence: with backbone.recurrence == 'recurrent' the minibatch arrives as a batch of EPISODES (leading (B,T)); the actor is re-applied along each trajectory via a per-episode scan (_actor_forward_recurrent) so the per-step hidden is recomputed under the current params, then everything is flattened to (M=B*T) and the rest of the loss is shape-identical to the feedforward path. The feedforward path keeps the flat (M,...) minibatch and the per-row vmap forward EXACTLY as before (byte-unchanged).

Selector (selector == 'on'): a NEW gated path mirroring the goal+role two-action PPO math EXACTLY, only the second action is the SKILL (over {disperse,flock,hold}) rather than the role. The forward is the selector variant (_actor_skill_forward_{ff,recurrent}) which returns the per-step skill-logits AND all three skills' (N,K) offset-logits; the loss (a) sources the OFFSET log-prob from the SELECTED skill's offset-logits (gather per the stored skill) — so the clipped-PPO ratio for the goal action starts at exactly 1 — and (b) adds the SKILL clipped-PG + entropy (a clone of the role-PG block). Roles are off when the selector is on (selector supersedes the role picker). With selector == 'off' (default) this whole branch is dead and the loss is byte-identical to v0.

Source code in experiments/ctde_v0/ppo.py

def loss_fn(actor, critic, batch, cfg: CTDEConfig, key):
    """Total loss = PPO(goal) [+ PPO(role)] + vf*value + beta*aux + degreeReg
    - ent*(goal entropy [+ role entropy]). The role terms are added ONLY when
    ``role_picker == 'expl_relay'`` (off -> identical to v0).

    Recurrence: with ``backbone.recurrence == 'recurrent'`` the minibatch arrives as
    a batch of EPISODES (leading (B,T)); the actor is re-applied along each trajectory
    via a per-episode scan (``_actor_forward_recurrent``) so the per-step hidden is
    recomputed under the current params, then everything is flattened to (M=B*T) and
    the rest of the loss is shape-identical to the feedforward path. The feedforward
    path keeps the flat (M,...) minibatch and the per-row vmap forward EXACTLY as
    before (byte-unchanged).

    Selector (``selector == 'on'``): a NEW gated path mirroring the goal+role two-action
    PPO math EXACTLY, only the second action is the SKILL (over {disperse,flock,hold})
    rather than the role. The forward is the selector variant
    (``_actor_skill_forward_{ff,recurrent}``) which returns the per-step skill-logits AND
    all three skills' (N,K) offset-logits; the loss (a) sources the OFFSET log-prob from the
    SELECTED skill's offset-logits (gather per the stored ``skill``) — so the clipped-PPO
    ratio for the goal action starts at exactly 1 — and (b) adds the SKILL clipped-PG +
    entropy (a clone of the role-PG block). Roles are off when the selector is on
    (selector supersedes the role picker). With ``selector == 'off'`` (default) this whole
    branch is dead and the loss is byte-identical to v0."""
    # both 'recurrent' (per-agent GRU) and 'gcrn' (recurrent message-passing) carry a
    # per-agent hidden state across the episode, so BOTH take the per-episode BPTT scan
    # forward + minibatch-over-episodes path (only the in-actor hidden update differs).
    recurrent = cfg.backbone.recurrence in ("recurrent", "gcrn")
    use_selector = cfg.selector == "on"
    obs = batch["obs"]                 # FF: (M,N,C,H,W) | REC: (B,T,N,C,H,W)
    adj = batch["adj"]                 # FF: (M,N,N)     | REC: (B,T,N,N)
    # the non-default message_content modes stored the per-step comm-graph distance in the
    # trajectory ("dist"); replay it alongside adj. None for 'learned' -> ignored forward.
    dist = batch.get("dist")           # FF: (M,N,N) | REC: (B,T,N,N) | None (learned)
    # the selector path stored per-step agent positions (for the scripted flock / hold
    # skills) and the sampled skill + its log-prob. None on the off-path (never read).
    pos = batch.get("position")        # FF: (M,N,2) | REC: (B,T,N,2) | None (selector off)
    use_roles = cfg.role_picker == "expl_relay" and not use_selector
    # the selector logits / offset-logits / sampled-skill produced by the forward, gated.
    skill_logits = offset_logits = None

    if recurrent:
        B, T = obs.shape[0], obs.shape[1]
        _f = lambda x: x.reshape((B * T,) + x.shape[2:])
        if use_selector:
            # selector forward along each episode (BPTT scan), flattened to (M=B*T,...).
            skill_logits, offset_logits, l2_hat, _z = _actor_skill_forward_recurrent(
                actor, obs, adj, pos, key, dist, remat=cfg.loss.remat)
            role_logits = None
        else:
            # actor forward along each episode (BPTT scan), flattened to (M=B*T,...). The
            # per-step targets/old-logps/adv arrive (B,T,...) and are flattened the SAME way.
            goal_logits, role_logits, l2_hat, _z = _actor_forward_recurrent(
                actor, obs, adj, key, dist, remat=cfg.loss.remat)
        central = _f(batch["central"]); goal = _f(batch["goal"])
        old_logp = _f(batch["goal_logp"]); gmask = _f(batch["goal_mask"])
        role = _f(batch["role"]); old_role_logp = _f(batch["role_logp"])
        adv = _f(batch["adv"]); ret = _f(batch["ret"])
        true_l2 = _f(batch["true_l2"]); degree = _f(batch["degree"])
        if use_selector:
            skill = _f(batch["skill"]); old_skill_logp = _f(batch["skill_logp"])
    else:
        central = batch["central"]         # (M,Cg,H,W)
        goal = batch["goal"]               # (M,N) sampled goal index
        old_logp = batch["goal_logp"]      # (M,N)
        gmask = batch["goal_mask"]         # (M,N,K)
        role = batch["role"]               # (M,N) sampled role index
        old_role_logp = batch["role_logp"] # (M,N)
        adv = batch["adv"]                 # (M,) team advantage
        ret = batch["ret"]                 # (M,) team return
        true_l2 = batch["true_l2"]         # (M,) aux target
        degree = batch["degree"]           # (M,N) per-node comm degree
        if use_selector:
            skill = batch["skill"]             # (M,N) sampled skill index
            old_skill_logp = batch["skill_logp"]  # (M,N)
            skill_logits, offset_logits, l2_hat, _z = _actor_skill_forward_ff(
                actor, obs, adj, pos, key, dist, remat=cfg.loss.remat)
            role_logits = None
        else:
            goal_logits, role_logits, l2_hat, _z = _actor_forward_ff(
                actor, obs, adj, key, dist, remat=cfg.loss.remat)
    # Selector: this step's goal_logits ARE the SELECTED skill's offset-logits — gather
    # offset_logits (M,3,N,K) at the stored skill (M,N) -> (M,N,K). The offset log-prob /
    # mask / PG below are then IDENTICAL to v0 (the goal action math is reused verbatim).
    if use_selector:
        goal_logits = jnp.take_along_axis(
            offset_logits, skill[:, None, :, None], axis=1)[:, 0]    # (M,N,K)
    # goal_logits (M,N,K); apply the SAME mask used at sample time.
    masked = jnp.where(gmask, goal_logits, _NEG)
    logp_all = jax.nn.log_softmax(masked, axis=-1)                  # (M,N,K)
    logp = jnp.take_along_axis(logp_all, goal[..., None], axis=-1)[..., 0]  # (M,N)
    probs = jnp.exp(logp_all)
    entropy = -(jnp.where(gmask, probs * logp_all, 0.0)).sum(-1)    # (M,N)

    # adv is (M,) shared team advantage -> (M,1) broadcast to all agents (v0/A0), OR
    # (M,N) per-agent difference-reward advantage (v3 credit=="agent"); whiten either way.
    _a = jax.lax.stop_gradient(adv)
    adv_b = _a if _a.ndim == 2 else _a[:, None]                     # (M,N) | (M,1)
    adv_norm = (adv_b - adv_b.mean()) / (adv_b.std() + EPS)
    clip = cfg.loss.ppo_clip
    goal_pg = _clipped_pg(logp, old_logp, adv_norm, clip)

    # role policy head (shares the team advantage; trained only when enabled).
    if use_roles:
        role_logp_all = jax.nn.log_softmax(role_logits, axis=-1)    # (M,N,R)
        role_logp = jnp.take_along_axis(
            role_logp_all, role[..., None], axis=-1)[..., 0]        # (M,N)
        role_probs = jnp.exp(role_logp_all)
        role_entropy = -(role_probs * role_logp_all).sum(-1)        # (M,N)
        role_pg = _clipped_pg(role_logp, old_role_logp, adv_norm, clip)
        policy_loss = goal_pg + role_pg
        role_ent = role_entropy.mean()
    else:
        policy_loss = goal_pg
        role_pg = jnp.zeros(())
        role_ent = jnp.zeros(())

    # SELECTOR policy head (the second PPO action when on; supersedes the role head). This
    # is a verbatim clone of the role-PG block above — a clipped-PG + entropy on the
    # categorical skill choice, sharing the same team advantage — only the logits/old-logp
    # come from the selector. Added to policy_loss + entropy bonus exactly like the role.
    if use_selector:
        skill_logp_all = jax.nn.log_softmax(skill_logits, axis=-1)  # (M,N,3)
        skill_logp = jnp.take_along_axis(
            skill_logp_all, skill[..., None], axis=-1)[..., 0]      # (M,N)
        skill_probs = jnp.exp(skill_logp_all)
        skill_entropy = -(skill_probs * skill_logp_all).sum(-1)     # (M,N)
        skill_pg = _clipped_pg(skill_logp, old_skill_logp, adv_norm, clip)
        policy_loss = policy_loss + skill_pg
        skill_ent = skill_entropy.mean()
    else:
        skill_pg = jnp.zeros(())
        skill_ent = jnp.zeros(())

    if cfg.critic_mode == "decentral":
        # DTE: team value = the actor's own per-agent value head (mean over agents); the
        # centralized critic is NOT read here (zero-grad in this mode). ``_z`` is the
        # post-compass belief/hidden the heads read, so value_head(_z) reproduces the
        # rollout's per-agent value exactly (consistency with the v_team source).
        v_agent = jax.vmap(jax.vmap(actor.value_head))(_z)[..., 0]    # (M,N)
        v_pred = v_agent.mean(-1)                                     # (M,) team value
    elif _deepsets_critic(cfg):
        # DeepSets central critic: per-agent obs stack (M,N,C,H,W) + STORED team scalars
        # [coverage, mean λ̂₂]. The old (rollout-time) l2_hat/coverage are used as fixed data
        # (not the recomputed differentiable l2_hat) so the critic never gradients the actor's
        # aux head, matching the rollout's team-scalar source.
        if recurrent:
            obs_c = _f(batch["obs"]); cov_c = _f(batch["coverage"]); l2_c = _f(batch["l2_hat"])
        else:
            obs_c = batch["obs"]; cov_c = batch["coverage"]; l2_c = batch["l2_hat"]
        ts = jnp.stack([cov_c.astype(jnp.float32),
                        l2_c.mean(-1).astype(jnp.float32)], axis=-1)  # (M,2)
        v_pred = jax.vmap(lambda o, t: critic(o, t, inference=False))(obs_c, ts)  # (M,)
    else:
        v_pred = jax.vmap(lambda c: critic(c, inference=False))(central)  # (M,)
    value_loss = jnp.mean((v_pred - jax.lax.stop_gradient(ret)) ** 2)

    aux_loss = _aux_loss(l2_hat, true_l2, cfg)

    # SizeShiftReg-style degree regularizer: variance across the batch of the
    # per-sample mean comm-degree (penalize drift of local structure).
    mean_deg = degree.mean(-1)                                      # (M,)
    degree_reg = jnp.var(mean_deg)

    ent = entropy.mean()
    reg = cfg.regularization
    total = (
        policy_loss
        + cfg.loss.vf_coef * value_loss
        + cfg.loss.aux_beta * aux_loss
        + reg.degree_reg * degree_reg
        - reg.entropy_coef * (ent + role_ent + skill_ent)
    )
    metrics = {
        "policy_loss": policy_loss, "goal_pg": goal_pg, "role_pg": role_pg,
        "skill_pg": skill_pg,
        "value_loss": value_loss, "aux_loss": aux_loss,
        "entropy": ent, "role_entropy": role_ent, "skill_entropy": skill_ent,
        "degree_reg": degree_reg,
    }
    return total, metrics

init_dual

init_dual(cfg)

Initial dual state: λ = mission_safety.lambda_init (scalar f32), integral and prev-error 0. Same shape regardless of mechanism (jit-stable).

Source code in experiments/ctde_v0/ppo.py

def init_dual(cfg: CTDEConfig) -> DualState:
    """Initial dual state: λ = ``mission_safety.lambda_init`` (scalar f32), integral
    and prev-error 0. Same shape regardless of mechanism (jit-stable)."""
    init = jnp.asarray(cfg.mission_safety.lambda_init, dtype=jnp.float32)
    z = jnp.zeros((), dtype=jnp.float32)
    return DualState(lam=init, integral=z, prev_v=z)

dual_update

dual_update(dual, violation, cfg)

One dual step from the realized connectivity violation v = relu(τ − mean_rollout(true λ₂)) ≥ 0. Pure / jit-safe.

• lagrangian — dual ASCENT: λ_next = relu(λ + lambda_lr · v).
• pid_lagrangian — PID (Stooke et al. 2020): integral += v; λ = relu(kp·v + ki·integral + kd·(v − prev_v)); carry prev_v = v.
• else — returned unchanged (action_mask / soft_lambda inert).

Source code in experiments/ctde_v0/ppo.py

def dual_update(dual: DualState, violation: jax.Array, cfg: CTDEConfig) -> DualState:
    """One dual step from the realized connectivity ``violation``
    ``v = relu(τ − mean_rollout(true λ₂)) ≥ 0``. Pure / jit-safe.

    * lagrangian      — dual ASCENT: ``λ_next = relu(λ + lambda_lr · v)``.
    * pid_lagrangian  — PID (Stooke et al. 2020): ``integral += v``;
        ``λ = relu(kp·v + ki·integral + kd·(v − prev_v))``; carry ``prev_v = v``.
    * else            — returned unchanged (action_mask / soft_lambda inert).
    """
    ms = cfg.mission_safety
    if ms.mechanism == "lagrangian":
        lam = jax.nn.relu(dual.lam + ms.lambda_lr * violation)
        return DualState(lam=lam, integral=dual.integral, prev_v=violation)
    if ms.mechanism == "pid_lagrangian":
        integral = dual.integral + violation
        deriv = violation - dual.prev_v
        lam = jax.nn.relu(ms.pid_kp * violation + ms.pid_ki * integral
                          + ms.pid_kd * deriv)
        return DualState(lam=lam, integral=integral, prev_v=violation)
    return dual                                                      # inert otherwise

init_state_from_checkpoint

init_state_from_checkpoint(env, cfg, ckpt_path, key)

Warm-start the train state from a previously-saved (actor, critic).

The scale-strategy entry point. The LPAC backbone + heads are scale-invariant by construction: the (actor, critic) parameter shapes depend only on the obs/central CHANNELS, width/depth/mp_rounds, the goal head K and n_roles — NOT on the grid size or agent count (those only set runtime tensor dims via same-padding conv + global-average-pool + the per-agent vmap). So a model saved at one rung (e.g. 16²/4) loads into a fresh skeleton built for the NEXT rung (e.g. 32²/10) with byte-identical param shapes — only --grid / --n-agents / --comm-r differ, none of which appear in the params.

We build a fresh Actor/Critic from the CURRENT cfg (so K / width / depth / mp_rounds / agg / message_content / explorer_tool / compass / recurrence all match the run you're launching), then eqx.tree_deserialise_leaves the saved params INTO that skeleton. Deserialise validates leaf shapes against the template, so a mismatched backbone (different width / channels / K) raises here rather than corrupting the run; we additionally assert leaf-count + per-leaf shape compatibility up front with a clear scale-strategy error message.

Optimizer state is NOT carried. Adam's moment estimates are tied to the previous rung's loss landscape (and we deliberately do not serialise opt_state in the deployable model.eqx snapshot), so we warm-start the POLICY and re-initialise a FRESH optimizer on the loaded params (opt.init(params)). The dual variable is likewise re-initialised from cfg (a per-run safety knob, not a learned weight). Net effect: same architecture, transplanted weights, clean Adam moments + clean dual — exactly the warm-start-ladder rung hand-off.

Source code in experiments/ctde_v0/ppo.py

def init_state_from_checkpoint(env, cfg: CTDEConfig, ckpt_path: str, key) -> TrainState:
    """Warm-start the train state from a previously-saved ``(actor, critic)``.

    The **scale-strategy** entry point. The LPAC backbone + heads are
    scale-invariant by construction: the (actor, critic) parameter shapes depend
    only on the obs/central CHANNELS, ``width``/``depth``/``mp_rounds``, the goal
    head ``K`` and ``n_roles`` — NOT on the grid size or agent count (those only
    set runtime tensor dims via same-padding conv + global-average-pool + the
    per-agent vmap). So a model saved at one rung (e.g. 16²/4) loads into a fresh
    skeleton built for the NEXT rung (e.g. 32²/10) with byte-identical param
    shapes — only ``--grid`` / ``--n-agents`` / ``--comm-r`` differ, none of which
    appear in the params.

    We build a fresh Actor/Critic from the CURRENT ``cfg`` (so K / width / depth /
    mp_rounds / agg / message_content / explorer_tool / compass / recurrence all
    match the run you're launching), then ``eqx.tree_deserialise_leaves`` the saved
    params INTO that skeleton. Deserialise validates leaf shapes against the
    template, so a mismatched backbone (different width / channels / K) raises here
    rather than corrupting the run; we additionally assert leaf-count + per-leaf
    shape compatibility up front with a clear scale-strategy error message.

    **Optimizer state is NOT carried.** Adam's moment estimates are tied to the
    previous rung's loss landscape (and we deliberately do not serialise opt_state
    in the deployable ``model.eqx`` snapshot), so we warm-start the POLICY and
    re-initialise a FRESH optimizer on the loaded params (``opt.init(params)``).
    The dual variable is likewise re-initialised from ``cfg`` (a per-run safety
    knob, not a learned weight). Net effect: same architecture, transplanted
    weights, clean Adam moments + clean dual — exactly the warm-start-ladder rung
    hand-off.
    """
    from . import checkpoint as _ckpt

    _fork_guard(cfg)
    ka, kc = jax.random.split(key)
    in_ch = env.obs.obs_channels
    cg = env.obs.central_channels
    # The bootstrap checkpoint is ALWAYS a single shared Actor (B-fork replicates AFTER
    # loading), so the LOAD template is a single Actor regardless of cfg.fork_groups.
    template_actor = _make_actor(in_ch, cfg, ka)
    template_critic = _make_critic(in_ch, cg, cfg, kc)

    # Shape-compatibility gate BEFORE deserialise, with a scale-strategy-aware
    # message (eqx would also catch this, but later and more cryptically).
    saved = _ckpt.load_model(ckpt_path, (template_actor, template_critic))
    _assert_param_shapes_match((saved[0], saved[1]),
                               (template_actor, template_critic), ckpt_path)
    actor, critic = saved

    # Arm-A curriculum symmetry-break: optionally perturb the loaded ACTOR before
    # training so "train small → scale up" doesn't transplant a frozen huddle. σ == 0
    # (default) leaves the warm-start byte-unchanged; the critic is always loaded clean
    # (it is re-fit at the new rung regardless).
    if cfg.warmstart_noise > 0:
        actor = _perturb_actor(actor, float(cfg.warmstart_noise),
                               jax.random.fold_in(key, 0x515E))

    # B-fork warm-start: replicate the single shared bootstrap into G sub-actors. Copy 0
    # is the bootstrap exactly; copies>0 are lightly perturbed (σ=0.02·RMS) so the groups
    # diverge under their own per-group gradients rather than staying tied by symmetry.
    if cfg.fork_groups > 1:
        fk = jax.random.split(jax.random.fold_in(key, 0xF02C), cfg.fork_groups)
        subs = [actor] + [_perturb_actor(actor, 0.02, fk[g])
                          for g in range(1, cfg.fork_groups)]
        actor = GroupedActor(subs)

    # Fresh optimizer on the LOADED params (opt_state intentionally NOT carried
    # across rungs — see docstring) + a fresh dual from cfg.
    opt = make_optimizer(cfg)
    params = (eqx.filter(actor, eqx.is_array), eqx.filter(critic, eqx.is_array))
    return TrainState(actor=actor, critic=critic, opt_state=opt.init(params),
                      dual=init_dual(cfg))

train_step

train_step(env, state, cfg, key, opt, stencil, mvplanner=None)

One PPO iteration: collect -> GAE -> ppo_epochs of minibatch updates -> dual update (adaptive mechanisms). The dual variable read at rollout time is the CURRENT state.dual.lam; it is updated AFTER the policy step from the realized connectivity violation and carried forward in the returned state.

Source code in experiments/ctde_v0/ppo.py

def train_step(env, state: TrainState, cfg: CTDEConfig, key, opt, stencil, mvplanner=None):
    """One PPO iteration: collect -> GAE -> ppo_epochs of minibatch updates ->
    dual update (adaptive mechanisms). The dual variable read at rollout time is
    the CURRENT ``state.dual.lam``; it is updated AFTER the policy step from the
    realized connectivity violation and carried forward in the returned state."""
    ck, pk = jax.random.split(key)
    traj = collect(env, state.actor, state.critic, cfg, stencil, ck, state.dual.lam, mvplanner)
    adv, ret = compute_advantages(traj, cfg)
    # v3 per-agent credit: swap the shared team advantage for a PER-AGENT one (B,T,N). ``ret``
    # stays team (the critic target is unchanged). Gated -> default (credit=="shared" AND
    # loss.credit=="shared") is byte-identical. Downstream _flatten_BT / _f collapse
    # (B,T[,N]) either way. The two schemes are mutually exclusive:
    #   cfg.credit == "agent"            -> each agent's own full composed reward (A2 axis).
    #   cfg.loss.credit == "difference"  -> the EXACT submodular difference reward D_i (the
    #                                       coverage-marginal per-agent reward "rew_agent_credit").
    if cfg.credit == "agent":
        adv = compute_agent_advantages(traj, cfg)
    elif cfg.loss.credit == "difference":
        adv = compute_agent_advantages(traj, cfg, reward_field="rew_agent_credit")

    # The KB adjacency the actor consumed at rollout time is stored in the traj
    # ("adj"), so the loss replays the actor on exactly the same comm graph.
    #
    # Recurrence changes ONLY how the minibatch axis is built. Feedforward (default):
    # flatten (B,T)->M and shuffle/minibatch over the M independent steps EXACTLY as
    # before (byte-unchanged). Recurrent: KEEP the (B,T,...) episode structure and
    # shuffle/minibatch over EPISODES (B) so each 100-step sequence stays intact for
    # the per-episode BPTT scan in loss_fn (the hidden resets per episode, so episodes
    # are the independent unit — never mix steps across episodes). The minibatch slicer
    # (_update_epoch) is identical in both cases; only the indexed leading axis differs
    # (M flat rows vs B episodes), and loss_fn detects which by the leading rank.
    fields = ["obs", "adj", "central", "goal", "goal_logp", "goal_mask",
              "role", "role_logp", "true_l2", "degree"]
    # the non-default message_content modes carry a per-step "dist" in the trajectory;
    # minibatch it alongside "adj" (gated so the 'learned' field set is unchanged).
    if cfg.backbone.message_content != "learned":
        fields.append("dist")
    # the selector carries the sampled skill + its log-prob + the per-step agent positions
    # (for the scripted flock / hold skills); minibatch them too (gated so the off-path
    # field set is byte-unchanged).
    if cfg.selector == "on":
        fields += ["skill", "skill_logp", "position"]
    # the DeepSets central critic reads the per-agent obs stack (already in "obs") plus two
    # STORED team scalars — coverage + l2_hat — so minibatch those alongside (gated so the
    # conv-critic field set is byte-unchanged).
    if _deepsets_critic(cfg):
        fields += ["coverage", "l2_hat"]
    if cfg.backbone.recurrence in ("recurrent", "gcrn"):  # both carry a per-episode hidden
        flat = {k: traj[k] for k in fields}          # keep (B,T,...) — minibatch episodes
        flat["adv"], flat["ret"] = adv, ret          # (B,T) per-episode advantage/return
        nperm = flat["obs"].shape[0]                 # B episodes
    else:
        flat = {k: _flatten_BT(traj[k]) for k in fields}   # (B*T,...) flat steps
        flat["adv"], flat["ret"] = _flatten_BT(adv), _flatten_BT(ret)
        nperm = flat["obs"].shape[0]                 # M = B*T steps

    def one_epoch(carry, ek):
        pkey, lkey = jax.random.split(ek)
        perm = jax.random.permutation(pkey, nperm)
        carry, metrics = _update_epoch(carry, flat, perm, lkey, cfg, opt)
        return carry, metrics

    epoch_keys = jax.random.split(pk, cfg.trainer.ppo_epochs)
    carry0 = (state.actor, state.critic, state.opt_state)
    carry, epoch_metrics = jax.lax.scan(one_epoch, carry0, epoch_keys)
    actor, critic, opt_state = carry
    last_metrics = jax.tree_util.tree_map(lambda x: x[-1], epoch_metrics)

    # ---- dual update (adaptive connectivity mechanisms; inert otherwise) ----
    # realized violation on this iter's rollout (a CTDE training-time signal);
    # dual_update advances λ (+ PID state). The aggregation pattern is the SAME for
    # both conn_signals — a rollout-mean shortfall — only the per-step quantity swaps
    # (I1c): global_lambda2 -> v = relu(τ − mean-over-rollout true λ₂) [byte-unchanged];
    # local_edge_margin -> v = mean_i p_i over the rollout (the aggregate per-agent
    # margin shortfall, == traj["margin_step"].mean()).
    lam_used = state.dual.lam                                         # λ this rollout used
    if cfg.mission_safety.conn_signal == "local_edge_margin":
        violation = traj["margin_step"].mean()
    else:
        violation = jax.nn.relu(cfg.constraint_threshold - traj["true_l2"].mean())
    dual = dual_update(state.dual, violation, cfg)
    state = TrainState(actor=actor, critic=critic, opt_state=opt_state, dual=dual)

    # ---- iteration diagnostics (on the freshly collected traj) ----
    ep_reward = traj["rew_team"].sum(axis=1).mean()
    coverage_pct = traj["coverage"][:, -1].mean()
    connectivity_pct = (traj["true_l2"] > cfg.connectivity.threshold).mean()
    # the REAL connectivity bar (λ₂ > real_threshold, default 0.5) logged ALONGSIDE the
    # trivial one — the honest 90/90 grade, never flattered by the loose 1e-3 floor.
    connectivity_real = (traj["true_l2"] > cfg.connectivity.real_threshold).mean()
    mean_lambda2 = traj["true_l2"].mean()

    # behavioural diversity (diagnostic; never enters training): SND = instantaneous
    # spread between agents' goal distributions (mean over steps of mean pairwise TV);
    # role_div = persistent spread of each agent's TIME-AVERAGED distribution (do agents
    # settle into distinct roles). Both in [0,1]; ≈0 for a homogeneous huddle, higher =
    # more differentiated. These are the yardstick the B-fork / B-dico arms move.
    gp = traj["goal_probs"]                                          # (B,T,N,K)
    snd = jax.vmap(jax.vmap(_pairwise_tv_mean))(gp).mean()           # mean over B,T
    role_div = jax.vmap(_pairwise_tv_mean)(gp.mean(axis=1)).mean()   # mean over B of time-avg

    # per-agent coverage (divide-vs-flood): redundancy = Σ_i cells_i / team-covered
    # (1 = perfect division of labour, →N = everyone covers the same ground = flood);
    # top_share = max_i cells_i / team-covered (the top agent's fraction — small = no
    # single agent can solo the mission); mean cells per agent for scale.
    cpa = traj["cells_per_agent"]                                    # (B,N)
    team_cov = traj["team_covered_cells"]                            # (B,)
    redundancy = (cpa.sum(-1) / team_cov).mean()                     # () in [1, N]
    top_agent_share = (cpa.max(-1) / team_cov).mean()               # () top agent fraction
    mean_cells_per_agent = cpa.mean()                               # () avg cells/agent

    # aux λ₂ accuracy = 1 - median rel-err (l2_hat vs true_l2) over CONNECTED steps
    l2_true_bt = traj["true_l2"]                       # (B,T)
    l2_hat_bt = traj["l2_hat"].mean(axis=-1)           # (B,T) mean over agents
    connected = l2_true_bt > cfg.connectivity.threshold
    rel = jnp.abs(l2_hat_bt - l2_true_bt) / jnp.maximum(jnp.abs(l2_true_bt), EPS)
    rel_connected = jnp.where(connected, rel, jnp.nan)
    median_rel = jnp.nanmedian(rel_connected)
    aux_acc = jnp.clip(1.0 - median_rel, 0.0, 1.0)

    # controller validity: fraction of emitted moves that were in the env action
    # mask (the greedy controller's guarantee; should be exactly 1.0).
    valid_frac = traj["move_valid"].astype(jnp.float32).mean()

    # role split: fraction of agent-steps assigned EXPLORER vs RELAY. When
    # role_picker is off every agent is an explorer (frac=1.0 by construction).
    explorer_frac = (traj["role"] == ROLE_EXPLORER).astype(jnp.float32).mean()

    logs = {
        "ep_reward": ep_reward,
        "coverage_pct": coverage_pct,
        "connectivity_pct": connectivity_pct,
        "connectivity_real": connectivity_real,
        "snd": snd,
        "role_div": role_div,
        "redundancy": redundancy,
        "top_agent_share": top_agent_share,
        "mean_cells_per_agent": mean_cells_per_agent,
        "mean_lambda2": mean_lambda2,
        "aux_loss": jnp.mean(last_metrics["aux_loss"]),
        "aux_acc": aux_acc,
        "median_rel_l2": median_rel,
        "policy_loss": jnp.mean(last_metrics["policy_loss"]),
        "value_loss": jnp.mean(last_metrics["value_loss"]),
        "entropy": jnp.mean(last_metrics["entropy"]),
        "role_entropy": jnp.mean(last_metrics["role_entropy"]),
        "degree_reg": jnp.mean(last_metrics["degree_reg"]),
        "ctrl_valid_frac": valid_frac,
        "explorer_frac": explorer_frac,
        "relay_frac": 1.0 - explorer_frac,
        # adaptive-mechanism diagnostics: λ the rollout USED, λ AFTER the dual step,
        # and the realized connectivity violation that drove it (0 / constant unless
        # mechanism is lagrangian / pid_lagrangian).
        "dual_lambda": lam_used,
        "dual_lambda_next": dual.lam,
        "dual_violation": violation,
    }
    # selector diagnostics (only when on; `cfg.selector` is static so the log key set is
    # stable per config): the per-skill USAGE fraction over all agent-steps + the
    # mode-usage entropy (mean over agents/steps of the per-agent selector entropy — high
    # = the team keeps mixing modes, low = it collapsed onto one skill).
    # difference-credit diagnostics (only on the "difference" path; the log key set is
    # static per config): the per-agent uniquely-provided new-coverage D_i, mean over all
    # agent-steps, and the per-episode total credited cells (Σ_i D_i summed over the horizon,
    # episode-mean) — the "how much coverage was non-redundant" signal.
    if cfg.loss.credit == "difference":
        dc = traj["diff_credit"]                                      # (B,T,N) raw D_i
        logs["diff_credit_mean"] = dc.mean()                          # () per-agent D_i mean
        logs["diff_credit_ep_total"] = dc.sum(axis=(1, 2)).mean()     # () unique cells/episode
    if cfg.selector == "on":
        skill_bt = traj["skill"]                                     # (B,T,N) sampled skill
        for m, name in ((0, "disperse"), (1, "flock"), (2, "hold")):
            logs[f"skill_frac_{name}"] = (skill_bt == m).astype(jnp.float32).mean()
        logs["skill_pg"] = jnp.mean(last_metrics["skill_pg"])
        logs["skill_entropy"] = jnp.mean(last_metrics["skill_entropy"])
        # mode-usage entropy from the realized skill choices (in nats), team-mean: how
        # spread the SAMPLED skill distribution is across {disperse,flock,hold}.
        counts = jnp.stack([(skill_bt == m).mean() for m in range(3)])  # (3,) usage probs
        p = counts / jnp.maximum(counts.sum(), EPS)
        logs["skill_usage_entropy"] = -(p * jnp.log(p + EPS)).sum()
    return state, logs

train

train(env, cfg, *, key=None, log_fn=None, init_from=None, mvplanner=None)

Full training loop over cfg.iters PPO iterations.

log_fn(it, host_logs) is called each iteration. Returns (TrainState, history).

init_from (the scale-strategy / warm-start dial): if a path is given, the train state is warm-started from that saved (actor, critic) snapshot via :func:init_state_from_checkpoint (scale-invariant cross-rung load; fresh optimizer + dual) instead of a random init. None (default) = random init, byte-identical to before this option existed (the init_from branch is never touched, so the RNG draw / param surface is unchanged).

Source code in experiments/ctde_v0/ppo.py

def train(env, cfg: CTDEConfig, *, key=None, log_fn=None, init_from: str | None = None,
          mvplanner=None):
    """Full training loop over ``cfg.iters`` PPO iterations.

    ``log_fn(it, host_logs)`` is called each iteration. Returns (TrainState, history).

    ``init_from`` (the scale-strategy / warm-start dial): if a path is given, the
    train state is warm-started from that saved ``(actor, critic)`` snapshot via
    :func:`init_state_from_checkpoint` (scale-invariant cross-rung load; fresh
    optimizer + dual) instead of a random init. ``None`` (default) = random init,
    byte-identical to before this option existed (the ``init_from`` branch is never
    touched, so the RNG draw / param surface is unchanged)."""
    if key is None:
        key = jax.random.PRNGKey(cfg.seed)
    opt = make_optimizer(cfg)
    stencil = make_stencil(cfg)
    if init_from is None:
        state = init_state(env, cfg, key)
    else:
        state = init_state_from_checkpoint(env, cfg, init_from, key)

    @eqx.filter_jit
    def jitted_step(state, k):
        return train_step(env, state, cfg, k, opt, stencil, mvplanner)

    history = []
    k = key
    for it in range(cfg.iters):
        k, sk = jax.random.split(k)
        state, logs = jitted_step(state, sk)
        host_logs = {kk: float(np.asarray(v)) for kk, v in logs.items()}
        history.append(host_logs)
        if log_fn is not None:
            log_fn(it, host_logs)
    return state, history