Skip to content

ctde_v0.run_es

run_es

MERL coexistence runner — ES evolves the selector, CTDE-gradient trains the executor.

Wires the injectable ES trainer (es.py) to the real selector + the PPO/CTDE executor. Each outer round: (a) grad_steps of ppo.train_step train the whole agent (the dense per-step CTDE signal — incl. the selector); (b) the gradient-current selector is taken as the ES mean (MERL "inject the learner"); (c) es.es_step refines ONLY the selector head against TEAM FITNESS (mean episode return via ppo.collect) on the CURRENT executor, using common random numbers (one fixed eval key per round) for a fair population compare; (d) the evolved selector is written back. This is the MERL / feudal-evolutionary loop (Khadka & Tumer 2019): ES + gradient share CTDE's centralized-training signal (team return / central critic) and touch disjoint-ish params, so they compose rather than fight.

Requires --selector on (ES evolves actor.selector_head). CPU smoke: JAX_PLATFORMS=cpu PYTHONPATH=.:../../../FiedlerValueEstimation /Users/bijanmehr/Project.Zymera/zymera_lab/.venv/bin/python -m ctde_v0.run_es --grid 10 --n-agents 4 --outer 3 --grad-steps 3 --pop 8 --flock scripted --rollouts 4

merl_train 

merl_train(env, cfg, *, key, n_outer, grad_steps, es_cfg, log_fn=None, init_from=None)

Run the MERL coexistence loop on env/cfg (which must have selector on). init_from (path to a prior model.eqx) warm-starts the (actor, critic) — the scale-ladder entry point (16²→24²→32²), carrying the ES-evolved selector up. Returns (final_state, history) where history is a list of per-round records.

Source code in experiments/ctde_v0/run_es.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
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
86
87
88
89
90
91
def merl_train(env, cfg: CTDEConfig, *, key, n_outer: int, grad_steps: int,
               es_cfg: es.ESConfig, log_fn=None, init_from=None):
    """Run the MERL coexistence loop on ``env``/``cfg`` (which must have selector on).
    ``init_from`` (path to a prior model.eqx) warm-starts the (actor, critic) — the
    scale-ladder entry point (16²→24²→32²), carrying the ES-evolved selector up.
    Returns ``(final_state, history)`` where history is a list of per-round records."""
    opt = ppo.make_optimizer(cfg)
    stencil = ppo.make_stencil(cfg)
    gstate = (ppo.init_state_from_checkpoint(env, cfg, init_from, key) if init_from
              else ppo.init_state(env, cfg, key))

    @eqx.filter_jit
    def eval_return(actor, critic, dual_lam, ekey):
        # team fitness for a given selector = mean episode return on the CURRENT executor.
        traj = ppo.collect(env, actor, critic, cfg, stencil, ekey, dual_lam)
        return traj["rew_team"].sum(axis=1).mean()

    history = []
    k = key
    for outer in range(n_outer):
        # (a) gradient on the executor (dense CTDE signal; trains the whole agent).
        glog = None
        for _ in range(grad_steps):
            k, sk = jax.random.split(k)
            gstate, glog = ppo.train_step(env, gstate, cfg, sk, opt, stencil)
        # (b) take the gradient-current selector as the ES mean (inject the learner).
        theta = es.module_theta(gstate.actor.selector_head)
        # (c) ES step on the selector — common random numbers (one eval key for the round).
        k, ek, pk = jax.random.split(k, 3)

        def eval_fn(th, _ek=ek, _g=gstate):
            actor = eqx.tree_at(
                lambda a: a.selector_head, _g.actor,
                es.set_module_theta(_g.actor.selector_head, th))
            return float(eval_return(actor, _g.critic, _g.dual.lam, _ek))

        theta, info = es.es_step(theta, eval_fn, es_cfg, pk)
        # (d) write the evolved selector back into the executor's actor.
        gstate = eqx.tree_at(
            lambda s: s.actor.selector_head, gstate,
            es.set_module_theta(gstate.actor.selector_head, theta))

        rec = {"outer": outer,
               "best": round(float(info.get("best_fitness", 0.0)), 3),
               "mean": round(float(info.get("mean_fitness", 0.0)), 3),
               "grad_cov": round(float(glog["coverage_pct"]) * 100, 1) if glog else 0.0,
               "grad_conn5": round(float(glog.get("connectivity_real", 0.0)) * 100, 1)
               if glog else 0.0}
        history.append(rec)
        if log_fn is not None:
            log_fn(outer, rec)
    return gstate, history