Skip to content

Opponent Model on joint_graph

src/opponent_model now assumes the world model comes from the exact joint_graph pipeline rather than the retired true_dynamics path.

Core assumptions

  • The opponent-model imagination stack consumes legal full joint actions from env_transition_graph.pkl.
  • Graph state lookup is observation-driven, so this path requires the graph to be observation-Markov: the stored node bits must uniquely identify a graph node.
  • If multiple graph nodes share the same observation bits, GraphWorldModelAdapter rejects the graph in strict_legal=True mode instead of guessing.

This used to exclude pursuit-style graphs when they relied on latent suffix bits not present in the public state. After the public Markov-state update, supported envs are expected to expose every future-relevant variable in state() / shared observation so their graph nodes remain uniquely recoverable from the stored bits.

Exact legality in imagination rollouts

The old rollout code approximated subset legality by taking the cartesian product of per-agent legal actions. That is too loose for an exact joint-action graph because a per-agent legal action can still be incompatible with the other agents' choices.

The current rollout path uses exact partial-joint legality:

  • GraphWorldModelAdapter precomputes the legal full joint actions available from each state.
  • Rollout code carries a partial joint action vector filled with UNASSIGNED_ACTION = -1.
  • When a fixed controller, fixed sampler, or controlled policy needs a legality mask, the adapter filters the legal full joint actions by the assignments already chosen and asks whether each candidate action still has at least one legal completion.

This keeps the rollout factorisation intact while ensuring every sampled partial choice can still be completed to a legal full joint action in the graph.

Cached legality fast path

The exact partial-joint legality rule is still the source of truth, but the rollout code no longer re-filters full legal joint-action lists for every row in every mask query.

  • GraphWorldModelAdapter now caches canonical legality masks keyed by: state_idx, the controlled env-index subset, and the full partial-joint assignment pattern.
  • Batch legality queries first group identical (state, partial_joint) requests, compute the exact mask once, and then broadcast it back to every matching row.
  • The cache stores masks in canonical base-A joint-action order, so rollout code can reuse them directly for the standard _joint_table(...) action tables used by soft-rollout PI.

This keeps semantics exact while removing the main Python/NumPy hotspot that had been dominating Matrix opponent-model wall time.

WM → OM handoff requirements

Opponent-modelling notebooks now rely on a shared prepare_om_training_inputs(...) helper before loading policies or starting rollout PI.

That helper requires:

  • wm/D.pkl for the real level-0 dataset,
  • wm/env_transition_graph.pkl for imagination rollouts,
  • matching public-state width between the OM dataset and the graph,
  • observation-Markov graph nodes that can be resolved uniquely from stored bits.

Failures in that handoff now stop the notebook immediately with a targeted error instead of letting the rollout loop discover the mismatch later.

Compatibility limits

  • Supported in this scope: environments where graph node bits are observation-complete, such as the matrix games and the observation-Markov gridworld graphs.
  • Unsupported in this scope: latent-key graphs that require extra hidden state to map an observation back to a unique graph node.
  • No heuristic disambiguation is attempted. If the graph is ambiguous, fail fast and either change the graph representation or pass richer state identifiers through the opponent-model data path.