ctde_v0.run_scale_ladder

run_scale_ladder

Scale-strategy experiment: single-point vs warm-start-ladder (vs multi-scale-joint).

Does CLIMBING the scale ladder with a warm start (train small -> transplant the weights -> keep training bigger) reach the top rung better/cheaper than training the top rung from scratch? The LPAC backbone is scale-invariant by construction (param shapes depend on channels / width / depth / mp_rounds / goal-K / n_roles, NOT grid size or agent count), so a model trained @16²/4 loads UNCHANGED into a @32²/10 run — which is exactly what makes the ladder possible.

Three strategies (this launcher runs (a) + (b); (c) is documented, not built — see MULTI-SCALE-JOINT below):

(a) single-point : train from scratch @32²/10 (the top rung only). (b) warm-start-ladder: train @16²/4 -> save -> --init-from train @24²/6 -> save -> --init-from train @32²/10. Each rung warm-starts the POLICY from the previous rung's saved (actor, critic); the optimizer + dual are re-initialised fresh per rung (see ppo.init_state_from_checkpoint). (c) multi-scale-joint: train on the rungs MIXED per-episode (each episode samples a rung). NOT BUILT — see the limitation note below.

It just orchestrates python -m ctde_v0.train_ctde SUBPROCESS calls with the right --grid/--n-agents/--comm-r/--run-dir/--ckpt/--init-from. Each rung writes its own run-dir (config.json + history.json + model.eqx), so results are inspectable and the ladder's hand-off (rung N's model.eqx -> rung N+1's --init-from) is explicit on disk.

RESUMABLE / shard-friendly (mirrors run_ctde_sweep.py): a rung whose run-dir already holds a finished model.eqx is SKIPPED, so a re-launch never redoes a completed rung (and the ladder's dependency is respected — a later rung waits for its predecessor's model.eqx). --only single|ladder runs one strategy; --shard i --nshards n splits the INDEPENDENT units of work across workers (the single-point run and the ladder are two units; within the ladder the rungs are sequential by data dependency, so the whole ladder is one unit).

# one worker, both strategies, the default 16->24->32 ladder:
JAX_PLATFORMS=cpu PYTHONPATH=.:../../../FiedlerValueEstimation     /Users/bijanmehr/Project.Zymera/zymera_lab/.venv/bin/python -u         ctde_v0/run_scale_ladder.py --out runs/scale_ladder

# just the warm-start ladder, custom rungs + iters:
... run_scale_ladder.py --only ladder --rungs 16x4x5,24x6x5,32x10x5 --iters 400

This experiment is CPU-only by construction (do NOT run on the GPU server); it only SHELLS OUT to train_ctde — keep the per-rung --iters/--rollouts in the train_ctde budget. ===>>> This launcher LAUNCHES REAL TRAINING; run it deliberately. <<<===

MULTI-SCALE-JOINT (c) — known limitation, NOT built here. The current trainer compiles ONE env at a fixed (grid, n_agents, comm_r): the grid H×W and the agent count N are static array dims baked into the jitted rollout (the env is a trace-time constant; see env_utils.build_env). Mixing rungs PER EPISODE would need either (i) one compiled rollout per rung shape (different N / H×W can't share a scan, so they cannot be vmapped together), or (ii) a single padded/masked MAX-shape env (pad to the largest grid + agent count and mask the slack) so every episode shares one shape and a per-episode key selects the active sub-rung. Both are a trainer change, out of scope for this launcher; the scale-invariant backbone makes either viable as a follow-up. Until then, the ladder (b) is the supported way to train one model across multiple scales.

run_single_point

run_single_point(out, top, *, iters, rollouts, seed, extra)

(a) train the TOP rung from scratch (no --init-from).

Source code in experiments/ctde_v0/run_scale_ladder.py

def run_single_point(out: str, top: tuple[int, int, int], *, iters: int,
                     rollouts: int, seed: int, extra: list[str]) -> dict:
    """(a) train the TOP rung from scratch (no --init-from)."""
    g, n, r = top
    run_dir = os.path.join(out, f"single_{_rung_tag(g, n, r)}")
    rec = {"strategy": "single-point", "rung": _rung_tag(g, n, r), "run_dir": run_dir}
    if _done(run_dir):
        print(f"[single-point {_rung_tag(g, n, r)}] model.eqx present -- skip",
              flush=True)
        rec["status"] = "skipped"
        return rec
    rc = _run(_train_cmd(grid=g, n_agents=n, comm_r=r, run_dir=run_dir, iters=iters,
                         rollouts=rollouts, seed=seed, extra=extra, init_from=None))
    rec["status"] = "ok" if rc == 0 else f"rc={rc}"
    return rec

run_ladder

run_ladder(out, rungs, *, iters, rollouts, seed, extra)

(b) warm-start ladder: rung[0] from scratch, each later rung --init-from the previous rung's model.eqx. Sequential by data dependency; resumable per rung.

Source code in experiments/ctde_v0/run_scale_ladder.py

def run_ladder(out: str, rungs: list[tuple[int, int, int]], *, iters: int,
               rollouts: int, seed: int, extra: list[str]) -> dict:
    """(b) warm-start ladder: rung[0] from scratch, each later rung --init-from the
    previous rung's model.eqx. Sequential by data dependency; resumable per rung."""
    rec: dict = {"strategy": "warm-start-ladder",
                 "rungs": [_rung_tag(*x) for x in rungs], "steps": []}
    prev_ckpt: str | None = None
    for i, (g, n, r) in enumerate(rungs):
        run_dir = os.path.join(out, f"ladder_{i}_{_rung_tag(g, n, r)}")
        ckpt = os.path.join(run_dir, "model.eqx")
        step = {"rung": _rung_tag(g, n, r), "run_dir": run_dir,
                "init_from": prev_ckpt}
        if _done(run_dir):
            print(f"[ladder rung {i} {_rung_tag(g, n, r)}] model.eqx present -- skip",
                  flush=True)
            step["status"] = "skipped"
        else:
            # a later rung needs its predecessor's ckpt; bail clearly if absent.
            if prev_ckpt is not None and not (os.path.exists(prev_ckpt)
                                              and os.path.getsize(prev_ckpt) > 0):
                step["status"] = "blocked: predecessor model.eqx missing"
                rec["steps"].append(step)
                print(f"[ladder rung {i}] BLOCKED: {prev_ckpt} missing "
                      f"(run the earlier rung first)", flush=True)
                break
            rc = _run(_train_cmd(grid=g, n_agents=n, comm_r=r, run_dir=run_dir,
                                 iters=iters, rollouts=rollouts, seed=seed,
                                 extra=extra, init_from=prev_ckpt))
            step["status"] = "ok" if rc == 0 else f"rc={rc}"
            if rc != 0:
                rec["steps"].append(step)
                break
        rec["steps"].append(step)
        prev_ckpt = ckpt
    return rec