Efficient Generation of Diverse Cooperative Agents with World Models

Yi Loo; Akshunn Trivedi; Malika Meghjani

cs.AIJun 9, 2025

Efficient Generation of Diverse Cooperative Agents with World Models

Authors: Yi Loo, Akshunn Trivedi, Malika Meghjani

Organizations: Singapore University of Technology and Design (SUTD)

Abstract

A major bottleneck in the training process for Zero-Shot Coordination (ZSC) agents is the generation of diverse partner agents via Population-based Training (PBT). Prominent PBT methods, such as Cross-play Minimization (XPM), can be very computationally expensive and sample inefficient, as the training objective requires sampling multiple types of trajectories, in particular the need of Mixed-Play (MP) trajectories to mitigate self-sabotage in cross-play. Each partner agent in the population is also trained from scratch, despite all of the partners in the population learning policies of the same coordination task. In this work, we propose that simulated trajectories from the dynamics model of an environment can drastically speed up the training process for XPM methods. We introduce XPM-WM, a framework for generating simulated trajectories for XPM via a learned World Model (WM). We show XPM with simulated trajectories removes the need to sample Mixed-Play trajectories while still mitigating any self-sabotage. The learned WM also bootstraps the training of subsequent agents in the population. We show that these two factors combined gives XPM-WM significantly better sample efficiency as compared to other XPM methods and scales linearly to the size of the population. We show our proposed method can effectively generate partners with diverse conventions that match the performance of previous methods. Finally, we show that our proposed training paradigm of using learned world models to generate diverse population can be extended to other PBT methods beyond XPM.

Figures & tables

Figure 1: Diagram of the training process of XPM-Sim (right) compared to CoMeDi. Our key insight is that by expanding the starting states of SP simulated trajectories, we remove the need for Mixed Play trajectories while still producing policies that do not self-sabotage, resulting in a more sample efficient method.

Figure 2: Diagram of the training process for XPM-WM. For the first phase, only the self-play (in blue) portion of the diagram is used and the agent is trained with just the objective SP. During the training of subsequent agents using XPM-Sim, the parameters of the partner agent is kept frozen while we continuously update the World Model.

	$1$	$2$	$3$	$4$	$5$	$6$	$7$	$8$	Average	SP $\uparrow$	XP $\downarrow$
LIPO	$0.0$	$1.0$	$1.0$	$0.33$	$0.5$	$0.25$	$0.98$	$0.65$	$0.59$	$2.12$	$-130.65$
CoMeDi	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.175$	$0.57$	$0.09$	$30.67$	$-114.01$
XPM-Sim (Ours)	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.09$	$0.0$	$\mathbf{0.011}$	$21.50$	$-114.50$

Table 1: Results of the MPPMR environment. The first 8 columns indicate the proportion of self-sabotage behaviors when in Cross-Play with preceding agents. E.g. Agent 4 will be paired with agents 1,2, and 3, with 1.0 indicating every episode containing self-sabotage behavior with 0 indicating None. SP indicates the average SP rewards across all 8 agents and XP indicates the average cross-play rewards.

Figure 3: Scaling by number of agents. We show the training cost in terms of environmental steps for all 5 Overcooked layouts as number agents $M$ in the population increases. LIPO and CoMeDi both have non-linear scaling in environmental steps required as $M$ increases due to the requirement to sample additional trajectories whereas XPM-WM has a linear scaling after the initial agent that is trained together with the world model.

Layout	LIPO		CoMeDi		XPM-WM (Ours)
Layout	SP $\uparrow$	XP $\downarrow$	SP $\uparrow$	XP $\downarrow$	SP $\uparrow$	XP $\downarrow$
CR	$8.99$	$2.40$	$8.29$	$1.58$	$10.23$	$0.43$
AA	$9.07$	$1.76$	$11.23$	$6.80$	$18.28$	$0.83$
C Ring	$7.77$	$0.11$	$9.45$	$2.74$	$10.81$	$0.03$
FC	$10.50$	$1.99$	$8.84$	$1.61$	$14.23$	$0.07$
CC	$8.47$	$0.36$	$8.94$	$1.26$	$8.25$	$0.72$

Table 2: The average number of dishes served in Self Play and Cross Play across all 5 Overcooked Environments for each method of XPM-based diverse agent population generation. Both self play and cross play scores are averaged over 20 episode of 400 steps.{CR = Cramped Room, AA = Asymmetric Advantages, C Ring = Coordination Ring, FC= Forced Coordination, CC=Counter Circuit}

	E-LIPO	E-CoMeDi	E-XPM-WM (Ours)
LIPO	$5.66\pm 0.07$	$5.34\pm 0.06$	$4.88\pm 0.07$
CoMeDi	$6.42\pm 0.07$	$6.25\pm 0.06$	$6.18\pm 0.06$
XPM-WM (Ours)	$4.88\pm 0.07$	$3.75\pm 0.07$	$6.62\pm 0.07$
Average	$5.66\pm 0.04$	$5.11\pm 0.04$	$\bm{5.80\pm 0.04}$

Table 3: Results for Ego agents averaged across 5 Overcooked layouts . For each entry we compute the mean number dishes served as well as the standard error over 40 XP episodes for each partner agent in holdout population. E-X denotes an Ego agent trained with partner agent generated using X method.

	E-LIPO	E-CoMeDi	E-XPM-WM (Ours)
LIPO	13.75 $\pm$ 0.43	12.54 $\pm$ 0.48	15.11 $\pm$ 0.41
CoMeDi	19.20 $\pm$ 0.25	18.69 $\pm$ 0.25	19.87 $\pm$ 0.24
XPM-WM (Ours)	7.71 $\pm$ 0.52	8.58 $\pm$ 0.51	10.17 $\pm$ 0.52
Average	13.55 $\pm$ 0.28	13.27 $\pm$ 0.28	$\bm{15.05\pm 0.27}$

Table 4: Results for XPM Ego agents on Overcooked v2 Asymmetric Advantages with center receipe layout. The evaluation framework is similar to that of Table 3 .

	E-FCP-WM	E-FCP	E-MEP-WM	E-MEP
CR	7.38 $\pm$ 0.15	7.47 $\pm$ 0.11	6.63 $\pm$ 0.22	5.70 $\pm$ 0.21
AA	8.41 $\pm$ 0.35	4.87 $\pm$ 0.28	4.10 $\pm$ 0.2	8.15 $\pm$ 0.29
C. Ring	6.82 $\pm$ 0.09	5.72 $\pm$ 0.11	3.29 $\pm$ 0.13	5.70 $\pm$ 0.18
FC	4.73 $\pm$ 0.26	8.00 $\pm$ 0.15	2.32 $\pm$ 0.17	6.13 $\pm$ 0.35
CC	3.54 $\pm$ 0.11	3.77 $\pm$ 0.09	1.79 $\pm$ 0.16	5.57 $\pm$ 0.16
Average	$\bm{6.18\pm 0.11}$	5.97 $\pm$ 0.08	3.38 $\pm$ 0.10	$\bm{6.25\pm 0.11}$

Table 5: Results of Ego agents trained with FCP-WM, MEP-WM evaluated on a held out Self-Play population generated using FCP and MEP respectively. Similar to 3 we compute the mean dishes served over 40 episodes as well as standard error.

Figure 4: A comparison of the user study scores for HiPT ego agent trained with XPM-WM, CoMeDi and LIPO. Error bars represent the standard deviation.

	Plating Preference		Onion Picking Preference		Movement Activity Preference
	High	Low	High	Low	High	Low
E-XPM-WM (Ours)	$0.74$	$-0.43$	$1.3$	$-1.33$	$0.17$	$-0.20$
E-LIPO	$1.76$	$-1.17$	$1.14$	$-0.4$	$1.13$	$-0.51$
E-CoMeDi	$1.72$	$-1.87$	$0.90$	$-0.44$	$0.30$	$-0.59$

Table 6: Comparison of difference in average number of dishes served with specific human conventions relative to all humans across all Overcooked environments.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: The initial state of Multi Particle Point Mass Rendezvous Environment. The 2 players control the red and green particles respectively. The four blue landmarks are spaced evenly in the environment. Both players receive a high reward if both particles are near one-another and to one of the landmarks. The environment episode terminates if one of the players go out of bounds of the 3x3 space.

Panel of Figure 6 — Figure 6: The Five original Overcooked layouts Top row from left to right: Cramped Room, Asymmetric Advantages, Coordination Ring, Forced Coordination and Counter Circuit. The Blue and Green hat chefs indicate the 2 different starting positions for each layout. Bottom row Asymmetric Advantages Recipes Center layout from Overcooked v2 The bowl at the center incicates which recipe to serve.

Event	Reward scaling
Player $i$ picks up ingredient	$1$
Player $i$ picks up a serving bowl	$1$
Player $i$ puts ingredient in pot	$1$
Player $i$ picks up soup from pot	$3$
Player $i$ delivers soup	$12$

Appendix

Table 7: Reward Events for the Overcooked AI environment

Event	Reward scaling
Player $i$ picks up a serving bowl	$3$
Player $i$ puts ingredient in pot	$3$
Player $i$ picks up soup from pot	$5$
Soup starts cooking in pot	$5$
Player $i$ delivers soup	$20$

Appendix

Table 8: Reward Events for the Overcooked v2 environment

hyperparameter	value
No. Actor MLP Layers	5
No. Critic MLP Layers	5
MLP Layer Size	512
activation	LeakyRELU
LIPO $\lambda_{XP}$	0.25
CoMeDi $\lambda_{XP}$	0.5

Appendix

Table 9: Hyperparameters for MPPMR Agents

hyperparameter	value
CNN Layers	3
CNN filter size	5x5, 3x3, 3x3
MLP Layers	2
No. Actor MLP Layers	1
No. Critic MLP Layers	1
MLP Layer Size	512

Appendix

Table 10: Hyperparameters for MAPPO (FCP, MEP, LIPO and CoMeDi) agents on Overcooked AI. XP coefficients, Training Steps, learning rate and Entropy coefficients are listed in order of layouts as in Figure 6 .

hyperparameter	value
World Model Specific Parameters
Encoder	3 x CNN, 3 x MLP
Decoder	3 x MLP, 3 x Transpose CNN
Transition Predictor	1 x GRU
Reward Predictor	5 x MLP
Continue Predictor	5 x MLP

Appendix

Table 11: Hyperparameters for XPM-WM on Overcooked AI Layout. Training steps are listed in order of layouts as in Figure 6 .

hyperparameter	value
CNN Layers	3
CNN filter size	5x5, 3x3, 3x3
MLP Layers	2
LSTM Layers	1
No. Actor MLP Layers	1
No. High Level Actor MLP Layers	1

Appendix

Table 12: Hyperparameters for HiPT Ego Agent.

Figure 7: Screenshot of an ongoing round of game in our user studies

Panel of Figure 8 — Figure 8: Trajectories of 2 sample agents under cross-play in the MPPMR environment.

Figure 9: Training Curves for all 5 Overcooked layouts. We plot the number of dishes served in each environment in each layout against the number of SP training steps (steps from XP and MP trajectories are excluded) averages across each subsequent agent in the population. The bold lines are the average dishes served across all subsequent agents and shaded area indicate the standard deviation.

	E-LIPO	E-CoMeDi	E-XPM-WM
LIPO	$6.49\pm 0.10$	$7.15\pm 0.10$	$5.21\pm 0.10$
CoMeDi	$8.20\pm 0.10$	$7.44\pm 0.12$	$7.38\pm 0.09$
XPM-WM	$6.35\pm 0.11$	$6.49\pm 3.07$	$8.52\pm 0.08$
Average	$7.01\pm 0.06$	$7.03\pm 0.06$	$\bm{7.04\pm 0.06}$

Appendix

Table 13: Ego Agent Results for Cramped Room

	E-LIPO	E-CoMeDi	E-XPM-WM
LIPO	$11.98\pm 0.15$	$8.05\pm 0.17$	$8.34\pm 0.17$
CoMeDi	$12.61\pm 0.12$	$10.11\pm 0.19$	$9.59\pm 0.16$
XPM-WM	$10.62\pm 0.18$	$6.47\pm 0.20$	$10.74\pm 0.24$
Average	$\bm{11.74\pm 0.09}$	$8.21\pm 0.11$	$9.56\pm 0.11$

Appendix

Table 14: Ego Agent Results for Asymmetric Advantages

	E-LIPO	E-CoMeDi	E-XPM-WM
LIPO	$3.99\pm 0.07$	$3.02\pm 0.09$	$3.01\pm 0.07$
CoMeDi	$5.32\pm 0.07$	$4.91\pm 0.10$	$4.41\pm 0.07$
XPM-WM	$4.14\pm 0.09$	$1.12\pm 0.06$	$4.69\pm 0.09$
Average	$\bm{4.48\pm 0.04}$	$3.02\pm 0.06$	$4.04\pm 0.05$

Appendix

Table 15: Ego Agent Results for Coordination Ring

	E-LIPO	E-CoMeDi	E-XPM-WM
LIPO	$3.75\pm 0.12$	$6.26\pm 0.11$	$2.87\pm 0.08$
CoMeDi	$1.58\pm 0.08$	$5.20\pm 0.09$	$3.70\pm 0.06$
XPM-WM	$1.30\pm 0.06$	$3.56\pm 0.06$	$4.81\pm 0.12$
Average	$2.21\pm 0.06$	$\bm{5.01\pm 0.06}$	$3.70\pm 0.06$

Appendix

Table 16: Ego Agent Results for Forced Coordination

	E-LIPO	E-CoMeDi	E-XPM-WM
LIPO	$2.09\pm 0.07$	$2.31\pm 0.06$	$4.36\pm 0.08$
CoMeDi	$4.41\pm 0.10$	$3.60\pm 0.11$	$6.08\pm 0.07$
XPM-WM	$2.01\pm 0.08$	$1.09\pm 0.07$	$4.36\pm 0.08$
Average	$2.75\pm 2.61$	$2.33\pm 2.72$	$\bm{4.66\pm 2.56}$

Appendix

Table 17: Ego Agent Results for Counter Circuit

Figure 10: A comparison of the training curves for XPM-WM with and without the use to vectorized event-based reward prediction on the Coordination Ring layout.

Figure 11: A comparison of the training curves for agents 2-8 for XPM-Sim with and without the shared encoder on the MPPMR environment. For XPM-Sim with shared encoder, we first train the first agent together with a VAE. The VAE is trained to reconstruct the joint observation. The actor and critic takes in the latent representation of the VAE. The shared encoder is then used to bootstrap training for subsequent agents and VAE is still continuously updated similar to XPM-WM.

Figure 12: A comparison of the training wall clock time for XPM-WM with CoMeDi and LIPO for all 8 agents on the Coordination Ring layout.

Figure 13: A comparison of the training wall clock time for XPM-WM with CoMeDi and LIPO for agents 2-8 (Phase 2 training only) on the Coordination Ring layout.

Cooperative MARL

Cooperative multi-agent reinforcement learning (MARL) trains agents to coordinate toward shared goals under partial information and decentralized control.

See all 247 papers

Shaping Zero-Shot Coordination via State Blocking
May 12, 2026Mingu Kang, Sunwoo Lee, Yonghyeon Jo +1Multi-Agent Coordination Cooperative MARL
Graduate School of Artificial Intelligence UNIST Ulsan, South Korea 44919
PopGen: Generating Semantically Diverse Populations for Zero-Shot Coordination
Oct 11, 2026Paul Teiletche, Kanghyun Ryu, Kyungjoon Park +3Zero-Shot Coordination Multi-Agent Coordination
EPFL · UC Berkeley · Yonsei University +2
Is Inter-Seed Cross-Play Enough? Evaluating the Robustness of Zero-Shot Coordination Algorithms to Implementation Details
Aug 4, 2026Maksymilian Wolski, Nicholas Hoernle, Johannes Forkel +1Multi-Agent Coordination AI Agent Evaluation
ProrokLab University of Cambridge United Kingdom · QuantCo · BOLD† University of Oxford United Kingdom

May 12, 2026J/K move · Enter open · S save

Shaping Zero-Shot Coordination via State Blocking

Mingu Kang, Sunwoo Lee, Yonghyeon Jo, Seungyul Han

Graduate School of Artificial Intelligence UNIST Ulsan, South Korea 44919

Multi-Agent Coordination Cooperative MARL Zero-Shot Coordination

Open paper AskPDF arXiv

Zero-shot coordination (ZSC) aims to enable agents to cooperate with independently trained partners without prior interaction, a key requirement for real-world multi-agent systems and human-AI collaboration. Existing approaches have largely emphasized increasing partner diversity during training, yet such strategies often fall short of achieving reliable generalization to unseen partners. We introduce State-Blocked Coordination (SBC), a simple yet effective framework that improves ZSC by inducing diverse interaction scenarios without direct environment modification. Specifically, SBC generates a family of virtual environments through state blocking, allowing agents to experience a wide range of suboptimal partner policies. Across multiple benchmarks, SBC demonstrates superior performance in zero-shot coordination, including strong generalization to human partners.

Abstract

Cooperative MARL

Explore similar work

Shaping Zero-Shot Coordination via State Blocking

PopGen: Generating Semantically Diverse Populations for Zero-Shot Coordination

Is Inter-Seed Cross-Play Enough? Evaluating the Robustness of Zero-Shot Coordination Algorithms to Implementation Details