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

Authors: Yang Zhang, Chenjia Bai, Bin Zhao, Junchi Yan, Xiu Li, Xuelong Li
Primary area: reinforcement learning
Keywords: multi-agent reinforcement learning, world models, learning in imagination

Abstract:
Learning a world model for model-free Reinforcement Learning (RL) agents can significantly improve the sample efficiency by learning policies in imagination. However, building a world model for Multi-Agent RL (MARL) can be particularly challenging due to the scalability issue in a centralized architecture arising from a large number of agents, and also the non-stationarity issue in a decentralized architecture stemming from the inter-dependency among agents. To address both challenges, we propose a novel world model for MARL that learns decentralized local dynamics for scalability, combined with a centralized representation aggregation from all agents. We cast the dynamics learning as an auto-regressive sequence modeling problem over discrete tokens by leveraging the expressive Transformer architecture, in order to model complex local dynamics across different agents and provide accurate and consistent long-term imaginations. As the first pioneering Transformer-based world model for multi-agent systems, we introduce a Perceiver Transformer as an effective solution to enable centralized representation aggregation within this context. Main results on Starcraft Multi-Agent Challenge (SMAC) and additional results on MAMujoco show that it outperforms strong model-free approaches and existing model-based methods in both sample efficiency and overall performance.

Decision: reject
Reviewer scores: 6, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1729515162028

thread_type:
official_review

<thread_object>
creator: reviewer_cbz9

<content>
summary:
The paper introduces MARIE (Multi-Agent auto-Regressive Imagination for Efficient learning), a Transformer-based architecture designed to improve sample efficiency through improving the accuracy of multi-agent world modelling. The authors aim to address challenges of world modelling in MARL, particularly the scalability and non-stationarity issues, by using decentralised local dynamics combined with centralised aggregation through a Perceiver Transformer. The architecture/algorithm is evaluated on SMAC and additional experiments are conducted on MAMujoco, showing improved sample efficiency and overall performance compared to existing model-free and model-based methods.

strengths:
- The idea is presented clearly, and the architecture and experimental setup are detailed well. The integration of decentralised and centralised components is relatively straightforward and understandable.

- The submission provides comprehensive implementation details and references all the open-source repositories used for baselines, making it likely reproducible. The authors also mention that code will be released after the review process, which supports transparency.

- The paper presents results that show improvements over baselines. The usage of SMAC and MAMujoco environments offers a broad view of the architecture’s capabilities i.e. both discrete and continuous action spaces.

- The introduction of a Perceiver Transformer for centralised aggregation in a multi-agent context is an interesting approach and could provide valuable insights for the community.

weaknesses:
- The experiments lack rigourous statistical testing, which is critical given the limited number of seeds (only four). This raises concerns about the reliability and significance of the results. Referring to rigorous RL evaluation protocols such as those outlined in [rliable_code](https://github.com/google-research/rliable), [rliable_paper](https://arxiv.org/pdf/2108.13264) and [marl_eval](https://proceedings.neurips.cc/paper_files/paper/2022/file/249f73e01f0a2bb6c8d971b565f159a7-Paper-Conference.pdf) (among others) would have strengthened the empirical claims. These evaluation protocols have become a common standard that the community should uphold. Without statistical validation, it's hard to confirm that the reported improvements are statistically significant.

- The use of SMAC, particularly SMAC v1, is problematic as it is an outdated benchmark with known serious flaws, see [SMACv2](https://arxiv.org/abs/2212.07489). The evaluation would benefit from using the updated SMAC v2 version, which addresses some of these issues and gives more credibility to the method. Furthermore, the absence of comparisons with MAMBA in environments beyond SMAC makes it difficult to comprehensively evaluate the advantages of MARIE over existing architectures.

- While the architecture is interesting, the novelty might be overstated. There are similarities between MARIE and existing methods, such as MAMBA, with no stark performance difference (at least without the statistical testing, i don't believe we can make a fair claim that the difference is stark). Additionally, a recent approach, see https://openreview.net/forum?id=f2bgGy7Af7,  using graph attention networks (GATv2) (which are essentially transformers in a sense) closely mirrors the methodology, questioning the novelty of MARIE's transformer-based aggregation. Not mentioning this in the related work section detracts from the contribution's novelty marginally. However, i will say this is the least important weak point considering that if the results were rigorously validated i would still believe this methodology is worth it for the community to see.

questions:
### Questions

1. Why were only four seeds used for evaluation, and are there plans to conduct more extensive statistical testing to validate the results? I understand that experiments can be expensive to run and lots of tasks were used thus I thoroughly recommend reading the RLiable paper to see how you can address this without increasing computational budget. I believe you could simply run the tests with existing results.

2. Have you considered using SMAC v2, or is there a rationale for continuing with SMAC v1 despite its known flaws?

3. Could you provide a direct comparison on environments beyond SMAC for methods like MAMBA and potentially CoDreamer?

4. 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?

### Suggestions:

- **Use of Benchmarks**: Including evaluations on more diverse environments, such as those in the updated SMAC version or other cooperative multi-agent benchmarks, would strengthen the paper’s claims.

- **Statistical Validation**: Incorporating more seeds and employing robust statistical methods would add credibility to the results.

In conclusion, I don’t feel confident enough that the results presented truly indicate a statistically significant performance improvement and that the architecture itself doesn’t provide enough of a difference to warrant acceptance without the solid empirical proof that it is a superior method. If a) the RLiable evaluation methodology is run and the results present statistically significant improvements and b) MAMBA is run on MAMujoco with similar conclusions to SMAC then i will be willing to raise my score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730719933194

thread_type:
official_review

<thread_object>
creator: reviewer_qblx

<content>
summary:
The authors propose Perceiver Transformer-based world model for MARL that addresses the actual problems of scalalability and the non-stationarity. The evaluations on SMAC and MAMujuco prove the model superiority over multiple baselines.

strengths:
- The proposed model combined decentralized dynamics modeling with centralized representation aggregation using Transformer sequence modeling.
- The paper is well-written and easy to follow.
- The authors provide the ablation results and the analysis of attention patterns to reveal the implicit decison-making features.

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

questions:
- The paper presentation would benefit from increasing the font size of figures in the main text.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730725621258

thread_type:
official_review

<thread_object>
creator: reviewer_ystm

<content>
summary:
This paper introduces a novel framework, **MARIE (Multi-Agent auto-Regressive Imagination for Efficient learning)**, which leverages a Transformer-based world model to enable sample-efficient policy learning in **Multi-Agent Reinforcement Learning (MARL)**. MARIE addresses two key challenges in MARL: scalability and **non-stationarity**. The approach combines decentralized local dynamics learning with centralized feature aggregation using a Perceiver Transformer. The authors evaluate the proposed method on the Starcraft Multi-Agent Challenge (SMAC) and **MAMuJoCo**, demonstrating improved sample efficiency and performance over existing model-free and model-based methods.

strengths:
1. The integration of decentralized local dynamics learning and centralized feature aggregation is well-motivated and effectively addresses key challenges in MARL, such as scalability and non-stationarity.
2. The use of the Perceiver Transformer for centralized representation aggregation is an innovative contribution that facilitates efficient global information sharing between agents while maintaining scalability.

weaknesses:
1. **Necessity of individual components**: The authors claim that this work is “the first pioneering Transformer-based world model for multi-agent systems,” but the underlying techniques—centralized feature aggregation, the Perceiver Transformer, and autoregressive modeling of discretized tokens—are already present in the literature. More ablation experiments to demonstrate the necessity of these components would strengthen the paper. It is necessary to investigate whether it is a kind of simple combinations of different techniques, or more reasonable design for MARL.
2. **Limited comparison to existing Transformer-based world models**: While the paper compares its method with model-free and some model-based MARL approaches, a more in-depth exploration of existing Transformer-based methods in MARL, or related architectures from single-agent RL that could directly be extended to MARL (e.g., IRIS, TWM and other methods), is lacking. It would be beneficial to further discuss why existing single-agent Transformer-based approaches cannot be directly adapted to MARL.

questions:
- 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.
- Since the policy relies on reconstructed observations, a deeper analysis of how errors in reconstructed observations impact final performance would be insightful.
- "The policies π are exclusively trained using imagined trajectories." Does this lead to wasted real experience collected during training?
- I am curious about the prediction accuracy for discounts at each step. As the horizon (H) increases, can the model accurately predict the end of the game, and how does this affect performance?
- Since MARIE separates model learning from policy learning, providing intuitive or experimental comparisons with methods that jointly learn the model and policy would increase the persuasiveness of the approach. For example, the following references could be useful:
  - [1] Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Ruslan Salakhutdinov. Mismatched no more: Joint model-policy optimization for model-based RL. Advances in Neural Information Processing Systems, 35:23230–23243, 2022.
  - [2] Raj Ghugare, Homanga Bharadhwaj, Benjamin Eysenbach, Sergey Levine, and Ruslan Salakhutdinov. Simplifying model-based RL: learning representations, latent-space models, and policies with one objective. arXiv preprint arXiv:2209.08466, 2022.
- In the first ablation experiment, regarding learning local dynamics instead of joint dynamics, the authors state that “the scalability issue is exacerbated by a growing number of agents.” However, in Figure 4, the performance of learning joint dynamics degrades more with 3 agents (3s_vs_3z) than with 5 agents (2s3z). This seems inconsistent with the authors' claim.
- In the third ablation study, it would be worth exploring the effect of not discretizing observations and instead learning a continuous embedding through a linear layer. Additionally, assessing the impact of policies that directly depend on the internal hidden states of the Transformer, rather than reconstructed observations, would be insightful. Intuitively, the policy input only needs to capture decision-relevant information, not a complete image reconstruction. Moreover, errors in reconstruction may negatively impact policy learning.
- In Figure 7, the authors claim that MARIE “has remarkably better error,” but the actual curves do not seem to support the term “remarkably.” Providing a corresponding performance comparison curve would make this claim more visually intuitive.
- Including a more comprehensive set of experimental results in Table 1 would enhance the paper.
- Can the authors provide more details on the computational cost of using a Perceiver Transformer for centralized aggregation? How does this affect MARIE's scalability as the number of agents increases?
- Could the authors clarify the role of intra-step autoregression and how it contributes to the overall performance of the model? A comparison between the Perceiver and other common aggregation methods would also be helpful.
- The paper shows strong performance on simulation benchmarks, but adding results or discussions on how MARIE could be applied to real-world scenarios would increase its impact and relevance.
- The Preliminary section could benefit from a brief introduction to the Perceiver and other aggregation techniques, making the paper more accessible to readers unfamiliar with these concepts.
- Are there any limitations of MARIE that might make it less effective in certain multi-agent settings, such as environments with highly heterogeneous agents or asymmetric observation spaces?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730738095999

thread_type:
official_review

<thread_object>
creator: reviewer_gjdy

<content>
summary:
Considering the inevitable challenges of both centralized and decentralized learning in developing a world model, this paper proposes MARIE (Multi-Agent auto-Regressive Imagination for Efficient learning), a Transformer-based approach that integrates both methods. The process is divided into three stages: the first involves collecting multi-agent trajectories, the second focuses on learning the world model from these experiences, and the third uses the world model for policy learning. The second stage, which centers on the learning of the world model, involves discretizing observations with the VQ-VAE method and using the learned discretized codes to construct a Transformer architecture for transition prediction. Additionally, the authors incorporate agent-wise aggregation information to mitigate non-stationarity. Experiments on SMAC and MAMujoco are conducted to validate the method's effectiveness.

strengths:
1.	In constructing the world model, the authors considered both centralized information and decentralized information.
2.	The overall logic of the paper is coherent and easy to understand.
3.	The paper conducted extensive experiments.

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

questions:
1.	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 $\{r_t^i \}_{i=1}^{N}$?
2.	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?
3.	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.
4.	Please explain in detail the role of learning $\gamma$ in the overall method, and why it cannot be replaced with a constant.
5.	How is the codebook $Z$ initialized, and does the initialization affect the learning outcomes under different conditions?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733060323205

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear reviewers and AC,

We sincerely appreciate your valuable time and effort spent reviewing our manuscript.

We propose MARIE, a sample-efficient model-based MARL algorithm driven by an effective Transformer-based multi-agent world model, which is built by casting the decentralized agent dynamics into sequence modeling over discrete tokens to solve the *scalability* issue, and incorporating it with a centralized agent-wise aggregation to solve the *non-stationarity* issue. As reviewers highlighted:
1. Our method, especially the integration of decentralized dynamics and centralized aggregation, is a notable strength, relatively straightforward and well-motivated (all);
2. Our method is demonstrated by extensive experiments and analysis (GJdy, qbLx, Cbz9);
3. The use of the Perceiver Transformer for centralized aggregation is innovative and provides valuable insights (ysTM, Cbz9).

We appreciate your constructive feedback on our manuscript. In response to the comments,we have carefully revised and enhanced the manuscript as follows:
1. Additional re-implementation, adaptation and evaluation of MAMBA evaluation of MAMBA on MAMujoco (Appendix E.1);
2. Additional experiments on SMACv2 (Appendix E.2);
3. Comparison with existing Transformer-based world models (Appendix E.4);
4. Comparison against MAMBA with different imagination horizon (Appendix E.5);
5. Additional Discussion on the related works reviewers provided (Appendix F);
6. Rigorous statistical evaluation by using RLiable (Appendix G);
7. Limitation discussion on using limited seeds (Section 5);
8. Increased font size of figures in the main text to enhance readability;
9. A follow-up conclusion in the caption Figure 3~6 to improve the presentation.

Besides, we have also provided the results of essential experiments including:
1. The comparison on computational cost and performance between different means of aggregation (Response to Reviewer ysTM);
2. The evaluation of the accuracy of MARIE's discount prediction under different imagination horizon (Response to Reviewer ysTM).

The above results would be also incorporated into the future revision right after having access to revising our manuscript. Lastly, in the current revised manuscript, these updates are temporarily highlighted in blue for your convenience to check. We sincerely believe that these updates could help us better deliver the benefits of the proposed MARIE to the multi-agent community. 

Thank you very much,

Authors.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces MARIE, a Transformer-based architecture designed to improve sample efficiency by enhancing the accuracy of multi-agent world modeling. The authors aim to address challenges of world modeling in multi-agent reinforcement learning (MARL), particularly the scalability and non-stationarity issues, by using decentralized local dynamics combined with centralized aggregation through a Perceiver Transformer. Reviewers believe that the method proposed by the authors is somewhat novel, the writing is good, and the experimental setup is very clear. Considering that none of the reviewers expressed clear support for the paper (scores were 6665) and were not positive about the rebuttal, AC encourages the authors to supplement the experiments based on the reviewers' opinions and carefully revise the paper, believing that its quality can be greatly improved.