Skip to content

World Modelling

World modelling turns environment experience into probabilities and rewards on top of an exact legal joint graph.

The Three World-Model Artifacts

Artifact Purpose
wm/legal_joint_graph.pkl Reachable public states, legal joint actions, and legal successor topology.
wm/world_model.pt Neural checkpoint predicting next-state bits and per-agent rewards from (state, joint_action).
wm/env_transition_graph.pkl A copy of the legal graph annotated with learned env_prob and rewards.

The legal graph cache also has wm/legal_joint_graph.metadata.pkl. Its 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

train_wm_mo.py collects policy rollouts into D.pkl, trains WorldModelMLP, and evaluates the checkpoint. 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.

The normal evaluator writes wm_report.{json,txt} with held-out next-state and per-agent reward metrics, state calibration bins, and uncertainty summaries. When a compatible exact graph already exists, the workflow can also write transition_tv_report.{json,txt} over every enumerated state/action group. This graph-exhaustive comparison measures the learned conditional transition kernel; it is not automatically a certificate about states outside the graph.

Most notebooks collect live-policy experience. Ice Duel additionally uses a coverage-aware mix of PPO and uniform-random live-environment transitions, then splits train/validation data after collection. This is an environment-specific calibration measure, not a change to the graph or shield interfaces.

From Bit Predictions to Edge Probabilities

The MLP emits independent next-state bit probabilities. For one legal (state, joint_action) group, the graph augmenter:

  1. scores each exact candidate successor under those bit probabilities,
  2. normalizes the scores across that candidate set,
  3. writes the normalized value to Edge.env_prob, and
  4. stores the predicted rewards in environment-agent order.

The learned model therefore chooses probabilities among exact candidates; it does not add new graph nodes or successors.

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/model problem. That distinction is important because shielding treats missing probability mass and missing graph coverage conservatively.