Skip to content

Training Pipeline

The complete learned-shield workflow is ordered because each stage consumes artifacts from the preceding stage.

train_wm_mo.py  ->  train_om_mo.py  ->  experiments_mo.py
 exact topology      level-0 data        baselines
 successor kernel    imagined levels     learned shield (OmSh)
                     one per role        TrueShielded oracle
                                         frozen evaluation

Stage 1: World Model

Environment train_wm_mo.py notebooks:

  1. instantiate the labelled environment,
  2. build or validate wm/legal_joint_graph.pkl,
  3. collect a replay buffer and policy history into wm/D.pkl,
  4. train and evaluate the nominal model selected by wm_variant (successor by default, or bit_mlp), and
  5. annotate legal edges in wm/env_transition_graph.pkl.

The successor model trains under a frozen protocol, so its reports are comparable across environments. The stage also records detailed wm_* timings and a component_wm total in the environment-root timing sidecar. See World Modelling for the two variants and their reports.

Stage 2: Opponent Model

train_om_mo.py loads the learned environment graph and world-model rollout data. It prepares the level-0 dataset, reports train/holdout metrics, trains the higher imagined reasoning levels, and saves om/iop_stack.pt. Several opponents are modelled as one joint team. With more than one opponent, the team uses an exact autoregressive factorization rather than independent per-agent models.

The OM stage must match the WM public-state width, action size, focal agent, and agent count. Graph-backed imagination rejects ambiguous or incomplete legal state/action coverage.

To shield a second role in Congestion or Markov Stag Hunt, run the stage again with --focal-idx=1; it writes om/agents/player_1/iop_stack.pt. Run role models one after another, because they share the environment's timing file.

Stage 3: Experiments

experiments_mo.py can run five algorithm stage families:

  • IPPO,
  • IPPO-Lagrangian,
  • CPO,
  • IPPO with the learned transition shield (OmSh), and
  • IPPO with the true-policy exact shield (TrueShielded, an oracle reference).

It also exposes an independent, off-by-default exact-anytime certificate audit. That diagnostic builds and validates an arbitrary-opponent certificate without initializing an RL trainer or consuming the policy-training step budget.

The three unshielded baselines need only the environment. OmSh needs wm/env_transition_graph.pkl and the opponent stack for every shielded role; its cached shield/transition_shield_bundle*.pkl is rebuilt when metadata changes. TrueShielded builds or loads true_shield/true_transition_graph.pkl and does not need the learned WM/OM artifacts. Once that exact graph exists, OmSh setup also writes learned-versus-exact diagnostics.

The learned-shield stage is itself configurable. shield_action_mode selects pure actions, pairwise mixtures, learned successor budgets, or the combined parameterization. shield_opponent_mode selects the default monotone floor or the credible-floor, all-level, and Bayesian-mixture ablations. shield_padding_radius pads the certificate, and shielded_agent_ids protects several roles at once. These are conditions of run_ippo_shielded, not separate run_* stages. Every condition starts its policy from fresh initialization while reusing the compatible pretrained WM/OM inputs.

All algorithm training stages export completion manifests, episode tables, update metrics, aggregate diagnostics, and checkpoints. With final_policy_eval_episodes > 0 they also evaluate the frozen final policy; see Final-Policy Evaluation. Tagged campaigns place these artifacts below experiments/runs/<tag>/. The exact-anytime audit instead writes its certificate and report below true_shield/exact_anytime/.

Local Example

uv run python notebooks/matrix/chicken/train_wm_mo.py --n-steps=25000
uv run python notebooks/matrix/chicken/train_om_mo.py
uv run python notebooks/matrix/chicken/experiments_mo.py \
  --run-ippo-shielded=true \
  --num-runs=1

For Gathering, Pursuit, and Congestion, pass the same --n-agents value to all three commands. --n-steps belongs to WM collection; --timesteps and --num-runs belong to experiments. For Pursuit with the default successor world model, also pass --max-risk=0.16 to the experiment stage (see Safety Budgets).

For production sweeps, use the dependency-aware launchers in Cluster Execution instead of manually timing stage submissions. Notebook defaults are not the campaign protocol; see Reproduce the Campaign Protocol.