Title: Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective
Venue: NeurIPS 2025
Accepted: yes
Keywords: Multi-Agent Model-based Reinforcement Learning

World Models

Learning in Imaginations

Diffusion Models
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, \textbf{D}iffusion-\textbf{I}nspired \textbf{M}ulti-\textbf{A}gent 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 \url{https://github.com/breez3young/DIMA}.


Peer reviews:
comments: SUMMARY:
This manuscript proposes to model the multi-agent transition systems using a sequential diffusion model, where each action of each agent is denoised step by step instead of concatenated together. The authors develop a world model for MARL using diffusion models and empirically compared with baselines over some multi-agent benchmarks.

STRENGTHS AND WEAKNESSES:
Strengths:

1. The paper is easy to follow and well-written.
2. The paper conducts extensive experiments and compared with some baselines.

Weaknesses:

1. One major concern I have is that it's not clear to me why the proposed approach has superior sample efficiency. I don't think the final returns upon convergence have big difference with other baselines (e.g., HAPPO, if you check the results therein). The authors doesn't explain where is the sample efficiency from and no ablation study can help understand it either. I'm also curious about the comparison on the training cost (e.g., time, resources) for each method.

2. The second concern I have is the motivation of denoising each step instead of using "conventional approaches". In the section 5.4, the joint case and sequential case has almost the same expected returns, sample efficiency, and variance. The difference is very marginal. However, there are downsides of using sequential training, like the authors mention, permutation invariance and condition-independent noise process. In this case, being simpler and using the joint case is preferred.

QUESTIONS:
Where is the proof of Thm. 2? It should be in the appendix but the readers can't be find it in the paper.

LIMITATIONS:
Please see weakness.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes DIMA as world model for multi-agent systems with diffusion techniques. \nDifferent from traditional joint action condition in multi-agent world models, DIMA applies sequential action conditions with permutation invariance, to reduce the modeling complexity. This design is novel, while I expect to see stronger evidence that it is better than joint action conditions in terms of both computation, prediction accuracy and policy performance. The experiments on MAMuJoCo and Bi-Dexhands show better efficiency and final returns over prior model-based and model-free MARL baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nThe paper is well written and easy to follow.\n\nThe design of sequential modeling for multi-agent action condition is novel, with some experiments showing comparable performances against joint action condition. This should save some modeling efforts due to the reduced complexity.\n\nExperimental advantages over prior methods are clear.\n\n\nWeakness:\n\nFor comparison of future trajectory prediction over different modeling methods (Fig. 6), please provide numerical results of observation and reward errors at each step over a long horizon. \n\nThe comparison of DIMA’s sequential modeling against the traditional joint agent action modeling, which is shown in Fig. 8, does not clearly tell the difference of the performances and the variances. The learning curves overall have high fluctuations, and it is not significant to see the lower variances by DIMA. Please prove this argument by more environments, and perhaps more runs to show with higher statistical confidence.\n\nQUESTIONS:\nCurrent world dynamics is modeled by diffusion, while reward and termination is modeled by transformer using MinGPT, is there any benefits for this separate modeling choice? and would it be possible to unify the modeling of different variables in one big diffusion model instead?\n\nLIMITATIONS:\nThe paper should also discuss the computational cost for DIMA over other model-based or model-free MARL methods, since diffusion model rollout for multi-agent system can be quite heavy and RL policies usually require a large number of samples.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- The paper is about modeling environment dynamics in MARL using a diffusion process.\n- The dynamic modeling problem is formulated as a diffusion-based conditional denoising process, which is claimed to reduce complexity without requiring additional communication mechanisms.\n - The key idea is to shift from modeling joint state-action transitions directly to modeling only the state transitions step by step, conditioning progressively on each agent’s action.\n - Experiments on MAMuJoCo and Bi-DexHands benchmarks demonstrate that DIMA outperforms existing model-based MARL baselines (such as MAMBA and MARIE) as well as some model-free methods (MAPPO, HAPPO, and HASAC).\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- The connection between sequential agent modeling and diffusion processes is conceptually interesting.\n - DIMA compresses the information space from \n∣S∣×∣S∣×∣A∣ to ∣S∣, which significantly reduces model complexity.\n-  Results on challenging tasks like Bi-DexHands and various MAMuJoCo partitions (Figures 4–5) show clear improvements in sample efficiency and stability over strong baselines.\n- The paper includes long-horizon prediction reconstructions (Figure 6), permutation invariance tests (Figure 7), and an ablation comparing sequential vs. joint agent modeling (Figure 8), which make the empirical analysis quite complete.\n\n**Weaknesses**\n\n- While I appreciate the idea of reducing the information space from \n∣S∣×∣S∣×∣A∣ to  ∣S∣, I find this relatively simple and not very innovative.\n- The paper mentions MADiff and MADiTS as relevant diffusion-based works but does not sufficiently compare with them beyond qualitative differences.\n- While MAMuJoCo and Bi-DexHands are good benchmarks, adding an additional environment (such as SMACv2 [1]) could demonstrate the generality of the approach for more discrete-action cooperative tasks.\n- Theorem 2 seems to be a key theoretical result in the paper, but its practical implication for the overall algorithm is not clear. I am wondering if the bound is trivial and does not actually help guide the main algorithm.\n- There are some more recent model-free MARL algorithms that are worth mentioning and possibly including in the comparison (e.g., IMAX-PPO [2] ) as well as some transformer-based MARL approaches (e.g. MARIE).\n\n**Refs**\n[1] Ellis, B., Cook, J., Moalla, S., Samvelyan, M., Sun, M., Mahajan, A., Foerster, J. N., & Whiteson, S. (2023). SMACv2: An improved benchmark for cooperative multi-agent reinforcement learning. In Advances in Neural Information Processing Systems\n\n[2]  Bui, T. V., Mai, T., & Nguyen, T. H. (2024). Mimicking To Dominate: Imitation Learning Strategies for Success in Multi-Agent Games. In Advances in Neural Information Processing Systems 37 (NeurIPS 2024)\n\nQUESTIONS:\n- Can you clarify the practical implication of Theorem 2 for the main algorithm?\n- How realistic is Assumption 1? Could you provide an example showing how this assumption holds in the multi-agent environments under consideration?\n\nPlease also respond to the specific points raised in the Weaknesses section above.\n\nLIMITATIONS:\n- Yes, the authors briefly mention a limitation regarding scalability. However, there might be other limitations, such as the need for sufficiently large and diverse datasets to accurately model the environment dynamics.\n-  I do not see any major potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a DIMA, a Diffusion-Inspired Multi-Agent world model to support training of multi-agent reinforcement learning policies. \n\nThe authors consider a Dec-POMDP (decentralized partially observable Markov decision process) where N agents interact in the same environment: given a state s_t, each agent i takes an action a^i_t given its local observation o^i_t of the state. Then, the environment moves from the current state s_t to the next state s_t+1, given the joint actions a_t, from all the agents. \nGiven the high complexity  of the joint action space A --- which increases exponentially with the number of agents and actions --- the authors propose a diffusion inspired module (DIMA) that scale linearly wrt the state space dimension. The main and novel idea is to leverage the denoising capability of diffusion models: starting from a noisy state representation s^0_t the module iteratively  applies the actions (a^i) from the agents, until all the actions  have been executed and the state have been denoised (thus we obtain the final state s_t). \n\nExperimental results on two benchmarks support the claims: DIMA outperforms state-of-art approaches, providing higher episodic returns with fewer steps.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n\n-\tThe paper is well written and easy to follow. I like Fig.1, it provides a clear understanding of the approach.\n\n-\tThe idea of using the diffusion denoising approach is novel. The denoising is used to predict the next state while handling the joint actions. \n\n-\tThe results are interesting, they show better performance on the two tested benchmarks.\n\n\nWeakness: \n\n-\tThe scalability of the approach is not sufficiently tested. The evaluated prediction/ imagination horizon is limited, with H=15. This does not provide insights on the scalability of DIME to longer horizon. This is particularly important, as a well-known issue in auto-regressive world models is the accumulation of compounding errors, which can lead to unrealistic predictions (out-of-distribution). \n\n-\tThe main paper lacks details or intuition about the architecture of the diffusion model. How does it handle (discrete?) finite actions and states?\n\n-\tThe experiments have a limited number of steps (10^5 or 10^6 steps). For example, the cited HASAC paper studied the algorithm using up to 1e7 steps. I understand the approach have faster convergence, but what are the performance at convergence? I believe this experiment would be interesting for the reader. \n\n-\tThe computational cost and training time of the world-agent are not investigated (also in comparison wrt similar solution).\n\nMinor and typos:\n\n- does the authors use “§” to refer to the appendix? If so, this should be stated explicitly at least once, for clarity.\n\nQUESTIONS:\nQ1: How does the model handle long prediciton horizions? What is the impact of an increasing state/action space? How does the model interpolate or generalize to previously unseen actions/states?\n\nQ2: what are the performance of the tested models using more steps / at convergence (e.g., 1e7 steps)? \n\nQ3: what is the computational cost of training the world model (e.g., time), compared to similar approaches?\n\nLIMITATIONS:\nYes - additional details would be helpful", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}