Skip to content

Opponent-Model Runtime Fast Path

The April 2026 OM runtime pass kept staircase PI semantics unchanged and focused on the shared runtime path used by both Matrix and Gridworld notebooks.

What changed

  • OM input preparation now goes through prepare_om_training_inputs(...).
  • In marimo notebooks, prepare_om_training_inputs must be imported and returned by the setup cell before any later cell can list it as an input. Matrix and Gridworld OM notebooks use this explicit reactive dependency because script execution otherwise fails with NameError before OM prep reaches dataset or graph validation.
  • That helper validates the WM → OM handoff once, before training:
  • wm/env_transition_graph.pkl must exist,
  • graph state width must match the OM dataset width,
  • every OM training state must map back into the graph,
  • observation-Markov violations still fail fast through GraphWorldModelAdapter.
  • Successful prep attaches resolved graph indices to OMBuffer, so training can sample graph-aligned states directly.

Main runtime win

The dominant hotspot was repeated legality-mask construction inside soft-rollout PI, not transition sampling.

  • GraphWorldModelAdapter now serves canonical legality masks from a cache keyed by state, controlled controller subset, and partial-joint assignment.
  • Batch legality queries collapse duplicate (state, partial_joint) requests before doing any exact filtering work.
  • soft_rollout_policy_improvement(...) uses those cached masks and prefers preindexed OM buffers when available.

What did not change

  • staircase level count,
  • rollout horizon,
  • n_rollouts,
  • reward aggregation,
  • soft target construction.

This means notebook hyperparameters remain directly comparable to earlier runs; the change is meant to reduce overhead in the shared rollout implementation rather than to weaken the imagination objective.