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
 learned dynamics    imagined levels     learned shield
                                          exact baseline

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,
  4. train and evaluate wm/world_model.pt, and
  5. annotate legal edges in wm/env_transition_graph.pkl.

The stage also records detailed wm_* timings and a component_wm total in the environment-root timing sidecar.

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. The maintained configuration uses one autoregressive joint team policy per level; it does not split the joint opponent into 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.

Stage 3: Experiments

experiments_mo.py can run five algorithm stage families:

  • IPPO,
  • IPPO-Lagrangian,
  • ICPO,
  • IPPO with the learned transition shield, and
  • IPPO with the true-policy exact shield.

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. Learned-shielded IPPO needs wm/env_transition_graph.pkl and om/iop_stack.pt; its cached shield/transition_shield_bundle.pkl is rebuilt when metadata changes. The true-policy baseline builds or loads true_shield/true_transition_graph.pkl and does not need the learned WM/OM artifacts.

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. Non-pure modes automatically use ParameterizedIPPO; they are 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. Tagged campaigns place these condition artifacts below experiments/runs/<tag>/ and can be summarized with scripts/summarize_evaluation.py. 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 production sweeps, use the dependency-aware launchers in Cluster Execution instead of manually timing stage submissions.