Skip to content

Training Pipeline

The main experiment pipeline has three stages.

train world model -> train opponent model -> run experiments

Each supported environment has notebooks following this pattern.

Stage 1: World Model

World-model notebooks are named:

notebooks/.../train_wm_mo.py

They usually:

  • instantiate the environment,
  • build or load the legal joint graph,
  • collect rollout data,
  • train WorldModelMLP,
  • save world_model.pt,
  • build env_transition_graph.pkl,
  • write timing and evaluation outputs.

Stage 2: Opponent Model

Opponent-model notebooks are named:

notebooks/.../train_om_mo.py

They usually:

  • load wm/env_transition_graph.pkl,
  • prepare opponent-model training inputs,
  • fit level 0 from observed behavior,
  • train higher levels with imagined graph rollouts,
  • save the opponent stack as om/iop_stack.pt.

Stage 3: Experiments

Experiment notebooks are named:

notebooks/.../experiments_mo.py

They compare baseline and shielded learners, such as:

  • IPPO,
  • IPPO Lagrangian,
  • ICPO,
  • IPPO with the learned shield,
  • IPPO with the exact true shield baseline when available.

Shielded runs depend on the world-model and opponent-model artifacts. Plain baseline runs only need the environment.

For a complete learned-shield run:

uv run python notebooks/matrix/chicken/train_wm_mo.py
uv run python notebooks/matrix/chicken/train_om_mo.py
uv run python notebooks/matrix/chicken/experiments_mo.py --run-ippo-shielded=true

Use smaller --timesteps, --num-runs, and notebook-specific CLI options when doing a smoke test.