Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective
Yang Zhang, Xinran Li, Jianing Ye, Shuang Qiu, Delin Qu, Xiu Li, Chongjie Zhang, Chenjia Bai
Abstract
World models have recently attracted growing interest in Multi-Agent Reinforcement Learning (MARL) due to their ability to improve sample efficiency for policy learning. However, accurately modeling environments in MARL is challenging due to the exponentially large joint action space and highly uncertain dynamics inherent in multi-agent systems. To address this, we reduce modeling complexity by shifting from jointly modeling the entire state-action transition dynamics to focusing on the state space alone at each timestep through sequential agent modeling. Specifically, our approach enables the model to progressively resolve uncertainty while capturing the structured dependencies among agents, providing a more accurate representation of how agents influence the state. Interestingly, this sequential revelation of agents' actions in a multi-agent system aligns with the reverse process in diffusion models-a class of powerful generative models known for their expressiveness and training stability compared to autoregressive or latent variable models. Leveraging this insight, we develop a flexible and robust world model for MARL using diffusion models. Our method, Diffusion-Inspired Multi-Agent world model (DIMA), achieves state-of-the-art performance across multiple multi-agent control benchmarks, significantly outperforming prior world models in terms of final return and sample efficiency, including MAMuJoCo and Bi-DexHands. DIMA establishes a new paradigm for constructing multi-agent world models, advancing the frontier of MARL research. Codes are open-sourced at https://github.com/breez3young/DIMA.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 1438f521-4ea2-4478-bf1a-be7c4e986f44Cited by top-tier papers2
- Empowering Multi-Robot Cooperation via Sequential World ModelsZijie Zhao, Honglei Guo, Shengqian Chen, Kaixuan Xu et al.ICLR 2026 · 16 citations
- Diffusing to Coordinate: Efficient Online Multi-Agent Diffusion PoliciesZhuoran Li, Hai Zhong, Xun Wang, Qingxin Xia et al.ICML 2026 · 2 citations
Builds on36
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 13,211 citations
- Elucidating the Design Space of Diffusion-Based Generative ModelsTero Karras, Miika Aittala, Timo Aila, Samuli LaineNeurIPS 2022 · 3,959 citations
- Dream to Control: Learning Behaviors by Latent ImaginationDanijar Hafner, Timothy P. Lillicrap, Jimmy Ba, Mohammad NorouziICLR 2020 · 1,852 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
Related papers
- Learning Disentangled Multi-Agent World Model for Decentralized ControlDi Xue, Jing Jiang, Shaowei Zhang, Wenhao Guo et al.ICML 2026 · 2 citations
- Diffusion for World Modeling: Visual Details Matter in AtariEloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto et al.NeurIPS 2024 · 359 citations
- EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence ModelingJia-Hua Lee, Bor-Jiun Lin, Wei-Fang Sun, Chun-Yi LeeNeurIPS 2025 · 4 citations
- Hieros: Hierarchical Imagination on Structured State Space Sequence World ModelsPaul Mattes, Rainer Schlosser, Ralf HerbrichICML 2024 · 8 citations
- Imitating Human Behaviour with Diffusion ModelsTim Pearce, Tabish Rashid, Anssi Kanervisto, David Bignell et al.ICLR 2023 · 23 citations
