Dear Reviewer n42q,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

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.

Response: Thanks for the comment.

1. We would like to clarify a potential misunderstanding about the "sample efficiency" and "final return" in the context of our work. "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. In 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. Illustrated 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. For 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. 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.

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

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.

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

Here, we also include the detailed final return results in Figure 8.

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

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.

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

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.

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.

Reference

[1] Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." NIPS (2021): 8780-8794.

Dear Reviewer Uzmx,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

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.

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.

Moreover, 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.

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.

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

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

To 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. Notably, 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.

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

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?

Response: Thanks for the question.

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

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

Reference

[1] Zhang, Yang, et al. "Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models." TMLR.

[2] Micheli, Vincent, Eloi Alonso, and François Fleuret. "Transformers are Sample-Efficient World Models." ICLR 2023.

[3] Robine, Jan, et al. "Transformer-based World Models Are Happy With 100k Interactions." ICLR 2023.

Dear Reviewer Y1VU,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

Weakness 1: The scalability of the approach is not sufficiently tested.

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

To 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:

We also evaluate policy learning performance under the extended imagination horizon, comparing DIMA with $H=15$ and $H=25$. The results are reported as follows.

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.

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?

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. Regarding 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. Given 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.

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.

Response: Thanks for the insightful comment.

1. 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.
2. 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.

The extended results are summarized in the table below (bold for the highest value, underlined for the second highest).

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

Furthermore, 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.

Weakness 4: The computational cost and training time of the world-agent are not investigated (also in comparison wrt similar solution).

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.

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?

Response: Thanks for the insightful questions. We address each of them in turn:

1. 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.
2. 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.
3. 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. Since 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.

Question 2: what are the performance of the tested models using more steps / at convergence ?

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.

Question 3: What is the computational cost of training the world model (e.g., time), compared to similar approaches?

Response: We refer the reviewer to our response to Weakness 4 for a detailed evaluation.

Question 4: Does the authors use “§” to refer to the appendix?

Response: Yes, we would add one sentence in the revised manuscript to clarify this symbol usage.

Reference

[1] Alonso, Eloi, et al. "Diffusion for world modeling: Visual details matter in atari." NIPS (2024): 58757-58791.

[2] Janner, Michael, et al. "When to trust your model: Model-based policy optimization." NIPS (2019).

Dear Reviewer DfzB,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

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.

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). This 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.

Weakness 2: The paper mentions MADiff and MADiTS as relevant diffusion-based works but does not sufficiently compare with them beyond qualitative differences.

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,

1. 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].
2. MADiTs used diffusion models for offline trajectory stitching, generating synthetic transitions to augment the original offline dataset.

Crucially, 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. Furthermore, 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.

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.

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.

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. In 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. SMACv2, 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. Thus, we chose not to include SMACv2 in our evaluation.

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.

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

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

Response: Thanks for the constructive suggestion. Regarding 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. The results under a low data regime (limiting real-environment samples to 1M for MAMuJoCo) is listed as follows. DIMA still significantly outperforms HATD3.

Question 1: Can you clarify the practical implication of Theorem 2 for the main algorithm?

Response: Thanks for the question. Please see the response to the Weakness 4.

Question 2: How realistic is Assumption 1? Could you provide an example showing how this assumption holds in the multi-agent environments under consideration?

Response: We appreciate this insightful question. We illustrate Assumption 1 with two examples from multi-agent firefighting scenarios.

First, 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.

Second, consider all three firefighters working together at the same location. When 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.

Limitation 1: However, there might be other limitations, such as the need for sufficiently large and diverse datasets to accurately model the environment dynamics.

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. And our experimental results show that DIMA can learn an accurate enough world model to effectively accelerate policy learning under a low-data regime. Compared with prior online methods like MARIE and MAMBA which also focus on learning in imaginations paradigm, DIMA achieves superior performance with fewer environment interactions. In contrast, offline model-based MARL algorithms like MADiff and MADiTs instead demand a sufficiently large dataset to deliver good performance.

Reference

[1] Zhu, Zhengbang, et al. "MADiff: Offline multi-agent learning with diffusion models." NIPS (2024): 4177-4206.

[2] Yuan, Lei, et al. "Efficient multi-agent offline coordination via diffusion-based trajectory stitching." ICLR 2025.

[3] Ajay, Anurag, et al. "Is Conditional Generative Modeling all you need for Decision Making?." ICLR 2023.

[4] Zhang, Yang, et al. "Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models." TMLR.

[5] Mai, Tien, and Thanh Nguyen. "Mimicking to dominate: Imitation learning strategies for success in multiagent games." NIPS (2024): 84669-84697.

[6] Zhong, Yifan, et al. "Heterogeneous-agent reinforcement learning." JMLR (2024): 1-67.

We would like to highlight that, as reviewers Y1VU, DfzB, and n42q have acknowledged, all their concerns have been addressed through our thorough efforts and additional experiments. Although Reviewer Uzmx still holds some concerns regarding our ablation study and has provided valuable suggestions for better evaluating the reduced modeling complexity, we appreciate and accept these suggestions and are committed to incorporating them in our further evaluations. More importantly, our proposed DIMA demonstrates a significant and clear improvement in sample efficiency compared to existing model-free and model-based MARL approaches.

We sincerely appreciate all reviewers’ constructive feedback, and we believe these improvements will help us better communicate the advantages of DIMA to the multi-agent and NIPS community.

Best regards,

Authors