World Modelling¶

World modelling turns environment experience into transition probabilities and rewards on top of an exact legal joint graph. The learned model never invents topology: it only assigns probabilities to successors the exact builder already enumerated.

Two Nominal Models¶

train_wm_mo.py trains one of two nominal models, selected by wm_variant:

Both produce the same handoff artifact, wm/env_transition_graph.pkl, with env_prob summing to one within each group. The opponent model, the shield, and the RL learners consume it unchanged, so the variant can be swapped without touching them.

Why the successor kernel is the default¶

The shield consumes a kernel Gamma(s' | s, a) over legal successors. The baseline recovers that kernel indirectly by scoring every candidate under independent bit probabilities and renormalizing. That has two weaknesses:

• Independence is the wrong inductive bias. Correlated effects such as joint slips, collisions, and simultaneous respawns cannot be represented by independent marginals. High held-out bit accuracy can coexist with large kernel error.
• Rewards cannot depend on the realized successor. The baseline's reward head sees only (s, a), so every stochastic successor of a group gets the same reward.

The successor model scores each legal candidate directly from the shared state-action features and the candidate's bits, including what changed (b(s') - b(s)). A masked softmax normalizes over the legal set, and a parallel head predicts r(s, a, s') per edge.

It was adopted after a predeclared Stage 1 comparison of five kernel estimators over six environments and ten paired seeds. The comparison scored exhaustive transition TV and oracle cross-entropy against the exact graph. The baseline remains available as wm_variant=bit_mlp and reproduces the published product-Bernoulli model exactly.

Artifacts¶

The legal graph identity includes environment type, public-state width, exact-builder source, seed, and builder arguments. A mismatch triggers a rebuild instead of silently reusing incompatible topology.

Data Collection and Training¶

The notebook collects live policy rollouts into D.pkl. Environment and action data use the public flattened state and the complete joint action. Hidden episode caps are converted into explicit truncation boundaries so advantage bootstrapping and replay episodes remain consistent.

For the successor model, buffered transitions are resolved against the legal graph and aggregated into per-group successor counts. This aggregation is exact for both the likelihood and the reward loss, and it makes one forward pass per visited group rather than per transition. A transition that cannot be resolved to a legal successor is an integration error, not probability mass: more than a small allowance fails the stage.

The successor protocol is frozen in src/world_model/successor/protocol.py, so the downstream campaign trains the same model the Stage 1 comparison measured. Two details matter when reading its reports:

• Checkpoint selection uses the full validation objective (row NLL plus scaled reward error). On deterministic rows the NLL is zero from the first pass, so NLL-only selection would keep an untrained reward head.
• Training has a floor of 2,000 optimizer updates. Exact aggregation makes one epoch a single update on small graphs such as the matrix games, so an epoch cap alone would stop before the reward head converges. On the stochastic gridworlds, early stopping already ends training well past this floor, so it has no effect there.

Evaluation¶

The two variants report different quantities because they predict different objects:

• Successor model. Row negative log-likelihood of the observed successor, categorical Brier score, and reward error at the realized successor. Bit-level accuracy and calibration are undefined for a categorical kernel and are omitted rather than substituted.
• Product-Bernoulli model. Held-out next-state bit metrics, per-agent reward metrics, state calibration bins, and uncertainty summaries.

When the exact graph true_shield/true_transition_graph.pkl already exists, the learned-shield setup also writes wm/transition_tv_report.{json,txt}: the total variation between learned and exact rows over every enumerated group. This measures the learned kernel that the shield actually consumes. It is not a certificate about states outside the graph.

Why the Separation Matters¶

A missing state or joint action is a legal-graph integration or stale-cache problem. Poor probabilities on covered edges are a data or model problem. The distinction matters because shielding treats missing probability mass and missing graph coverage conservatively, while kernel error on covered rows is what padding and robustness margins account for.

See the maintainer notes on the successor kernel and the Stage 1 variant study for the gate, metrics, and protocol amendments.