Skip to content

Opponent Model API

Opponent-model APIs cover supervised level-0 training, imagined higher-level policies, graph-world adaptation, checkpoint loading, and action encoding.

Buffers, Networks, and Preparation

src.opponent_model.buffer.OMBuffer

add(state, action)

Add either a single (state, action) or a batch. - state: shape (state_dim,) or (B, state_dim) - action: shape () / (,) or (B,) / (B,1)

src.opponent_model.om_mlp.OpponentMLP

Bases: Module

Opponent model: p(a_opp | s_flat). - input: state_flat_mh ∈ R^{cells * n_types} - output: logits over 5 actions

forward(x)

Returns logits

action_probs(state_flat_oh)

Returns probabilities

src.opponent_model.preparation.OMPreparedInputs dataclass

src.opponent_model.preparation.prepare_om_training_inputs(export_dir, *, action_size, n_agents, state_dim, device, focal_idx=0, strict_legal=True, seed=0)

Load the real-experience OM dataset and the world-model transition graph, then verify that every sampled OM state can be mapped back into the graph used for imagination rollouts.

src.opponent_model.utils.encode_joint_action(actions, action_size)

Base-A encoding of a list of discrete actions.

src.opponent_model.utils.extract_level0_dataset(buffer, *, action_size, n_opps, skip_inactive=True)

Extract states and joint-opponent action labels for level-0 training.

Returns:

Type Description

A pair (S_list, A_list) where S_list contains states and

A_list contains joint-opponent action indices in

[0, action_size**n_opps).

Imagined Opponent Stack

src.opponent_model.imagination.iop.MBOMConfig dataclass

src.opponent_model.imagination.iop.Level0SplitMetrics dataclass

src.opponent_model.imagination.iop.Level0FitMetrics dataclass

src.opponent_model.imagination.iop.ImaginedOpponent

update_mixture_with_observation(s, a_obs)

Bayes filter over levels using observed (possibly JOINT) opponent action index. s: [B, D] a_obs: [B] long, in [0, A_joint)

src.opponent_model.imagination.iop.print_level0_fit_metrics(metrics, *, title='Level-0 Fit Metrics', console=None)

Soft-Rollout Policy Improvement

src.opponent_model.imagination.soft_rollout_pi.SoftRolloutPIConfig dataclass

steps property writable

Backward-compatible alias for root_state_samples.

src.opponent_model.imagination.soft_rollout_pi.make_masked_sampler_from_logits(*, logits_fn, dyn, agent_idx_env)

src.opponent_model.imagination.soft_rollout_pi.soft_rollout_policy_improvement(*, dyn, policy, states_source, device, controlled_indices, reward_env_indices=None, fixed_samplers, fixed_joint_controllers=(), cfg=SoftRolloutPIConfig(), log_context=None)

Soft rollout PI for a controller over controlled_indices. The action space is the JOINT action space of those indices: A_joint = action_size**len(controlled_indices).

Budget/accounting: - cfg.root_state_samples counts sampled root states drawn from states_source. - Legacy alias: cfg.steps. - It does NOT count real environment steps. - It also does NOT directly count imagined transitions inside the world-model graph. - Actual imagined-transition compute is much larger because each root state may evaluate multiple candidate actions, and each candidate runs cfg.n_rollouts rollouts of length cfg.horizon.

per-agent samplers for all remaining env indices NOT covered by:
  • controlled_indices (this policy)
  • fixed_joint_controllers (coordinated fixed policies)

src.opponent_model.imagination.imagination_loop.StaircasePIConfig dataclass

src.opponent_model.imagination.imagination_loop.staircase_train_pi(*, iop, dyn, agent_policy, states_source, device, cfg=StaircasePIConfig())

Staircase PI (IOP stack), multi-agent version where: - focal agent is dyn.focal_idx - all other agents are treated as a single JOINT "team" opponent.

For m = 0..levels-2: (1) Improve agent policy as BR to TEAM at level m (or mixture) (2) Improve TEAM opponent model at level (m+1) as BR to the updated agent

Graph Adapter and Checkpoints

src.opponent_model.GraphWorldModelAdapter

Imagination dynamics from an exact legal joint-action graph whose edges carry

env_prob : P_env(s' | s, a_joint), normalized within each legal successor group rewards : E[r | s, a_joint] in environment agent order

This opponent-model path assumes graph states are observation-Markov: the stored node bits must uniquely identify a graph state. If two graph nodes share the same stored observation bits, observation-only opponent-model buffers cannot resolve which graph state they came from.

src.opponent_model.checkpoints.load_iop_stack(path, *, device=None)