Learning Disentangled Multi-Agent World Model for Decentralized Control

World models enable learning policies via latent imagination, offering benefits such as history compression and sample efficiency. The primary challenge in applying world models to multi-agent tasks is that modeling multi-agent dynamics in latent space requires integrating information from different agents, often creating spurious correlations between their latent states. Existing methods either reconstruct the observation for each agent or employ communication to maintain correlation during execution, failing to learn disentangled latent states that are crucial for effective decentralized control. To address this, we present the Disentangled Multi-Agent World Model (DMAWM). It facilitates learning decentralized policies in the latent space through a novel architecture comprising independent agent modules and a shared environment module. During real-environment execution, agent modules independently process local information to form a factorized latent representation. The environment module is then trained to mirror the factorized structure generated by the agent modules, effectively disentangling individual latent states from the interaction dynamics. Consequently, imaginary rollouts generated by the environment module more faithfully simulate decentralized execution dynamics, facilitating the transfer of policies from imagination to decentralized execution. Empirically, DMAWM outperforms existing model-based and model-free approaches in convergence speed and final performance, with additional visualization demonstrating its efficacy in capturing agent interactions.

% Lay summary for ICML 2026 (max 10 sentences / 200 words) % Guidelines: https://blog.icml.cc/2026/05/07/icml-2026-lay-summaries/ Model-based agents often build a compact "mental model" of the world and practice inside their imagination rather than in the real environment. However, when several agents maintain their own mental models with respect to a shared environment, they need to consider teammates' mental models during imagination. This creates the problem of "entanglement", where their imagined views of the world secretly rely on the mental models of their teammates during centralized imagination, which is unavailable once the agents act in the real world. We introduce DMAWM, a framework that learns disentangled multi-agent world models, which could effectively capture the interactions among agents during imagination while ensuring that each agent's mental model is faithful to the real-environment dynamics. Multi-agent teams trained with DMAWM can then be deployed directly in real environments in a fully decentralized manner, requiring no communication or access to teammates' mental models in real time. Across standard cooperative benchmarks, including SMAC, SMACv2, and vision-based Melting Pot, DMAWM learns faster and reaches higher final performance than existing model-based and model-free approaches. These results show that imagination-based learning can be sample-efficient while still supporting fully decentralized multi-agent execution.