MA-WAM: Multi-Agent World-Action Model for Test-Time Planning

Guowei Zou; Haitao Wang; Guoxin Wang; Beiwen Zhang; Zhiquan Chen; Guojie Wang; Hejun Wu

cs.AISep 25, 2026

MA-WAM: Multi-Agent World-Action Model for Test-Time Planning

Authors: Guowei Zou, Haitao Wang, Guoxin Wang, Beiwen Zhang, Zhiquan Chen, Guojie Wang, Hejun Wu

Organizations: Sun Yat-sen University

Abstract

Multi-agent cooperative tasks require different agents to execute a joint action simultaneously, and each agent's action affects both the observations and responses of the other agents. Hence, a world model is needed to predict the team return resulting from the joint actions of all agents. A naive extension directly applies a single-agent world model to each agent's action when predicting the team return step by step. However, such an extension fails to capture the dependencies among the simultaneous actions of multiple agents. We propose Multi-Agent World-Action Model (MA-WAM), a test-time planning framework that enables a frozen multi-agent flow policy to evaluate futures of candidate joint actions. To our knowledge, MA-WAM is the first test-time world-model planner for multi-agent flow policies. MA-WAM predicts the consequences of each joint action according to cross-agent dependencies and enables efficient candidate scoring. Across 30 offline multi-agent reinforcement learning (MARL) settings on MAMuJoCo, SMAC, and MPE, MA-WAM achieves mean relative gains of 22.0% over direct execution and 25.6% over uniform action selection. Under the standard evaluation protocol on an A100 GPU, MA-WAM adds 12.1 ms, accounting for 2.5% of the measured generation-and-scoring time.

Figures & tables

Figure 1: Multi-Agent World-Action Model (MA-WAM) overview. At each environment step, a frozen policy proposes fixed-length joint-action sequences from the current observation context. RWM rolls out and scores the candidates by predicted cumulative return. The planner executes the first joint action of the highest-scoring candidate and replans from the next observation. Both models remain frozen during deployment.

Task	Quality	BC	MA-ICQ	MA-TD3+BC	MA-CQL	OMAR	MADiff	MA-SfBC	DOM2	MA-WAM (ours)
Spread	Expert	35.0 $\pm$ 2.6	104.0 $\pm$ 3.4	108.3 $\pm$ 3.9	98.2 $\pm$ 5.2	114.9 $\pm$ 2.6	95.0 $\pm$ 5.3	87.5 $\pm$ 7.3	88.7 $\pm$ 6.3	118.3 $\pm$ 1.5
	Md-Replay	10.0 $\pm$ 3.8	13.6 $\pm$ 5.7	15.4 $\pm$ 5.6	31.4 $\pm$ 7.2	37.9 $\pm$ 6.1	30.3 $\pm$ 2.5	8.2 $\pm$ 4.6	63.1 $\pm$ 9.5	61.6 $\pm$ 3.0
	Medium	31.6 $\pm$ 4.8	29.3 $\pm$ 5.5	39.4 $\pm$ 3.6	34.1 $\pm$ 7.2	47.9 $\pm$ 18.9	64.9 $\pm$ 7.7	51.6 $\pm$ 14.2	78.6 $\pm$ 8.1	86.7 $\pm$ 2.8
	Random	-0.5 $\pm$ 3.2	6.3 $\pm$ 3.5	9.8 $\pm$ 4.9	24.0 $\pm$ 9.8	34.4 $\pm$ 5.3	6.9 $\pm$ 3.1	5.1 $\pm$ 3.9	37.4 $\pm$ 11.3	66.2 $\pm$ 3.2
Tag	Expert	40.0 $\pm$ 9.6	113.0 $\pm$ 14.4	115.2 $\pm$ 12.8	119.3 $\pm$ 14.0	123.9 $\pm$ 10.5	103.0 $\pm$ 12.0	77.4 $\pm$ 13.9	98.2 $\pm$ 14.4	132.9 $\pm$ 5.5
	Md-Replay	0.9 $\pm$ 1.4	34.5 $\pm$ 27.8	28.7 $\pm$ 20.9	41.7 $\pm$ 15.3	47.1 $\pm$ 15.3	53.9 $\pm$ 11.4	12.7 $\pm$ 7.3	68.2 $\pm$ 16.7	77.4 $\pm$ 5.0

Table 1: Main results across MPE, SMAC, and MA-MuJoCo. Entries are episode returns; MPE uses the normalized score scale, whereas SMAC and MA-MuJoCo use native returns. MA-WAM (ours) reports mean

\pm

standard deviation across seeds, while published baselines retain their source-paper uncertainty conventions. Method abbreviations are expanded in Appendix A.2 .

Figure 2: Reactive vs. planning per setting with three denoising steps per candidate. Bars show means across the evaluation seeds, and error bars show standard deviations across these seeds. Table 2 (in Appendix) reports the underlying results. Panel (d) tallies

26

positive,

0

zero, and

4

negative

\Delta

values.

Figure 3: Two-factor view of when planning helps. Settings are grouped by reactive-policy headroom and $H{=}8$ ranking reliability according to whether each factor is above or below its global median. Each cell reports the mean relative gain and its setting count. Colors use a signed square-root scale to keep the large-gain outlier from compressing smaller gains.

Figure 4: World-model ranking reliability across horizons and return alignment at

H{=}8

. The top panel reports benchmark-family medians and interquartile ranges for Spearman correlation,

8

-way Top-2 accuracy, pairwise accuracy, and normalized selection regret. The bottom panel plots standardized predicted versus realized returns. The dotted line marks

H{=}8

, and the diagonal marks perfect agreement. Per-setting results are in Appendix A.5 .

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Qualitative illustration of the candidate-selection problem on three MPE tasks. The top, middle, and bottom bands show Spread in rows (a)–(h), Tag in rows (i)–(p), and World in rows (q)–(x), respectively. Each band shows eight candidate trajectories from one shared initial state at five time fractions from

0

100\%

. Rows (a), (i), and (q), marked by green borders, replay real expert-dataset episodes and serve as reference high-quality candidates. The seven remaining rows have red borders and show synthetically perturbed variants with progressively weaker but still coherent behavior, from a near-miss to a clearly late or off-course pursuit. The panels illustrate best-of-

M

selection; later experiments evaluate world-model rollouts.

MPE (normalized score)					SMAC (episode return)
Task	Quality	Reactive	Planning	$\Delta$	Task	Quality	Reactive	Planning	$\Delta$
Spread	Expert	111.6	118.2	+6.6	3m	Good	19.0	19.5	+0.5
	Md-Replay	34.8	55.7	+20.9		Medium	13.4	13.8	+0.4
	Medium	54.0	84.7	+30.7		Poor	10.0	11.0	+0.9
	Random	37.4	64.4	+27.0	2s3z	Good	19.9	19.5	$-0.4$
Tag	Expert	130.8	132.0	+1.1		Medium	17.3	17.2	$-0.1$

Appendix

Table 2: Per-setting planning gain with three denoising steps per candidate. Reported scores are means under the common evaluation protocol; MPE uses normalized score, while SMAC and MA-MuJoCo use raw return.

\Delta

is Planning minus Reactive, computed before display rounding, and bold marks a positive

\Delta

. Planning improves the mean on

26

30

settings.

Benchmark	Settings	Positive/Zero/ Negative $\Delta$	Mean $\Delta$	Agg. rel. gain	Mean setting-wise gain	Median setting-wise gain
MPE	12	11/0/1	$+10.81$	$+14.42\%$	$+46.27\%$	$+19.16\%$
SMAC	12	10/0/2	$+0.41$	$+2.75\%$	$+3.29\%$	$+2.70\%$
MA-MuJoCo	6	5/0/1	$+125.00$	$+7.68\%$	$+10.70\%$	$+14.04\%$

Appendix

Table 3: Benchmark-level seed-averaged gain summary with three denoising steps per candidate. Mean

\Delta

is in each benchmark’s native unit. Aggregate relative gain is computed after summing reactive and planning scores within each benchmark family. Mean and median setting-wise gains first compute a relative gain for each task–quality setting and then aggregate those percentages within the benchmark family.

MPE (normalized score)					SMAC (episode return)
Task	Quality	Random $M{=}8$	RWM $M{=}8$	$\Delta$	Task	Quality	Random $M{=}8$	RWM $M{=}8$	$\Delta$
Spread	Expert	111.9	118.2	+6.3	3m	Good	19.5	19.5	+0.0
	Md-Replay	34.2	55.7	+21.4		Medium	14.9	13.8	-1.1
	Medium	51.5	84.7	+33.2		Poor	8.9	11.0	+2.1
	Random	37.1	64.4	+27.4	2s3z	Good	19.8	19.5	-0.3
Tag	Expert	133.7	132.0	-1.8		Medium	16.7	17.2	+0.5

Appendix

Table 5: Full

30

-setting equal-candidate selector control with three denoising steps per candidate. Each entry is a mean under the common evaluation protocol. Both arms use

M{=}8

candidates; Random selects uniformly, whereas RWM ranks candidates by predicted return. MPE uses OMAR normalized score; SMAC and MA-MuJoCo use native returns. Within each block,

\Delta

is RWM minus Random, computed before display rounding; bold marks a positive difference. RWM exceeds Random in mean on

23

settings.

Panel of Figure 7 — Figure 7: Two-factor diagnostic for planning gains, split into one panel per simulation environment. Each point is a task–quality setting with an available $H{=}8$ score-calibration diagnostic. The $x$ -axis is a benchmark-normalized headroom proxy for the reactive policy, and the $y$ -axis is $8$ -way Top-2 ranking reliability of the world-model score. Top-2 is the fraction of groups in which the realized-best candidate appears among the two highest-scoring candidates. Here, $H$ denotes the world-model scoring horizon. Marker shapes are reused independently within each panel and denote task, marker fill denotes dataset quality, and color indicates the relative planning gain. Full per-setting values are in Table 2 . Settings with both high reactive-policy headroom and high ranking reliability tend to have the largest positive gains.

MPE
Setting	Top-1 $\uparrow$		Top-2 $\uparrow$		Spearman $\uparrow$		PairAcc $\uparrow$		NormReg $\downarrow$
	RWM	Chance	RWM	Chance	RWM	Chance	RWM	Chance	RWM	Rand.
Spread-Medium	0.422	0.125	0.703	0.250	0.584	0.000	0.737	0.500	0.208	0.555
Tag-Medium	0.250	0.125	0.422	0.250	0.240	0.000	0.590	0.500	0.506	0.534
World-Medium	0.266	0.125	0.438	0.250	0.254	0.000	0.602	0.500	0.455	0.572
MAMuJoCo

Appendix

Table 6: Same-state counterfactual candidate-ranking diagnostic on MPE and MAMuJoCo Medium tasks using

64

anchor states per task. RWM is compared with the analytical chance references for Top-1, Top-2, and Spearman rank correlation, and with the

0.5

chance reference for PairAcc (pairwise ranking accuracy). NormReg (normalized selection regret) compares RWM selection with empirical uniform selection from the same candidate sets.

Figure 9: Per-setting world-model score calibration at

H{=}8

on MPE , MA-MuJoCo , and SMAC . Here,

H

denotes the world-model scoring horizon. The first row contains the MPE and MA-MuJoCo tasks, and the second row contains the SMAC tasks. Each square panel overlays its dataset settings using the inset colors and markers. Realized and predicted returns are standardized independently within each task–dataset setting. For readability,

120

of the

2{,}000

points per setting are displayed; Pearson

r

and Spearman

\rho

are computed from all

2{,}000

points. The grey line is

y{=}x

Figure 10: Per-setting score-reliability distributions at horizons

H\in\{1,4,8\}

for Spearman rank correlation, pairwise ranking accuracy, and normalized selection regret. Here,

H

denotes the world-model scoring horizon. Boxes summarize per-setting values within each benchmark family, and the dashed line marks the chance pairwise accuracy of

0.5

	RTX 3090				A100
$M$	Gen. (ms)	Score (ms)	Total (ms)	Share	Gen. (ms)	Score (ms)	Total (ms)	Share
1	55.9	10.1	66.0	15.3%	64.4	11.1	75.5	14.8%
2	106.3	12.8	119.1	10.7%	120.5	12.5	132.9	9.4%
4	209.0	13.3	222.3	6.0%	234.0	12.2	246.2	4.9%
8	414.6	13.4	428.0	3.1%	461.6	12.1	473.7	2.5%
16	827.9	13.3	841.2	1.6%	931.0	12.6	943.6	1.3%

Appendix

Table 7: Candidate-count profiling on an RTX 3090 and an A100. Results are averaged equally over MPE Spread-Medium, MAMuJoCo 2Ant-Medium, and SMAC 3m-Medium. Candidates are generated sequentially and scored in one batched RWM pass. We use four parallel environments,

K{=}3

, and a requested rollout cap

H{=}8

; the effective action horizons are

8/8/3

in the listed task order. Each entry is the mean of three repeats, each with 10 warm-up and 30 timed steps per setting. Total is Gen. plus Score.

Panel of Figure 12 — Figure 12: Representative dynamics-expert diagnostics on MAMuJoCo Medium datasets using $4{,}096$ sampled transitions. Panels (a) and (b) show the mean combine-probability mass assigned by each agent to each dynamics expert after summing over the expert’s four slots, thereby quantifying agent-specific use of the shared experts. Panels (c) and (d) report the increase in next-observation mean squared error (MSE) after setting one expert’s four slot outputs to zero without retraining. The unequal increases measure the trained model’s sensitivity to each expert.

Panel of Figure 13 — Figure 13: Aggregate routing diagnostics across all $30$ task–quality settings. Panel (a) uses corrected normalized mutual information (NMI) to compare context and agent-identity association. Panel (b) uses Jensen–Shannon divergence (JS divergence) to compare routing variation across transition contexts and agent identities. Panel (c) summarizes expert utilization and switching.

Figure 15: Selector comparison on

10

representative task–quality settings. Each panel uses the native score scale of one benchmark family. The five vertical bars in each subplot correspond to RWM, the monolithic scorer, the current-action Q-ranker, the direct-return predictor, and Independent-WM. Error bars show standard deviations across the evaluation seeds. A dark outline marks the largest displayed value in a group.

Figure 16: Coordination routing in the Routed World Model (RWM). (a) Soft-routed dynamics [ 37 ] : agent tokens are dispatched to $4$ slots for each of $8$ experts and softly combined after expert processing. (b) Top- $2$ reward routing [ 39 ] : all $4$ expert outputs are computed, but only the selected two receive nonzero mixture weights per token. Solid arrows mark retained routes, and dashed arrows mark routes excluded from the mixture.

Symbol	Meaning
$n$	Number of agents
$o_{t}^{i}$	Observation token of agent $i$ at time $t$
$a_{t}^{i}$	Single-agent action of agent $i$ ; discrete actions use class indices for execution and one-hot encoding as RWM input
$s_{t}$	Current joint observation $(o_{t}^{1},\ldots,o_{t}^{n})$
$c_{t}$	Joint-observation context ending at $s_{t}$ , including configured history
$\mathbf{a}_{t}$	Joint action $(a_{t}^{1},\ldots,a_{t}^{n})$

Appendix

Table 9: Notation for the flow policy, world model, and planner.

Multi-Agent Planning

Multi-Agent Planning coordinates multiple decision-makers to assign tasks, plan actions and paths, and resolve conflicts under shared constraints.

See all 115 papers

MA-JEPA: Joint-Embedding World Models for Multi-Agent Reinforcement Learning
Sep 27, 2026Brandon Gary Kaplowitz, Osaze James Obahor, Christian Schroeder de WittGame-Playing Agents World Model Learning
Department of Engineering Science, University of Oxford
MultiWorld: Scalable Multi-Agent Multi-View Video World Models
Apr 20, 2026Haoyu Wu, Jiwen Yu, Yingtian Zou +1Multi-View Consistency Action-Conditioned World Models
The University of Hong Kong · Sreal AI
DualWAM: Dual-System World Action Models for Asynchronous Global Planning and Local Refinement
Sep 21, 2026Yixin Zheng, Jiangran Lyu, Yuntian Deng +6Long-Horizon Robotic Manipulation World Models for Robotics
Institute of Automation, Chinese Academy of Sciences · Galbot · Peking University +2

Sep 27, 2026J/K move · Enter open · S save

MA-JEPA: Joint-Embedding World Models for Multi-Agent Reinforcement Learning

Brandon Gary Kaplowitz, Osaze James Obahor, Christian Schroeder de Witt

Department of Engineering Science, University of Oxford

Game-Playing Agents World Model Learning Joint-Embedding Predictive Architecture Centralized Training with Decentralized Execution Multi-Agent Reinforcement Learning Model-Based RL

Open paper AskPDF arXiv

World models improve sample efficiency by training policies on imagined trajectories, but their usefulness depends on learning representations that capture the information needed for future control. We study whether self-supervised joint-embedding prediction (JEPA) can provide this learning signal for multi-agent reinforcement learning. We introduce MA-JEPA, a stochastic world model that replaces observation reconstruction with prediction of target representations, enabling model-based multi-agent reinforcement learning with centralized training and decentralized execution. A categorical latent state and a causal Transformer are trained with posterior and action-conditioned dynamics prediction objectives and are then used for actor-critic learning from latent imagination. A training-only joint predictor conditions on all agents' local states and actions to predict each agent's next local observation embedding. These predictions are passed through the same local posterior used during real interaction with a centralized critic that is used only for value learning, with execution remaining decentralized. Our experiments show that this architecture performs strongly on SMAC, matching or exceeding the strongest reported comparator mean win rate on four of eight evaluated maps.

Abstract

Multi-Agent Planning

Explore similar work

MA-JEPA: Joint-Embedding World Models for Multi-Agent Reinforcement Learning

MultiWorld: Scalable Multi-Agent Multi-View Video World Models

DualWAM: Dual-System World Action Models for Asynchronous Global Planning and Local Refinement