Skip to content

Opponent-Model Training Semantics

This note makes the opponent-modelling training loop explicit, especially the meaning of the soft-rollout sample budget.

Does opponent modelling use a neural network?

Yes.

  • The learned opponent policy is OpponentMLP, a PyTorch nn.Module with two hidden ReLU layers and a final linear logits head.
  • ImaginedOpponent builds one OpponentMLP per reasoning level in self.layers.
  • Level 0 is trained by supervised maximum likelihood on real (state, opponent action) data from the OMBuffer.
  • Higher levels are also neural nets: each level is warm-started from the previous one, then improved with soft rollout policy improvement inside the world-model graph.

What data is level 0 trained on?

extract_level0_dataset(...) converts real transitions into:

  • state: the focal-agent state representation.
  • action: the opponent team action, encoded as a single joint-action class.

Those samples are stored in OMBuffer, and ImaginedOpponent.fit_level0(...) trains level 0 with cross-entropy on that dataset.

What does SoftRolloutPIConfig.root_state_samples mean?

root_state_samples is a training-sample budget for soft rollout policy improvement.

It counts sampled root states drawn from states_source, usually from OMBuffer. It does not count:

  • real environment steps,
  • world-model rollout transitions,
  • optimizer iterations.

Concretely, the loop does:

  1. Draw b = min(batch, remaining_steps) root states from states_source.
  2. For those states, estimate action values by imagined rollouts in the graph world model.
  3. Do one optimizer update.
  4. Increase the counter by b.

So with batch=512 and root_state_samples=100_000, the intended budget is about 100k sampled training states, which is about ceil(100000 / 512) = 196 optimizer updates.

What is the hidden compute behind one counted step?

One counted root-state sample can fan out into many imagined transitions.

For each sampled state, the code may:

  • evaluate multiple candidate actions for the controlled policy,
  • for each candidate, run n_rollouts simulated rollouts,
  • each rollout advances the graph world model for horizon imagined transitions.

So the dominant imagination compute is roughly:

sampled_root_states x evaluated_actions_per_state x n_rollouts x horizon

This is why root_state_samples should be read as a root-state training budget, not as environment time.

One subtlety

If sampled states are not present in the legal graph, they can be filtered before training on that iteration. The root-state sample counter still advances by the requested batch size b, so the effective number of valid states used for optimization can be slightly smaller than cfg.root_state_samples.

Current notebooks now validate this earlier during OM input preparation. prepare_om_training_inputs(...) checks that the saved wm/env_transition_graph.pkl exists, that graph-state width matches the OM dataset width, and that every extracted OM state can be mapped back into the graph. When that validation succeeds, the helper stores the resolved graph indices directly on OMBuffer, so soft-rollout PI can sample preindexed states instead of re-hashing observation bits every batch.

OMBuffer mutation invariants

OMBuffer is a fixed-capacity ring. A batch larger than its capacity replaces the ring with the newest capacity samples from that batch. Any non-empty write clears graph_state_idx, because those cached indices describe the old sample contents; callers must attach graph indices only after the real dataset has finished loading. Capacity and state width must both be positive.

What does the soft-rollout PI progress bar mean?

Inside a single soft_rollout_policy_improvement(...) call, progress is shown with one phase-local tqdm bar rather than repeated printed lines.

A typical description looks like:

soft-rollout-PI 1/4 | agent opp=L0

and the bar postfix carries metrics such as:

local=300/3000 global=76800/3840000 (2.0%) roots=76800/768000 valid=256/256 policy_ce=0.6862 entropy=0.6930 objective=0.6827

  • soft-rollout-PI 1/4: phase 1 out of 4 in the full staircase schedule.
  • agent opp=L0: the focal agent is being improved against opponent level 0.
  • team level=L1: for team phases, the opponent/team policy being trained becomes level 1.
  • The bar total is the local phase budget root_state_samples, so the visible progress is for the current PI call only.
  • local=i/n: successful optimizer-update count inside the current PI phase. n = ceil(root_state_samples / batch).
  • global=x/y: cumulative root-state budget consumed across the whole staircase run, not environment steps.
  • roots=x/y: root-state budget consumed inside just the current PI phase.
  • valid=v/b: how many sampled states survived graph-index filtering in that update, out of the requested batch size b.
  • policy_ce: cross-entropy from the rollout-induced soft target policy to the current policy logits.
  • entropy: entropy of the current masked policy.
  • objective: the optimized scalar, equal to policy_ce - entropy_coef * entropy.

Visible refreshes are intentionally sparse for notebook output. The display cadence is:

max(print_every, ceil(root_state_samples / 10))

So the bar updates at roughly 10% milestones unless the caller asks for an even sparser interval. print_every=0 disables the progress bar entirely.

Why not redefine it as env steps?

That would be misleading here.

  • The loop does not interact with the real environment.
  • One sampled root state can branch into many imagined graph transitions.
  • The amount of imagined rollout work per sampled state depends on legality filtering, how many candidate actions are evaluated, n_rollouts, and horizon.

So a rename is the clean fix. Reusing the old name for a different unit would make old experiments harder to compare and easier to misunderstand.