Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models

Dear Reviewer Cbz9,

We sincerely appreciate your precious time and constructive comments. Below, we provide a detailed point-by-point response, and we hope our clarifications and revision adequately address your concerns.

W1 & Q1: Lack of standardized performance evaluation protocol given the limited number of seeds.

Response: Thanks for the valuable comment! Upon the reviewer's suggestion, we summarize in Figure 13 the win rate with stratified bootstrap confidence intervals for mean, median, and interquartile mean (IQM). For finer comparisons, we also provide probabilities of improvement in Figure 14. Please refer to the revision for further details.

There are mainly two reasons why we use four random seeds for each experiment on a single scenario. First, the GPU resources we can use are quite limited. Meanwhile, as shown in Table 13, our algorithm has a demand for a relatively high computational cost and computational time. Second, we aim to align with the evaluation protocols used in related works, such as IRIS (single-agent model-based RL), MBVD, and MAZero (multi-agent model-based RL), which also use 4-5 seeds for their experiments. We acknowledge the importance of more random seed to get a rigourous statistical validation, and would be happy to test MARIE for more runs if sufficient computational resources were available.

Part of W2 & Q2: The use of SMAC, particularly SMAC v1, is problematic. The evaluation would benefit from using the updated SMAC v2 version.

Response: Thanks for the valuable comment. To address the reviewer's concern, we conduct additional experiments on 3 scenarios of SMACv2 -- zerg_5_vs_5, protoss_5_vs_5 and terran_5_vs_5. Without modifying any hyperparameters, MARIE consistently demonstrates higher sample efficiency compared to other baselines in a low data regime. Please refer to Appendix E.2 in the revision for further details.

Regarding the use of SMACv1 in the main experiments, as noted in the SMACv2 paper, certain maps in SMACv1 allow open-loop policies, which only utilize agent IDs and timesteps as input, to achieve non-trivial performance. However, not all maps suffer from this issue. For example, corridor, one of the maps we select, is not affected by this limitation. The significant improvement on the superhard scenario corridor can be a reasonable evidence to demonstrate the effectiveness of our algorithm. Furthermore, our policy is implemented using a 3-layer MLP, only taking the reconstructed observation as input, excluding agent IDs and timesteps. Based on these considerations, we believe that the results of the main experiments on SMACv1 are reliable and effectively validate the superiority of our algorithm.

Part of W2 & Q3: A direct comparison on environments beyond SMAC for methods like MAMBA needs to be provided.

Response: Thanks for the comment. Upon the reviewer's suggestion, we include MAMBA as a model-based baseline for comparisons on 3 chosen scenarios of MAMujoco. Since MAMBA was originally originally designed for domains with discrete action space, significant effort was required to adapt and evaluate it on MAMujoco, which features continuous action space. Notably, MAMBA fails to enhance policy learning in the Walker2d-v2-2x3 scenario and remains exceptionally time-consuming. Consequently, we report its results only for 1 million environment steps in this scenario. Our MARIE consistently shows superior sample efficiency and achieves the best performance in 2 of 3 scenarios with limited 2M environment steps. For the performance difference between MAMBA and MARIE in HalfCheetah-v2-3x2, we hypothesize that MAMBA's policy learning benefits significantly from using the internal recurrent features of the world model as inputs in this scenario, while the policy in our method only takes the reconstructed observation as input in order to support fast deployment in the environment without the participation of the world model. Please refer to Appendix E.1 in the revision for further details.

Q4: Given that the difference in compounding error of the world models between MAMBA and MARIE get worse over time, would the performance gap between the two be reduced if using a smaller imagination horizon for training and does this possibly bridge the gap?

Response: Thanks for the question. We report the performance of MARIE and MAMBA on 3s_vs_5z with different imagination horizon. And the result shows that the performance gap between the two is not related to the choice of imagination horizon. Interestingly, a larger imagination horizon may help policy learning in imagination. We hypothesize that longer imagined trajectories help alleviating shortsighted behaviours in policy learning.

Please do not hesitate to contact us if you need any further clarification or experiments.

Thanks for your further feedback. Below, we would like to address your concerns step by step.

Q1: SMACv2 Comparisons.

Response: Since MAMBA was originally originally designed for domains with discrete action space and only conducted the main experiment on SMAC, it would take considerable time and effort to evaluate it on SMACv2. Moreover, due to our limited GPU resources, we were unable to provide MAMBA's experimental results on SMACv2 in our initial response, as we wanted to allow sufficient time for the reviewer to give further feedback. Now, we manage to include the results of MAMBA on SMACv2 and update the figure for plotting the results of all algorithms on SMACv2 in the revision. We refer the reviewer to Figure 11 for details. While MARIE shows competitive performance to MAMBA on zerg_5_vs_5, MARIE is superior to MAMBA in terms of sample efficiency and final performance in the rest 2 scenarios.

Q2: MAMujoco Results.

Response: First, what we want to emphasize again is that MAMBA was originally designed for domains with discrete action space and only conducted the main experiment on SMAC. To ensure that MAMBA could successfully evaluate on the MAMujoco benchmark, we spent substantial time and effort on its implementation and tuning.

1. Specifically, across all the experiments on MAMujoco, the actor and critic of MAMBA were the same as those of MARIE.
2. All parameters, except those related to the difference on world model and corresponding algorithm pipeline, shared between the two methods. For example, the imagination horizon $H$ is set to 15 for both methods, which is also the optimal choice provided in the official repository of MAMBA in the SMAC experiments.
3. Furthermore, both approaches were trained on the same server equipped with 1 AMD EPYC 7C13 64-Core Processor and 4 NVIDIA Geforce RTX 4090, ensuring computational fairness.

The averaged training time for a run within 2 million environment steps is reported as follows:

The poor performance of MAMBA on the Walker2d scenario can be attributed to the higher difficulty of this task, which includes a healthy_reward bonus (we refer the reviewer to Walker2d documentation). This bonus is absent in the HalfCheetah scenario (refer to HalfCheetah documentation for this notable difference).

MAMBA fails to capture the presence/absence of bonus in the reward function across different scenarios, leading to poor performance of the policy learned in the imaginations of the world model of MAMBA itself. Its performance is even worse than that of model-free algorithms, further highlighting that MAMBA’s world model cannot generalize to complex continuous action space case. Another evidence that support this claim is that during the evaluation of MAMBA, the averaged episode length is 1000 (i.e., the default maximum episode length in MAMujoco) in HalfCheetah without the healthy_reward bonus while it becomes 60~100 in Walker2d. In contrast, the averaged episode length would gradually converge to 1000 across all three scenarios during the evaluation of MARIE. We believe that the significant performance gap in the more challenging Walker2d scenario demonstrates that MARIE’s improvements are non-trivial. Additionally, on HalfCheetah, MAMBA exhibits a lack of robustness and consistency, performing worse than MARIE in one case but better in the other.

Considering the overall performance across all three scenarios, especially the pronounced gap in the more difficult Walker2d environment with the healthy_reward bonus, we respectfully disagree with the reviewer’s statement: "Given that MAMBA has been shown to perform comparably in certain environments (as evidenced in MAMujoco) ...".

Q3: Imagination Horizon Experiment.

Response: We would like to point out that imagination horizon $H = 15$ is the recommended optimal choice provided in the official MAMBA benchmark. We refer the reviewer to the official link for details.

Additionally, we added the experiment of evaluating MAMBA on 3s_vs_5z with imagination horizon $H = 8$. The comparison between MARIE and MAMBA with different imagination horizon is reported as follows.

Dear Reviewer qbLx,

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

W1: The paper presentation could be improved with captioning figures of experimental results with short conclusions

Response: Thanks for the valuable comment. We have added a corresponding conclusion in Figure 3~6. Please see the revision for details.

Q1: The paper presentation would benefit from increasing the font size of figures in the main text.

Response: Thanks for the valuable comment. In response to your suggestion, we have redrawn Figure 1~2 in the main text and increased the font size to enhance readability.

Please do not hesitate to contact us if you need any further clarification or experiments.

Dear Reviewer qbLx,

Thank you for your continued feedback. We are grateful that our efforts have addressed your concerns.

We believe that your comments have notably improved the quality of our paper. Additionally, we would like to bring your attention to new experiments conducted during the discussion period. Following the insightful suggestions of Reviewer ysTM, we have conducted more comprehensive experiments, including an analysis of how errors in reconstructed observations impact final performance, an evaluation of the accuracy of discount prediction and the final performance with an increasing imagination horizon, as well as a comparison between different aggregation methods. Furthermore, thanks to the valuable feedback from Reviewer Cbz9, we have provided a statistically rigorous validation of our main experiment on SMAC and evaluated MARIE on an additional benchmark, SMACv2, to further test MARIE's superiority. As a result of these revisions, among other improvements, Reviewer ysTM and Cbz9 increased their scores from 5 to 6 and from 3 to 5, respectively.

If you have any remaining questions during the discussion period, we would be happy to answer them. And we are dedicated to earning a higher score from you.

Dear Reviewer ysTM,

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

W1: Necessity of individual components.

Response: Thanks for the constructive comment.

• First, we would like to clarify that our work is pioneering in the sense that it introduces the first effective transformer-based world model within the multi-agent context, enabling multi-agent world models to enjoy the benefit of recent progress on powerful generative models.
• Second, casting local dynamics learning as autoregressive modeling over discretized tokens is inspired by IRIS[1]. Compared with TWM[2] which used autoregressive modeling over continuous embedding, IRIS attained superior performance to TWM on Atari 100K, which demonstrates that autoregressive modeling over discrete tokens can further unleash the power of Transformer on dynamics learning. It motivates us to adopt a similar manner to build the Transformer-based multi-agent world model.
• For the Perceiver Transforme, we conduct additional experiment to compare our proposed aggregation via Perceiver with the common aggregation with self-attention. The difference on the final performance and the FLOPs shows that our agent-wise aggregation with Perceiver offers a more computationally efficient solution compared to self-attention aggregation. See our response to Q10 for details.
• For centralized feature aggregation, we show that the world model struggles to capture the underlying multi-agent dynamics in the trajectories and further hinders the policy learning without the agent-wise information aggregation in the second ablation experiment.

W2: Limited comparison to existing Transformer-based world models.

Response: Thank you for the valuable comment. Existing Transformer-based world models are primarily designed for single-agent scenarios, but they can be naturally adapted to multi-agent settings, modeling either independently local dynamics or joint dynamics. Based on the reviewer's suggestion, we have included IRIS as a Transformer-based world model baseline. Notably, the "Centralized Manner" and "MARIE w/o aggregation" variants from our ablation experiments correspond to IRIS baseline variants under different deployment strategies. Here we would like to clarify that these IRIS baseline variants also uses the same actor-critic method as MARIE during learning in imaginations phase, which can form a more fair comparison. To deliver a more intuitive demonstration, we additionally plot a new figure in Appendix E.4. Shown by the result, without incorporating CTDE principle, the learning of single-agent world model would be disrupted by the scalability and non-stationarity issues, validating our motivation in the introduction.

Q1: Given that the global state is known, could it be directly used as the aggregated global feature? If analysis or experiments were performed to validate the effectiveness of the current agent-wise aggregation, it would be more convincing.

Response: Thanks for the insightful question. In real-world applications, access to the global state, as available in simulations, is often impractical. Therefore, we opted not to incorporate the global state into the training of the world model, including using it directly as the aggregated global feature. Furthermore, incorporating the global state during training would necessitate its estimation during inference. Without global state estimation, long-term continuous imagination within the world model would require querying the simulation environment for intermediate global states, which is infeasible. Additionally, estimating the global state would introduce more prediction errors.

Q2: Since the policy relies on reconstructed observations, a deeper analysis of how errors in reconstructed observations impact final performance would be insightful.

Response: Thanks for the question. The policy uses reconstructed observations as input to avoid a input distribution shift between collecting experience from real environments and learning in imagination. Here, we conduct additional experiment with setting the number of observation token to 8 which results in more errors in reconstructed observations. We report the result within 50 thousand (50000) environment steps in the following table:

As the reconstructed observation contains more errors, the performance of the policy learned in imaginations also significantly degrades.

Dear Reviewer GJdy,

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

W1: The learning results of the world model depend on the supervisory signals, specifically the trajectories generated by a superior policy used as labels. In complex scenarios, without trajectories produced by an optimal policy, it may be difficult to learn a complete dynamic transition.

Response: Thanks for the comment. While it is possible to train a world model from a pre-collected offline dataset, our algorithm adopts a Dreamer[1]-like scheme (i.e., learning in imaginations) to train and leverage the world model, which runs the following three parts iteratively:

1. collect experience by executing the policy;
2. learn the world model from the collected experience;
3. learning the policy with the rollouts from the world model.

Thus, our algorithm doesn't need a superior policy to generate trajectories for world model training since the learning of world model is coupled with a randomly initialized policy online learning inside it. If the world model can provides trajectories that are sufficiently accurate compared to reality, the policy can progressively converge to an optimal one.

Q1: Considering that in SMAC all agents share the same environment reward, is the $r_t^i$ predicted by the "Shared Transformer" the same for each agent? If they are different, how is the team reward used to train the agents during the imagination phase? Is it averaged from $\sum_{i=1}^Nr_t^i/N$?

Response: Thanks for the question. Yes, the label of predicting $r_t^i$ for each agent is the same, i.e., the ground-truth team reward in the SMAC. Here, we suppose that using the team reward as a shared label for each agent's reward prediction can be considered as a simple yet potentially inappropriate credit assignment where each agent makes equal contribution and gets their individual reward $r_{\text{team}} / N$ which is further magnified by a factor of constant $N$. Hence, we use the individual predicted reward $\hat{r}_{t}^i$ instead of the averaged reward when training each agent during the imagination phase.

Q2: Based on Question 1, did the authors consider having each agent learn a different reward while learning the world model, in order to address the credit assignment problem in MARL through the world model learning process?

Response: Thanks for the constructive comment. As shown in response to Q1, we suppose currently the learning w.r.t. the reward prediction in our algorithm is implicitly a simple credit assignment. In terms of learning a different reward for each agent, we have not yet considered using different individual reward labels. But if the current team reward can be decomposed into different reward labels according to the actual individual contribution, it is intuitively promising for further enhancing policy learning in imaginations of the world model. Generally speaking, credit assignment can be approached through non-learning-based methods (e.g., manually assigning rewards based on each agent's specific actions) or learning-based methods (e.g., counterfactual inference via the value function as in COMA[2]). Non-learning-based approaches tend to introduce excessive human priors, making the learning world model overly domain-specific and lacking generalization capability. On the other hand, jointly training learnable credit assignment and the world model may face the challenge from co-training. We would leave this to the future work.

Q3: In the training process of the world model, how are the trajectories used as labels obtained? Also, please discuss what should be done in complex scenarios when there is no good initial policy to generate the trajectories.

Response: Thanks for the question. The trajectories for world model training is collected with a policy which is randomly initialized and iteratively trained inside the learned world model. The learning pipeline is stated in Line 164-166. In complex scenarios, we can still deploy the same procedure to improve the sample efficiency of policy learning.

Q4: Please explain in detail the role of learning $\gamma$ in the overall method, and why it cannot be replaced with a constant.

Response: Thanks for the question. Regarding learning discount $\gamma$ prediction, the labels $\gamma = 0.99$ for time steps within an episode and are set to zero for terminal time steps. Given an imagination horizon $H$, the predicted discount sequence $\hat{\gamma}_{1:H}^{i}$ is used to weigh the loss terms of the actor and critic of agent $i$ by the cumulative predicted discount factors to softly account for the possibility of episode ends while computing the cumulative return of the imaginations.