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 ¶

OMBuffer(
    capacity: int,
    state_dim: int,
    device: device | None = None,
)

add ¶

add(
    state: Union[Tensor, ndarray, list],
    action: Union[int, Tensor, ndarray, list],
) -> None

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 ¶

OpponentMLP(
    in_dim: int, n_actions: int = 5, hidden: int = 128
)

Bases: Module

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

forward ¶

forward(x: Tensor) -> torch.Tensor

Returns logits

action_probs ¶

action_probs(state_flat_oh) -> torch.Tensor

Returns probabilities

src.opponent_model.preparation.OMPreparedInputs `dataclass` ¶

OMPreparedInputs(
    graph: Graph,
    dyn: GraphWorldModelAdapter,
    om_buffer: OMBuffer,
    n_timesteps: int,
    dataset_size: int,
)

src.opponent_model.preparation.prepare_om_training_inputs ¶

prepare_om_training_inputs(
    export_dir: str | Path,
    *,
    action_size: int,
    n_agents: int,
    state_dim: int,
    device: device,
    focal_idx: int = 0,
    buffer_focal_idx: int = 0,
    strict_legal: bool = True,
    seed: int = 0,
) -> OMPreparedInputs

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 ¶

encode_joint_action(
    actions: Sequence[int], action_size: int
) -> int

Base-A encoding of a list of discrete actions.

src.opponent_model.utils.extract_level0_dataset ¶

extract_level0_dataset(
    buffer: Iterable[Transition],
    *,
    action_size: int,
    n_opps: int,
    skip_inactive: bool = True,
    focal_idx: int = 0,
    buffer_focal_idx: int = 0,
)

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

BufferWrapper stores one focal action followed by the remaining agents' actions in environment order. focal_idx may select a different focal role when the stored observation is a shared/global state. This lets two role-conditioned opponent models reuse one pretraining buffer without pretending that a model of player 1's actions is also a model of player 0.

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` ¶

MBOMConfig(
    n_actions_opp: int,
    architecture: str = "auto",
    action_size_opp: int | None = None,
    n_opponents: int = 1,
    dense_chunk_size: int = 65536,
    n_levels: int = 3,
    hidden: int = 128,
    lr_om: float = 0.0003,
    batch_size: int = 1024,
    dense_confusion_max_actions: int = 512,
    bayes_beta: float = 0.05,
    temp_tau: float = 1.25,
    device: Optional[Device] = None,
)

src.opponent_model.imagination.iop.Level0SplitMetrics `dataclass` ¶

Level0SplitMetrics(
    size: int,
    cross_entropy: float,
    accuracy: float,
    error_rate: float,
    multiclass_brier: float = float("nan"),
    expected_calibration_error: float = float("nan"),
    mean_confidence: float = float("nan"),
    action_counts: tuple[int, ...] = (),
    confusion_matrix: tuple[tuple[int, ...], ...] = (),
    confusion_entries: tuple[
        tuple[int, int, int], ...
    ] = (),
    confusion_shape: tuple[int, int] = (0, 0),
)

src.opponent_model.imagination.iop.Level0FitMetrics `dataclass` ¶

Level0FitMetrics(
    holdout_fraction: float,
    train: Level0SplitMetrics,
    holdout: Level0SplitMetrics,
)

src.opponent_model.imagination.iop.ImaginedOpponent ¶

ImaginedOpponent(
    state_dim: int,
    cfg: MBOMConfig,
    device: Optional[Device] = None,
)

update_mixture_with_observation ¶

update_mixture_with_observation(
    s: Tensor, a_obs: Tensor
) -> MixtureUpdateMetrics

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 ¶

print_level0_fit_metrics(
    metrics: Level0FitMetrics,
    *,
    title: str = "Level-0 Fit Metrics",
    console=None,
    export_dir: str | Path | None = None,
)

Soft-Rollout Policy Improvement¶

src.opponent_model.imagination.soft_rollout_pi.SoftRolloutPIConfig `dataclass` ¶

SoftRolloutPIConfig(
    root_state_samples: int = 100000,
    batch: int = 512,
    horizon: int = 5,
    n_rollouts: int = 8,
    gamma: float = 0.99,
    tau: float = 1.0,
    lr: float = 0.0003,
    clip_grad: float = 1.0,
    entropy_coef: float = 0.0,
    eval_topk: int = 0,
    eval_rand: int = 0,
    print_every: int = 2000,
)

src.opponent_model.imagination.soft_rollout_pi.make_masked_sampler_from_logits ¶

make_masked_sampler_from_logits(
    *,
    logits_fn: Callable[[Tensor], Tensor],
    dyn: GraphWorldModelAdapter,
    agent_idx_env: int,
) -> ActionSampler

src.opponent_model.imagination.soft_rollout_pi.soft_rollout_policy_improvement ¶

soft_rollout_policy_improvement(
    *,
    dyn: GraphWorldModelAdapter,
    policy: Module,
    states_source: OMBuffer,
    device: device,
    controlled_indices: Sequence[int],
    reward_env_indices: Optional[Sequence[int]] = None,
    fixed_samplers: Dict[int, ActionSampler],
    fixed_joint_controllers: Sequence[
        FixedJointController
    ] = (),
    cfg: SoftRolloutPIConfig = SoftRolloutPIConfig(),
    log_context: SoftRolloutPILogContext | None = None,
) -> list[dict[str, Any]]

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. - 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` ¶

StaircasePIConfig(
    levels: int = 3,
    agent_cfg: SoftRolloutPIConfig = SoftRolloutPIConfig(),
    opp_cfg: SoftRolloutPIConfig = SoftRolloutPIConfig(),
    agent_vs: str = "level",
    use_active_level: bool = False,
)

src.opponent_model.imagination.imagination_loop.staircase_train_pi ¶

staircase_train_pi(
    *,
    iop: ImaginedOpponent,
    dyn: GraphWorldModelAdapter,
    agent_policy,
    states_source: OMBuffer,
    device: device,
    cfg: StaircasePIConfig = StaircasePIConfig(),
    export_dir: str | Path | None = None,
    diagnostic_state_sample: int = 8192,
) -> list[dict[str, Any]]

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 ¶

GraphWorldModelAdapter(
    graph: Graph,
    *,
    action_size: int,
    n_agents: int,
    focal_idx: int = 0,
    strict_legal: bool = True,
    seed: Optional[int] = None,
    legal_mask_cache_size: int | None = 100000,
)

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 ¶

load_iop_stack(
    path: str | Path, *, device: Optional[device] = None
) -> ImaginedOpponent

Runtime Mixture¶

src.opponent_model.mixing.temp_softmax ¶

temp_softmax(
    x: Tensor, tau: float = 1.0, dim: int = -1
) -> torch.Tensor

High-temp softmax; tau>1 makes weights smoother, tau<1 sharper.

src.opponent_model.mixing.bayes_update ¶

bayes_update(
    prior_alpha: Tensor,
    per_level_probs: Tensor,
    eps: float = 1e-12,
) -> torch.Tensor

Posterior p(m|a^o) ∝ π_m(a^o|s) * p(m).

src.opponent_model.mixing.ema_update ¶

ema_update(
    old: Tensor, new: Tensor, beta: float
) -> torch.Tensor

(1-beta)old + betanew, kept normalized.

src.opponent_model.mixing.mix_action_probs ¶

mix_action_probs(
    per_level_action_probs: Tensor, alpha: Tensor
) -> torch.Tensor

Return ∑_m α_m π_m(a|s) → [B, A].