Successor world-model campaign (Stage 2A)¶

The downstream reward–safety campaign that follows the Stage 1 gate (wm-variant-study.md): the published algorithm protocol, rerun with the masked categorical successor kernel as the nominal world model.

Two launchers, two stages¶

The pipeline has two entry points, and they are not interchangeable.

pipeline_env.sh routes wm_variant to the WM stage correctly, but its algorithm stage inherits the notebook defaults — 3,000,000 timesteps, no experiment_tag, no frozen-policy evaluation. The published campaign (paper_final_corrected_2026_08_v2) was produced by corrected_final_campaign.sh, which passes 250,000 timesteps (1,000,000 for Markov Stag Hunt), seed_offset=2200, bundles=4, experiment_tag, and 1000 frozen-policy evaluation episodes at horizons 1/10/50/100/200. Its algorithm runs are the only ones either aggregator (summarize_final_policy_campaign.py, summarize_evaluation.py) can find.

So: pipeline_env.sh for WM and OM, then corrected_final_campaign.sh for the algorithms. The first Stage 2A launch used pipeline_env.sh end to end and its algorithm runs had to be discarded — they would have compared a 3M-step successor campaign against a 250k-step baseline.

corrected_final_campaign.sh takes two environment variables beyond the published defaults:

• PURSUIT_MAX_RISK (default 0.11). Pursuit-Evasion's budget is the first clean hundredth above the learned shield's robust reset-state demand, which depends on the world model: 0.102817 under the product-Bernoulli model, 0.151488 under the successor kernel. A successor campaign passes 0.16; at 0.11 the shield correctly refuses to deploy. Since 6 Oct 2026 the Pursuit notebook itself defaults to max_risk=0.16, the successor-kernel budget; this script still defaults to 0.11 so that it reproduces the published product-Bernoulli campaign.
• ONLY_ENVIRONMENTS (default: all five). A space-separated subset, so environments whose opponent model is ready can start while another's trains.

Congestion is a sixth environment with its own protocol¶

The paper's final tables have six environments; corrected_final_campaign.sh covers five. Congestion's published row comes from convention_merge_gate_confirm_v1_joint (historical path gridworlds/convention_merge/2): both roles shielded, 64-step horizon, p = 0.10, 2000 evaluation episodes, trainer seeds 5300-5304, evaluation seed block 4,200,000, run on amd24. Its experiment notebook's defaults match everything except the evaluation seed block, whose default is 1,000,000, so it runs through pipeline_env.sh with that one override:

OMSH_SLURM_WM_MEMORY=8G bash pipeline_env.sh congestion amd24 \
  wm_variant=successor om_focal_indices=0,1 \
  seed_offset=5300 experiment_tag=<tag> shielded_agent_ids=player_0,player_1 \
  final_policy_eval_seed_offset=4200000

om_focal_indices=0,1 trains an opponent model per role, sequentially; shielded_agent_ids must name the same roles. Omitting Congestion would leave one paper row on the old world model.

The first successor Congestion run missed the seed-block override: every manifest field matched August's apart from final_policy_eval_seed_offset (1,000,000 against 4,200,000). The results are valid, just an independent evaluation sample, and they ran on amd48 rather than amd24, so the model-free rows matched August only within noise rather than bit for bit. The algorithm stage was rerun on amd24 with the override (jobs 8538-8541). The first run is kept at runs/successor_campaign_v1_evalseed1000000_amd48. Always diff the complete manifest config against the baseline campaign's, not only the flags you remember passing.

When you copy a DRY_RUN=1 line by hand, remember that pipeline_env.sh prints its commands with printf %q. The \, in player_0\,player_1 is display quoting, not part of the argument.

Exact post-processing of the model-free baselines¶

The paper tables' \(U_\infty\) columns come from <stem>_safety.json beside each run's frozen-policy reports. The algorithm jobs do not write these. They come from scripts/evaluate_final_campaign_baselines.py, one job per environment once IPPO, IPPO-Lagrangian and CPO are complete:

sbatch -p amd48 --mem <M> py_cpu.sh omsh/scripts/evaluate_final_campaign_baselines.py \
  venv=omsh tag=<tag> environment-name=<env> \
  horizons=1,10,50,100,200 evaluation-seed-offset=2200000 episodes=1000

Congestion uses horizons=1,2,10,64 evaluation-seed-offset=4200000 episodes=2000. The seed offset must equal the live reports' block, so the script reuses them and computes only the exact graph quantities. Memory follows the exact graph: 64G Markov Stag Hunt, 32G Gathering, 16G Pursuit, 8G otherwise. Without these files the summarizer still runs, but the compact table's third number and the all-conditions \(U_\infty\) columns come out empty.

Memory¶

Reservations are per stage, and the launch commands must set each one that matters. The WM stage needs OMSH_SLURM_WM_MEMORY sized to the graph (64G Markov Stag Hunt, 32G Gathering, 16G Pursuit, 8G otherwise). The OM stage needs OMSH_SLURM_OM_MEMORY=64G on Markov Stag Hunt, whose 1.5 GB annotated graph plus the OM's own grouped index does not fit the 4 GB default; Gathering (536 MB) and Pursuit (91 MB) are fine at 4 GB.

The budget-ablation experiment notebooks (notebooks/ablations/<env>/experiments_mo.py) train three runs of 3M steps, and their end-of-run plotting holds the whole history in memory. Pursuit's history is 410 MB, and its plotting peaked at 9.65 GB, so give these notebooks at least 16G. At the 4 GB default, 8510 finished training and wrote every result, then died while plotting. Markov Stag Hunt's ran at 64G.

A marimo kernel that dies does not end the job: it stays in RUNNING indefinitely, holding a slot while its afterok dependents wait forever. Check logs for kernel was stopped unexpectedly rather than trusting squeue.

Errno 122 Disk quota exceeded at wandb.init, seconds into an algorithm job, is not this project's storage. ~/.bashrc on the cluster points WANDB_DIR and WANDB_CACHE_DIR at /vol/gpudata/oja24-cac/tmp/, a personal cache volume that also holds the uv Python installs; when it is over quota, only notebooks with use_wandb=True (Congestion and Markov Stag Hunt by default) fail. Pass use_wandb=False — wandb only mirrors logs, results come from the exported histories — or free that volume. corrected_final_campaign.sh does not forward extra arguments, so for a subset rerun build the env.sh call by hand and diff its DRY_RUN=1 output against the script's.

To resume from the OM stage, do not rerun pipeline_env.sh — passing it a train_om_mo.py path still submits the whole chain, and the WM notebook seeds its own buffer, so the rerun would train a different world model. Submit the OM job directly and chain the algorithm jobs to it, copying the lines from a DRY_RUN=1.

Reading the result¶

scripts/compare_final_policy_campaigns.py compares two campaigns' final_policy_campaign.json. Its design rests on one structural fact: only the shielded condition consumes a world model. IPPO, IPPO-Lagrangian and CPO are model-free — their training cells read no world-model, opponent-model or shield artefact — and the campaign reuses the published seeds. The comparison therefore reports two things:

• a reproduction check on the three model-free conditions, which should match the published run up to CPU nondeterminism. The reward difference is expressed in units of the baseline's across-seed standard deviation, so seed noise is not mistaken for a discrepancy. A real mismatch would mean something other than the world model changed since August — environment or training code.
• the world-model effect, read off the shielded condition alone.

A risk-threshold difference between the campaigns is printed in bold with a caveat rather than hidden: Pursuit's shielded row compares 0.11 against 0.16 and is not a like-for-like budget.

Why Pursuit's budget moved¶

The expectation was that a more accurate kernel would relax the shield. It did the opposite: the product-Bernoulli kernel was optimistic about risk, and the successor kernel — exhaustive transition TV 0.058 mean and 0.691 max, against the published deployed model's 0.129 and 0.817 — exposes that. (Stage 1's freshly trained bit_mlp sat at 0.293; compare deployed models with deployed models.) Stage 2B's attribution found the same direction of error one level down on the same environment, where the maximum-likelihood count centre also understated risk (0.72833 estimated against 0.88632 structural). A kernel wrong in the optimistic direction buys feasibility the true dynamics do not support.

Ablations under the successor kernel¶

The paper's appendix ablations were rerun with the successor kernel. Results are in the audit report, reports/09/2026-09-23 successor world-model audit.md, and the dated goal.md entries.

Parameterized replacement (successor_replacement_ablation_v1). The shielded condition is rerun with shield_action_mode=egalitarian_mixed or learned_budget_pure, with every other argument equal to the primary campaign's, including the trainer and evaluation seeds. So each variant seed pairs with the primary seed that shares its initialisation. Congestion is excluded because its coordinated shield implements pure replacement only. scripts/summarize_replacement_ablation.py pairs variant with primary and writes markdown, JSON and the paper's replacement_ablation.tex.

Promoting a variant needs a disjoint, predeclared screen. The final seeds may not select a configuration. When a variant looks better on them, as mixed replacement did in Markov Stag Hunt, the procedure is:

1. Write the promotion rule into goal.md and push it before launching anything.
2. Run both the primary and the variant on fresh trainer seeds and a fresh evaluation-seed block, one job per (seed, condition). Tag each seed <prefix><seed>, so both conditions' run0 land in the same directory. Single-seed jobs with num_runs=1 must not share a tag: every one of them writes run0.
3. Apply the rule with scripts/decide_replacement_screen.py, which only reads the reports and writes a decision JSON. Do not revisit the rule afterwards.

The Markov Stag Hunt screen (seeds 7200-7204, evaluation block 7,200,000, jobs 8567-8576) kept pure replacement: one mixed seed reached U₂₀₀ = 0.201 against p = 0.20. The "reward gain" that motivated it was an artefact of the pre-2026-09-24 evaluator, which sampled the frozen pure policy outside the shield's mask (see ../reinforcement-learning/final-policy-safety-evaluation.md). Evaluated as trained, pure earns more than mixed (0.605 against 0.585 on the screen seeds), so the rule fails on both counts. Each Markov Stag Hunt shielded job peaks near 40 GB, so the ten screen jobs ran at 48G, spread across amd48 and amd24.

Private-reward inference. See ../opponent-model/private-reward-inference-ablation.md. It needs four algorithm-stage settings the notebook defaults get wrong, and a live-only evaluation with a CPU-only JAX environment. Under the successor kernel it reproduces August bit for bit, because both games are deterministic.

Budget ablation. The notebooks under notebooks/ablations/<env>/ keep their own protocol: 3 runs of 3M steps, p = 0.20. A weak world model is the successor kernel trained on 75k steps for 6 epochs with an optimiser-step floor of 200. The strong-WM cells symlink the campaign's successor model into exports/ablations/<env>/strong_wm_weak_om/wm/. Give the experiment notebooks at least 16G for their plotting step (see Memory). The successor kernel reverses August's feasibility in the finished cells. Pursuit weak WM + strong OM becomes infeasible, with reset demand 0.964. The cause is the weak kernel, not the opponent model: over the exact kernel the same opponent model gives 0.103. The weak kernel is diffuse. It puts 0.32 of a row's mass on the exact modal successor, against 0.84 for the campaign model, so the focal agent cannot steer and the risk compounds. Its one-step unsafe mass is nearly exact. Both strong-WM + weak-OM cells become feasible. ablations_budgets.tex in the paper is hand-maintained from those runs' logs and diagnostics.

Plotting defects fixed during the campaign. - The with-prefix learning curves prepended the world-model history to seed 0 only, and truncated after prepending. The shielded curve was therefore mostly data collection, averaged with unprefixed seeds. See prefix_and_truncate_histories in src/visualisation/common.py. - The coordinated Congestion shield took IPPO's colour. See series_color in src/visualisation/styling.py.

Running jobs plot with the code they pulled at start, so regenerate figures with just_vis=True after a plotting fix.