World-model variant study (Stage 1)¶

A model-level comparison of transition-kernel estimators, run before any downstream campaign. A failure here is a useful result: it says the more structurally appropriate model did not materially improve kernel estimation, without paying for the full RL campaign.

The arms and the adoption gate are predeclared in goal.md before the sweep runs. This note describes the machinery; it must not be used to redefine the criteria afterwards.

Why a standalone harness¶

The WM notebooks hard-code dir = f"exports/.../{N_AGENTS}/wm" and pass no seed to collect_rollouts_ppo, so they cannot be fanned out per variant without colliding. scripts/wm_variant_study.py is also the better experiment: every arm sees the same rollout buffer and the same fitting protocol, so a difference between arms is representational rather than procedural, and results pair cleanly by (environment, seed).

The notebooks stay untouched until the gate passes.

The five arms¶

\(q\) and \(\alpha\) are fixed per arm so no estimator is defined after seeing results.

Protocol¶

PPO buffer
├── train subset       -> all learned and count arms fit here
└── validation subset  -> epoch selection / early stopping only

held-out random buffer -> evaluation only
oracle graph           -> final metrics only, never tuning

A representation comparison is only valid if the training procedure is held constant. The published train_world_model_mlp has no validation split, so giving only the successor model best-epoch selection would confound representation with procedure. It therefore gained opt-in val_indices and early_stopping_patience, off by default so the published path is unchanged, and both neural arms use the same rule on the same split.

Scoring¶

Each arm is scored from its annotated graph, not from the model object. That is exactly the kernel the shield will consume, and it is the only way to put the product-Bernoulli baseline on the same footing as the categorical arms, since its row distribution exists only after candidate renormalisation.

• Exhaustive transition TV reuses export_world_model_transition_tv unchanged, so numbers stay comparable with reports/08/2026-08-16 reward-safety and guarantee audit.md §12.1.
• Row NLL, oracle cross-entropy, KL, categorical Brier, successor-conditioned reward MAE and oracle_row_mass come from src/world_model/successor/eval.py. See ../world-model/successor-kernel.md.
• Bit and cell accuracy remain available as debugging columns only.

Environments and seeds¶

Six environments have a tractable exact oracle. Bertrand, Chicken and Congestion-2 are deterministic — every row is a singleton, so TV is structurally zero for every arm — and act as regression guards. The signal is Gathering-2, MSH-2 and Pursuit-2.

Ten seeds per environment. The gate rests on bootstrap CIs over paired differences, and five paired observations per environment gives unstable intervals. If the stochastic environments prove expensive, the fallback is a five-seed pilot expanded to ten on those three before any adoption decision — never deciding on the pilot.

Caching¶

The legal graph and the oracle graph are seed-independent and are cached once per environment under a FileLock, not per cell. Ten concurrent seeds would otherwise rebuild them ten times, and those builds dominate cost on the larger environments.

Running it¶

# local smoke check on the smallest graph
uv run python scripts/wm_variant_study.py --env bertrand --seeds 1 \
  --train_steps 4000 --holdout_steps 2000 --epochs 8 --hidden 64

# inspect the fan-out before submitting
DRY_RUN=1 bash notebooks/experiments/scripts/wm_variants.sh

# ten seeds across six environments on cpucluster: 60 jobs
bash notebooks/experiments/scripts/wm_variants.sh

# stochastic environments only, larger budget
STOCHASTIC_ONLY=1 bash notebooks/experiments/scripts/wm_variants.sh train_steps=250000

# apply the predeclared gate
uv run python scripts/summarize_wm_variants.py --tag wm_variant_stage1

The cluster Python runner requires strict key=value arguments and translates each into --key value. The launcher therefore rejects CLI-flag syntax in its pass-through arguments, and the study script accepts underscore spellings (--seed_offset, --train_steps, …) with --allow_unresolved taking an explicit boolean value rather than being a bare store_true.

Results land under <output_root>/<env>/<agents>/seed<N>/summary.json, with wm_variant_gate.json, wm_variant_summary.md and wm_variant_comparison.tex written by the aggregator alongside them. The .tex is the paper's Stage 1 table; it is generated rather than maintained by hand so it cannot drift from the summaries.

Note that --output_root must already include the tag. The study writes <output_root>/<env>/<agents>/seed<N> and does not append --tag to the path — the tag labels the run, not the directory — so passing a bare reports root scatters cells beside the other tags' directories, where the next aggregation over that root will mix two sweeps.

The gate¶

Applied by scripts/summarize_wm_variants.py. Paired by (env, seed) on identical buffers, over the three stochastic environments, with 95% bootstrap CIs on the paired differences.

• Primary criteria: mean TV, p95 TV, oracle cross-entropy.
• Adopt an arm over bit_mlp if: paired mean-TV relative reduction \(\ge 20\%\) in at least 2 of 3 environments, with no environment worse than \(+5\%\) relative; paired p95-TV no worse than \(+5\%\) relative anywhere; paired oracle cross-entropy improved in all 3, with the CI excluding zero in at least 2.
• Max TV is a safety-oriented non-regression diagnostic: flagged if worse by more than \(10\%\) relative, but a single worst-row outlier does not by itself fail an arm.
• Deterministic environments: TV \(\le 10^{-6}\) for every arm.
• All four non-baseline arms are judged on the same footing. If a count model wins, that is the result — adopt and report it rather than forcing adoption of the neural arm. A null result is reported as-is.

Two protocol amendments, and how to reproduce either¶

The harness defaults reproduce the original 2026-09-20 gate, not the shipped protocol. That is deliberate — the gate verdict is immutable and must stay reproducible — but it means a probe of the current model has to ask for it:

Selection moved to the full validation objective because the row NLL is identically zero from the first pass on a singleton-support row, so selecting on it restores a checkpoint whose reward head is untrained. The optimiser-step floor followed because exact count aggregation makes an epoch a single update on a low-group-count environment, so the epoch cap ended training before patience ever fired. Neither is sufficient alone, which was measured rather than assumed: running the floor while leaving selection at row_nll recovers the original reward errors to within seed noise — 3.227, 1.438 and 1.331 over three seeds against 3.229, 1.427 and 1.321 over ten — because selection still pins at update 1 however long training runs. The 2,000 updates are taken and then discarded.

The whole 60-cell matrix was rerun with objective selection (wm_variant_stage1_objective, the paper's table) and reached the same verdict with the same passing arms; both verdicts stand, the later one alongside rather than replacing the earlier. That rerun predates the step floor: its archived summary.json configs have selection_metric: objective and no min_optimizer_steps key (checked 6 Nov 2026), so it ran with floor 0. The CX3 re-run repeats exactly that (cx3_arxiv_independent.sh).

min_optimizer_steps and max_optimizer_steps are pinned config keys in the aggregator alongside selection_metric, so a tree mixing budgets is reported as two protocols rather than one experiment.

What follows¶

Stage 2A (the downstream RL campaign) runs only if the gate passes; a null result leaves bit_mlp as the published nominal model. Stage 2B, the certified ambiguity-set evaluation, runs regardless — see ../shielding/anytime-ambiguity-sets.md.