Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective
Authors: Yang Zhang
Xinran Li
Jianing Ye
Shuang Qiu
Delin Qu
Xiu Li
Chongjie Zhang
Chenjia Bai
Venue: NeurIPS 2025 poster ()
Keywords: Multi-Agent Model-based Reinforcement Learning
World Models
Learning in Imaginations
Diffusion Models
Primary area: reinforcement_learning
TL;DR: We introduce a sample-efficient Diffusion-Inspired Multi-Agent world model (DIMA) for multi-agent control environments--MAMuJoCo and Bi-DexHands

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}.

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
TTinnhua University, "The Hong Kong University of Science and Technology, Washington University in St. Louis, *City University of Hong Kong, FFudan University, IIntiiute of Artificial Intelligence (TeleAI), China Telecom, 7 Shenzhen Research Institute of Northwestern Polytechnical University 2 z-yang2 @mails. tsinghua. edu. en, baicj baicj@chinatelecom cn
Abstract
World models have recently attracted growing interest in Multi-Agent Reinforce- ment Learning (MARL) due to their ability 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 complex- ity 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 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 latent variable models. Leveraging this insight, we develop a flexible and robust world model for MARL using diffusion models. Our method, Diffusion-Insiee Inspired Multi-Agent world model (DIMA), achieves state-of-the-as performance across multiple multi-agent control benchmarks, significantly outperforming prior world models terms of final return and sample efficiency, including MAMuloCo and Bi-DexHands. DIMA establishes new paradigm for constructing multi-agent world models, advancing the frontier MARL research. Codes are open-sourced at http:///gtttubbc.oooorreee/o///p
Introduction
Learning accurate world models to capture environmental dynamics is crucial for effective decision- making. In the realm of model-based reinforcement learning (MBRL), such models play a pivotal role by enabling policy training through learning in imagination 2 3 4) facilitating planning with look-ahead search [5] 66, combining both approaches 7 g. While MBRL has achieved significant success single-agent settings, extending these methodologies to multi-agent scenarios presents unique challenges, necessitating new approaches for multi-agent world modeling.
In multi-agent settings, where multiple agents simultaneously interact within a shared environment, two primary challenges emerge. First, the joint action space grows exponentially with the number of
Work done during the internship at TeleAI. Corresponding Author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

agents [0], making it computationally expensive to directly handle joint dynamics. Second, the complex interdependencies among agents [TT] make difficult to accurately capture how individual actions impact global state transitions. Current multi-agent world modeling approaches face fundamental tradeoff. On one end of the spectrum, centralized modeling schemes directly capture full joint dynamics but incur computational costs that scale exponentially with the number of agents. On the other end, decentralized approaches [22 [4] model individual agent dynamics separately and rely additional mechanisms, such as sophisticated communication aggregation modules, to recover the global state. However, this misalignment between decentralized model structure and the global Markov decision process (MDP) impose inherent limitations on model accuracy and those communication or aggregation modules do not have explicit signal for supervision, further hindering the training. This tradeoff motivates fundamental rethinking of the world model structure: Can we develop a centralized modeling scheme that maintains global consistency without auxiliary components in decentralized methods, while keeping computational complexity manageable as the number of agents increases?
To address this challenge, we adopt sequen- tial agent modeling perspective that processes agents' actions incrementally, as illustrated in Figure I Specifically, consider a multi-agent system timestep with global state sq. When
Enviromment Timestep (t)
Model Rollout
Decoising Time
a2-1
1.
Policy Sample
t Da
2
all agents' actions a. are unknown, the next state 4 remains highly uncertain. As agents' actions are progressively revealed, this uncer- tainty gradually decreases. This sequential un- certainty reduction process bears striking simi- larity to the reverse process in diffusion models 15 16 [7). where generation is framed as iter- ative denoising from noise clean samples.
st
S/+1
41-1
1o
at
Dg
Do
Temporal Context
at-3
RSL-2
Inspired by this conceptual similarity and the recent success of diffusion models in image- based world modeling [18 19 20), we propose Iiffusion-Inspired Multi-Agen world model (DIMA), which reformulates multi-agent dy- namics prediction modified conditional de- noising process. Despite employing a central- ized modeling scheme, DIMA achieves com- putational complexity that scales linearly with the state space dimensionality, regardless of the number of agents. We summarize our contribu- tions as follows:
Illustraiion of the DIMA world model. From the temporal perspective, each environmen- tal timestep is modeled as a complete denoising process, analogous to diffusion models. Within each timestep, we further consider an agent-wise perspective, where the introduction of each individ- ual agent's 's action information represents a single denoising step, progressively reducing uncertainty about the next state.
Figure 1:
We leverage the connection between sequential agent modeling and diffusion processes to reformu- late multi-agent dynamics prediction as a conditional denoising process. This enables centralized modeling scheme that reduces complexity without additional communication mechanisms.
We propose DIMA, centralized multi-agent world model tailored for model-based MARL, and derive its corresponding evidence lower bound (ELBO), providing theoretical insights. We then instantiate DIMA within the EDM training framework 21 and integrate it into the learning-in- imagination paradigm for policy optimization.
We evaluate DIMA on challenging continuous MARL benchmarks, including MAMuJoCo 22 and Bi-DexHands 23 in low-data regimes. Experimental results show that DIMA consistently improves the prediction accuracy of environment dynamics and outperforms both model-fre free and strong model-based MARL baselines in terms of sample efficiency and overall performance.
2
Preliminaries
2.1
Multi-Agent Systems as Dec-POMDP
We focus on fully cooperative multi-agent systems where all agents share team reward signal. We formulate the system as a decentralized partially observable Markov decision process (Dec-POMDP) M. which can be described by a tuple A, S, N [I....,n} denotes a set

---

of agents, S is the finite global state space, A = IT AA' is the product of finite action spaces of all agents, i.e. the joint action space, P: S x A Ax [0. 1] is the global transition probability function, R : S x A + R is the shared reward function, s2 - St is the product of finite observation spaces of all agents, i.e., the joint observation space, O = {Oi E M is the set of observing functions of all agents. Oi :S - 52' maps global states to the observations for agent i, and y is the discount factor. Given global state at timestep ł, agent i is restricted to obtaining solely its local observation of = (s+). takes an action oj drawn from its policy x' =-(oog)) based on the history of its local observations o'se which together with other agents' actions gives joint action a, = (a) af)) E A, equivalently drawn from a joint policy ग T(-1og)) I x² (-loss)) Then the agents receive a shared reward ry=- R[s,, az), and the environment moves to next state st+-1 with probability P(s++1|5, az). The aim of all agents is to learn a joint policy π that maximizes the expected discounted return [T) = Ess.ag.... .p...... |Et= 7 R(ss, a,)]. Note that recent approaches [2 [4] build the multi-agent world models via modeling P(o or, as) in the joint observation and action space, which mismatches with the transition formulation in Dec-POMDD However, DIMA is trained to recover the well-defined global state transition 1 18, as) according to the proposed multi-agent dynamics formulation.
2.2 Score-based Diffusion Models
In this work, we directly utilize the unified framework and the accompanying practical design choice of diffusion models introduced by Karras al. (2T)
Notation. Let us consider a diffusion process {x" ] indexed by a continuous time variable T E [0, 77, with corresponding marginals {p" hrejo,T. and boundary conditions p² Pdata and pT = Pprior, where Pptoo is usually pure Gaussian distribution in practical implementation. For clarity, we use the superscript 7 denote the diffusion process timestep and the subscript to denote the trajectory timestep.
ODE Expression. Song et al. 17 models the forward and reverse diffusion processes with stochastic differential equations (SDEs) which describe how the desired distribution of sample x evolves over time Assuming the stochasticity only comes from the initial sample xJ of prior distribution Pprior. Karras et al. [21] expresses diffusion models via its corresponding probability flow ordinary differential equation (ODE) 117 which continuously increases reduces the noise level of the image when moving forward or backward in time, respectively. The defining characteristic of the probability flow ODE is that evolving a sample xT ~ (x" = p(x: (TTa) from time To To (either forward or backward in time) yields sample x ~ (x) = p(x; (T.)). where (T) is schedule that defines the desired noise level at time T. It is described by
dx= (7)) (x)
where the dot denotes time derivative. marginals (p" re along the process. Equipped with the score function, we can thus smoothly mold random noise into data for sample generation, or diffuse data point into random noise.
Vx log (x) is the score function 24 associated with the
Denoising Score Matching. By using the score matching objective (24), we can evaluate the score function easily. Specifically, Do r) is parameterized denoiser function that minimizes the expected Le denoising error for samples x" drawn from Pdata for every ((T),
(())- = Exo (x)xx <pix (x0)
- Do |
- x²1]2,
where x p=x=|x''" denotes that x is obtained by applying Gaussian noise of scale (T) clean sample x". Then the score estimation can be given by Vx log (x) (D,(x; T) (T)) at any given time т. Thanks to this estimation, we can solve the ODE by numerical integration, i.e., taking finite steps over discrete time intervals with the help of various ODE solvers.
3 Methodology
In the following, we first elucidate our proposed formulation of modeling multi-agent dynamics from a diffusion-inspile perspective in $3.1. Based on this formulation, we derive the corresponding ELBO and score matching objective for implementing the diffusion model that incorporates such a perspective. Then, we describe the behavior learning process within the world model in $3.2
3

---

3.1 Modeling Multi-Agent Dynamics from a Diffusion-Inspired Perspective
Given a dataset {(01, $1. d1,71........................ .......... OT, ST,, ar, rr, rr.); containing all collected episodes, the aim of the multi-agent world model is to precisely predict how the next state is like based on an action intervention, i.e., recovering the unknown ground truth environment dynamics P(s (1,11ss,,)))
Diftusion-Inspired Formulation. Supposing there are n. agents 2 n)
and n noise levels (t) 0i on 0, the noisy sample st11 is corrupted
fon..... 0,, 01 that satisfy max 5.
from the clean next state 8t+1 Se+1 by adding noise of the corresponding level o;. Note that
(0)
here we use the superscript to denote the diffusion process timestep except for the action notation
Following the definition in [25]. we start by defining similar conditional Markovian forward
diffusion process 9.
((v/+1) p(8t-1);
(k+1)) @( 1s St, am = g(
(k+1) (8t41);
(3)
():6))) ((*+11 1 1641 S1, ae 1:m) = (( H1 141 86. ")
(k+1) (8)
where q denotes the unconditional forward diffusion process. While the conditional forward diffusion process @ is conditioned on the control signal (st, a!!"), we can prove that it behaves exactly like the unconditional one q. The following equations hold,
(5) The detailed proof referred to Dhariwal and Nichol 25 Since the above equations suggest that the forward diffusion process is independent of the control signal (st, "?""), we can now fully focus
(($
(k+ (1)) 1s = @( "eett"1ss, |144))
(k+1))
a!""), ((s (1):()))
(522)) = (((++1))) 1+11)
(1):(m)
on describing our formulation via the conditional reverse diffusion process.
To describe how the predicted next state gets sharpened progressively with sequentially given action of each agent, we have specify the conditioning order. Without loss of generality, adopt the descending order of agent id (n, n - ...... 1) as the conditioning order. Formally, we make the following assumption terms of the global state transition.
Assumption 1 (Diffusion-Ine Inspired Decomposition of Multi-Agent Dynamics). In our diffusion- inspired formulation with the descending order of agent (n, n - 1,...,1) as the conditioning order, the global state transition P(s 8t, a. yields the next state in manner akin rypical reverse diffusion process, i.e., satisfying
(1)()) (i a! = p(((a)) P(s S ( |s,,
p(
(K) 1211 a4, st),
(6)
(n)) where $:41 is corrupted with the noise of maximum level o. practically indistinguishable from pure Gaussian noise.
Under the assumption, we have the following new form of Evidence Lower Bound (ELBO) on the log ((s141|s a. ".
Theorem 2 (ELBO under the Diffusion-Inspiine Formulation). Under Assumption I the log-
likelihood of the multi-agent global state transition (i.e., the evidence of the transition) is lower bounded as follows.
log P(s al 2 E qi [log (( (1)
at st)]
(KL( (
prior motching
torm
D (k- 1 S lplss 1541
( (K- (k-1)
denoising matching term
The detailed proof is deferred to A The denoising matching term in Eq. secretly reveals
(k-1) (.1)))_.0 (k) &
that we can learn parameterized denoising intermediate step po(s (841 a st) that matches

---

the tractable ground-truth denoising intermediate step g( ((8441")004 (k) (0)
1 st+1), thereby realizing the
formulation we propose. When we utilize Gaussian noise for corruption in the forward diffusion process, the denoising matching term can be simplified
variant of Eq. (
( - Σ k=1 [DD(s Ok, 8t, at) - s1+11 given the order (n....,,,,1 1)
(8)
However, there are still two properties to be incorporated into Eq. (8) (i) Permutation Invariance. Note that our formulation merely provides a novel perspective for modeling multi dynamics, rather than changing the underlying mechanism of global state transitions. In other words, regardless of how the conditioning order of a!"h is specified, the next state should remain unchanged given the same current state and joint action, i.e., exhibiting permulation invariance. Therefore, for any possible order (i1, i2...., in) uniformly sampled from the whole permutation set perm 2........ n2), we should optimize an expectation of Eq. (8) over the whole permutation set. (ii) Condition- Independent Noising Process. According to Eqs. (2)-(5). the conditional forward diffusion process independent of the conditions. It allows us randomly sample the noise levels (o1,...,0,, Jn. with the predefined continuous- timeuuuuus--eueeeeee noise scheduler (T) in $2.
Putting the above two together, we finally derive the optimization objective of DIMA,
(() = Eq (ro (ro PompPom(.......
Σ k=1 [ (8) ++1 0k. St, at, s
E. UnIBB--Uiifommm....... [II Do (T), 8t, at) s 11
(9)
where k ~ Uniform (1, ......n.. n indicates that the agent index k is uniformly sampled from the set (1.2....,.11]
Comparison with Conventional Approaches. We present concise illustration to highlight the fundamental difference between our DIMA and recent diffusion- based methods 26 27 in modeling multi-agent dynamics. As shown in Figure 2 recent methods attempt to inject the entire joint action information into the progres- sively denoised next state every intermediate step, whereas DIMA incorporates only a single agent's s action at each step. Denoting the state space size as |S] and the individual action space size as [,A], DIMA compresses the relevant in- formation from a |S| X |A| x |S] space into a |S| space for each intermediate state transi- tion pa ps (k=11) $
Conventional Flattened Dystamics
2. Pe.p
DEMA
A
PMOST 114 a2.n)
Figure 2: Comparison between conventional flat- tened multi-agent modeling and DIMA's sequential agent modeling. Light gray indicates clean states; dark gray indicates noisy states.
| st). In contrast, exist- ing methods must handle significantly higher cost due to compressing information from (S] |A|l |S| space into (S]. This simple qualitative analysis demonstrates that despite model- ing multi-agent dynamics in a centralized manner, DIMA enjoys a linear complexity in modeling difficulty with respect to the number of agents.
Practical Implementation. Inspired by the success of DIAMOND 18) a powerful single-agent diffusion- world model, we adopt similar design choice and employ the EDM framework (2] to effectively train the desired diffuusion model. Specifically, the denoiser Dg is reparameterized using the EDM preconditioners as follows:
De(5 sc, at) = chip + C. FolC Cuviser St, St,a)) a.))
(10)
where Fg is the neural network. These preconditioners (<aiip Cout: CT Coose)) are detailed in B In addition, we incorporate two practical techniques to further improve DIMA's predictive performance: (i) we maintain running mean and standard deviation of global states to normalize the state before training, ensuring stable dynamics ranges; (ii) we augment the model input with fixed window of past l global states and joint actions to provide richer temporal context for next-state prediction.
3.2 Learning Behaviors in Imagination
To support reinforcement learning with imagined rollouts, we pair DIMA with two necessary compo- nents. The first is reward and termination model fo where reward prediction and termination pre-
5

---

diction are framed as scalar regression and binary classification tasks, respectively. Motivated by the advanced sequence modeling capability of Transformer [28], we employ Transformer architecture as the backhone. As illustrated in Figure 3 the model takes sequences of (.... 8t, atth 8841, (.4).....) as input and predicts reward and terminttoo at each timestep via two separate 3-layer multilayer perceptron (MLP) heads on top of the shared output embedding. The model is built upon MinGPT implementation [29]. The second component is a special auto-encoder 9 (o!-"|ss)
(t) that encodes the global state 8t into a compact latent space and decodes it into the joint observation o !: We implement this using a simple yet effective VQ-VAE [30] with Finite Scalar Quantization [3I. We adopt an actor-critic framework to learn the behavior policy of each agent, where the actor and critic are parameterized by two 3-layer MIPs, (af and V²(+s), respectively.
22-2
Yz-2
Pt-1
Y--1
Ft
Figure 3: Overview of the reward and termination model. DIMA ad- dresses reward and termination pre- diction from a global perspective using a transformer architecture to capture temporal correlations. Both functions share the same backbone with separate prediction heads.
MinGPT
2
1
CC-1
St
state emb.
joint action emb.
valid output emb.
Thanks to DIMA'S global state transition predictions,
leverage oracle information from the global state to train a centralized critic, which in turn guides the optimization of decentralized agent actors. This naturally aligns with the centralized training with decentralized execution (CTDE) paradigm commonly used in model-free MARL [32 33
34]- Moreover, this provides a clear advantage over recent model-based MARL methods that also rely learning in imagination [4]. As these methods typically model local observation dynamics for scalability, they lose the benefits of accessing oracle global state, which our approach fully exploits. Here we train the actor and critic with MAPPO [4] A-return [I] is used as the target to update the value function. The details of behavior learning objectives and algorithmic description are presented in HCIand $F respectively.
As we evaluate DIMA under the learning in imagination paradigm, our approach iteratively exccutes cycle that comprises three steps: (i) collecting experience by executing the policy, (ii) updating the world model with the collected experience, and (iii) learning the policy through imagined rollouts within the learned world model. Note that throughout the whole procedure, the historical experiences stored in the replay buffer are only used for training the world model, while the policy is optimized through unlimited imagined trajectories generated by the world model.
4 Related Works
Diffuusio Model for RL. Diffusion models 16.35 have been applied in reinforcement learning (RL) for its strong generation capability. Specifically, they are capable of modeling complex action distributions in online RL [36 37
38 39], offline policy learning [40] 42), and imitation learning Other works also adopt diffusion models as planners to generate state-action sequences Recently, diffusion policies are also used as an action expert to combine with LLMs and lllaaaagn-aattiiooo model 44 50 For dynamics modeling, diffusion models have been employed as alternatives to autoregressive models learn the complex transition function of MDPs |85119 19). while they are limited in the single-agent domain. MADiff (26) learns the distribution of the whole trajectory in the offline multi-agent settings, without modeling the step-wise transitions.
43
44 45
obtain
Other works [52 53 54 treat diffusion-based world modeling as a video generation problem without taking actions as a condition, limiting their abilities. In contrast, our proposed DIMA predicts future states based on sequential action conditions, which effectively builds the multi-agent world model.
Multi-Agent RL. In cooperative MARL, agents coordinate maximize a joint reward function. Centralized Training with Decentralized Execution (CTDE) (55 is foundational framework that leverages the global state of agents during training to facilitate policy learning while relying on partial information during execution. CTDE framework serves the basis for both based [32 B3 [6]
and policy-based MARL methods (57
34
58 Additionally, some works reformulate MARL as
a sequential decision-making problem
55 60
61), offering insight into sequential denoising in
diffusion-basee dynamics. Model-based MARI. has gained significant attention for its ability to explicitly model the underlying MDPs in multi-agent environments. Notable examples include

---

MAZero 13,, which adapts MuZero-style planning with MCTS, and Dreamer- based methods These approaches have demonstrated the potential of model-based methods to improve coordination in multi-agent systems.
12 62 14
which leverage learning in imaginations for multi-agent setups 2 63
More recently, diffusion models have been introduced into MARL to enhance coordination and trajectory modeling, motivated by their advanced modeling capabilities. MADiff [26] first introduces diffusion models in MARL through offline trajectory learning via attention- diffusion. Subse- quent works have extended the use of diffusion models in MARL. Specifically, DoF [64] investigates offline MARL by factorizing a centralized diffusion model into multiple sub-models, aligning with the CTDE framework. Similarly, MADiTS (27] explores diffusion- based data augmentation by stitching high-quality coordination segments together. While effective, these methods primarily use diffusion models as goal-conditioned trajectory generators, failing to account for the underlying multi-agent dynamics. Our proposed DIMA addresses this research gap by constructing that explicitly captures the multi-agent dynamics. By leveraging the strengths of diffusion- inspired modeling, DIMA assists policy training and improves the overall performance of MARL.
effective world model
5 Experiments
5.1 Experiments Setup
Environments. We evaluate our method on two widely-used multi-agent continuous control bench- marks requiring heterogeneous--geny cooperation: Multi-Agent MuJoCo (MAMuJoCo) [22 and Bimanual Dexterous Hands (Bi-DexHands) [23]. MAMuJoCo extends MuJoCo 665 multi-agent settings by partitioning a robot into agents controlling different degrees of freedom (DoFs), requir- ing coordination for coherent movement. We use seven agent- peett-pattiioning settings: HalfCheetah [2x3, 3x2, and 6x1]; Walker [2x3 and 3x2]; and Ant [2x4 and 4x2]. HinDexHands features dual Shadow hands (26 Dol's each) performing precise bimanual manipulation. We evaluate on four tasks: ShudowIHandPen
and Shaddw/tandBotleeCup To highlight the sample efficiency learning in imaginations, we adopt low-data regime [66]. limiting real-environment samples to IM for MAMuJoCo and 300k for Bi-DexHands, adjusted for their different episode lengths. In model- based MARL where policies are learned in imaginations, performance direetly reflects the accuracy of the world model, enabling transparent evaluation.
BhnddowlandDoooOOpeee
We compare DIMA against two strong model-based baselines with the same policy learning paradigm as ours MAMBA [12 and MARIE [4]. MAMBA extends DreamerV2 (39 to the multi-agent context and establishes an effective Recurrent State Space Model (RSSM)-based world model. MARIE incorporates Transformer-ban -based autoregressive world modeling 3 with CTDE principle and demonstrates remarkable sample efficiency on the benchmark with discrete action space. We also compare DIMA with strong model-free baselines, including two on-policy algorithms MAPPO 34] and HAPPO (58), and an off- off-policy algorithm HASAC 667 68). HASAC is a heterogencous-
Baselines.
extension of SAC 69 which is well known for its high sample efficiency. Each algorithm is evaluated using 4 random seeds per scenario. For each random seed, we report the averaged episode return across 10 evaluation episodes at fixed intervals of environment steps. To ensure a fair comparison, we restrict the imagination horizon H = 15 for all model- algorithms. Results of MARIE would not be reported in Bi-DexHaand due to severe out-of of-memory issues under our available computational resources.
5.2 Main Results
DIMA consistently outperforms all evaluated baselines across a wide range of multi-agent continuous control tasks, achieving superior sample efficiency and higher final returns. As shown in Figure 4 and DIMA exhibits rapid and consistent policy convergence across all chosen MAMuJoCo and Bi- Bi-DexHands tasks, while other model-based baselines fail to demonstrate such stable learning behavior. This highlights the advantage of our approach in leveraging an effective world modeling formulation that is better aligned with the global state transitions of the environment. MARIE and MAMBA suffer from a mismatch with the true global transition dynamics inherent in Dec-POMDPss due to their integration of local dynamics modeling with the CTDE principle. This discrepancy potentially imposes an inherent limitation on model accuracy, particularly in environments like MAMuJoCo where inter-agent dependencies strongly correlated. Although
7

---

MAPPO
HAPPO
HASAC
MAMBA
MARIE
DIMA (Ours)
HalfCheetah
HalfCheetah
Episode Returns
5000
EBisode Returns
6000
6000
4000
5000
5000
4000
4000
3000
3000
3000
2000
2000
2000
1000
1000
1000
0
-.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
-0.0
0.2
0.4
0.6
0.8
1.0
HalfCheetah 6x1
Walker2d-2x3
Walker2d-34 3xx
6000
5000
Eo.sode Retums
5000
Episode Returns
4000
Eoisode Returns
4000
4000
3000
2000
2000
1000
1000
-0.0
0.2
0.4
0.6
0.8
1.0
-0.0
0.2
0.4
0.6
0.8
1.0
-0.0
0.2
0.4
0.6
0.8
1.0
Environment Steps
3et
Environment Steps
1ce
Environment Steps
3et
Figure 4: Curves of averaged episode returns for all methods in MAMuJoCo.
MAPPO
HAPPO
HASAC
MAMBA
DIMA [Ours)
SlademmtandBottkCCp
ShadowHaanDDorrree puntnward
Shatio Door Outnewa
StaddowHarddaa
250
30-
Epiiede fettrre
250
Epiiock Barturte
300
Epieod Betuarr
25
150
3
200
20
N
100
Epiund
100
15
50
10
-50
-0.0
0.6
122
1.8
2.4
3.0
-0.0
0.5
2.4
3,0
-0.0
0.5
18
2.4
3,0
-0.0
0.5
18
2.4
3,0
Environmant Stops
Ervironmant Stops
Environment Stops
Ervironment Stops
Figure 5: Curves averaged episode returns for all methods in Bi-DexHands.
enjoying world modeling complexity at a linear rate, such architectural misalignment limits their scalability highly coupled settings. Interestingly, MARIE and MAMBA perform comparably to-or even worse than sample-inefficien on-policy model-free methods HАPPO and MAPPO (c.g.. in HalfCheetah [3x2, 6x1]), whereas DIMA consistently demonstrates superior performance. This performance gain reflects the accuracy and robustness of DIMA in enabling more precise and reliable imaginations for policy optimization.
similar trend is observed in the Bi-DexHands benchmark, characterized by the control of two dexterous hands, each with 26 DoFs (i.e., af TR:26). Benefiting from the expressiveness of the diffusion model, DIMA is able to more accurately capture especially sophisticated and contact-rich dynamics. By learning denoising generative process under our formulation, DIMA enables more faithful representations of the underlying transition distribution and leads to more stable, coherent imagined trajectories for downstream policy learning, compared to RSSM-based models. As result, DIMA substantially improves the learning of dexterous manipulation policies in scenarios requiring fine-grained, high-precision coordination. The numerical results are further provided in
5.3 Model Analysis
DIMA demonstrates substantially more accurate and stable long-horizon predictions than existing multi-agent world models. To better evaluate the model capabilities among MAMBA. MARIE and our DIMA, we visualize their imagined trajectories alongside the ground truth (GT) on Ant [2x4] task. As visualized Figure DIMA generates consistent imagined trajectory that closely aligns with the ground truth (GT) across the full prediction horizon H 15, maintaining coherent agent structures and motion patterns. In contrast, MARIE and MAMBA both exhibit significant degradation as the horizon extends, and suffer from varying degrees of distortions. These issues are especially pronounced at challenging future timesteps such as 4 and = 12, highlighted by red bounding boxes. The qualitative results underscore DIMA's superior modeling capability and stability in capturing complex continuous multi-agent control dynamics, which critical for generating reliable imagined rollouts that support sample-efficient policy learning.

---

1w4 1-n 1-10
Figure 6: Reconstructions of long-term predictions of different multi-agent world models. We qualitatively compare the reconstruction quality of different multi- agent world models. Each model perform forward imagination over horizon of 15.
DIMA scales robustly to longer horizons with significantly lower compounding error. DIMA demonstrates superior scalability and robustly mitigates compounding errors, a capability we validated by testing prediction horizons (H 25) significantly longer than those seen during training (H = 15). This setup provides a rigorous -uitt--distributiin (OOD) generalization test for auto- regressive world models. The results in Table 1 are conclusive: the Ant [2x4] benchmark, DIMA not only exhibits the lowest L1 accumulated observation and reward errors at the training horizon (H 15) but significantly widens this performance gap against MARIE and MAMBA the unseen H 25 horizon. This finding underscores DIMA'S effectiveness in long-range predictive accuracy. Downstream policy performance using this extended imagination horizon detailed in E.
Table I: Accumulated observation and reward errors at extended prediction horizons (H 25) Ant [2x4]. Results are averaged over 100 trajectory segments. DIMA exhibits the lowest compounding
еπгоr
Obs Accumulation Errors
Rew Accumulation Errors
Methods
= 15
25
H = 15
25
DIMA
2.42 3.64
4.32 1.44
15.08
31.07 + 17.63
MARIE
6.6212.68
21.64 38.90 244.88
47.66 32.002
MAMBA
3.98
6.8512.51
71.67=41.77
Obs Erors Diffrent Orders
Acconting
Mean Curmulative Errors
1.75
Descending
Randem
2.30
Awcding Doverding Raasers Odkrr
Orier
1.25
2.50
0.75
0.30
0.25
0.50
12
15
Horizons
Figure 7: Visualization of long-term predictions with different conditioning orders, together with the accompanying cumulative observation errors curve.
DIMA effectively preserves permutation invariance over long-horizon multi-agent predictions. To validate whether and how DIMA exhibits the desired permulation invariance property elaborated in $3.1 we evaluate its unrolling behavior under different conditioning orders of agent actions. As shown in Figure 7 (left), we generate imagined rollouts from the same initial state and joint action set, but vary the conditioning order using three representative orders: random, ascending and descending w.r.t. agent ids. DIMA produces visually consistent rollouts across different orders until notable visual differences emerges at t 12, highlighted by the red bounding box. This demonstrates that DIMA maintains this consistency effectively up to a long prediction horizon least H = 10. To further quantify the consistency, we plot the mean cumulative observation errors over prediction horizons under each order, as depicted in Figure (right). The resulting curves seems quite aligned, with no significant deviation among the three conditions, indicating that DIMA exhibits the permutation invariance property within a considerably long horizon via optimizing Eq. (.. Details of this experiment setup are provided in

---

5.4 Ablation Study
Our proposed formulation improves sample efficiency and stability in low-data regimes. To evaluate the core contribution of DIMA's agent-wise sequential modeling, we compare it against a "Joint" modeling baseline, which adopts the conventional centralized approach of conditioning on the full joint action at at every denoising step. The primary benefit of our sequential formulation is the reduction of modeling complexity from an exponential to linear dependency on the number of agents. This reduction is particularly impactful in low-data regimes, where simpler model can learn more effectively. We conducted experiments on several Bi-DexHands tasks with independent runs. As shown in T on the Shudow/HandBottee
task, our sequential approach (DIMA) achieves higher returns and lower variance at 100k and 150k steps. As the data budget increases (200k-300k), the performance of the joint model catches up, and both methods converge to a similar performance. This aligns with our hypothesis: sequential modeling provides a significant sample efficiency boost when data is scarce. This benefit is even more pronounced in more complex tasks. Table 3 shows that on DoorOpenOunt advantage and reduced variance over the Joint model even at the 300k step limit. This demonstrates that for harder tasks, the reduced modeling complexity of sequential modeling remains beneficial for longer, leading more stable and effective policy learning. This addresses the key request from the review process to validate the benefit of sequential modeling.
and DoorOpen/nward, DIMA (Sequential) maintains a clear performance
Table 2: Ablation study on ShadowHanddoote under varying data budgets (8 runs). Sequential modeling shows superior performance and lower variance in lower-data regimes.
comparing sequential (DIMA) vs. joint modeling
Method
100K Steps
150K Steps
200K Steps
250K Steps
300K Steps
Joint
234.1+20.6
238.6
246. 246.7
243.7+18.2
255.2+7.0 .0
Sequential (Ours)
251.8117.3
248.2
246.3
251.9
249.2
Table 3: Ablation study on complex Bi-DexHands tasks 300k steps (8 runs). The advantage of sequential modeling persists in more challenging environments.
Method
DoorOpenOutwarr @ 300K steps DoorOpenlnward @ 300K steps
Joint
302.5
235. 236.1
Sequential (Ours)
352.44440.
290.33:300
Sequential Modeling Retains Full Predictive Accuracy with Reduced Complexity. We conducted a direct empirical comparison of prediction error between sequential and joint modeling, independent of downstream policy optimization, to validate our design choice. We utilized IM-step replay dataset from HASAC, training both models on the first 500k transitions and evaluating their accumulated LI observation errors held-out set of the final 500k transitions (averaged over 100 segments). As shown in Table from $E.4 our sequential model achieves predictive accuracy that is statistically on par with the more complex joint modeling approach. Across three challenging Bi-DexHands tasks and at both H = 15 and extended H = 20 horizons, the error metrics are statistically indistinguishable when considering the standard deviations. It provides direct, quantitative evidence that the reduced modeling complexity (and associated benefits, e.g., computational efficiency) of our sequential approach comes no cost raw predictive capability. This result confirms that the added complexity of joint modeling is unnecessary for achieving high-fidelity predictions in these environments.
6 Conclusion
This paper presented a multi-agent world model motivated by the conceptual similarity between the progressive denoising process and the incremental reduction of uncertainty in predicting the global next state in MARL. Then, we propose DIMA that models multi-agent dynamics from a centralized perspective while achieving reduced complexity, seamlessly aligning the world model with the underlying MDPs to obtain more accurate predictions. To validate the efficacy of DIMA, we training scheme and conducted extensive experiments on
integrated it into the learningiin-imagnnale the MAMuJoCo and Bi-DexHands benchmarks. The results demonstrated DIMA's superior accuracy and robustness in predicting environment dynamics, as well as its ability to enhance sample efficiency
10

---

and overall performance. Despite its effectiveness, DIMA may encounter scalability challenges when applied large-scale multi-agent systems with hundreds of agents. To address this, we plan to explore grouping techniques to further extend DIMA's applicability and scalability in future work.
Acknowledgments
We would like thank Liyuan Mao for his insightful discussions and comments, and Jiaqi Peng for his generous help. This work is supported by the National Natural Science Foundation of China (Grant fo.62306242), the Young Elite Scientists Sponsorship Program by CAST (Grant No. 2024QNRC001). the Yangfan Project of the Shanghai (Grant No.2 1462000, and the General Research Fund (GRF 16209124).
References
[1] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream control: Learning behaviors by latent imagination. In International Conference on Learning Representa- tions, 2020. URL httpttts:://perrraatieeeee http://www.....ooorraaaaatttttA.
[2] Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021. URL https://openreviex .ttt:/oorummdddooabyyyy...
[3] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL
https://opeureviee..neee/ www.::eeeeeettoooa..dddddddeuuae
[4] Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer- world models are happy with 100k interactions. In The Eleventh International Conference on Learning Representations, 2023. URL /ttp:tttt:://peereevie
[5] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-60) 2020.
[6] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Picter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in neural information processing systems, 34:25476-25488, 2021.
[7] Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In International Conference on Machine Learning, pages 8387-8406. PMLR, 2022.
[8] Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for continuous control. In The Twelfth International Conference on Learning Representations, 2024. URL /Rtpppsp////peerreieie
[9] Pablo Hernandez-Leal, Bilal Kartal, and Matthew E. Taylor. A very condensed survey and critique of multiagent deep reinforcement learning. In Proceedings of the 19th International Conference Autonomous Agents and MultiAgent Systems, 2020.
[10] Thanh Thi Nguyen, Ngoc Duy Nguyen, and Saeid Nahavandi. Deep reinforcement learning for multiagent systems: A review of challenges, solutions, and applications. IEEE Transactions on Cybernetics, 50:3826-3839, 2020. doi: 10. 109/TCYB.BBBB... 2.2.0...227277777
[11] Frans A Olichoek, Christopher Amato, et al. A concise introduction to decentralized POMDPs, volume 1. Springer, 2016.
[12] Vladimir Egorov and Alexei Shpilman. Scalable multi-agent model-based reinforcement learning. In Proceedings of the 2 International Conference Autonomous Agents and Multiagent Systems, pages 381-390, 2022.
[13] Qihan Liu, Jianing Ye, Xiaoteng Ma, Jun Yang. Bin Liang, and Chongjie Zhang. Efficient multi-agent reinforcement learning by planning. In The Twelfth International Conference on Learning Representations, 2024. URL httt::///ppeeeerrreeeeeffrrrrrrn
11

---

[14] Yang Zhang, Chenjia Bai, Bin Zhao, Junchi Yan, Xiu Li, and Xuelong Li. Decentralized transformers with centralized aggregation are sample-efficient agent world models. arXiv preprint arXiv:2406. 15836, 2024.
[15] Jascha Sih--Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-2265. pmlr, 2015.
[16] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-685 2020.
[17] Yang Song. Jascha Sohl-Dicksteni Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-hased generative modeling through stochastic differential equations. arXiv preprint arXiv:201 I. 13456, 2020.
[18] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos J Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. Advances in Neural Information Processing Systems, 37:58757-58791 2024.
[19] Zihan Ding, Amy Zhang, Yuandong Tian, and Qinging Zheng. Diffusion world model: Fu- ture modeling beyond step-by-step rollout for offline reinforcement learning. arXiv preprint arXiv:2402.035700 2024.
[20] Dani Valevski, Yaniv Leviathan, Moab Arar, and Shlomi Fruchter. Diffusion models are real- time game engines. In The Thirteenth International Conference on Learning Representations, 2025. URL /ttL:tttss/:/pporvviiii
wwt/ forum?id P8pgeekn
[21] Tero Karras, Mīik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors. Advances in Neural Information Processing Systems, 2022. URL
httts::/ http://oooo...i...////o/oo
[22] Bei Peng, Tabish Rashid, Christian Schroeder de Witt, Plexre-Alexandre Kamienny, Philip Torr, Wendelin Böhmer, and Shimon Whiteson. Facmac: Factored multi-agent centralised policy gradients. Advances in Neural Information Processing Systems, 34:12208- 2021.
[23] Yuanpei Chen, Yaodong Yang. Tianhao Wu, Shengjie Wang. Xidong Feng, Jiechuan Jiang, Zongqing Lu, Stephen Marcus McAleer, Hao Dong, and Song-Chun Zhu. Towards human-level bimanual dexterous manipulation with reinforcement learning. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022. URL http://pprrrrieeeeeeeeffoorrrmmp
[24] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[25] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis.
Advances in neural information processing systems, 34:8780-8794, 2021.
[26] Zhengbang Zhu, Minghuan Liu, Liyuan Mao, Bingyi Kang, Minkai Xu, Yong Yu, Stefano Ermon, and Weinan Zhang. Madiff: Offline multi-agent learning with diffusion models. Advances in Neural Information Processing Systems, 37:4177 4206, 2024.
[27] Lei Yuan, Yuqi Bian, Lihe Li, Ziqian Zhang, Cong Guan, and Yang Yu. Efficient multi-agent offline coordination via diffusion- based trajectory stitching. In The Thirteenth International Conference on Learning Representations, 2025. URL https:: /ttttp:///ppnneviee wttpeeeiwwwwwewweeeeeee um! id-EpnZEZYDUT
[28] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[29] Andrej Karpathy. mingpt: A minimal pytorch re-implemennatioon of the openai gpt (generative pretrained transformer) training. attps://githhe cow /ttp//aaaaaahh/iinpTr 2020.
[30] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in
neural information processing systems, 30, 2017.
[31] Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantization: VQ-VAE made simple. In The Twelfth International Conference on Learning
Representations, 2024. URL
12

---

[32]
Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojeiech Marian Czarnecki, Vinicius Zambaldi, Max Jaderberg, Mare Lanctot, Nicolas Sonnerat, Joel Z. Leibo, Karl Tuyls, and Thore Graepel. Value-decompositton networks for cooperative multi-agent learning, 2017. URL https://
1etp:/110..000 05296
[33]
Tabish Rashid, Mikayel Samvelyan, Christian Schroeder de Witt. Gregory Farquhar, Jakob Foerster, and Shimon Whiteson. Qmix: Monotonic value function factorisation for deep multi-agent reinforcement learning, 2018. Utt.httpss//arziii org abs 1000 11485
[34]
Chao Yu, Akash Velu, Eugene Vinitsky, Jiaxuan Gao, Yu Wang, Alexandre Bayen, and Yi Wu. The surprising effectiveness of ppo cooperative. multi-agent games, 2022. URL https
//arxiv. wwwg/bb////aa 2103. 01955
[35]
Ling Yang. Zhilong Zhang, Yang Song, Shenda Hong. Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming- Hasann Yang. Diffusion models: A comprehensive survey of methods and .0099)
applications, 2024. URL https://arxiii org/
[36] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies an expressive policy class for offline reinforeement learning. In The Eleventh International Conference on Learning Representations. 2023. URL /tttp:://ppenreetee
[37] Linjiajie Fang, Ruoxue Liu, Jing Zhang. Wenjia Wang, and Bingyi Jing. Diffusion actor-critic: Formulating constrained policy iteration as diffusion noise regression for offline reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025, URL
https: petps://ont http///ppneeeireeeee
ceet/forum?idee forum??d:=1dv:A000K
[38] Wenhao Li. Efficient planning with latent diffusion. In The Twelfth International Conference on Learning Representations, 2024. URL
[39] Anurag Ajay, Yilun Du Abhi Gupta, Joshua Tenenbaum. Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211. 15657, 2022.
[40] Long Yang, Zhixiong Huang, Fenghao Lei, Yucun Zhong, Yiming Yang, Cong Fang, Shiting Wen, Binbin Zhou, and Zhouchen Lin. Policy representation via diffusion probability model for reinforcement learning. arXiv preprint arXiv:2305. /3122, 2023.
[41] Haitong Ma, Tianyi Chen, Kai Wang, Na Li, and Bo Dai. Soft diffusion actor-critic: Efficient online reinforcement learning for diffusion policy, 2025. URLhttps: http://axxvvv.//// /bs//
2502.00361
[42] Shutong Ding, Ke Hu, Zhenhao Zhang, Kan Ren, Weinan Zhang, Jingyi Yu, Jingya Wang, and Ye Shi. Diffusion-basee reinforcementt learning via w-weighted variational policy optimization,
2024. URL StRLsstts:://// //axxi http://///aarivooo//aab//////
16173.
[43] Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion.
The International Journal of Robotics Research, page 02783649241223668. 2023.
[44] Xiang Li, Varun Belagali, Jinghuan Shang, and Michael S. Ryoo. Crossway diffusion: visuomotor policy via self-supervised learning, 2024. URL
Improving
https://arxiv. org/ abs 01849
[45] Po-Chen Ko, Jiayuan Mao, Yilun Du Shao-Hur Muoo-Hua Sun, and Josh Tenenbaum. Learning to act
from actionless videos through dense correspondences. ArXiv, abs/2310.085577 2023. URL
httts::/app..semmtt www.seeeaaaannhoo.ocoeeo
[46] Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Plamning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
[47] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.0m 04333000..44...... 2023.
[48] Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. In Neural Information Processing Systems, 2023.
[49] Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, et al. T0: A lasiio-languaae-acttoon flow model for general robot control, 2024. URL https://arxiv. wtg/abs/2410.44166 2024.
13

---

[50] Songming Liu, Lingxuan Wu, Bangguo Li, Hengkai Tan, Huayu Chen, Zhengyi Wang, Ke Xu, Hang Su, and Jun Zhu. Rdt- lb: a diffusion foundation model for bimanual manipulation. arXiv preprint arXiv:2410.07864, 2024.
[51] Marc Rigter, Jun Yamada, and Ingmar Posner. World models via policy-guided trajectory diffusion, 2024. URL (tRts:hts:://// /ata:v/aavivvi/// www. www.ooogg////aaee.eee.oeee 08533
[52] Yilun Du, Sherry Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Joshua B. Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. In shirty-seventh Conference on Neural Information Processing Systems, 2023. URL https
//openreview.o (ttt//forrmm2--booBBMRRcc
[53] Sherry Yang. Yilun Du, Seyed Kamyar Seyed Ghasemipour, Jonathan Tompson, Leslie Pack Kaelbling, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. In The Twelfth International Conference on Learning Representations, 2024. URL https / //openreview wwtp:/eeeenoondsaf..ooooooo
[54] Siyuan Zhou, Yilun Du, Jiaben Chen, YANDONG LI, Dit- Yeung, and Chuang Gan. Robodreamer: Learning compositional world models for robot imagination. In Forty-first International Conference on Machine Learning, 2024. URL https://opeereviiee net/ KHj
[55] Kyunghwan Son, Daewoo Kim, Wan Ju Kang, David Earl Hostallero, and Yung Yi. Qtran: Learning to factorize with transformation for cooperative multi-agent reinforcement learning. org/ abs/abas///9 1905/ 05408.
2019. BtRL:https:///
[56] Jianhao Wang, Zhizhou Ren, Terry Liu, Yang Yu, and Chongjie Zhang. Qplex: Duplex dueling multi-agent q-learning, 2021. URL htt//axxvvv.rrrg hwtp://bs20000.0000 01062
[57] Jakob Foerster, Gregory Farquhar, Triantafyllos Afouras, Nantas Nardelli, and Shimon Whiteson. Counterfactual multi-agent policy gradients, 2024. URL https:////a wrw.:/.gggbaaa.b. abs/ 1705
08926
[58] Jakub Grudzien Kuba, Ruiqing Chen, Muning Wen, Ying Wen, Fanglei Sun, Jun Wang, and Yaodong Yang. Trust region policy optimisation in multi-agent reinforcement learning. In International Conference on Learning Representations, 2022.
het//oorum.ddEECCFXUTT.
[59] Dimitri Bertsekas. Multiagent rollout algorithms and reinforcement learning, 2020. URL
https: fitpss//axte htt:://axxvv.v.eov org/ 19900 00120
[60] Jianing Ye, Chenghao Li, Jianhao Wang, and Chongjie Zhang. Towards global optimality in cooperative marl with the transformation and distillation framework, 2023. URL https
//arxiv.org/ wwwp:/g/g//aa/////// aba/2207.2.. 221121177.111
[61] Zijie Zhao, Honglei Guo, Shengqian Chen, Kaixuan Xu, Bo Jiang, Yuanheng Zhu, and Dongbin Zhao. Empowering multi-robot cooperation via sequential world models. arXiv preprint arXiv:2509.1300 2025.
[62] Zifan Wu, Chao Yu, Chen Chen, Jianye Hao, and Hankz Hankui Zhuo Models as agents: Optimizing multi-step predictions of interactive local models in model-based multi-agent reinforcement learning, 2023. URL httpst/t:: //axxiv org/ abs/ /20///30.... 2
[63] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2024. URL
URL https: //axxiv. org/ abs/23/1.....0.////// 04104.
[64] Chao Li, Ziwei Deng, Chenxing Lin, Wenqi Chen, Yongquan Fu, Weiquan Liu, Chenglu Wen, Cheng Wang, and Siqi Shen. Dof: A diffusion factorization framework for offline multi-agent reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025. URL httptt:://ppprrrteees wwwpppprrrvieeeeeeeooooioooe www.ormmmmndnnnnndn111
[65] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033. IEBE, 2012.
[66] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020. URI. https //openreview wwt.wvwwwww...eeeeeee het:/forumiid-SIxcPDDHe
14

---

[67] Yifan Zhong, Jakub Grudzien Kuba, Xidong Feng, Siyi Hu. Jiaming Ji, and Yaodong Yang. Heterogeneous-ageen agent reinforcement learning. Journal of Machine Learning Research, 25(32):
1-67, 2024. URL http://jmlr wwwpppppp..vv....//...
html
[68] Jiarong Liu, Yifan Zhong, Siyi Hu, Haobo Fu, QIANG FU, Xiaojun Chang, and Yaodong Yang. Maximum entropy heterogeneous-ageen reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024. URL
id=tmqOhBC4a5
[69] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference machine learning, pages 1861-1870. Pmir, 2018.
[70] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 1xXv:1666600415,, 2016.
15

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The paper's contributions and scope concretely claimed in the abstract and introduction. Our key contributions are summarized in
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitation of our approach at the end of 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
16

---

We propose formal and concrete assumption in $3.1 and provide detailed proof of our Theorem
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We provide the training details and hyperparameters in G
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
17

---

Answer: [Yes]
Justification: First, we use online Reinforcement Learning setup, which only requires available simulators. The evaluated environments in our experiments reported in $5.1 As for the code of our approach, we promise we would release our code when this work gets accepted.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
seelicc/guies//// http///euumiissoonn..... ooloyyyooyyyyooooooooo for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ htips.cc/public//uiic. htth///oeeeeeiiiseeeeeee coeeinsssoooooiiii
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity. the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: We provide the training details and hyperparameters in G And we describe the chosen tasks and evaluation protocol in $5.1
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We plot the mean episode return curves in $5.2 with shaded areas represent- ing the deviation of episode returns. And we report the number of seeds we use in our experiments in $5.1
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
18

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We report the computational resources we used
G
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: We read the NeurIPS Code of Ethics, and make sure that our research conducted in the paper conform with the NeurlPS Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the societal impacts in
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
19

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite works of related assets.
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
20

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
21

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: Our method does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https://nen) /ttp://////upp.ppeeee. /t//Confereence//222222 for what should or should be described.
22

---

A Proof of Theorem 2
Proof. Given the log-likelihood of the global state transition, we have the following:
log P(8 ) log / p(( 8 (1)((0)
18: af
ds
(1)(())
log /
p( 841 81, a
(1)(m))
(1) (1))t))
14(5+11 88,, at s141) (2)(6) ds
(
(1)
sec ag 1.n s(+1)
p(5141, 1441
st ) ((s1t1 (1)(() 1st, at 8t+1).
(1)(())
= =log EngEE (11 sc-1)
p(s 8411 (8, at d:n)
(1)(())
2 E (200) 20022100 seat me44)
log [($141 18t, .041)))+1).
(11)
m041)
(1)=(m)
where the last inequality results from Jensen's 's inequality. Under the definition and property of the conditional Markovian forward diffusion process á in Eqs. 2 and Assumption we can rewrite
Eq. II as follows,
p(s
p(s si+11 St, at))
(8-11) (k)
log P(s 1:N
E
log
(12)
((
m041)
Th= (( 1 84+1
(k) ((:-1))
where we denote 8141 == 8- Then, RHS Eq.
(0)
can be further simplified,
p(5
(m) IT p( (k-1) (k)
1141 81, ((,
E (2)(0))))
log
(0 (e44s)
Il (($22115144 11
(k)
p(
(n) (0) | (::11 St: at)
n
p p( 1 s1 St, at)
(k-1) (k)
log
log
((84118824)
@ (
(k) (1-1))
8141
p(s (n) (0)
s11 si+ St, at)
p p
15 18
a
E (B (941)
log
log
((5 (s
(1) (0)
ak 1)
1
(0) 10( 1.
k=2
2 (
p(s (n)
)p(s (0)
8411 St, at)
4(s (11 (41
(0)
p( (k- 1) (k)
s411 8t, at)
i+1
E
(0 (1:0)) (1441)
log
log
log
(1 (1)
(
((s (a) (0)
(111
( |
8411 S 8(11)
(k) (0)
p(s (n)
(p(())((()))))))) 1 Sc, at)
(0) (1)
p
(k- 1) (k)
1+1 8c, at)
=E ()100 (8441)
log
log
((s (n) (1))
(( 1 $411
(k) (0)
km
Therefore, the evidence of dynamics transition can be bounded as follows:
log P( I.N) > E ((o (1) 12 [log
(0) [log] [log p(s 19 St, af)l DKL (
aall (2)) [KL( (k- 1s 11
(K-1))
1 at s1))
(13)
As shown by 25 the conditional forward diffusion process q behaves identically to the unconditional one q. Therefore, we can substitute the q with the q Eq. [13). concluding our proof.
B EDM Preconditioners and Noise Scheduler
To keep input and output signal magnitudes fixed the same scale and avoid large variance in gradient magnitudes on per-sample basis, Karras al. 21 introduced the following preconditioners for
23

---

normalization and re-scaling output to stabilize and improve the training dynamics of the network:
e [TT))
(14)
duna
Cout
(15)
(T(T)) + data
Cuvise 1og(((7))
(16)
Skip
data + (7)
(17)
where data 0.5 our experiment hyperparameter setup. The noise scheduler for training the diffusion model follows the same design in [21], described as follows:
(T) log(o) ~ ((PPmean, P
(18)
where Pmman = -0.4 and Pad 1.2.
C
Behavior Learning Details
B4] to train both the actor and critic inside Inspired by the success of MARIE [14], we adopt MAPPO
the imaginations of DIMA. key distinction from MARIE is that our model explicitly predicts the global state, enabling seamless integration with CTDE techniques well as actor-critic architectures commonly used in model-free MARL.. Therefore, we implement both the actor & and critic & with two 3-layer MLPs together with ReLU activation and Layer Normalization, respectively, Similar to off-the-shelf CTDE model-free MARL algorithms, we adopt actor parameter sharing across agents.
Critic loss function. We utilize A-return in DreamerV1 [ which employs an exponentially weighted average of different k-steps TD targets to balance bias and variance the regression target for the critic. Given an imagined trajectory {š,, 6!.". aa fr, 30 over all agents, A-return is calculated recursively as,
{
(1 (1-A)) ()-)) + (1+4)) if 6<H
(a() = ft + +
(19)
if 1=
The objective of the critic & is minimize the mean squared difference Le with X-returns over imagined trajectories, as
Σ (Ve(+)-ss(Vv(55) sV(((())))
Le = Exp
(20)
where sg(-) denotes the soopggradient operation. We optimize the critic loss with respect to the critic parameters & using the Adam optimizer.
Actor loss function. The objective for the actor T5. (-1ot) = (- is to output actions that maximize the prediction of long-term future rewards made by the critic. To incorporate intermediate rewards more directly, we train the actor to maximize the same A-return that was computed for training the critic. In terms of the non-stationarity issue in multi-agent scenarios, we adopt PPO updates, which introduce importance sampling for actor learning. The actor loss function for agent i is:
C -
-E I
Σ min (l(()) 1=0
+ +1HH(
(21)
where r(())= = T///
Voia the policy ratio and A, = sg(A(()) - Ve())) is the advantage. Unlike MAPPO, we choose not design agent- speent- global states, as such designs are overly hand-crafted and inject task-specific human priors, which undermines the generality and soundness of the approach. Instead, we retain the environment's original agent-agnostic global state shared among all agents, and feed it into the value function Ve- As result, the estimated advantage function A, is also shared across all agents during actor updates. We optimize the actor loss with respect to the actor parameters sb using the Adam optimizer. Overall hyperparameters are shown in Table 4
24

---

Table 4: Behaviour learning hyperparameters.
Hyperparameter
Value
Common Imagination horizon (H)
15
À
0.95
Clipping parameter e
0.1
MAMuJoCo
Discount factor y
0.99
0.001
Bi-DexHands
Discount factor ?
0.95
0.01
D Illustrations of Experimental Environments
ShadowHandPor
Shadow Cam
Shadow Deeeawaannnn
HalfCheetah
6-agent partitioning
(4)
MAMuJoCo
Figure 8: Illustrations of the experimental environments our work. Left: Visualizations of three Bi-DexHands tasks: removing pen cap, opening bottle cap, and opening door inwards. Right: Visualization of 6-agent partitioning w.r.t. HalfCheetah MAMuJoCo.
25

---

E
Additional Results
E.1 Additional Experiments: Detailed Returns of All Methods on MAMuJoCo and Bi-DexHands
Table 5: Comparison of final episode returns across MAMuJoCo and Bi-DexHands benchmarks. We report the mean final episode return and standard deviation over 4 random seeds. DIMA consistently outperforms all baselines across all chosen tasks on both MAMuJoCo and Bi-DexHands. The best result per task is highlighted in bold and shaded in blue color, while the second-bes best is underlined.
Methods
Tasks
Steps
DIMA (Ours)
MARIE
MAMBA
HASAC
НАРPP
MΑPPO
MAMuJoCo
Ant-2x4
4881
4471 #553
131 4476
1844
1716-449
859
Ant-4x2
4766
1173
16184931
850g128
1917
253
854+41
HalfCheetah-2x3
IM
6370
4045 #227
2813 1813
249911061
2628
899
3196175
HalfCheetah-33x
6175-222
2380 1144
3029
28722+800
3402 -317
2936
HalfCheetah-6x1
5643 1163
1738 1213
1848
2044
2939
11113
925
Walker2d-2x3
3329 1106
2822
124
1135+:210
1007 =282
752
Walker2d-3x2
4084
0044:40
466
958
932
1004 480
Bi-DexHands
BottleCap
259.9144
203.815.2
210.9
100.7
104.0
DoorOpenlnward
300K
290.4
225.0
246.3
30.7
65.846.9
DoorOpenOutward
367. 367.14100 19..4
177.4443.1
221.9-7.3
58.8
96.448.5
BottleCap
24.4:11.4
1.340.4
0.0.0.0
0.0.0.0
0.0
Table 6: Policy learning performance (final return) of DIMA with different imagination horizons (H 15 H 25).
Scenarios
DIMA w/ H = 15
DIMA w/ H = 25
Ant [4x2] (4-agent)
4766
4328+1055
HalfChoetah [6x1] (6-agent)
56431163
6310
E.2 Additional Experiments: Policy Learning with Longer Imagination Horizons
We investigated the downstream impact on policy learning when utilizing the extended imagination horizon (H 25) during the whole training process. The final policy returns presented in T The results are environment-dependenn and provide critical insights. On HalfCheetah [6x 1}, training with longer horizon (H 25) yields substantial performance improvement. This strongly suggests that when DIMA's long-range predictions are stable, the policy optimizer can successfully leverage these extended rollouts discover more complex and far-sighted strategies, leading to superior returns. Conversely, on Ant [4x2], we observe a performance degradation accompanied by significantly higher variance. This indicates that for this specific task, while comparatively lower than baselines (as shown in Table []. the accumulated prediction errors in H = 25 rollouts are still sufficient introduce noise that misleads the policy learning. This phenomenon underscores both the practical potential of DIMA in scalable model-basee MARL (demonstrated on HalfCheetah) and highlights key challenge: ensuring that the predictive accuracy is robust enough to provide stable gradient signal for policy learning across all task types.
E.3 Additional Experiments: Comparison on Training Compute
We measure the training time and GPU memory usage all evaluated model-based MARL methods, including proposed DIMA. All experiments were conducted using single NVIDIA RTX 4090
26

---

GPU to ensure fair comparison across methods and tasks. As shown in the table below, DIMA substantially more efficient than MARIE both training time and GPU memory usage, and is comparable to or slightly more efficient than MAMBA.
Table 7: Comparison on consumed computational resources over 2 test scenarios.
Methods
Training Time
Usage of GPU Mem
Anf-2x4
DIMA
1d 19h
3.10 GB
MARIE
3d 17h
14.33 GB
MAMBA
1d 1h
2.74 GB
Ant-4x2
DIMA
1d 19h
3.15 GB
MARIE
3d 17h
3.10 GB
MAMBA
1d 1h
4.37 GB
E.4
Ablation: Sequential vs. Joint Modeling Prediction Accuracy
Table 8: Ablation study comparing the cumulative L1 observation errors of sequential vs. joint modeling. Models were trained on 500k transitions and evaluated on 500k held-out set. Sequential modeling achieves statistically indistinguishable prediction accuracy. validating its design.
Task
Method
Obs L1 Error @ H 15
Obs L1 Error @ H 20
Sequential (Ours)
5.33344..23
7. 7.081=0.325
DoorOpenOutward
Joint
5.345.10.267
7.092
Sequential (Ours) Joint
5.563
7.447
DnwrOpenInward
5.565+40.22
7.453
Sequential (Ours) Joint
6.667 11.47417764 7664
8.936 8.936_2.328
Pen
6.676+1.762
8.947 22.322
E.5
Experiment Details:
Imagination Evaluation across Different Conditioning Orders
To evaluate DIMA's imagination under different conditioning orders on the 2- 2-agent Ant [2x4] task, collect 10 episodes by using the final policy induced by algorithm, and randomly sample 100 trajectory segments to form our trajectory segment dataset. For each segment, we generate imagined rollouts using DIMA with different action conditioning orders.
As the EDM framework decouples inference- time sampling from training, the number of denoising steps need not match the number of agents. Thus, set the number of denoising steps equal to 4, i.c., twice the number agents. Letting the agent set be {1, 2]. we consider three conditioning orders: (i) random order: (2, 1, 1,2), (ii) ascending order w.r.t. agent id: (1, 2, and (iii) descending order w.r.t. agent id: 22 1.)
To provide quantitative evaluation in Figure (right), we compute the L1 error per observation dimension at each timestep between the 100 sampled trajectory segments and their corresponding imagined rollouts, and accumulate the errors over the prediction horizon. All observation LI errors are averaged across 2 agents.
Overview of DIMA with Learning in Imaginations
We summarize the overall training procedure of DIMA paired with learning in imaginations in Algorithm B below. We denote D the replay databuffer which stores data collected from the real environment.
27

---

G Training Details and Hyperparameters
G.1 Model Architecture Details
State decoder. To enable decentralized execution of policies trained within DIMA's imagination rollouts, each agent must make decisions based solely on its local observation rather than the shared global state. To support such policy learning, we introduce a necessary state decoder that maps the global state &t into the corresponding joint local observations o!!
Due to our online model-based MARL setup, the state decoder must be continually updated under non-stationary data distribution which also shifts continually. Using vanilla MLP as the state decoder this setting may lead to issues such as overfitting or mode collapse. To mitigate these risks, we incorporate additional regularization into the decoder design by adopting Vector Quantized Variational Autoencoder (VQ-VAE) 30), which enforces a compact latent codebook representation via vector quantization. Among various VQ- VQ-V variants, we choose Finite Scalar Quantization (FSQ) [3] as our final implementation as it removes any auxiliary losses and achieves remarkably high codebook utilization,, which indicates its strong and effective regularization.
Our implementation is based on the open-source repository: https: thub com wwt.moccc///c/oec.ccoooooo hwt/draiidddr//// vector- cestorr-quatiin quantize serranttzz-ppyrrcch We simply build the encoder E,g and decoder D.p as MLPs to deal with continuous non-vision global states and joint local observations. The decoder designed with the same hyperparameters as the encoder. The loss function for learning the autoencoder is follows:
( Dp) = of
||lo 1m (E)((( + sg (f
E (s:))))
(22)
where f is a bounding function such that i-th channel/entry in & = round( (E, takes one of L, unique values (here f : + [ tanh(z) for i-th channel/entry) and round is the operation to map real-valued inputs to the nearest integers. Therefore, we have an implicit codebook C with |C| = - Li. After training the VAE, our state decoder can be expressed by gup (op (of:"(ss) = (( The hyperparameters are listed in Table
Diffusion model for dynamics modeling. We use the 1-D variant adapted from the U-Net 2D in DIAMOND [18] as the backbone of the diffusion model Dg. To predict the next state 8t+1. the diffusion model Do is initially conditioned on the current global st, joint action af!" and the diffusion prediction accuracy,
time To improve next-global-stae
empirically augment the temporal context by additionally incorporating the last global states and joint actions, extending it from & and 1m to 5;-2:t 2:t and amtn Note that the effect of sequential denoising is confined to the joint action a? conditioning the current timestep 1, and does not extend the past joint actions.
Inspired by the success of DIAMOND, we directly adopt the same conditioning mechanism in DIAMOND, and use temporal stacking for global state conditioning and adaptive group normalization for joint action and diffusion time conditioning. The hyperparameters are listed in Table 9
Transformer as reward and termination model. The Transformer for predicting the reward and termination is built upon the implementation of minGPT 29 Given a fixed imagination horizon H, it first takes a sequence of length 2H composed of global states and joint actions (..., St, I:n ...... and encodes every single global state and joint action into d, ddddmensiionll embedding tensor via 2 separate encoding functions. Then the sequence tensor of shape 2H x de is forwarded through fixed Transformer blocks. Finally, the Transformer predicts reward and termination via two separate 3-layer MLP heads (relsgg,, a Iin) and fo(Ti) a lin), respectively. In general, the loss function is described by
Ls E Σ"= H -lg fo (relsss, a a
log fa(") a
(23)
But in practice, we optimize the reward prediction with a smooth LI loss function and the termination prediction with cross-entropy loss function. The hyperparameters are listed in Table 9
G.2 Computational Resources Used for Training
All our experiments including the evaluation of chosen baselines are run on a machine with a single NVIDIA RTX 4090 GPU, 24-core CPU, and 256GB RAM.
28

---

G.3 Baseline Implementation Details
In our experiments, we reran and evaluated all baseline methods. To ensure faimesss for comparisons.
we followed the optimal hyperparameters providd in their official implementations, listed below:
MARIE: httppp///ggttub..oo////eeeeeaaap
MAMBA: http:://gtthuuo.cmm/b-aaaaaaaaa
HASAC, HAPPO and MAPPO: https://githut: wtt.://w..coo///--MA.R.A
G.4 DIMA hyperparameters
We list the hyperparameters DIMA paired with learning in imaginations in Table
H Broader Impact
Our work introduces DIMA, diffusion innsiired multi-agent world model that significantly improves sample efficiency in cooperative multi-agent control environments. By enabling more faithful imagined rollouts, DIMA can accelerate the development of complex autonomous systems such as multiple real robots coordination, traffic management, and energy-efficient buildings-thered reducing real-world trial costs. However, these capabilities also carry potential risks: misuse of high- fidelity neural simulators for adversarial planning could exacerbate privacy and security concerns.
29

---

Algorithm 1: DIMA paired with learning imaginations
Procedure training_loop():
for epochs do collect_experiont (steps_collect)
for steps_ state_ decoder do
L cicaate__ttte_deccteee ()
for steps_ dlifuusionmmodee do
update_ seaattteddffusiioomdeeeee ()
for steps_ prrps_reward_rrr crrward_en__aae monde do
update__eward___ mondell ()
for steps_actorccrilicc do update__ctor_criticc)
Procedure colleett_eperriecceen
80: 00% f- anv.reset! ()
for | 0
do
Sample at ~ T (f) V agent
St+1, 4:n 0-11 Tt,
satp
DU (s++1, 0411; 1 To 7t)
if then
86-1, 1+1 env reset
Inn
Procedure update_ sttteetttttteeeeeeeeeeeeeeeeeett decoder ():
Sample state-observation pair (st, of D
Compute Crso Eq. (22)
Update State Decoder 9ç
Procedure
Sample sequence $ --t1::... St. a, s6-1) D
Sample log(a) [MPP Mean: add and get = a since (TT) =
Sample 8T41 ~ [/(=-1100
Sample chosen agent id Uniform (1, ,.
(0) Do D0(8 81 1+ Compute *+1
1:m d-L+1
temporal context
Compute loss C(0) = | 8141 in Eq. 9 Update Diffusion Model Dg
Procedure reward_ end mod_model()) ():
Sample sequence (8t, atin l:m Tt, .t, 8t+ 11 11 a, 1:n for to + H do
L Compute F; fol (P: a ai and 9 fol a 22r)
8- Compute L. = Σ!# i=t
CrossEntropy 7)) + SmoothLI (f;, corresponding to
Eq. (23
Update Reward and Termination Model fs
Procedure update_acao poodate_atto___e wwt.aacctttticcaoeeeoaaaaaa critiic)::
Set the joint action condition order p (i1, ,,,
Sample starting point (st-L+1, al 1+ a. L+1 1....... 8t, o. ~ D of the imagination
Let 6.:n = 1.:11
for
to
do
Sample a: ~ (o( V agent
Sample the reward f; and the termination 2/2 with fo
Sample the next global state $44 by iteratively denoising with Do and p
Sample the next joint observation state alm with 9,
Update actor TH. and critic Ve via Le and Ct over imaginations {s,, 6!:n a In fr, 11+H--
1
30

---

Table 9: Architecture details.
Hyperparameter
Value
State Decoder (gpp
MLP layers
3
Hidden size
512 GELU [70]
Activation
FSQ Levels L;
[8, 6, 5]
Diffusion Model (Dg)
Global state conditioning mechanism Joint action conditioning mechanism
Temporal stacking
Adaptive Group Normalization Adaptive Group Normalization
Diffusion time conditioning mechanism
Residual blocks layers
[2, 2, 2]
Residual blocks channels Residual blocks conditioning dimension
[64. 64, 64] 256
Reward and Termination Model (fz)
Embedding dimension de Transformer block layers
256
6
Attention heads Embedding dropout
4
0.1
Attention dropout Residual dropout
0.1
0.1
31

---

Table 10: Hyperparameters for DIMA.
Hyperparameter
Value
Batch size for State Decoder training
256
Batch size for Diffusion Model training
64
Batch size for Reward and Termination Model training
128
Optimizer for State Decoder
Adam W
Optimizer for Diffusion Model
Adam W
Optimizer for Reward and Termination Model
Adam
Optimizer for Actor & critic
Adam
Learning rate for State Decoder Learning rate for Diffusion Model
0.0003
0.0001
Learning rate for Reward and Termination Model Learning rate for Actor & critic
0.0001
0.0005
Gradient clipping for State Decoder
10
Gradient clipping for Diffusion Model
1
Gradient clipping for Reward and Termination Model
10
Gradient clipping for Actor & critic
10
Weight decay for State Decoder
0.01
Weight decay for Diffusion Model
0.01
Weight decay for Reward and Termination Model À for X-return computation Discount factor Entropy coefficient
0.01
0.95
see Table
Table
Buffer size (transitions)
2.5 x 10
Training steps per epoch
200
Training steps per epoch for policy learning
4
200 in MAMuJoCo
Sampling Environment steps per epoch
500 in Bi-DexHands
PPO epochs
5
PPO Clipping parameter
0.1
Number imagined rollouts
600
Imagination horizon H
15
Diffusion sampling solver
Euler
[2 I [ <
Number of denoising steps
i > 2
32


Reviews:
[{"review_id": "P9zKn7Mmat", "reviewer": "Reviewer_Y1VU", "summary": "The 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.", "questions": "Q1: 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?", "limitations": "Yes - additional details would be helpful", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \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.", "was_revised": true, "final_justification": "The paper provides a novel perspective and approach to support training of multi-agent reinforcement learning policies.", "rebuttal": "Dear Reviewer Y1VU,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**Weakness 1**: The scalability of the approach is not sufficiently tested.\n\n**Response**: We fully agree that evaluating the scalability of the world model—especially under longer imagination horizons—is essential, due to the well-known issue of compounding errors in predicting future observation with an auto-regressive manner. These errors can lead to significant deviations from the true dynamics over time, and addressing them is key to ensuring reliable imagined rollouts for downstream policy learning.\nAs we provide a novel perspective to reduce the modeling complexity on learning multi-agent dynamics, it can implicitly enable DIMA to more accurately capture the underlying dynamics while demanding less data.\n\nTo further assess the scalability of DIMA, we conduct additional evaluations with a longer prediction horizon of $H=25$, measuring accumulated observation reconstruction errors over time. Here we still select the scenario *Ant-v2-2x4* for the experiment. The results are summarized below:\n| Methods | Obs Accumulation Errors @ $H=15$ | Obs Accumulation Errors @ $H=25$ |\n| :-: |:-:|:-:|\n| DIMA |**2.42** ± 0.93|**4.32** ± 1.44|\n| MARIE |3.54 ± 1.60|6.62 ± 2.68|\n| MAMBA |3.98 ± 1.54|6.85 ± 2.51|\n\nWe also evaluate policy learning performance under the extended imagination horizon, comparing DIMA with $H=15$ and $H=25$. The results are reported as follows.\n| Scenarios | DIMA w/ $H=15$ | DIMA w/ $H=25$ |\n| :-: |:-:|:-:|\n| Ant 4x2 (4-agent) |4766 ± 450| 4328 ± 1058 |\n| HalfCheetah 6x1 (6-agent) |5643 ± 163| 6310 ± 335 |\n\n**In conclusion**, these results suggest that DIMA not only exhibits lower compounding error trends over time compared to other baselines, but also sustains effective policy learning under a longer imagination horizon, thereby demonstrating its scalability.\n\n**Weakness 2**: The main paper lacks details or intuition about the architecture of the diffusion model. How does it handle (discrete?) finite actions and states?\n\n**Response**: Thank you for the constructive feedback. We provide detailed descriptions of our diffusion model architecture in Section G of the appendix (available in the supplementary materials). Briefly, our model adopts a 1D variant of the 2D U-Net architecture used in DIAMOND [1]. The global state is first pre-normalized by a running mean and standard deviation computed throughout training, and corrupted by adding noise for inputting the diffusion models.\nRegarding the joint actions, in order to achieve a sequentially conditioned reverse diffusion process, we employ an agent-wise mask to mask out the unused agents, and use a Transformer encoder to extract the feature from the action information in temporal context $a_{t-k:t-1}^{1:n}$ and current action $a_t^i$ of selected agent $i$. These actions are first encoded into embeddings via an MLP, and then processed by the Transformer encoder. The resulting representation is used as a conditioning input to the diffusion model.\nGiven that all evaluated environments feature continuous actions, we employ an MLP to project the raw action inputs into embeddings. For tasks with discrete action spaces, an embedding layer would be a natural alternative. Regarding the discrete state, no modification is required for our method.\n\n**Weakness 3**: The 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.\n\n**Response**: Thanks for the insightful comment.\n1. The focus of our paper is on evaluating sample efficiency brought by a learned world model. **To this end, we intentionally adopt a low-data regime to assess how effectively different methods can learn under constrained data budgets. This setting is practically relevant for scenarios where data collection is expensive or limited, and where fast convergence is crucial.** Note that the \"sample efficiency\" within the context of RL refers to the ability of a policy to achieve higher performance with fewer samples (transitions) during training, indicating more effective sample utilization. \n2. To address the reviewer’s concern regarding final performance at convergence, we additionally run HASAC for up to $10^7$ steps on selected scenarios. Specifically, we focus on the MAMuJoCo suite (e.g., Ant 2×4, Ant 4×2, and HalfCheetah 6×1), where performance continues to improve beyond 1M steps. In contrast, for the Bi-DexHands environments, HASAC exhibits early convergence to suboptimal policies around 300K steps (see Figure 5), making further training less beneficial.\n\nThe extended results are summarized in the table below (bold for the highest value, underlined for the second highest).\n| Scenarios | HASAC @ 1M steps | HASAC @ 5M steps | HASAC @ 10M steps | DIMA @ 1M steps |\n| :-: |:-:|:-:|:-:|:-:|\n| Ant 2x4 |1344 ± 282|4830 ± 757|**5253** ± 617| $\\underline{4881}$ ± 756|\n| Ant 4x2  |850 ± 126| $\\underline{5023}$ ± 753|**5338** ± 667|4766 ± 450|\n| HalfCheetah 6x1 |2044 ± 110|5320 ± 2120|**6419** ± 2564| $\\underline{5643}$ ± 163|\n\nAccording to the theoretical analysis in [2], policies optimized under an approximate dynamics model would exhibit a certain suboptimality compared to those trained with true dynamics. Although our policy is learned in a learned dynamics model rather than the ground-truth environment, its performance remains competitive. Meanwhile, as shown in Figures 4 and 5, the policy learning curve in DIMA continues to improve and has not yet converged, indicating potential for further gains.\n\nFurthermore, when comparing DIMA's performance at 1M steps to HASAC's performance at 5M steps, we observe that DIMA achieves higher returns with 5x fewer samples, highlighting its strong sample efficiency.\n\n**Weakness 4**: The computational cost and training time of the world-agent are not investigated (also in comparison wrt similar solution).\n\n**Response**: We measure the training time and GPU memory usage of all evaluated model-based MARL methods, including our proposed DIMA. All experiments were conducted using a single NVIDIA RTX 4090 GPU to ensure a fair comparison across methods and tasks.\n\n| Methods | Training Time | Usage of GPU Mem |\n| :-: |:-:|:-:|\n| DIMA in Ant 2x4  | ~1d 19h |3.10 GB|\n| MARIE in Ant 2x4 | ~3d 17h |14.33 GB|\n| MAMBA in Ant 2x4 | ~1d 1h  |2.74 GB|\n| DIMA in Ant 4x2  | ~1d 19h |3.15 GB|\n| MARIE in Ant 4x2 | ~3d 17h |15.73 GB|\n| MAMBA in Ant 4x2 | ~1d 1h  |4.37 GB|\n\n**Question 1**: 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\n**Response**: Thanks for the insightful questions. We address each of them in turn:\n1. As our DIMA predicts the next global state based on a fixed window of past global states, it can predict the future over an infinite horizon in an auto-regressive manner.\n2. With an increasing state and action space, learning an accurate world model becomes more challenging especially in continuous control settings. For example, in the MAMuJoCo suite, each agent has a 2–4 dimensional action space, while in Bi-DexHands, each agent’s action space is 26-dimensional—over 6× larger. This makes modeling the transition dynamics significantly harder. Despite this, DIMA remains robust and performs well in Bi-DexHands, whereas other methods (e.g., MARIE) fail to scale and encounter severe out-of-memory (OOM) issues due to its computational overhead.\n3. We refer the reviewer to our **response to Weakness 1** for a detailed evaluation. Briefly, we tried to analyze the prediction error accumulation by sampling trajectories using policies from all methods.\nSince learning the world model is tied to a progressively improving policy both in MARIE, MAMBA and DIMA, we separately use their final policies to sample 5 episodes for fairness. We then compute L1 errors per observation between 100 trajectory segments randomly sampled from all 15 episodes and their imagined counterpart. As most of these trajectories are not induced by DIMA’s own policy, they can test how well DIMA generalizes to unseen cases. Results indicate that DIMA significantly outperforms the baselines not only at the training horizon of $H=15$, but also when evaluated beyond the training horizon.\n\n\n**Question 2**: what are the performance of the tested models using more steps / at convergence ?\n\n**Response**: We report the performance of MAPPO, HAPPO and HASAC at 1e7 steps as follows (bold for the highest value, underlined for the second highest). But we would like to emphasize again that **our work focuses on the low data regime**. As shown in Figure 4, our method has not yet fully converged within the original data regime. And we are extending the training of our method to 1e7 steps and would report it as soon as possible.\n\n| Scenarios | MAPPO @ 10M steps | HAPPO @ 10M steps | HASAC @ 10M steps | DIMA @ 1M steps |\n| :-: |:-:|:-:|:-:|:-:|\n| Ant 2x4  |4597 ± 433|4823 ± 775| **5253** ± 617| $\\underline{4881}$ ± 756|\n| Ant 4x2  |3375 ± 425|3792 ± 1166| **5338** ± 667| $\\underline{4766}$ ± 450|\n| HalfCheetah 6x1 |**7141** ± 180|5453 ± 217|$\\underline{6419}$ ± 2564|5643 ± 163|\n\n**Question 3**: What is the computational cost of training the world model (e.g., time), compared to similar approaches?\n\n**Response**: We refer the reviewer to our **response to Weakness 4** for a detailed evaluation.\n\n**Question 4**: Does the authors use “§” to refer to the appendix?\n\n**Response**: Yes, we would add one sentence in the revised manuscript to clarify this symbol usage.\n\n\n**Reference**\n\n[1] Alonso, Eloi, et al. \"Diffusion for world modeling: Visual details matter in atari.\" NIPS (2024): 58757-58791.\n\n[2] Janner, Michael, et al. \"When to trust your model: Model-based policy optimization.\" NIPS (2019)."}, {"review_id": "t8apF7T7FM", "reviewer": "Reviewer_DfzB", "summary": "- 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).", "questions": "- 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.", "limitations": "- 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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**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)", "was_revised": true, "final_justification": "The authors have addressed most of my concerns. I believe the paper makes a meaningful contribution, and I maintain my evaluation as a borderline acceptance.", "rebuttal": "Dear Reviewer DfzB,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**Weakness 1**: While I appreciate the idea of reducing the information space from ∣S∣×∣S∣×∣A∣ to ∣S∣, I find this relatively simple and not very innovative.\n\n**Response**: Thank you for the comment. While our approach may appear relatively simple, we believe that its strength lies in this simplicity. By reducing the modeling space from $|\\mathcal{S} | \\times |\\mathcal{A}|^n \\times |\\mathcal{S}|$ to $|\\mathcal{S} | \\times |\\mathcal{A}| \\times |\\mathcal{S}|$, we are able to significantly lower the complexity of multi-agent dynamics modeling without sacrificing performance (see our ablation study in Section 5.4).\nThis is achieved through a principled formulation of multi-agent prediction as a sequentially conditional denoising process, which enables centralized modeling without requiring explicit agent communication. Despite its simplicity, DIMA yields strong empirical gains across challenging benchmarks and achieves higher sample efficiency than prior model-free and model-based methods. We view this as an instance of a simple yet effective design, which we believe is valuable for the community.\n\n**Weakness 2**: The paper mentions MADiff and MADiTS as relevant diffusion-based works but does not sufficiently compare with them beyond qualitative differences.\n\n**Response**: Thanks for the constructive comment. As stated in Section 4 of our manuscript, MADiff [1] and MADiTs [2] differ from our method not only in their design but also in the specific research questions they aim to address. Consequently, they cannot serve as suitable baselines for our work. Specifically,\n1. MADiff leveraged a diffusion model as a goal-conditioned trajectory generator and generated actions via an inverse dynamics model, which can be viewed as a multi-agent extension of DecisionDiffuser [3].\n2. MADiTs used diffusion models for offline trajectory stitching, generating synthetic transitions to augment the original offline dataset.\n\nCrucially, both MADiff and MADiTs are **offline model-based MARL algorithms** that rely on high-quality offline data (e.g., near-optimal trajectories) to perform well. In contrast, **DIMA operates in the online model-based MARL setting**, where the data buffer is initialized empty and the diffusion-based world model is learned entirely from scratch during training, without any access to pre-collected expert demonstrations.\nFurthermore, MADiff and MADiTs use diffusion models to model goal-conditioned trajectory distributions, i.e., $p(o_t^{1:n}, o_{t+1}^{1:n}, ..., o_{t+H}^{1:n} | g)$, where the diffusion model serves primarily as a planner conditioned on a high-level goal $g$. In contrast, DIMA directly models the multi-agent transition dynamics, i.e., $p(s_{t+1} | s_{t-k:t}, a_{t-k:t}^{1:n})$ and $p(o_t^{1:n} | s_t)$, which is essential for agents' interactions in imagined environment.\n\n**Due to these fundamental differences in training setting (offline vs. online), modeling objective (planning vs. dynamics modeling), and usage of the diffusion model, MADiff and MADiTs are not directly comparable to DIMA under our evaluation protocol focused on learning in imagination.**\n\n**Weakness 3**: While MAMuJoCo and Bi-DexHands are good benchmarks, adding an additional environment (such as SMACv2) could demonstrate the generality of the approach for more discrete-action cooperative tasks.\n\n**Response**: Thanks for the suggestion. The primary contribution of the proposed DIMA lies in introducing a novel multi-agent world model that achieves improved accuracy and reduced complexity. In light of this, we conducted extensive experiments comparing DIMA against baselines across two continuous control benchmarks with 10 scenarios. To the best of our knowledge, prior model-based MARL methods quite struggle in these continuous control benchmarks. Since the continuous action space is more complex than the discrete one, these continuous control tasks can be considered much more difficult.\nIn these experiments, DIMA consistently outperformed the baselines. The selected benchmarks span a variety of tasks, with the number of agents ranging from 2 to 6. This broad selection of tasks sufficiently demonstrates DIMA’s superior performance and general applicability. \nSMACv2, on the other hand, primarily focuses on environments with random agent types. When we decompose the original dynamics transition into a sequentially-conditioned structured transition, we also implicitly assume a fixed agent type for each agent ID throughout training, making SMACv2 incompatible with DIMA’s formulation.\nThus, we chose not to include SMACv2 in our evaluation.\n\n**Weakness 4**: 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\n**Response**: Thanks for raising this important point. Theorem 2 provides a theoretical justification for our training objective in the case of a sequentially conditioned diffusion-based world model under Assumption 1. Specifically, the denoisng term in Eq. 7 derives a denoising loss that guarantees the model learns the correct conditional transition distribution.\nWithout this result, it would be unclear how to formulate a principled objective to enforce the desired temporally sequential conditioning structure. Therefore, Theorem 2 plays an essential role in grounding our training procedure and ensuring the model’s behavior aligns with the intended design.\n\n**Weakness 5**: There are some more recent model-free MARL algorithms that are worth mentioning and possibly including in the comparison (e.g., IMAX-PPO) as well as some transformer-based MARL approaches (e.g. MARIE).\n\n**Response**: Thanks for the constructive suggestion.\nRegarding MARIE [4], we have already included it as a baseline in our experiments. As for IMAX-PPO [5], we were unable to find any publicly available implementation. As an alternative, we included HATD3 [6], an off-policy model-free MARL algorithm that can leverage experiences collected from different policies.\nThe results under a low data regime (limiting real-environment samples to 1M for MAMuJoCo) is listed as follows. DIMA still significantly outperforms HATD3.\n| Scenarios | HATD3 @ 1M steps | DIMA @ 1M steps |\n| :-: |:-:|:-:|\n| Ant 2x4  | 1751 ± 444 | **4881** ± 756 |\n| Ant 4x2  | 1639 ± 346 | **4766** ± 450 |\n| HalfCheetah 2x3 | 2405 ± 1119 | **6370** ± 121 |\n| HalfCheetah 3x2 | 2406 ± 1255 | **6175** ± 212 |\n| HalfCheetah 6x1 | 1823 ± 770  | **5643** ± 163 |\n\n**Question 1**: Can you clarify the practical implication of Theorem 2 for the main algorithm?\n\n**Response**: Thanks for the question. Please see the **response to the Weakness 4**.\n\n**Question 2**: How realistic is Assumption 1? Could you provide an example showing how this assumption holds in the multi-agent environments under consideration?\n\n**Response**: We appreciate this insightful question. We illustrate Assumption 1 with two examples from multi-agent firefighting scenarios.\n\nFirst, consider three firefighters independently extinguishing fires at different locations. Each denoising step receives one agent’s extinguishing fire action and removes the uncertainty corresponding to that specific fire (i.e., part of the global state). The remaining uncertainty depends on the actions of the other agents, which are revealed in later steps. This aligns with the case that each agent’s action partially denoises the global state in a dimension-wise manner.\n\nSecond, consider all three firefighters working together at the same location. \nWhen an agent take extinguishing fire action, it would slightly reduces the fire, progressively denoising the shared state. The final outcome depends on the combined effect of all actions, as the denoising steps cumulatively reduce uncertainty of the whole global state.\n\n**Limitation 1**: However, there might be other limitations, such as the need for sufficiently large and diverse datasets to accurately model the environment dynamics.\n\n**Response**: Thanks for the comment. As mentioned earlier, **DIMA operates in an online model-based MARL setting**, where the data buffer is initialized empty and the diffusion-based world model is learned entirely from scratch—without access to any pre-collected expert demonstrations.\nAnd our experimental results show that DIMA can learn an accurate enough world model to effectively accelerate policy learning under a low-data regime.\nCompared with prior online methods like MARIE and MAMBA which also focus on learning in imaginations paradigm, DIMA achieves superior performance with fewer environment interactions.\nIn contrast, offline model-based MARL algorithms like MADiff and MADiTs instead demand a sufficiently large dataset to deliver good performance.\n\n**Reference**\n\n[1] Zhu, Zhengbang, et al. \"MADiff: Offline multi-agent learning with diffusion models.\" NIPS (2024): 4177-4206.\n\n[2] Yuan, Lei, et al. \"Efficient multi-agent offline coordination via diffusion-based trajectory stitching.\" ICLR 2025.\n\n[3] Ajay, Anurag, et al. \"Is Conditional Generative Modeling all you need for Decision Making?.\" ICLR 2023.\n\n[4] Zhang, Yang, et al. \"Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models.\" TMLR.\n\n[5] Mai, Tien, and Thanh Nguyen. \"Mimicking to dominate: Imitation learning strategies for success in multiagent games.\" NIPS (2024): 84669-84697.\n\n[6] Zhong, Yifan, et al. \"Heterogeneous-agent reinforcement learning.\" JMLR (2024): 1-67."}, {"review_id": "IVr1RZ5wV7", "reviewer": "Reviewer_Uzmx", "summary": "The 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.", "questions": "Current 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?", "limitations": "The 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.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "Thanks for the rebuttal and additional results. My concern on baseline is addressed. However, I still do not see statistically significant results that joint modeling is better than sequential modeling (which is the main technical contribution of the work), given the large fluctuation itself in MBRL process, therefore I keep my current score.", "rebuttal": "Dear Reviewer Uzmx,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**Weakness 1**: For 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\n**Response**: Thanks for the constructive suggestion. Beyond the horizon $H = 15$ pre-defined during the training, We conduct additional evaluations with a longer prediction horizon of $H=25$, measuring accumulated observation and reward reconstruction errors over time. Here we still select the scenario *Ant-v2-2x4* for the experiment. The results are summarized below, and suggest that **DIMA exhibits lower compounding error trends over time compared to other baselines**.\n\n| Methods | Obs Accumulation Errors @ $H=15$ | Obs Accumulation Errors @ $H=25$ | Rew Accumulation Errors @ $H=15$ | Rew Accumulation Errors @ $H=25$ |\n| :-: |:-:|:-:|:-:|:-:|\n| DIMA |**2.42** ± 0.93|**4.32** ± 1.44|**15.01** ± 9.88|**31.07** ± 17.63|\n| MARIE |3.54 ± 1.60|6.62 ± 2.68|21.64 ± 17.51|47.66 ± 32.02|\n| MAMBA |3.98 ± 1.54|6.85 ± 2.51|38.90 ± 24.48|71.67 ± 41.78|\n\nMoreover, we would like to describe the detailed experiment setup of this error analysis. Since learning the world model is tied to a progressively improving policy both in MARIE, MAMBA and DIMA, we separately used their final policies to sample 5 episodes for fairness. We then computed L1 errors per observation between 100 trajectory segments randomly sampled from all 15 episodes and their imagined counterpart.\n\n\n**Weakness 2**: The 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. Please prove this argument by more environments, and perhaps more runs to show with higher statistical confidence.\n\n**Response**: Thanks for the constructive suggestion. To further clarify the performance differences between sequential modeling and joint modeling, we conducted additional experiments on the ShadowHandBottleCap task in the Bi-DexHands environment.\nMotivated by this comment, we took a closer look at the core question that our ablation study aims to answer—specifically what sequential modeling is fundamentally contributing to.\n\nThe primary benefit lies in decomposing the exponentially expanded space of the multi-agent dynamics into a linearly scaling space, significantly reducing modeling complexity. This facilitates more efficient learning of the world model under more limited data, which in turn accelerates policy convergence and improves sample efficiency.\n\nTo better highlight this, we calculate the averaged performance **under varying data budget** to test whether the reduced modeling complexity introduced by sequential modeling translates into improved policy performance. \nNotably, the result on ShadowHandBottleCap is computed by averaging 8 independent runs, which is twice the number of runs used in the original ablation study reported in the main paper.\n\n| Methods | BottleCap @ 100K steps | BottleCap @ 150K steps | BottleCap @ 200K steps |\n| :-: |:-:|:-:|:-:|\n| joint |234.1 ± 20.6|238.6 ± 22.9|246.7 ± 10.9|\n| sequential | **251.8** ± 17.3| **248.2** ± 11.6|246.3 ± 14.6|\n\nWe observe that under low-data regimes (e.g., 100K steps and 150K steps), sequential modeling outperforms joint modeling in both average return and variance. As the data budget decreases, the performance gap widens. This trend demonstrates that in extremely low-data regimes, sequential modeling benefits from reduced modeling complexity, enabling more accurate world model learning and accelerating policy learning. Our method's ability to maintain strong performance even with severely limited data underscores its effectiveness in resource-constrained scenarios, which is the primary focus of our work.\n\n**Question 1**: Current 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\n**Response**: Thanks for the question.\n1. We intentionally model the environment dynamics (i.e., state transitions) using a diffusion model, while predicting reward and termination separately via using a transformer-based architecture. This design choice stems from the inherent differences in the learning objectives of these variables. Note that reward prediction is typically learned using regression losses such as L2 loss, and termination prediction is often optimized via negative log-likelihood or binary cross-entropy loss. If all variables were modeled together in a single diffusion process, one would be forced to use a denoising loss, which is not well-aligned with the learning objectives for reward and termination. Therefore, we separate these components to better match their respective loss functions and learning dynamics.\n\n2. Our choice of using a transformer (MinGPT) for modeling reward and termination is motivated by prior successes in Transformer-based world models. Notably, transformer-based world models such as MARIE [1], IRIS [2], and TWM [3] have demonstrated strong performance in capturing temporal dependencies in reward and termination signals. Inspired by these works, we adopt a similar design for these components.\n\n**Reference**\n\n[1] Zhang, Yang, et al. \"Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models.\" TMLR.\n\n[2] Micheli, Vincent, Eloi Alonso, and François Fleuret. \"Transformers are Sample-Efficient World Models.\" ICLR 2023.\n\n[3] Robine, Jan, et al. \"Transformer-based World Models Are Happy With 100k Interactions.\" ICLR 2023."}, {"review_id": "5p72Ug9iI1", "reviewer": "Reviewer_n42q", "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.", "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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is easy to follow and well-written.\n2. The paper conducts extensive experiments and compared with some baselines.\n\nWeaknesses:\n\n1. 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.\n\n2. 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.", "was_revised": true, "final_justification": "The reviewer has addressed my concerns on sample efficiency and comparison with previous joint modeling-based approaches.", "rebuttal": "Dear Reviewer n42q,\n\nWe sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.\n\n**Weakness 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.\n\n**Response**: Thanks for the comment. \n1. We would like to clarify a potential misunderstanding about the \"sample efficiency\" and \"final return\" in the context of our work.\n\"Sample efficiency\" refers to the ability of a policy to achieve higher performance with fewer samples (transitions) during training, indicating more effective sample utilization. Thus, in the low-data regime where policy training relies on a very limited number of samples, higher performance directly reflects better sample efficiency, as it demonstrates the model's ability to learn effective policies with limited data.\nIn such context, \"final return\" refers to the performance at the end of training under this limited data budget, not the asymptotic return after sufficiently enough steps. Therefore, a higher final return under strict interaction constraints is itself an indicator of better sample efficiency.\nIllustrated in Figure 4, Figure 5 and Table 2 in Section E of the appendix, DIMA consistently outperforms all selected baselines in the low data regime, demonstrating superior sample efficiency.\nFor convenience, we provide numerical results below (with bold for the highest value and underlined for the second highest). HASAC is a strong off-policy model-free MARL algorithm which can also leverage the experience collected by different policy.\n**When comparing DIMA’s performance at 1M steps to HASAC at 5M steps, DIMA achieves comparable or better returns with 5x fewer samples, demonstrating strong sample efficiency**. Moreover, DIMA shows more stable performance, as indicated by lower variance across runs.\n\n| Scenarios | HASAC @ 1M steps | HASAC @ 5M steps |  DIMA @ 1M steps |\n| :-: |:-:|:-:|:-:|\n| Ant 2x4 |1344 ± 282| $\\underline{4830}$ ± 757| **4881** ± 756|\n| Ant 4x2  |850 ± 126| **5023** ± 753| $\\underline{4766}$ ± 450|\n| HalfCheetah 6x1 |2044 ± 110| $\\underline{5320}$ ± 2120| **5643** ± 163|\n\n\n2. We measure the training time and GPU memory usage of all evaluated model-based MARL methods, including our proposed DIMA. All experiments were conducted using a single NVIDIA RTX 4090 GPU to ensure a fair comparison across methods and tasks. As shown in the table below, DIMA is substantially more efficient than MARIE in both training time and GPU memory usage, and is comparable to or slightly more efficient than MAMBA.\n\n| Methods | Training Time | Usage of GPU Mem |\n| :-: |:-:|:-:|\n| DIMA in Ant 2x4  | ~1d 19h |3.10 GB|\n| MARIE in Ant 2x4 | ~3d 17h |14.33 GB|\n| MAMBA in Ant 2x4 | ~1d 1h  |2.74 GB|\n| DIMA in Ant 4x2  | ~1d 19h |3.15 GB|\n| MARIE in Ant 4x2 | ~3d 17h |15.73 GB|\n| MAMBA in Ant 4x2 | ~1d 1h  |4.37 GB|\n\n**Weakness 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.\n\n**Response**: We appreciate the reviewer’s comment. The core question addressed by our ablation is what advantages sequential modeling offers over joint modeling in multi-agent dynamics learning. The key benefit of sequential modeling lies in decomposing the exponentially large joint state and action space into a compositionally structured, linearly scaling space. This decomposition theoretically reduces the complexity of modeling multi-agent interactions, thereby enabling more sample-efficient learning of the world model, especially under the condition where limited data is available.\nWe performed additional experiments on the ShadowHandBottleCap task in the Bi-DexHands environment, to examine whether our proposed sequential modeling brings reduced modeling complexity and further facilitates more efficient policy learning.\nThe result on ShadowHandBottleCap is computed by averaging 8 independent runs, which is twice the number of runs used in the original ablation study reported in the main paper.\n\n| Methods | BottleCap @ 100K steps | BottleCap @ 150K steps | BottleCap @ 200K steps |\n| :-: |:-:|:-:|:-:|\n| joint |234.1 ± 20.6|238.6 ± 22.9|246.7 ± 10.9|\n| sequential |**251.8** ± 17.3|**248.2** ± 11.6|246.3 ± 14.6|\n\nHere, we also include the detailed final return results in Figure 8.\n\n| Methods | DoorOpenOutward @ 300K steps | DoorOpenInward @ 300K steps |\n| :-: |:-:|:-:|\n| joint |302.5 ± 76.9|235.1 ± 68.1|\n| sequential |**352.4** ($\\uparrow$ 49.9) ± 40.5|**290.3** ($\\uparrow$ 55.2) ± 30.4|\n\nThe overall results consistently favor the sequential modeling across these evaluated tasks. Notably, in the more data-constrained setting (BottleCap @ 100K and 150K steps), sequential modeling yields higher sample efficiency.\n\n**As for the potential downsides of our proposed formulation**, we would like to clarify that permutation invariance is not strictly required. Using a fixed order during both training and inference can be enough in practice. However, this can introduce implicit bias, as the predicted next state may become coupled with a fixed order. In contrast, the ground-truth environment dynamics exhibits order-invariance by nature.\n\nTo avoid this mismatch, we intentionally design DIMA to be permutation-invariant. In addition, the condition independence of the forward process is directly borrowed from the definition and derivation of conditional diffusion models in [1]. These design choices lead to a new objective (Eq. 9) as simple as the original one (Eq. 1), which we see as a strength rather than a limitation.\n\n\n**Question 1**: Where is the proof of Thm. 2? It should be in the appendix but the readers can't be find it in the paper.\n\n**Response**: The proof of Theorem 2 is provided in the **appendix PDF included in the supplementary material**, which was submitted alongside the main paper. We kindly ask the reviewer to refer to this file for the full derivation.\n\n**Reference**\n\n[1] Dhariwal, Prafulla, and Alexander Nichol. \"Diffusion models beat gans on image synthesis.\" NIPS (2021): 8780-8794."}]

Author rebuttal:

Meta-review: This paper introduces DIMA, a novel diffusion-inspired world model for multi-agent reinforcement learning. The core innovation lies in its sequential prediction of agent actions and state transitions, which reduces modeling complexity by aligning with the reverse steps of a diffusion process. 

The paper was praised for its exceptional clarity and the strong empirical performance of DIMA on established MARL benchmarks. Reviewers unanimously agreed on the novelty of its sequential, diffusion-inspired modeling approach. The primary concerns raised during review involved the method's scalability over longer prediction horizons, the depth of comparison to relevant prior works, and the need for more rigorous ablation studies to conclusively validate the benefits of sequential modeling over joint alternatives. The authors' rebuttal was exceptionally thorough and directly addressed the majority of these concerns, successfully satisfying three of the four reviewers. However, one reviewer (R3) remained unsatisfied, particularly regarding the sufficiency of the ablation studies. The final scores were three "Weak Accept" and one "Weak Reject." 

Given the paper's clear novelty, strong results, and the authors' comprehensive response, I recommend Acceptance. To fully address the residual concerns and solidify the paper's contribution, I strongly urge the authors to incorporate the following into their camera-ready version: additional experiments comparing observation and reward prediction errors for sequential vs. joint modeling. The addition will significantly enhance the paper's rigor and reproducibility for the community.
Decision: Accept (poster)