Adaptive opponent modeling for
adversarial co-training in MARL.

Two-team independent Q-learning.

Building blocks for adaptive opponent modeling in adversarial co-training: learn an uncertainty-aware model of the opponent, then exploit it inside a planner.

Environment

JaxMARL's MPE_simple_tag_v3. 3 predators vs 1 prey, 2D arena, 5 discrete actions, 25-step episodes. +10 / −10 per capture. Fully observable.

Two-team IQL

One TrainState per team, shared weights within team. Avoids gradient crosstalk from stock JaxMARL's single-net vmap.

The three variants

Architecture — the two-headed MLP

obs  ──▶  Dense(128) ──▶ ReLU ──▶ Dense(128) ──▶ ReLU ──┬──▶ Dense(action_dim=5)             ──▶  Q-values
                                                              │
                                                              └──▶ Dense(opp_n × action_dim)     ──▶  opp-action logits
                                                                    │
                                                                    └─ CE vs opp actions from replay  (the auxiliary loss, weight 0.5)

Shared trunk encodes opponent-predictive features. At training time the opp head is an auxiliary regularizer; at eval time (OA-plan) it becomes the opponent model the planner samples from.

Primary result — 3×3 cross-play tournament

Each cell: paired by seed, 100 eval episodes, averaged over 3 seeds.

Captures per episode — read the numbers directly

Column-wise: OA-plan prey is uniformly lowest (1.57, 1.15, 1.07) — ~2-capture drop. Row-wise: OA-plan pred is worst across the board.

Mechanism — opponent-head accuracy asymmetry

The prey predicts 3 predator actions per step; the predator predicts 1 prey action. 3× the CE gradient → better opponent model → better planner. The +7 pp accuracy gap drives the tournament asymmetry.

Summary

Setup

SimpleTagStaticMPE: obstacles pinned at (±0.5, ±0.5), plus a per-step directional bonus for the prey:

bonus = coef · dir_sign · sign(v_prev × v_new) · 1[r_in ≤ ‖v_new‖ ≤ r_out]

v_prev = prey_pos_t     − obstacle
v_new  = prey_pos_{t+1} − obstacle
dir_sign = +1 (CCW) or −1 (CW)
coef = 0.1,  r_in = 0.25,  r_out = 0.6

Direction-only (not speed). Annulus gate prevents farming on top of obstacles. Prey-only bonus; predator rewards untouched.

Four conditions

Behavior mining

64 episodes per (variant × seed). Track prey 2D position and signed angular displacement around the nearest obstacle.

The numbers behind the histograms

Rollouts

baseline Prey darts near cover, uses obstacles as shields, no stereotyped motion.

ccw — counter-clockwise Prey orbits counter-clockwise around the obstacles. Predators pre-position on the forward orbit arc.

cw — clockwise Mirror image of CCW. Predators exploit the same structure, opposite sense.

Each GIF is 150 steps = 6 auto-reset episodes chained together, greedy rollouts from seed 0 of the corresponding static variant.

Per-trajectory clustering

Can a single trajectory identify which policy produced it? Short answer: no, at shape_coef = 0.1. Two datasets tested: Reading A (900 labelled trajectories, 9 checkpoints) and Reading B (900 unlabelled, single policy).

1. Hand-crafted + occupancy features

2. MLP-VAE on trajectories

Flax MLP-VAE, 8-dim latent, KL-annealed, 8k Adam steps on Reading A.

3. Predator VAEs

Same VAE on each predator slot. Latent traversals are more interpretable (goal-directed chase trajectories), but still track arena geometry, not policy.

Predator clustering numbers, side-by-side with the prey

Reading A: predator ARI ~2× prey's, still at-chance. Reading B: predators are multi-modal (BIC drops 700+ nats), prey are not. Modes track which obstacle the chase converges on.

Implications & fixes

The VAE latent encodes trajectory shape, not policy identity. Two paths forward:

• Supervise the latent (SupCon or CE classification on checkpoint labels)
• Strengthen the data — Exp 4's resource environment makes strategy differences large enough that an unsupervised VAE should separate them

Design

4 collectable resources, prey-visible only. +5.0 reward per collection. Predator obs unchanged.

Placement modes

Circle

Radius 0.6 around origin. Optimal collection = loop.

Corners

At (±0.8, ±0.8). Requires long diagonal dashes.

Random

50/50 coin-flip per reset. Bimodal trajectory distribution by construction.

Trained policy rollouts

Same initial state (random placement, seed 0). Red = predators, blue = prey, green diamonds = resources.

Individual algorithm rollouts (click to enlarge):

Arena layout references

Information asymmetry

Predators: 16-d (no resource info). Prey: 26-d (base + 12 resource dims). The prey's strategy is shaped by resources the predator can't observe — an opponent-aware predator that infers the strategy type should outperform one that doesn't.

Implementation

ResourceState dataclass adds resource_pos and collected. CTRolloutManager pads to 26-d + 4-d one-hot = 30-d training obs. Predators zero-padded in resource dims.

Algorithm variants

IQL

Independent Q-Learning (MLP). Off-policy, replay buffer.

OA-IQL

+ opponent-action auxiliary head (CE, coef 0.5). MC planning at eval (K=5, H=3).

MAPPO

Centralized critic + per-agent actor. On-policy PPO, 32 envs, 128-step rollouts.

Training results — random placement

2M timesteps, 3 seeds, random placement. X-axis = env timesteps (comparable across algorithms).

MAPPO ~3× IQL on pred returns. Centralized critic coordinates pursuit. Resource collection comparable across all three (~0.4–0.5/step) — return gap is tag/evasion, not resources.

Trajectory datasets

300 MAPPO trajectories (3 seeds × 100 eps, T=50). Includes positions, resources, and collection flags.

Does the prey trajectory carry its strategy? — measuring the precondition

The proposal's Part 1 wants an encoder mapping the opponent's trajectory to a latent strategy z. That is only possible if the strategy is in the trajectory. The hypothesis was that resource placement (circle-loop vs corner-dash) would make it so. We tested this directly: label each first-episode prey trajectory by its placement (recoverable from resource geometry, never shown to the encoder), then ask whether a supervised linear probe on the raw trajectory can predict placement, and whether the unsupervised VAE latent recovers it.

Is it an incentive problem? — 4× the collect reward

The natural fix is to make the prey care more about resources. We retrained IQL with collect_reward raised from 5 to 20. The prey collected nearly twice as many resources per episode (0.68 → 1.28), confirming the incentive landed — yet placement stayed at chance (probe 0.49). Collecting more resources does not mean executing a layout-distinctive route.

Deliverable 1 — different placements produce different strategies

The two specialists trace visibly different routes. The circle-prey loops a diamond between the four ring resources; the corners-prey patrols a square around the perimeter. Pooling occupancy over a trajectory, a linear probe tells the two apart at 0.98 (chance 0.50) for IQL and 0.96 for MAPPO.

Deliverable 2 — a VAE latent that encodes the strategy

A single short trajectory is dominated by evasion noise, so the encoder needs both the right representation (the trajectory's 2D occupancy histogram, which exposes where the prey goes) and enough of it. Sweeping the observation length, the unsupervised VAE latent goes from chance to a clean two-cluster split: at L=100 steps the latent recovers placement at ARI 0.87 with a supervised latent-probe of 0.99.

This is the result the Exp 2 / Exp 4 VAEs could not reach. The difference is regime (the prey now expresses strategy) and representation (occupancy, not the raw evasion-dominated position sequence). MAPPO's prey is more dominated, so its latent encodes placement supervised (probe 0.83) but does not separate cleanly unsupervised — cleaner strategy expression gives a cleaner latent.

Deliverable 3 — predators are brittle to an unseen opponent

Take MAPPO predators co-trained against the circle-prey and drop in the corners-prey they never saw (and vice-versa). Captures fall 43% — from 1.59 to 0.91 per episode. Read column-wise (same prey and env, only the predator swapped): the co-trained predator catches the circle-prey 1.90×, the unfamiliar predator only 1.11; for the corners-prey, 1.29 vs 0.72. The centralized critic learns a layout-specific interception that does not transfer.

Part 1 (built) — a latent-conditioned opponent model

The advisor's next step: once the VAE is good, clone a policy conditioned on its latent. We log (obs, action) from both specialists, encode each trajectory to z with the validated occupancy-VAE, and behaviour-clone a single prey policy π(a | obs, z). Conditioning on zcircle vs zcorners makes the same policy trace the diamond or the square. Crucially, swapping z alone — with the environment and predators held fixed — steers the route at 0.85 (chance 0.50): the latent controls behaviour, not just the layout.

Three findings, one ladder

(1) Knowing the intent is worth +51%. An oracle predator that observes the hidden corner catches the prey 4.05 times/episode versus 2.68 for the strong intent-blind hedger. The opponent's hidden strategy genuinely matters.

(2) A naive model is brittle. That same oracle predator, fed a random intent, collapses to 1.42 — worse than not modeling at all. A model that ignores its own uncertainty is a liability.

(3) Uncertainty-aware inference wins. The belief predator, fed the encoder's online-sharpening posterior, reaches 2.82 — it beats the intent-blind baseline, by hedging while the belief is flat and committing once it sharpens. Reacting to the belief recovers only a sliver of the oracle's edge (2.82 vs 2.68 against an oracle of 4.05); the win comes from planning on it, below. Hard inference at a fixed step (2.56) doesn't help, because it gambles on one possibly-wrong guess.

Part 2 (built) — planning with the belief

The belief predator above is reactive. Now the predators plan: at each step, for every candidate joint action, they roll the true simulator a few steps ahead with the prey's moves sampled from its policy under an intent drawn from the belief, and pick the action that best intercepts the imagined future prey. Belief-conditioned opponent sampling inside a lookahead — the proposal's Part 2.

Sharpening Part 1 — predict, don't reconstruct (JEPA vs VAE)

The encoder above used labels. A Part 1 that scales should recover the opponent's strategy self-supervised, and the objective decides whether it can. Same prey trajectories, same 2-D latent, two self-supervised encoders: a VAE that reconstructs the trajectory (generative), and a JEPA that predicts the representation of the future window through an EMA target encoder (predictive, no reconstruction — LeCun's principle, cf. V-JEPA 2). JEPA keeps only the predictable structure (where the prey is heading = its strategy) and discards the evasion noise the VAE drowns in.

Because it predicts where the prey is heading instead of waiting for it to arrive, JEPA is also an anytime encoder — it reads the strategy from ~2× fewer observed steps than the VAE (JEPA at 11 steps matches the VAE at 20), so the planner can commit sooner.

And it closes the loop without labels: a self-supervised JEPA belief (encoder + a readout to predicted arrival position, then a softmax over the four known corners — no intent labels anywhere) drives the planner to 4.08 captures/episode, matching the supervised-belief planner (4.31, overlapping error bars) and beating the oracle reactive predator (4.05).

Here is that label-free belief running in the game. Left: intent-blind predators hedge. Right: JEPA-belief predators — the × marks where the JEPA model predicts the prey is heading, and it locks onto the prey's hidden corner (the star) as the pack closes in; the bar strip is the JEPA belief sharpening, with no labels.

Where JEPA stops helping — the world model (an honest negative)

The obvious next step is to use the same predictive principle for the planner's dynamics — a JEPA latent world model (encode → predict next latent → decode) the planner rolls through, fixing the state-space model that fell down in the earlier learned-model planner. It doesn't work: the latent compression makes the dynamics ~3× less accurate (5-step position RMSE 0.24 vs 0.07), and the planner collapses to 0.57 captures — worse than the state-space model (2.53) and the reactive baseline.

The second strategy axis — circle vs corners (new)

The same Part-1 question on the resource-placement axis: can an unsupervised encoder tell which specialist it is watching — the prey that orbits a resource circle, or the one that dashes between corners? Three findings (1,200 episodes, 3 encoder seeds, mopa package):

• A raw trajectory window carries no signal — for anyone. A supervised probe on absolute coordinates of one episode sits at chance (0.54). One episode is dominated by evasion; this is a property of the data, and it retroactively explains the earlier exp4/exp5 failures.
• Occupancy carries it. The same probe on a histogram of where the prey spends time climbs 0.71→0.88→0.93 over 25→50→100 observed steps.
• Predict-vs-reconstruct has a boundary condition. With 2-D latents JEPA again beats the VAE (0.63 vs 0.48 at 50 steps); with 4-D the VAE wins (0.81 vs 0.67) — occupancy histograms have already discarded the evasion noise, which is where JEPA's edge comes from. Neither objective forms unsupervised clusters (ARI ≈ 0, same as the earlier occupancy-VAE): the placement is a linear direction in the latent, not separated modes — the honest open problem on this axis.

And the behaviour-cloning test the meeting asked for — π(a|s) vs π(a|s,z), watch episode 1, predict the predator's episode-2 actions, episode-level split: conditioning on the true placement clearly helps (0.800 vs 0.776 held-out action accuracy), so the hypothesized effect is real; a single-episode unsupervised latent doesn't move it (0.772) — exactly as the probe curve predicts, since one episode of observation gives it almost no signal. The oracle–baseline gap is the prize a better opponent latent captures — and the follow-up below shows an unsupervised latent does close it once it sees enough.

What that metric looks like in the game: a held-out episode replays, and at every step each BC variant predicts each predator's actual action — green ring = predicted it, red = missed. On episodes where the prey's route matters, knowing the placement is the difference between guessing and following (0.60 → 0.83 circle, 0.60 → 0.92 corners).

Follow-up — the unsupervised latent does close the gap

Two checks turn the single-episode null into a positive scaling result. First, is behaviour cloning even a faithful clone? Deployed against the same prey, vanilla BC — trained from a hand-built state feature (all agent positions + a velocity proxy + predator id, strictly less than the policy network sees) — recovers 86% of the MAPPO expert's capture-edge over random (1.22 vs 1.35 captures/episode, random 0.40; 3 checkpoint seeds, two placements), at a held-out action match of 0.80 (episode-level split). The policy class is not the bottleneck; the residual 14% is the headroom a strategy signal could close.

Second, does conditioning help once the latent carries the strategy? The placement is fixed across all of a specialist's episodes, so we can give the encoder more observation. Reading the VAE latent from L steps before a held-out episode (steps 125–150 of a six-episode rollout), its placement probe climbs from 0.60 at L=25 to 0.78 at L=125, and latent-conditioned BC tracks it: from no gain at L=25 (0.764 vs vanilla 0.766) it overtakes vanilla by L=50 and reaches the oracle (true-placement) ceiling by L=125 (0.795 vs oracle 0.792). The JEPA latent, whose probe stays flat (≈0.6) at this 4-D latent size, gives no gain at any L. So π(a|s,z) > π(a|s) holds — exactly to the degree the latent recovers the strategy. Encoder quality translates directly into control benefit.

OA-IQL nudges both teams up ~2 return points. Opp-head accuracy: 0.53 (pred) vs 0.58 (prey) — same asymmetry as Exp 3.

Rollouts (single seed)

IQL baseline60-step greedy rollout. Predators close in; prey darts between corners.

OA-IQL60-step greedy rollout. Similar overall play; the quantitative effect shows up in the return numbers above.

Three implementation notes worth recording.

1. Two-team split

Per-team TrainState, shared within-team. Avoids gradient cancellation from stock JaxMARL's single-net vmap.

2. Use wrapped_step in the planner

Must preserve obs padding + one-hot. Using env.step_env silently returns NaN.

3. Pass LogEnvState directly

Don't unwrap to env_state.env_state — the log wrapper unwraps internally.

Hyperparameters

IQL / OA-IQL (Exp 1–4, off-policy):

Total env-steps

2 000 000 per training run

Parallel envs / train

8

Episode length

25 steps (Exp 2) / 32 at eval (Exp 3 tournament)

Discount γ

0.9

ε schedule

1.0 → 0.05 linear over 10 % of training

Learning rate

0.005, linear decay

Target update

every 200 steps, hard copy (τ = 1)

Buffer size / batch

5000 / 32

Hidden size

64 (RNN, Exp 1 & 2) / 128 (MLP, Exp 3 & 4)

Opp-aux coef

0.5

Planner K / H / γ

5 / 3 / 0.9

Seeds

3 (Exp 2–4), 1 (Exp 1)

MAPPO (Exp 4, on-policy):

Total env-steps

2 000 000

Parallel envs

32

Rollout length

128 steps

Hidden size

128

Learning rate

0.0003, linear anneal

Discount γ / GAE λ

0.99 / 0.95

PPO clip ε

0.2

Entropy coef

0.01

Value-fn coef

0.5

Update epochs / minibatches

4 / 4

Max grad norm

0.5

Seeds

3

Reproducing the results.

# Environment
python3.11 -m venv .venv && source .venv/bin/activate
pip install -e JaxMARL/ "flax==0.10.2" hydra-core flashbax wandb matplotlib distrax optax

# Exp 4 — resource collection (3 algos × 3 placements + tournament + traj dataset)
python src/smoke_test_resources.py                                             # 9 assertions

# training — IQL
python src/iql_teams_mlp.py alg=ql_teams_resources_circle  NUM_SEEDS=3
python src/iql_teams_mlp.py alg=ql_teams_resources_corners NUM_SEEDS=3
python src/iql_teams_mlp.py alg=ql_teams_resources_random  NUM_SEEDS=3

# training — OA-IQL
python src/iql_teams_oa_mlp.py alg=ql_teams_oa_resources_circle  NUM_SEEDS=3
python src/iql_teams_oa_mlp.py alg=ql_teams_oa_resources_corners NUM_SEEDS=3
python src/iql_teams_oa_mlp.py alg=ql_teams_oa_resources_random  NUM_SEEDS=3

# training — MAPPO
python src/mappo_teams_mlp.py alg=mappo_teams_resources_circle  NUM_SEEDS=3
python src/mappo_teams_mlp.py alg=mappo_teams_resources_corners NUM_SEEDS=3
python src/mappo_teams_mlp.py alg=mappo_teams_resources_random  NUM_SEEDS=3
python src/mappo_teams_mlp.py alg=mappo_teams_simple_tag        NUM_SEEDS=3   # vanilla baseline

# evaluation — 4×4 cross-play tournament
python src/tournament_resources.py --placement random  --seeds 3 --eps 100 --K 5 --H 3
python src/tournament_resources.py --placement circle  --seeds 3 --eps 100
python src/tournament_resources.py --placement corners --seeds 3 --eps 100

# evaluation — trajectory datasets for VAE pipeline
python src/generate_trajectory_dataset_resources.py --algorithm iql
python src/generate_trajectory_dataset_resources.py --algorithm oa_iql
python src/generate_trajectory_dataset_resources.py --algorithm mappo

# evaluation — VAE + multimodality
python src/train_traj_vae.py
python src/verify_multimodality.py
python src/verify_vae_modes.py

# Exp 3 — ~5 min training + ~50 s tournament (MLP, CPU)
python src/iql_teams_mlp.py    alg=ql_teams_mlp_simple_tag    NUM_SEEDS=3   # ~2 min
python src/iql_teams_oa_mlp.py alg=ql_teams_oa_mlp_simple_tag NUM_SEEDS=3   # ~2 min
python src/tournament.py --seeds 3 --eps 100 --K 5 --H 3                    # ~50 s
python src/plot_exp3.py

# Exp 2 — ~24 min total (3 × ~8 min RNN trainings)
python src/smoke_test_static.py
python src/iql_teams.py alg=ql_teams_static_baseline NUM_SEEDS=3
python src/iql_teams.py alg=ql_teams_static_ccw      NUM_SEEDS=3
python src/iql_teams.py alg=ql_teams_static_cw       NUM_SEEDS=3
python src/exp2_behavior_mining.py
python src/compare_static_plots.py \
  --static_baseline logs/MPE_simple_tag_v3/iql_teams_static_baseline_MPE_simple_tag_v3_seed0_metrics.npz \
  --static_cw       logs/MPE_simple_tag_v3/iql_teams_static_cw_MPE_simple_tag_v3_seed0_metrics.npz \
  --static_ccw      logs/MPE_simple_tag_v3/iql_teams_static_ccw_MPE_simple_tag_v3_seed0_metrics.npz

# Exp 1 — archived, 1 seed each, ~3 min
python src/iql_teams.py    alg=ql_teams_simple_tag    NUM_SEEDS=1
python src/iql_teams_oa.py alg=ql_teams_oa_simple_tag NUM_SEEDS=1
python src/compare_plots.py \
  --baseline logs/MPE_simple_tag_v3/iql_teams_MPE_simple_tag_v3_seed0_metrics.npz \
  --oa       logs/MPE_simple_tag_v3/iql_teams_oa_MPE_simple_tag_v3_seed0_metrics.npz