Skip to content

WM/OM Ablation Notebooks

The ablation notebooks live under notebooks/ablations/<env_name>/ and mirror the standard three-stage workflow:

  1. train_wm_mo.py
  2. train_om_mo.py
  3. experiments_mo.py

They currently cover:

  • pursuit, fixed to 2 agents
  • markov_stag_hunt

Conditions

Two ablation conditions are supported by --ablation:

  • weak_wm_strong_om
  • strong_wm_weak_om

strong means the checked-in default budget from the corresponding canonical environment notebook. weak means a 10% training-budget variant while keeping the model architecture and most optimization settings unchanged.

Budget Definition

Weak world model:

  • n_steps=75_000 instead of 750_000
  • epochs=6 instead of 60
  • hidden width, batch size, learning rate, graph build, and reporting settings otherwise match the canonical notebook for that environment.

Weak opponent model:

  • level-0 fit epochs: 3 instead of 30
  • pursuit soft-rollout PI: agent_updates=200, opp_updates=400
  • markov_stag_hunt soft-rollout PI: agent_updates=300, opp_updates=450
  • OM architecture, rollout horizon, rollout count, batch sizes, temperatures, and learning rates otherwise match the canonical notebook.

The notebooks also expose CLI overrides such as --n-steps, --epochs, --eval-steps, --level0-epochs, --agent-updates, and --opp-updates for smoke runs. Those overrides are intended for testing/debugging, not for the reported ablation setting.

Export Layout

All ablation artifacts are rooted at:

exports/ablations/<env_name>/<ablation>/

Each condition root is shaped like a normal experiment root:

  • wm/ stores D.pkl, models.pkl, policy checkpoints, world_model.pt, env_transition_graph.pkl, wm_history.pkl, and wm_report.txt
  • om/ stores iop_stack.pt and policy_after_imagination.pt
  • shield/ stores the learned transition-shield cache
  • experiments/histories/ippo_shielded_histories.pkl stores ablation shielded training histories
  • experiments/*.png stores cumulative, post-WM cumulative, episodic, post-WM episodic, and shield telemetry plots

The learned shield setup deliberately points at the ablation condition root, so it consumes that condition's wm/env_transition_graph.pkl and om/iop_stack.pt.

Comparison Sources

The ablation experiment notebooks train only the selected learned IPPO_Shielded condition. They load canonical histories for comparison only:

  • pursuit baseline: exports/gridworlds/pursuit/2
  • markov_stag_hunt baseline: exports/gridworlds/markov_stag_hunt

If the other ablation condition has already been run, its learned shielded history is included in plots automatically.

Expected Run Order

For each environment and condition:

uv run python notebooks/ablations/<env_name>/train_wm_mo.py --ablation=<condition>
uv run python notebooks/ablations/<env_name>/train_om_mo.py --ablation=<condition>
uv run python notebooks/ablations/<env_name>/experiments_mo.py --ablation=<condition>

For a quick script smoke run, reduce the experiment budget:

uv run python notebooks/ablations/pursuit/experiments_mo.py \
  --ablation=weak_wm_strong_om \
  --num-runs=1 \
  --timesteps=64 \
  --rollouts=64