Skip to content

t5lab.train

train

Isolated minimal MAPPO train step + loop for Phase 1a.

DTE (team value = mean of the per-agent value head), two-action PG (goal offset + MVProp move), GAE + clipped-PG + value MSE + entropy, AdamW. 16 rollouts, never reduced. Reuses ctde_v0's _gae / make_optimizer / make_stencil / goal_targets read-only. No dual/selector/roles.

Verified config paths: rollouts_per_iter/iters top-level; gamma/gae_lambda/clip/ppo_epochs/ minibatches on cfg.trainer; vf_coef on cfg.loss; entropy on cfg.regularization; optax pattern matches ctde_v0 (params = eqx.filter(actor, is_array); opt.update(grads, opt_state, params)).

train_step 

train_step(env, actor, opt, opt_state, cfg, key, stencil, controller='mvprop')

One PPO iteration: 16 rollouts -> GAE -> ppo_epochs x minibatches clipped-PG.

Source code in experiments/t5lab/train.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def train_step(env, actor, opt, opt_state, cfg, key, stencil, controller="mvprop"):
    """One PPO iteration: 16 rollouts -> GAE -> ppo_epochs x minibatches clipped-PG."""
    t = cfg.trainer
    B = int(cfg.rollouts_per_iter)                           # 16, never reduced
    rk, tk = jax.random.split(key)
    seeds = jax.random.split(rk, B)
    traj = jax.vmap(lambda s: rollout(env, actor, cfg, s, stencil, controller))(seeds)   # (B,T,...)
    adv, ret = _advantages(traj, t.gamma, t.gae_lambda)
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)
    old_logp = (traj["goal_logp"] + traj["move_logp"]).sum(-1)               # (B,T) joint over agents

    grad_fn = eqx.filter_value_and_grad(_ppo_loss, has_aux=True)
    logs = {}
    for _ep in range(int(t.ppo_epochs)):
        tk, pk, lk = jax.random.split(tk, 3)
        for mb in jnp.array_split(jax.random.permutation(pk, B), int(t.minibatches)):
            sub = jax.tree_util.tree_map(lambda x: x[mb], traj)
            (loss, logs), grad = grad_fn(
                actor, sub, adv[mb], ret[mb], old_logp[mb], stencil,
                t.clip, cfg.loss.vf_coef, cfg.regularization.entropy_coef, lk, controller)
            params = eqx.filter(actor, eqx.is_array)
            updates, opt_state = opt.update(grad, opt_state, params)
            actor = eqx.apply_updates(actor, updates)
    metrics = dict(ret=ret.mean(), **logs)
    return actor, opt_state, metrics

train 

train(env, actor, cfg, key, controller='mvprop')

Full loop over cfg.iters. Returns (actor, history).

Source code in experiments/t5lab/train.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def train(env, actor, cfg, key, controller="mvprop"):
    """Full loop over cfg.iters. Returns (actor, history)."""
    opt = make_optimizer(cfg)
    opt_state = opt.init(eqx.filter(actor, eqx.is_array))
    stencil = make_stencil(cfg)
    hist = []
    k = key
    step = eqx.filter_jit(train_step)
    for it in range(int(cfg.iters)):
        k, sk = jax.random.split(k)
        actor, opt_state, m = step(env, actor, opt, opt_state, cfg, sk, stencil, controller)
        row = {kk: float(v) for kk, v in m.items()}
        hist.append(row)
        if it % 10 == 0 or it == int(cfg.iters) - 1:
            print(f"[it {it}] ret={row['ret']:.2f} pg={row['pg']:.3f} "
                  f"vloss={row['vloss']:.1f} ent={row['ent']:.2f}", flush=True)
    return actor, hist