Masked categorical successor kernel¶

The shield consumes a transition kernel \(\Gamma(s' \mid s, a)\) over an already-enumerated legal successor set. src/world_model/successor/ learns that object directly, instead of predicting independent next-state bits and recovering a kernel from them afterwards.

Why the baseline representation is a mismatch¶

WorldModelMLP predicts one Bernoulli per (cell, channel) of the next state. apply_world_model_env_and_rewards then scores each enumerated candidate under a product-Bernoulli likelihood and renormalises across the candidate set:

Two consequences show up in the published diagnostics:

1. Independence is the wrong inductive bias. Renormalising over legal states introduces dependence after the fact, but the relative scores of legal successors remain constrained by independent marginals. Correlated stochastic effects — joint slip, collisions, simultaneous respawn — are exactly what the product form cannot represent. reports/08/2026-08-16 reward-safety and guarantee audit.md §12.1 records held-out bit accuracy \(\ge 0.996\) next to exhaustive kernel TV of \(0.113\) (MSH) and \(0.129\) (Pursuit) with maxima of \(0.859\) and \(0.817\); see also reward-loss-scaling.md, where cell accuracy moved \(0.996 \to 0.991\) while TV moved \(0.134 \to 0.284\). Bit accuracy and the shield's actual input are decoupled.
2. Rewards cannot depend on the realised successor. The reward head is \(\hat r(s,a)\), so every stochastic successor of a group receives an identical reward vector. Where a slip into a collision or an interception determines the payoff, the realised \(s'\) is the reward signal.

The model¶

SuccessorEdgeModel computes a trunk \(h = f_\theta(x)\) on the existing encode_state_action feature vector, so the input contract is shared with the baseline. Each candidate is embedded as \(\psi = g_\theta([\,b(s'),\ b(s') - b(s)\,])\) and scored by \(e_\theta = \mathrm{MLP}([\,h,\ \psi,\ h \odot \psi\,])\), with a parallel reward head \(r_\theta(s,a,s')\).

The delta channel \(b(s') - b(s)\) is load-bearing: it tells the scorer what changed, which is where correlated effects live. Values lie in \(\{-1, 0, 1\}\).

There is no sink logit¶

The premise of the method is that the legal successor set is known and complete, so genuine environment samples never target a sink. A learned sink logit would receive no legitimate positive examples, would let resolution failures leak into environment probability, and would break the deterministic-environment regression guard, since a single-successor row could then carry nonzero sink mass.

Three layers stay separate:

• the nominal model sums to one over known legal successors;
• src/shield/ambiguity represents uncertainty over the probabilities on that certified support;
• probability reaches the unsafe sink only through existing graph deficits or a separately justified missing-support bound.

The ambiguity set does not itself create sink mass. See ../shielding/anytime-ambiguity-sets.md.

Dataset: aggregation is exact, not an approximation¶

build_successor_groups reads Succ(s,a) from the legal joint graph, which every train_wm_mo.py already builds before world-model training, so no new enumeration is needed. Candidate sets come from _group_edges; buffered states resolve to node keys through state_key_from_bits at thr=0.5, the same convention the shield uses.

resolve_transitions hard-errors on an unresolved row by default, following prepare_om_training_inputs. allow_unresolved=True drops and records them; a resolution failure is never converted into probability mass.

Rows then aggregate to per-group counts. This is exact:

and the reward MSE against per-edge mean rewards differs from the per-sample MSE only by a term independent of the parameters, since \(\sum_i (r_\theta - r_i)^2 = N(r_\theta - \bar r)^2 + \sum_i (\bar r - r_i)^2\). Both gradients are therefore preserved, while the number of forward passes drops from one per transition to one per group.

The aggregated table is the empirical transition-count kernel, so src/world_model/counts and the confidence sequences consume the same object.

Memory layout¶

Candidates are stored as integer indices into a single node_bits array, never as a padded [N, K, state_width] tensor. A padded per-transition layout for Pursuit would be tens of gigabytes; the indexed layout is megabytes. Batches are padded to the batch-local maximum fan-out and masked with \(-\infty\) before the softmax, so a padded row is numerically identical to the same row batched alone — pinned by test_padded_row_is_identical_to_the_unpadded_row.

Annotation¶

apply_successor_model_env_and_rewards is the counterpart to apply_world_model_env_and_rewards. It differs in three ways:

• one batched forward per group, rather than a batch-size-1 forward per group through predict_probs_bits_rewards (which also calls model.eval() on every invocation);
• env_prob from the masked softmax, still summing to one within a group, so the structural contract seen by the shield, GraphWorldModelAdapter and export_world_model_transition_tv is unchanged;
• rewards per edge, which makes the probability-weighted reward averaging in induce_focal_mdp meaningful instead of a no-op.

The output artefact is the same wm/env_transition_graph.pkl, so every downstream consumer is untouched.

Anything implementing SuccessorKernel can be passed as kernel. NeuralSuccessorKernel wraps the trained model, CountSuccessorModel the empirical table, and GraphSuccessorKernel reads rows back out of an already annotated graph — the last is how the product-Bernoulli baseline is scored on the same footing as the categorical arms, since it has no native row distribution of its own.

Metrics¶

src/world_model/successor/eval.py scores the successor distribution itself: row NLL, oracle cross-entropy and KL against an exact graph, categorical Brier, successor-conditioned reward MAE, and oracle_row_mass.

Calibration is NLL plus categorical Brier. The binary 15-bin ECE used for independent bits has no direct multiclass meaning, so a reliability view is reported as an explicitly named top-label reliability (confidence of the argmax against how often the argmax was realised) rather than relabelled as ECE.

zero_probability_rate is surfaced because a count MLE can assign exactly zero to a legal-but-unobserved successor. That is a real weakness of that arm and should be visible rather than absorbed by the probability floor.

oracle_row_mass quantifies the deficit that guard_successors deliberately preserves rather than renormalising away. Per-agent slip distributions there sum to one, so the deficit is expected to be structurally absent; the diagnostic confirms that rather than assuming it.

Checkpoint selection, and why it is not the row NLL¶

train_successor_world_model takes selection_metric, which decides what the validation split selects on when restoring the best epoch:

• "row_nll" — the validation row NLL alone. This is what the 2026-09-20 gate ran under, and it stays the function default so that run is reproducible.
• "objective" — row NLL plus reward_weight times the scaled reward MSE, the same quantity the training split minimises. This is the frozen protocol's value since the 2026-09-21 amendment.

The reason is that the two terms converge on completely different timescales. The row NLL is a log-loss over a handful of candidates and falls fast; on a singleton-support row it is identically zero from epoch 1, so NLL-based selection there is vacuous and early_stopping_patience fires at epoch 11 with the epoch-1 weights restored. The reward head is a regression over raw environment units and is nowhere near converged at that point. Stage 1 measured the consequence directly: successor reward MAE of 3.229 on Bertrand against the baseline's 0.0001, with mean selected epoch 1.0, and the same story on Chicken and Congestion. Only Pursuit, whose kernel is genuinely hard, ran long enough (mean best epoch 44.9) for the reward head to converge — and there the successor model beat the baseline threefold.

This matters beyond the diagnostics because apply_successor_model_env_and_rewards writes r_theta(s,a,s') onto every edge, and those rewards are what induce_focal_mdp averages and what the opponent-model imagination rollouts score. A model that predicts the kernel perfectly and the reward badly is not a usable nominal model.

evaluate_counts_terms returns both validation terms from a single forward pass, so selecting on the sum costs nothing over selecting on the NLL. The report records selection_metric, validation_reward_mse and validation_objective alongside validation_nll, so which rule produced a checkpoint is always recoverable from the artefact.

An alternative considered and rejected: weighting the NLL term down, or warming up the reward term. Both change the objective, which would mean the adopted model is no longer the one the gate scored. Changing only the selection rule leaves the optimisation identical and fixes the actual defect, which was that the stopping criterion ignored half the loss.

The optimisation budget is counted in updates, not epochs¶

Aggregation is exact, and that exactness has a cost worth stating plainly: the dataset collapses to the number of visited state-action groups, so the gradient-step budget scales with the graph, not with the buffer. An epoch is therefore a completely different amount of optimisation in different environments. Measured on the immutable Stage 1 tree:

The diagnosis is not that sixty epochs is too few. Where a pass is hundreds of updates, patience-based early stopping fires on its own in every cell — the budget is never the binding constraint. Where a pass is a single update, the epoch cap ends training first and early stopping never gets the chance. Bertrand and Chicken selected at epoch 57.9 and 57.7 of 60, still descending when the cap hit, and shipped a reward MAE of 0.2166 against the baseline's 0.0002. The product-Bernoulli baseline is unaffected because it trains on individual transitions, so its 750,000-row buffer is about 58,000 updates however small the graph is.

train_successor_world_model therefore takes min_optimizer_steps and max_optimizer_steps. Training continues past epochs — with early stopping suppressed — until the floor is met, after which patience governs exactly as before. The floor is on optimisation, not on selection: the best checkpoint is restored either way, so raising it can only ever select a checkpoint at least as good on the validation objective. Both default to off, reproducing the epoch-budgeted runs.

Why the floor is 2,000¶

The value is bounded from both sides by Stage 1 data alone, which is what makes it predeclared rather than tuned.

Sufficiency. On a twenty-group singleton-support row set reproducing Bertrand's shape, at the frozen settings, the reward head converges at update 333: the reward MAE falls from 0.6368 at the sixty-update cap to 0.0, and the selected step stays 333 whether the floor is 500, 1,000, 2,000, 4,000 or 8,000. Raising the floor past that buys nothing but runtime, which grows at roughly 70 ms per update on a twenty-group set because validation runs once per pass. 2,000 is six times the observed convergence point.

Inertness. The floor must not bind where patience already governs. That is what confines the amendment to environments whose Stage 1 total variation was structurally zero for every arm, and so what leaves the gate verdict standing. Measured over the complete 60-cell re-verification, and per seed rather than on the mean — each figure is the fewest updates any single seed of that environment took before patience stopped it:

The floor binds on exactly the three environments it was meant for and on no seed of any stochastic one, where the smallest figure is more than three times the floor. The gate's three deciding environments are therefore untouched by it, as a matter of measurement rather than of argument.

MAX_OPTIMIZER_STEPS = 200000 is a safety rail rather than an operating point: the sixty-epoch budget on the largest environment is 31,500 updates. A cap below the floor is rejected as a contradictory budget rather than silently truncating.

The report records optimizer_steps, best_optimizer_step, steps_per_epoch, min_steps_floor_extended and max_optimizer_steps_reached, so how much optimisation a checkpoint actually received is recoverable from the artefact. summarize_wm_variants.py pins the budget alongside selection_metric: a tree mixing epoch-budgeted cells with floored ones is two protocols, not one experiment.

An alternative considered and rejected: expressing early_stopping_patience in updates as well. It is the more uniform rule, but at any single value it perturbs the stochastic environments — Pursuit-Evasion's patience is 1,310 updates and Gathering's is 5,250 — and changing a gate environment is exactly what the amendment must not do. Rescuing the environments where the cap bound, and leaving alone the ones where patience already governed, is the narrower and better-justified change.

Status¶

Implemented and unit-tested. The predeclared Stage 1 gate in goal.md adopted this model on 2026-09-20 over 60 cells; see ../experiments/wm-variant-study.md for the harness and the verdict. WorldModelMLP remains in place as the baseline arm. The settings the gate validated are frozen in src/world_model/successor/protocol.py, which also records the two amendments made since and why: checkpoint selection on the full validation objective, and the optimiser-step floor above. Both are justified from Stage 1 data that was complete and immutable before either was made, and neither was tuned against a Stage 2A outcome.