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

## Abstract

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.
Results on Starcraft Multi-Agent Challenge (SMAC) show that it outperforms strong model-free approaches and existing model-based methods in both sample efficiency and overall performance.

## 1 Introduction

Multi-Agent Reinforcement Learning (MARL) has made remarkable progress, which was driven largely by model-free algorithms . However, due to the complexity of multi-agent systems arising from large state-action space and partial observability, such algorithms usually demand extensive interactions to learn coordinative behaviors . A promising solution is building a world model that approximates the environment, which has exhibited its superior sample efficiency compared to model-free approaches in single-agent RL ; ; ; ; ; .
However, extending the design of world model in single-agent domain to the multi-agent context encounters significant challenges due to the unique biases and characteristics inherent to multi-agent environments.

The challenges primarily stem from two different means for multi-agent dynamics learning: *centralized* and *decentralized*.
Learning a world model to approximate the *centralized* dynamics encapsulates the inter-dependency between agents but struggles to be scalable to an increasing number of agents, which leads to the exponential surge in spatial complexity ; .
Conversely, applying a *decentralized* world model to approximating the local dynamics of each agent mitigates the scalability issue yet incurs non-stationarity, as unexpected interventions from other agents may occur in each agent’s individual environment .
Furthermore, beyond these unique challenges inherent to multi-agent scenarios, existing model-based MARL approaches ; ; excessively neglect the fact that the policy learned in imaginations of the world model heavily relies on the quality of imagined trajectories .
It thereby necessitates accurate long-term prediction, especially with respect to the non-stationary local dynamics.
Inspired by the capability of Transformer  in modeling complex sequences and long-term dependency ; ; , we seek to construct a Transformer-based world model within the multi-agent context for *decentralized* local dynamics together with *centralized* feature aggregation, combining the benefits of two distinctive designs.

In this paper, we introduce MARIE (Multi-Agent auto-Regressive Imagination for Efficient learning), the first Transformer-based multi-agent world model for sample-efficient policy learning. Specifically, the highlights of this paper are:

- 1.
To tackle the inherent challenges within the multi-agent context, we build an effective world model via scalable *decentralized* dynamics modeling and essential *centralized* representation aggregating, which mirrors the principle of Centralized Training and Decentralized Execution.
- 2.
To enable accurate and consistent long-term imaginations from the non-stationary local dynamics, we cast the *decentralized* dynamics learning as sequence modeling over discrete tokens by leveraging highly expressive Transformer architecture as the backbone. In particular, we successfully present the first Transformer-based world model for multi-agent systems.
- 3.
While it remains open for how to effectively enable *centralized* representation with the Transformer as the backbone, we achieve it by innovatively introducing a Perceiver Transformer for efficient global information aggregation across all agents.
- 4.
Experiments on the Starcraft Multi-Agent Challenge (SMAC) benchmark in low data regime show MARIE outperforms both model-free and existing model-based MARL methods w.r.t. both sample efficiency and overall performance and demonstrate the effectiveness of MARIE.

## 2 Related Works and Preliminaries

Multi-Agent Reinforcement Learning. In a model-free setting, a typical approach for cooperative MARL is centralized training and decentralized execution (CTDE), which tackles the scalability and non-stationarity issues in MARL. During the training phase, it leverages global information to facilitate agents’ policy learning; while during the execution phase, it blinds itself and has only access to the partial observation around each agent for multi-agent decision-making. Model-free MARL methods with this paradigm can be divided into 2 categories: value-based ; ; ; and policy-based ; ; ; ; ; ; ; ; ; . In contrast to model-free approaches, model-based MARL algorithms remain fairly understudied. MAMBPO incorporates MBPO-style techniques into multi-agent policy learning under the CTDE framework. Tesseract introduces the tensorised Bellman equation and evaluates the Q-value function using Dynamic Programming (DP) together with an estimated environment model. Similar to our setting where agents learn inside of an approximate world model, MAMBA integrates the backbone proposed in DreamerV2  with an attention mechanism across agents to sustain an effective world model in environments with an arbitrary number of agents, which leads to notably superior sample efficiency to existing model-free approaches.
In terms of model-based algorithm coupled with planning, MAZero expands the MCTS planning-based Muzero framework to the model-based MARL settings.
However, learning-based or planning-based policies in these two approaches are both overly coupled with their world models, downgrading their inference efficiency and further limiting expansion in combinations with other popular model-free approaches.

Learning behaviors within the imagination of world models. The Dyna architecture first emphasizes the utility of an estimated dynamics model in facilitating the training of the value function and policy. Inspired by the cognitive system of human beings, the concept of world model is initially introduced by composing a variational Auto-Encoder (VAE) and a recurrent network to mimic the complete environmental dynamics, then an artificial agent is trained entirely inside the hallucinated imagination generated by the world model. SimPLe  shows that a PPO policy  learned in a predictive model deliverer a super-human performance in Atari domains. Dreamer  builts the world model upon a Recurrent State Space Model (RSSM) that combines the deterministic latent state with the stochastic latent state to allow the model to not only capture multiple futures but also remember information over multi-steps. DreamerV2 further demonstrates the advantage of discrete latent states over Gaussian states. For MARL, MAMBA extends DreamerV2 to multi-agent contexts by using RSSM, underscoring the potential of multi-agent learning in the imagination of world models. Recently, motivated by the success of the Transformer , TransDreamer and TWM explored variants of DreamerV2, wherein the backbones of the world model were substituted with Transformers. Instead of incorporating deterministic and stochastic latent states, IRIS applies the Transformer to directly modeling sequences of observation tokens and actions of single-agent RL and achieves impressive results on Atari-100k. In contrast, the proposed MARIE concentrates on establishing effective Transformer-based world models in multi-agent contexts with shared dynamics and global representations.

Preliminaries. We focus on fully cooperative multi-agent systems where all agents share a team reward signal. We formulate the system as a decentralized partially observable Markov decision process (Dec-POMDP) , which can be described by a tuple $(\mathcal{N},\mathcal{S},\boldsymbol{\mathcal{A}},P,R,\boldsymbol{\Omega},
\boldsymbol{\mathcal{O}},\gamma)$. $\mathcal{N}=\{1,...,n\}$ denotes a set of agents, $\mathcal{S}$ is the finite global state space, $\boldsymbol{\mathcal{A}}=\prod_{i=1}^{n}\mathcal{A}^{i}$ is the product of finite actions spaces of all agents, i.e., the joint action space, $P:\mathcal{S}\times\boldsymbol{\mathcal{A}}\times\mathcal{S}\rightarrow[0,1]$ is the global transition probability function, $R:\mathcal{S}\times\boldsymbol{\mathcal{A}}\rightarrow\mathbb{R}$ is the shared reward function, $\boldsymbol{\Omega}=\prod_{i=1}^{n}\Omega^{i}$ is the product of finite observation spaces of all agents, i.e., the joint observation space, $\boldsymbol{\mathcal{O}}=\{\mathcal{O}^{i},i\in\mathcal{N}\}$ is the set of observing functions of all agents. $\mathcal{O}^{i}:\mathcal{S}\rightarrow\Omega^{i}$ maps global states to the observations for agent $i$, and $\gamma$ is the discount factor. Given a global state $s_{t}$ at timestep $t$, agent $i$ is restricted to obtaining solely its local observation $o_{t}^{i}=\mathcal{O}^{i}(s_{t})$, takes an action $a_{t}^{i}$ drawn from its policy $\pi^{i}(\cdot|o_{\leq t}^{i})$ based on the history of its local observations $o_{\leq t}^{i}$, which together with other agents’ actions gives a joint action $\boldsymbol{a}_{t}=(a_{t}^{1},...,a_{t}^{n})\in\boldsymbol{\mathcal{A}}$, equivalently drawn from a joint policy $\boldsymbol{\pi}(\cdot|\boldsymbol{o}_{\leq t})=\prod_{i=1}^{n}\pi^{i}(\cdot|o
_{\leq t}^{i})$. Then the agents receive a shared reward $r_{t}=R(s_{t},\boldsymbol{a}_{t})$, and the environment moves to next state $s_{t+1}$ with probability $P(s_{t+1}|s_{t},\boldsymbol{a}_{t})$. The aim of all agents is to learn a joint policy $\boldsymbol{\pi}$ that maximizes the expected discounted return $J(\boldsymbol{\pi})=\mathbb{E}_{s_{0},\boldsymbol{a}_{0},...\sim\boldsymbol{
\pi}}\left[\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},\boldsymbol{a}_{t})\right]$.

Figure: Figure 1: Overview of the proposed world model architecture in MARIE. VQ-VAE (*left*) maps local observations $o^{i}$ of each agent $i$ into discrete latent codes $(x^{i}_{1},...,x^{i}_{K})$, where $(E,D,\mathcal{Z})$ is shared across all agents. Together with discrete actions, this process forms local discrete sequences $(...,x^{i}_{t,1},...,x^{i}_{t,K},a^{i}_{t},...)$ of each agent. Then the Perceiver (*right*) performs aggregation of joint discrete sequences of all agents $(x^{1}_{t,1},...,x^{1}_{t,K},a^{1}_{t},...,x^{n}_{t,1},...,x^{n}_{t,K},a^{n}_{ t})$ independently at each timestep $t$, and inserts the aggregated global representations $(e_{t}^{1},e_{t}^{2},...,e_{t}^{n})$ into original local discrete sequences respectively. The resulting sequences $(...,x^{i}_{t,1},...,x^{i}_{t,K},a^{i}_{t},e^{i}_{t}...)$ contain rich information between transitions in local dynamics and are fed into the shared Transformer (*middle*), which learns observation token predictions in an autoregressive manner. Predictions of individual reward $r_{t}^{i}$ and discount $\gamma_{t}^{i}$ at timestep $t$ are computed based on all historical sequence $(x^{i}_{\leq t,1},...,x^{i}_{\leq t,K},a^{i}_{\leq t},e^{i}_{\leq t})$.
Refer to caption: x1.png

## 3 Methodologies

Our approach comprises three typical parts: (i) collecting experience by executing the policy, (ii) learning the world model from the collected experience, and (iii) learning the policy via imagination inside the world model. Throughout the process, the historical experiences stored in the replay buffer are used for training the world model only, while policies are learned from unlimited imagined trajectories from the world model. In the following, we first describe three core components of our world model in §3.1 and §3.2, and give an overview of the proposed world model in Fig. 1. Then we describe the policy-learning process inside the world model in §3.3. The comprehensive details of the model architecture and hyperparameter are provided in §A.

### 3.1 Discretizing Observation

We consider a trajectory $\tau^{i}$ of agent $i$ consists of $T$ local observations and actions, as

$$ $\displaystyle\tau^{i}=(o^{i}_{1},a^{i}_{1},\ldots,o^{i}_{t},a^{i}_{t},\ldots,o ^{i}_{T},a^{i}_{T}).$ $$

To utilize the expressive Transformer architecture, we need to express the trajectory into a discrete token sequence for modeling. Accounting for continuous observations, a prevalent but naive approach for discretization involves discretizing the scalar into one of $m$ fixed-width bins in each dimension independently .
However, with a higher dimension of the observation, the Transformer encounters higher computational complexity, which necessitates an approach that uses a discrete codebook of learned representations.
To this end, we employ the idea from neural discrete representation learning , and learn a Vector Quantised-Variational AutoEncoder (VQ-VAE) to play a role that resembles the tokenizer in Natural Language Processing ; . The VQ-VAE is composed of an encoder $E$, a decoder $D$, and a codebook $\mathcal{Z}$. We define the discrete codebook $\mathcal{Z}=\{z_{j}\}^{N}_{j=1}\subset\mathbb{R}^{n_{z}}$, where $N$ is the size of the codebook and $n_{z}$ is the dimension of codes.
The encoder $E$ takes an observation $o^{i}\in\mathbb{R}^{n_{\rm obs}}$ as input and outputs a $K$ $n_{z}$-dimensional latents $\hat{z}^{i}\in\mathbb{R}^{K\times n_{z}}$ reshaped from the direct outputs of encoder.
Subsequently, the tokens $\{x^{i}_{k}\}^{K}_{k=1}\in\{0,1,...,N-1\}^{K}$ for representing $o^{i}$ is obtained by a nearest neighbour look-up using the codebook $\mathcal{Z}$ where $x_{k}^{i}=\mathop{\arg\min}_{j}\|\hat{z}^{i}_{k}-z_{j}\|$.
Then the decoder $D:\{0,1,...,N-1\}^{K}\rightarrow\mathbb{R}^{n_{\rm obs}}$ converts $K$ tokens back into an reconstructed observation $\hat{o}^{i}$.
By learning this discrete codebook, we compress the redundant information via a succinct sequence of tokens, which helps improve sequence modeling. See §4.2 for a discussion.

### 3.2 Modeling Local Dynamics with Global Representations

Here, we consider discrete actions like those in SMAC, and the continuous actions can also be discretized by splitting the value in each dimension into fixed bins ; .
Therefore, a trajectory $\tau^{i}$ of agent $i$ can be treated as a sequence of tokens,

$$ $\tau^{i}=(\ldots,o_{t}^{i},a_{t}^{i},\ldots)=(\ldots,x_{t,1}^{i},x_{t,2}^{i}, \ldots,x_{t,K}^{i},a_{t}^{i},\ldots)$ (1) $$

where $x_{t,j}^{i}$ is the $j$-th token of the observation of agent $i$ at timestep $t$.
Given arbitrary sequences of observation and action tokens in Eq. (1), we try to learn over discrete multimodal tokens.

The world model consists of a tokenizer to discrete the local observation, a Transformer to learn the local dynamics, an agent-wise representation aggregation module, and predictors for the reward and discount.
The Transformer $\phi$ predicts the future local observation $\{\hat{x}_{t+1,j}^{i}\}^{K}_{j=1}$, the future individual reward $\hat{r}_{t}^{i}$ and discount $\hat{\gamma}_{t}^{i}$, based on the agent’s individual historical observation-action history $(x_{\leq t,\boldsymbol{\cdot}}^{i},a_{\leq t}^{i})$ and aggregated global feature $e_{t}^{i}$ of the agent.
The modules are shown in Eqs. (2)–(5).

$$ Transition: $\displaystyle\hat{x}_{t+1,\boldsymbol{\cdot}}^{i}$ $\displaystyle\sim p_{\phi}(\hat{x}_{t+1,\boldsymbol{\cdot}}^{i}|x_{\leq t, \boldsymbol{\cdot}}^{i},a_{\leq t}^{i},e_{\leq t}^{i})\mbox{\ with \ }\hat{x}_ {t+1,k}^{i}\sim p_{\phi}(\hat{x}_{t+1,k}^{i}|x_{\leq t,\boldsymbol{\cdot}}^{i} ,a_{\leq t}^{i},e_{\leq t}^{i},x_{t+1,<k}^{i})$ (2) Reward: $\displaystyle\hat{r}_{t}^{i}$ $\displaystyle\sim p_{\phi}(\hat{r}_{t}^{i}|x_{\leq t,\boldsymbol{\cdot}}^{i},a _{\leq t}^{i},e_{\leq t}^{i})$ (3) Discount: $\displaystyle\hat{\gamma}_{t}^{i}$ $\displaystyle\sim p_{\phi}(\hat{\gamma}_{t}^{i}|x_{\leq t,\boldsymbol{\cdot}}^ {i},a_{\leq t}^{i},e_{\leq t}^{i})$ (4) Aggregation: $\displaystyle(e_{t}^{1}$ $\displaystyle,e_{t}^{2},...,e_{t}^{n})=f_{\theta}(x_{t,1}^{1},x_{t,2}^{1},..., x_{t,K}^{1},a_{t}^{1},...,x_{t,1}^{n},x_{t,2}^{n},...,x_{t,K}^{n},a_{t}^{n})$ (5) $$

Transition Prediction.
In the transition prediction in Eq. (2), the $k$-th observation token is additionally conditioned on the tokens that were already predicted $x_{t+1,<k}^{i}\triangleq(x_{t+1,1}^{i},x_{t+1,2}^{i},...,x_{t+1,k-1}^{i})$, ensuring the autoregressive token prediction to facilitate modeling over the trajectory sequence.

Discount Prediction.
The discount predictor outputs a Bernoulli likelihood and lets us estimate the probability of an individual agent’s episode ending when learning behaviors from model predictions.

Reward Prediction.
Since the rewards distribution is widespread and diverse in multi-agent scenarios , we discrete the reward and learn a reward predictor via discrete regression instead of a mean-squared error.
Specifically, the reward predictor outputs an estimated categorical distribution $\mathcal{R}$ over $M$ buckets centered at $\{b_{m}\}^{M}_{m=1}$ through the softmax function. The estimated reward is represented as the expected value of this distribution, as

$$ $\displaystyle\hat{r}_{t}^{i}=\mathbb{E}\left[\mathcal{R}(x_{\leq t,\boldsymbol {\cdot}}^{i},a_{\leq t}^{i},e_{\leq t}^{i})\right],\quad\mathcal{R}(x_{\leq t, \boldsymbol{\cdot}}^{i},a_{\leq t}^{i},e_{\leq t}^{i})=\sum\nolimits_{m=1}^{M} p_{\phi}(b_{m}|x_{\leq t,\boldsymbol{\cdot}}^{i},a_{\leq t}^{i},e_{\leq t}^{i} )\cdot\delta_{b_{m}}$ (6) $$

Then we learn the reward predictor by using the cross-entropy loss function.
The target distribution $\mathcal{R}_{\rm tar}=\sum\nolimits_{m=1}^{M}p_{m}\cdot\delta_{b_{m}}$ can be calculated via HL-Gauss , which performs label smoothing for the target categorical distribution using a Gaussian distribution. The details are given in §B.

Agent-wise Aggregation.
Due to the partial environment, the non-stationarity issue stems from the sophisticated agent-wise inter-dependency on local observations generation.
To address it, we introduce a Perceiver to perform agent-wise representation aggregation which plays a similar role to communication.
To sustain the decentralized manner in transition prediction, we hope every agent can possess its own inner perception of the situations that all agents are in. Nonetheless, with discrete representation for local observation, the observation-action pair of agent $i$ at timestep $t$ is projected into a sequence $(x_{t,1}^{i},x_{t,2}^{i},...,x_{t,K}^{i},a_{t}^{i})$ of length $K+1$. It leads to a sequence of length $n(K+1)$ that linearly scales with the number of agents, which represents the joint observation-action pair of all agents.
Therefore, we choose the Perceiver as the agent-wise representation aggregation module, which excels at dealing with the case that the size of inputs and outputs scales linearly. Equipped with a flexible querying mechanism and self-attention mechanism, the Perceiver aggregates the joint representation sequence $(x_{t,1}^{1},x_{t,2}^{1},...,x_{t,K}^{1},a_{t}^{1},...,x_{t,1}^{n},x_{t,2}^{n}
,...,x_{t,K}^{n},a_{t}^{n})$ of length $n(K+1)$ into a sequence of $n$ feature vectors $(e_{t}^{1},e_{t}^{2},...,e_{t}^{n})$, where each feature vector serves as an intrinsic global abstraction of the environmental contexts perceived from each agent’s viewpoint.

The world model $\phi$ is trained with trajectory segments of a fixed horizon $H$ sampled from the replay buffer $\mathcal{D}$ in a self-supervised manner. The transition predictor, discount predictor, and reward predictor are optimized to maximize the log-likelihood of their corresponding targets:

$$ $\displaystyle\mathcal{L}_{\rm Dyn}(\phi,\theta)=\mathbb{E}_{i\sim\mathcal{N}} \mathbb{E}_{\tau^{i}\sim\mathcal{D}}\Bigl{[}\sum\nolimits_{t=1}^{H}$ $\displaystyle-\underbrace{\log p_{\phi}(r_{t}^{i}|x_{\leq t,\boldsymbol{\cdot} }^{i},a_{\leq t}^{i},e_{t}^{i})}_{\small\mbox{reward loss}}-\underbrace{\log p _{\phi}(\gamma_{t}^{i}|x_{\leq t,\boldsymbol{\cdot}}^{i},a_{\leq t}^{i},e_{t}^ {i})}_{\small\mbox{discount loss}}$ $\displaystyle-\underbrace{\Bigl{(}\sum\nolimits_{k=1}^{K}\log p_{\phi}(x_{t+1, k}^{i}|x_{\leq t,\boldsymbol{\cdot}}^{i},a_{\leq t}^{i},e_{t}^{i},x_{t+1,<k}^{ i})\Bigr{)}}_{\small\mbox{transition loss}}\Bigr{]}$ (7) $\displaystyle\mbox{where }(e_{t}^{1},e_{t}^{2},...,e_{t}^{n})=f_{\theta}$ $\displaystyle(x_{t,1}^{1},x_{t,2}^{1},...,x_{t,K}^{1},a_{t}^{1},...,x_{t,1}^{n },x_{t,2}^{n},...,x_{t,K}^{n},a_{t}^{n}),\forall t.$ $$

We jointly minimize this loss function in Eq. (3.2) with respect to the model parameters of local dynamics (i.e., $\phi$) and global representation (i.e., $\theta$) using the Adam optimizer .

Figure: Figure 2: Imagination procedure in MARIE. We unroll the imagination of all agents $\{1,...,n\}$ in parallel. Initially, each agent’s observation is derived from a joint observation sampled from a replay buffer. A policy, depicted in red arrows, generates actions based on reconstructed observations. Then, the Perceiver integrates joint actions and observations into global representations from each agent, appending them to each agent’s local sequence. The Transformer then predicts individual rewards and discounts, depicted by green and purple arrows respectively, while generating next observation tokens for each agent in an autoregressive manner, shown by blue arrows. This parallel imagination iterates for $H$ steps. The policies $\pi^{1:n}_{\psi}$ are exclusively trained using imagined trajectories.
Refer to caption: x2.png

### 3.3 Learning Behaviours in Imagination

We utilize the Actor-Critic framework to learn the behavior of each agent, where the actor and critic are parameterized by $\psi$ and $\xi$, respectively.
In the following, we take agent $i$ as an exemplar case for clarity and omit the superscript for denoting the index of the agent to avoid potential confusion.
Benefited from the shared local dynamics, the local trajectories of all agents are imagined in parallel, as illustrated in Fig. 2.
At timestep $t$, the actor takes a reconstructed observation $\hat{o}_{t}$ as input, and samples an action $a_{t}\sim\pi_{\psi}(a_{t}|\hat{o}_{t})$. The world model then predicts the individual reward $\hat{r}_{t}$, individual discount $\hat{\gamma}_{t}$ and next local observation $\hat{o}_{t+1}$. Starting from initial observations sampled from the replay buffer, this imagination procedure is rolled out for $H$ steps.
To stimulate long-horizon behavior learning, the critic accounts for rewards beyond the fixed imagination horizon and estimates the individual expected return $V_{\xi}(\hat{o}_{t})\simeq\mathbb{E}_{\pi_{\psi}}[\sum_{l\geq t}\gamma^{l-t}
\hat{r}_{l}]$.

In our approach, we train the actor and critic in a MAPPO-like manner.
Unlike other CTDE model-free approaches that require a global oracle state from the environment, we cannot obtain the oracle state from the world model, and only the predicted observations of each agent are available. To approximate the oracle information in critic training, we enhance each agent’s critic with the capability to access the observations of other agents.
Since the actor and critic only rely on the reconstructed observations, decoupling from the inner hidden states of the Transformer-based world model, we allow fast inference in the environment without the participation of the world model.
It is important for the deployment of policies learned with data-efficient imagination in real-world applications.
$\lambda$-target in Dreamer is used to updated the value function.
The details of behavior learning objectives and algorithmic description of MARIE are presented in §C and §G, respectively.

## 4 Experiments

We consider the most common benchmark – StarCraftII Multi-Agent Challenge (SMAC) for evaluating our method.
To highlight the sample efficiency brought by model-based imagination, we adopt a low data regime that resembles a similar setting in single-agent Atari domain .

### 4.1 Evaluations on SMAC

StarCraftII Multi-Agent Challenge. SMAC , a suite of cooperative multi-agent environments based on StarCraft II, consists of a set of StarCraft II scenarios. Each scenario depicts a confrontation between two armies of units, one of which is controlled by the built-in game AI and the other by our algorithm.
The initial position, number, and type of units in each army varies from scenario to scenario, as does the presence or absence of elevated or impassable terrain. And the goal is to win the game within the pre-specified time limit. SMAC emphasizes mastering micromanagement techniques across multiple agents to achieve effective coordination and overcome adversaries. This necessitates both sufficient exploration and appropriate credit assignment for each agent’s action. Another notable property of SMAC is that not all actions are accessible during decision-making of each agent, which requires world models to possess an in-depth comprehension of the underlying game mechanics so as to consistently provide valid available action mask estimation within the imagination horizon.

Experimental Setup.
We choose 13 representative scenarios from SMAC that includes three levels of difficulty – *Easy*, *Hard*, and *SuperHard*.
Specific chosen scenarios can be found in Table 1. In terms of different levels of difficulty, we adopt a similar setting akin to that in and restrict the number of samples from the real environment to 100k for *Easy* scenarios, 200k for *Hard* scenarios and 400k for *SuperHard* scenarios, to establish a low data regime in SMAC.
We compare MARIE with three strong model-free baselines – MAPPO , QMIX and QPLEX , and two strong model-based baselines – MBVD and MAMBA on SMAC benchmark. Specially, as a recently proposed multi-agent variant of DreamerV2 , MAMBA achieves state-of-the-art sample efficiency in various SMAC scenarios via learning in imagination. For each random seed, we compute the win rate across 10 evaluation games at fixed intervals of environmental steps.

Figure: Figure 3: Curves of evaluation win rate for methods in 8 chosen SMAC maps. See Table 1 for win rates. Y axis: win rate; X axis: number of steps taken in the real environment.
Refer to caption: x3.png

**Table 1: Median evaluation win rate and standard deviation on 13 SMAC maps for different methods over 4 random seeds. We bold the values of the maximum.**
| Maps | Difficulty | Steps | Methods |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MARIE (Ours) | MAMBA | MAPPO | QMIX | QPLEX | MBVD |  |  |  |
| 1c3s5z | Easy | 100K | 84.1(9.4) | 76.8(15.3) | 22.7(11.0) | 44.1(29.2) | 67.2(7.4) | 62.7(11.4) |
| 2m_vs_1z | 100.0(7.9) | 95.0(2.3) | 86.7(3.2) | 74.1(14.8) | 88.4(10.8) | 33.9(24.5) |  |  |
| 2s_vs_1sc | 100.0(7.1) | 99.1(7.1) | 100.0(0.0) | 0.0(0.0) | 7.1(19.5) | 0.3(14.8) |  |  |
| 2s3z | 81.8(9.3) | 69.1(12.7) | 31.3(12.9) | 37.7(15.5) | 51.1(8.4) | 52.3(4.1) |  |  |
| 3m | 99.5(0.4) | 86.4(7.1) | 84.4(12.8) | 53.7(22.7) | 89.0(6.9) | 72.4(6.9) |  |  |
| 3s_vs_3z | 99.5(1.5) | 92.3(10.1) | 0.8(1.3) | 0.0(0.0) | 0.0(0.0) | 0.0(0.0) |  |  |
| 3s_vs_4z | 63.6(24.9) | 27.7(12.3) | 0.0(0.0) | 0.0(0.0) | 0.0(0.0) | 0.0(0.0) |  |  |
| 8m | 89.1(3.9) | 65.0(7.7) | 77.3(19.5) | 74.5(12.8) | 83.0(6.4) | 74.7(9.7) |  |  |
| MMM | 25.0(3.4) | 45.0(27.6) | 4.7(4.5) | 25.0(17.3) | 88.4(35.1) | 20.4(2.1) |  |  |
| so_many_baneling | 97.7(5.9) | 93.6(4.1) | 43.8(15.0) | 22.7(8.9) | 31.2(6.1) | 12.6(10.4) |  |  |
| 3s_vs_5z | Hard | 200K | 88.6(38.8) | 10.5(14.0) | 0.0(0.0) | 0.0(0.0) | 0.0(0.0) | 0.0(0.0) |
| 2c_vs_64zg | 23.6(14.3) | 7.7(8.7) | 3.1(10.2) | 0.5(0.5) | 0.0(0.1) | 0.0(0.4) |  |  |
| corridor | SuperHard | 400K | 47.1(32.2) | 21.1(15.2) | 0.0(0.7) | 0.0(0.0) | 0.0(0.0) | 0.0(0.0) |

Main Results. Overall, we find MARIE achieves significantly better sample efficiency and a higher win rate compared with other strong baselines. We report the median win rates over four seeds in Table 1 and provide additional learning curves of several chosen scenarios, shown as Fig. 3.
As presented in Table 1 and Fig. 3, MARIE demonstrates superior performance and sample efficiency across almost all scenarios.
The improvements in sample efficiency and performance become particularly pronounced with increasing difficulty of scenarios, especially compared to MAMBA that adopts RSSM as the backbone for the world model.
We attribute such results to the model capability of the Transformer in local dynamics modeling and global feature aggregation.
Benefiting from more powerful strength in modeling sequences, the Transformer-based world model can generate more accurate and consistent imaginations than those relying on the recurrent backbone, which facilitates better policy learning within the imagination of the world model.
While the scenarios become harder, e.g. *3s_vs_5z*, our world model can address the challenge of learning more intricate underlying dynamics and further large quantities of accurate imaginations, thereby significantly outperforming other baselines on these scenarios. For *MMM* where the model-free baseline performs better, we hypothesize that the agent mainly relies on short-term behaviors for cooperation, without requiring long-term predictions for better performance.
Moreover, a special scenario *2c_vs_64zg* deserves attention, which features only 2 agents but with a considerably large action space of up to 70 discrete actions for each agent. It is easy for the world model to generate ridiculous estimated available action masks without understanding the mechanics behind this scenario, further leading to invalid or even erroneous policy learning in the imaginations of the world model. The performance gap on *2c_vs_64zg* proves that our Transformer-based world model has higher prediction accuracy and a deeper understanding of the underlying mechanics.

### 4.2 Ablation Studies

Incorporating CTDE principle with the design of the world model makes MARIE scalable and robust to different number of agents.
We compare our method with a *centralized* variant of our method, wherein the world model learns the joint dynamics of all agents together over the joint trajectory $\tau=(\ldots,o_{t}^{1},o_{t}^{2},\ldots,o_{t}^{n},a_{t}^{1},a_{t}^{2},\ldots,a
_{t}^{n},\ldots)$.
Given that $\tau$ already contains the joint observations and actions, we disable the aggregation module in this *centralized* variant. As illustrated in Figure 4, our comparisons span scenarios involving 2 to 7 agents.
When the number of agents is small enough, reducing the multi-agent system to a single-agent one over the joint observation and action space would not cause a prominent scalability issue, as indicated by the result in *2s_vs_1sc*.
However, the scalability issue is exacerbated by a growing number of agents.
In scenarios featuring more than 3 agents, the sample efficiency of the *centralized* variant encounters a significant drop, suffering from the exponential surge in spatial complexity of the joint observation-action space.
Furthermore, with equal prediction horizons, the parameter amounts in the *centralized* variant is increased by a factor of 4 or larger. And to achieve the same number of environment steps, the *centralized* variant demands over twice the original computational time.
Instead, with decentralized local dynamics and aggregated global features, MARIE delivers stable and superior sample efficiency.

Figure: Figure 4: Ablation on what manner to integrate into the design of the world model. *Decentralized* Manner denotes the standard implementation of MARIE, while *Centralized* Manner denotes that the world model is designed for learning the joint dynamics of all agents over the joint trajectory.
Refer to caption: x4.png

Agent-wise aggregation helps MARIE capture the sophisticated inter-dependency on the generation of each agent’s local observation.
To study the influence of agent-wise aggregation, we conduct ablation experiments on the aggregation module over scenarios where the number of agents gradually increases.
As shown in Fig. 5, in the 3-agents scenario (e.g., *3s_vs_3z*), the correlation among each agent’s local observation tends to be negligible. Therefore, the nearly independent generation of each agent’s local observation without any aggregated global feature still leads to performance comparable to that of standard implementation.
But as more agents get involved, the inter-dependency becomes dominant. Lacking the global features derived from agent-wise aggregation, the shared Transformer struggles to infer accurate future local observations, thus hindering policy learning in the imaginations of the world model and resulting in notable degradation in the win rate evaluation.

Figure: Figure 5: Comparisons between MARIE with and without the usage of the aggregation module.
Refer to caption: x5.png

VQ-VAE encapsulates local observations within a succinct sequence of tokens, promoting the learning of the Transformer-based world model and effectively improving algorithm performance.
Compared to VQ-VAE that discretizes each observation to $K$ tokens from $\mathcal{Z}$, perhaps a more naive tokenizer is projecting the value in each dimension into one of $m$ fixed-width bins , resulting in a $n_{obs}$-long token sequence for each observation, which we term *Bins Discretization*. We set the number of bins $m$ equal to the size of codebook $|\mathcal{Z}|$ and compare these two types of tokenizers in different environments with various $n_{\rm obs}$.
As shown in Fig. 6, the performance of the two tokenizers are comparable only in *2s_vs_1sc* where $n_{\rm obs}$ is close to 16. Even worse, *Bins Discretization* experiences a pronounced decline as $n_{\rm obs}$ increases in more complex environments (e.g., *3s_vs_4z*) under identical training durations.
We hypothesize that for a single local observation, a $n_{\rm obs}$-token-long verbose sequence yielded by *Bins Discretization* contains more redundant information compared to VQ-VAE that learns a more compact tokenizer through reconstruction
This not only renders the token sequences of *Bins Discretization* obscure and challenging to comprehend, but also results in an increase in model parameter amounts, being more computationally costly. Due to these two factors, *Bins Discretization* exhibits a notably slow convergence.
Meanwhile, the result in *2m_vs_1z* indicates *Bins Discretization* may ignore the correlation of different dimensions, which would be helpful in sequence modeling.

Figure: Figure 6: Ablation on the type of discretization for local observations. *Tokenizer* denotes the standard implementation of MARIE; *Bins Discretization* denotes the variant of MARIE where the $n_{obs}$-dimensional observation discretization is performed by projecting the value into one of $m$ fixed-width bins in each dimension independently. X-axis: cumulative run time of algorithms in the same platform.
Refer to caption: x6.png

### 4.3 Model Analysis

Error Accumulation.
A quantitative evaluation of the model’s accumulated error versus prediction horizon is provided in Fig. 7. Since learning the world model is tied to a progressively improving policy both in MARIE and MAMBA, we separately use their final policies to sample 10 episodes for fairness. We then compute $L_{1}$ errors per observation dimension between 1000 trajectory segments randomly sampled from these 20 episodes and their imagined counterpart. The result in Fig. 7 suggests architecture differences play a large role in the world model’s long-horizon accuracy. This also provides additional evidence that policy learning can benefit from accurate long-term imaginations, explaining MARIE’s notable performance in the *3s_vs_5z* scenario.
More precisely, lower generalization error between the estimated dynamics and true dynamics brings a tighter bound between optimal policies derived from these two dynamics according to theoretical results .

Attention Patterns.
During model prediction, we delve into the attention maps inside the shared Transformer and the cross attention maps in the Perceiver. Interestingly, we observe two distinct attention patterns involved in the local dynamics prediction. One exhibits a Markovian pattern wherein the observation prediction lays its focus mostly on the previous transition, while the other is regularly striated wherein the model attends to specific tokens in multiple prior transitions. During the agent-wise aggregation, we also identify two distinct patterns – *individuality* and *commonality* among agents. Such diverse patterns in the Transformer and Perceiver may be pivotal for achieving accurate and consistent imaginations of the sophisticated local dynamics. We refer to §D for further details and visualization results.

Figure: Figure 7: Compounding model errors. We compare the imagination accuracy of MARIE to that of MAMBA over the course of a planning horizon in *3s_vs_5z* scenario. MARIE has remarkably better error compounding with respect to prediction horizon than MAMBA.
Refer to caption: x7.png

## 5 Conclusion

We have introduced a model-based multi-agent algorithm – MARIE, which utilizes a shared Transformer as local dynamic model and a Perceiver as a global agent-wise aggregation module to construct a world model within the multi-agent context.
By providing long-term imaginations with policy learning, it significantly boosts the sample efficiency and improves final performance compared state-of-the-art model-free methods and existing model-based methods with same learning paradigm, in the low data regime.
As the first Transformer-based multi-agent world model for sample-efficient policy learning, we open a new avenue for combining the powerful strength of the Transformer with sample-efficient MARL.
Considering the notorious sample inefficiency in multi-agent scenarios, it holds important promise for application in many realistic multi-robot systems, wherein collecting tremendous samples for optimal policy learning is costly and impractical due to the safety.
While it has the great potential to bright the future towards achieving smarter multi-agent systems, there still exist limitations in MARIE. For instance, it would suffer from much slower inference speed when used with a very long prediction horizons, due to the auto-regressive property.

## Appendix A World Models Details and Hyperparameters

### A.1 Observation Tokenizer

Our tokenizer for local observation discretization is based on the implementation(^1^11Code can be found in [https://github.com/lucidrains/vector-quantize-pytorch](https://github.com/lucidrains/vector-quantize-pytorch)) of a vanilla VQ-VAE [].
Faced with continuous non-vision observation, we build the encoder and decoder as Multi-Layer Perceptrons (MLPs). The decoder is designed with the same hyperparameters as the ones of the encoder.
The hyperparameters are listed as Table 2. During the phase of collecting experience from the external environment, each agent takes the reconstructed observations processed by the VQ-VAE as input instead to avoid the distribution shift between policy learning and policy execution.

For training this vanilla VQ-VAE, we use a straight-through estimator to enable gradient backpropagation through the non-differentiable quantization operation in the quantization of VQ-VAE. The loss function for learning the autoencoder is as follows:

$$ $\mathcal{L}_{\rm VQ-VAE}(E,D,\mathcal{Z})=\mathbb{E}_{i\sim\mathcal{N}}\mathbb {E}_{o^{i}}\left[\|o^{i}-\hat{o}^{i}\|^{2}+\|{\rm sg}[E(o^{i})]-z_{q}^{i}\|^{2 }+\beta\|{\rm sg}[z_{q}^{i}]-E(o^{i})\|^{2}\right]$ (8) $$

where $\mathcal{N}=\{1,2,...,n\}$ denotes the set of agents, ${\rm sg}[\cdot]$ denotes the stop-gradient operation and $\beta$ is the coefficient of the commitment loss $\|{\rm sg}[z_{q}^{i}]-E(o^{i})\|^{2}$.
In practice, we found the codebook $\mathcal{Z}$ can suffer from codebook collapse when learning from scratch. Thus, we adopt the Exponential Moving Averages (EMA) [] technique to alleviate this problem.

**Table 2: VQVAE hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| Encoder&Decoder |  |
| Layers | 3 |
| Hidden size | 512 |
| Activation | GELU[] |
| Codebook |  |
| Codebook size ($N$) | 512 |
| Tokens per observation ($K$) | 16 |
| Code dimension | 128 |
| Coef. of commitment loss ($\beta$) | 10.0 |

### A.2 Transformer

The shared Transformer serving as the local dynamics model is based on the implementation of minGPT []. Given a fixed imagination horizon $H$, it first takes a token sequence of length $H(K+1)$ composed of observation tokens and action tokens, and embeds it into a $H(K+1)\times D$ tensor via separate embedding tables for observations and actions. Then, the aggregated feature tensor, returned by the agent-wise aggregation module, is inserted after the action embedding tensor at every timestep, forming a final embedding tensor of shape $H(K+2)\times D$. This tensor is forwarded through fixed Transformer blocks. Here, we adopt GPT2-like blocks [] as the basic blocks.
The hyperparameters are listed as Table 3. To enable training across all environments on a single NVIDIA RTX 3090 GPU, we adapt imagination horizon $H$ based on the number of agents.

**Table 3: Transformer hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| Imagination horizon ($H$) | {15, 8, 5} |
| Embedding dimension | 256 |
| Layers | 10 |
| Attention heads | 4 |
| Weight decay | 0.01 |
| Embedding dropout | 0.1 |
| Attention dropout | 0.1 |
| Residual dropout | 0.1 |

### A.3 Perceiver

The Perceiver [] is based on the open-source implementation(^2^22Code can be found in [https://github.com/lucidrains/perceiver-pytorch](https://github.com/lucidrains/perceiver-pytorch)).
By aligning the length of the latent querying array with the number of agents $n$, we obtain the intrinsic global representation feature corresponding to each individual agent.
We further dive into the process of agent-wise representation aggregation: (i) the embedding tensor of shape $(K+1)\times D$ at each timestep, mentioned in Appendix A.2, is concatenated with others from all agents, thereby getting a $n(K+1)\times D$ sequence for the joint observation-action pair at the current timestep; (ii) through the cross-attention mechanism with the latent querying array, the original sequence is compressed from length $n(K+1)$ to $n$; (iii) the compressed sequence is then forwarded through a standard transformer with bidirectional attention inside the Perceiver. The hyperparameters are listed as Table 4.

**Table 4: Perceiver hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| Length of latent querying | $n$ (number of agents) |
| Cross attention heads | 8 |
| Inner Transformer layers | 2 |
| Transformer attention heads | 8 |
| Dimension per attention head | 64 |
| Embedding dropout | 0.1 |
| Attention dropout | 0.1 |
| Residual dropout | 0.1 |

## Appendix B Reward Labels in HL-Gauss

HL-Gauss method [] leverages Histogram Losses introduced by []. Firstly, we can define the random variable $Y$ with probability density $f_{Y}$ and cumulative distribution function $F_{Y}$ whose expectation is $r_{t}^{i}$. Given $M$ buckets of width $\varsigma=(r_{\rm max}-r_{\rm min})/M$ centered at $\{b_{m}\}^{M}_{m=1}$, the distribution of $Y$ projected onto the histogram with these buckets can be calculated via integrating over the interval $[b_{m}-\varsigma/2,b_{m}+\varsigma/2]$,

$$ $\displaystyle p_{m}(b_{m})$ $\displaystyle=\int_{b_{m}-\varsigma/2}^{b_{m}+\varsigma/2}f_{Y}(y)\,dy$ $\displaystyle=F_{Y}(b_{m}+\varsigma/2)-F_{Y}(b_{m}-\varsigma/2)$ (9) $$

To stabilize classification learning over potential diverse reward distribution, we can model the distribution of $Y$ as a Gaussian distribution $\mathcal{N}(\mu=r_{t}^{i},\sigma^{2})$, which smooths the target distribution computed by Eq. (B).
The hyperparameter $\sigma$ is used to control the degree of label smoothing, and is typically set to be $0.75\varsigma$, as advised by []. We refer to [] for more details.

## Appendix C Behaviour Learning Details

In MARIE, we use MAPPO-like [] actor and critic, where the actor and critic should have been 3-layer MLPs. However, unlike other CTDE model-free approaches, whose critic takes additional global oracle states from the environment in the training phase, our world model hardly provides related predictions in the imagined trajectories. To alleviate this issue, we augment the critic with an attention mechanism and provide it all reconstructed observations $\boldsymbol{\hat{o}}_{t}$ of all agents.
Therefore, the actor $\psi$ remains a 3-layer MLP with ReLU activation, while the critic $\xi$ is enhanced with an extra layer of self-attention, built on top of the original 3-layer MLP, i.e., we overwrite the critic $V_{\xi}^{i}(\boldsymbol{\hat{o}}_{t})\simeq\mathbb{E}_{\pi_{\psi}^{i}}(\sum
\nolimits_{l\geq t}\gamma^{l-t}\hat{r}^{i}_{l})$ for agent $i$.
Similar to off-the-shelf CTDE model-free approaches, we adopt parameter sharing across agents.

#### Critic loss function

We utilize $\lambda$-return in Dreamer [], which employs an exponentially-weighted average of different $k$-steps TD targets to balance bias and variance as the regression target for the critic.
Given an imagined trajectory $\{\hat{o}_{\tau}^{i},a_{\tau}^{i},\hat{r}_{\tau}^{i},\hat{\gamma}_{\tau}^{i}\}
^{H}_{t=1}$ for agent $i$, $\lambda$-return is calculated recursively as,

$$ $V_{\lambda}^{i}(\boldsymbol{\hat{o}}_{t})=\left\{\begin{array}[]{lcl}\hat{r}_{ t}^{i}+\hat{\gamma}_{t}^{i}\left[(1-\lambda)V_{\xi}^{i}(\boldsymbol{\hat{o}}_{ t})+\lambda V_{\lambda}^{i}(\boldsymbol{\hat{o}}_{t+1})\right]&\mbox{if}&t<H\\ V_{\xi}^{i}(\boldsymbol{\hat{o}}_{t})&\mbox{if}&t=H\end{array}\right.$ (10) $$

The objective of the critic $\xi$ is to minimize the mean squared difference $\mathcal{L}_{\xi}^{i}$ with $\lambda$-returns over imagined trajectories for each agent $i$, as

$$ $\mathcal{L}_{\xi}^{i}=\mathbb{E}_{\pi_{\psi}^{i}}\left[\sum\nolimits_{t=1}^{H- 1}\left(V_{\xi}^{i}(\boldsymbol{\hat{o}}_{t})-{\rm sg}(V_{\lambda}^{i}( \boldsymbol{\hat{o}}_{t}))\right)^{2}\right]$ (11) $$

where ${\rm sg}(\cdot)$ donotes the stop-gradient operation. We optimize the critic loss with respect to the critic parameters $\xi$ using the Adam optimizer.

#### Actor loss function

The objective for the action model $\pi_{\psi}(\cdot|\hat{o}_{t}^{i})$ 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 $\lambda$-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 important sampling for actor learning. The actor loss function for agent $i$ is:

$$ $\displaystyle\mathcal{L}_{\psi}^{i}=-\mathbb{E}_{p_{\phi},\pi_{\psi_{\rm old}} ^{i}}\Big{[}\sum_{t=0}^{H-1}\mathop{\min}\Big{(}r_{t}^{i}(\psi)A_{t}^{i},{\rm clip }(r_{t}^{i}(\psi),1-\epsilon,1+\epsilon)A_{t}^{i}\Big{)}+\eta\mathcal{H}(\pi_{ \psi}^{i}(\cdot|\hat{o}_{t}^{i}))\Big{]}$ (12) $$

where $r_{t}^{i}(\psi)=\pi_{\psi}^{i}/\pi_{\psi_{\rm old}}^{i}$ is the policy ratio and $A_{t}^{i}={\rm sg}(V_{\lambda}^{i}(\boldsymbol{\hat{o}}_{t})-V_{\xi}^{i}(
\boldsymbol{\hat{o}}_{t}))$ is the advantage. We optimize the actor loss with respect to the actor parameters $\psi$ using the Adam optimizer.
In the discount prediction of MARIE, we set its learning target $\gamma$ to be 0.99. Overall hyperparameters are shown in Table 5.

**Table 5: Behaviour learning hyperparameters.**
| Hyperparameter | Value |
| --- | --- |
| Imagination Horizon ($H$) | {15, 8, 5} |
| Predicted discount label $\gamma$ | 0.99 |
| $\lambda$ | 0.95 |
| $\eta$ | 0.001 |
| Clipping parameter $\epsilon$ | 0.2 |

## Appendix D Extended Analysis on Attention Patterns

To provide qualitative analysis of our world model, we select typical scenarios – *3s_vs_5z* where our method achieves the most significant improvement compared to other baselines for visualizing attention maps inside the Transformer.
For the sake of simple and clear visualization, we set the imagination horizon $H$ as 5.
In terms of cross-attention maps in the aggregation module, we select a scenario *2s3z* including 5 agents for visualization.
Visualization results are depicted as Fig. 8 and Fig. 9.

The prediction of local dynamics entails two distinct attention patterns. The left one in Fig. 8 can be interpreted as a Markovian pattern, in which the observation prediction lays its focus on the previous transition.
In contrast, the right one is regularly striated, with the model attending to specific tokens in multiple prior observations.
In terms of the agent-wise aggregation, we also identify two distinct patterns: *individuality* and *commonality*. The top one in Fig. 9 illustrates that each agent flexibly attends to different tokens according to their specific needs. In contrast, the bottom one exhibits consistent attention allocation across all agents, with attention highlighted in nearly identical positions.
The diverse patterns in the Transformer and Perceiver may be the key to accurate and consistent imagination.

Figure: Figure 8: Attention patterns in the Transformer. We observe two distinct types of attention weights during the prediction of local dynamics. In the first one (*left*), the next observation prediction is primarily dependent on the last transition, which means the world model has learned the Markov property corresponding to Dec-POMDPs. The second type (*right*) exhibits a regularly striated pattern, where the next observation prediction hinges overwhelmingly on the same dimension of multiple previous timesteps. The above attention weights are produced by a sixth-layer and ninth-layer attention head during imaginations on the *3s_vs_5z* scenario.
Refer to caption: x8.png

Figure: Figure 9: Cross attention patterns in the Perceiver. We observe the *individuality* and *commonality* in the agent-wise aggregation. The top part of the figure represents the *individuality*, where agents adjust their attentions over the whole joint token sequence at timestep $t$ flexibly according to their own needs. In contrast, the bottom exhibits the *commonality*, where every agent’s attention over the joint token sequence is emphasized in the similar positions of the sequence. The cross attention weights mentioned above are produced by the first and sixth head of the cross attention within the Perceiver, during the agent-wise aggregation on the *2s3z* scenario.
Refer to caption: x9.png

## Appendix E Baseline Implementation Details

MAMBA [] is evaluated based on the open-source implementation: [https://github.com/jbr-ai-labs/mamba](https://github.com/jbr-ai-labs/mamba) with the hyperparameters in Table 6.

**Table 6: Hyperparameters for MAMBA in SMAC environments.**
| Hyperparameter | Value |
| --- | --- |
| Batch size | 256 |
| $\lambda$ for $\lambda$-return computation | 0.95 |
| Entropy coefficient | 0.001 |
| Entropy annealing | 0.99998 |
| Number of policy updates | 4 |
| Epochs per policy update | 5 |
| Clipping parameter $\epsilon$ | 0.2 |
| Actor Learning rate | 0.0005 |
| Critic Learning rate | 0.0005 |
| Discount factor $\gamma$ | 0.99 |
| Model Learning rate | 0.0002 |
| Number of model training epochs | 60 |
| Number of imagined rollouts | 800 |
| Sequence length | 20 |
| Imagination horizon $H$ | 15 |
| Buffer size | $2.5\times 10^{5}$ |
| Number of categoricals | 32 |
| Number of classes | 32 |
| KL balancing entropy weight | 0.2 |
| KL balancing cross entropy weight | 0.8 |
| Gradient clipping | 100 |
| Collected trajectories between updates | 1 |
| Hidden size | 256 |

MAPPO [] is evaluated based on the open-source implementation: [https://github.com/marlbenchmark/on-policy](https://github.com/marlbenchmark/on-policy) with the common hyperparameters in Table 7.

**Table 7: Common hyperparameters for MAPPO in SMAC environments.**
| Hyperparameter | Value |
| --- | --- |
| Batch size | num envs × buffer length × num agents |
| Mini batch size | batch size / mini-batch |
| Recurrent data chunk length | 10 |
| GAE $\lambda$ | 0.95 |
| Discount factor $\gamma$ | 0.99 |
| Value loss | huber loss |
| Huber delta | 10.0 |
| Optimizer | Adam |
| Optimizer learning rate | 0.0005 |
| Optimizer epsilon | $1\times 10^{-5}$ |
| Weight decay | 0.0 |
| Gradient clipping | 10 |
| Network initialization | orthogonal |
| Use reward normalization | True |
| Use feature normalization | True |

QMIX [] is evaluated based on the open-source implementation: [https://github.com/oxwhirl/pymarl](https://github.com/oxwhirl/pymarl) with the hyperparameters in Table 8.

**Table 8: Hyperparameters for QMIX in SMAC environments.**
| Hyperparameter | Value |
| --- | --- |
| Batch size | 32 episodes |
| Buffer size | 5000 episodes |
| Epsilon in epsilon-greedy | $1.0\rightarrow 0.05$ |
| Epsilon anneal time | 50000 timesteps |
| Train interval | 1 episode |
| Discount factor $\gamma$ | 0.99 |
| Optimizer | RMSProp |
| RMSProp $\alpha$ | 0.99 |
| RMSProp $\epsilon$ | $10^{-5}$ |
| Gradient clipping | 10 |

QPLEX [] is evaluated based on the open-source implementation: [https://github.com/wjh720/QPLEX](https://github.com/wjh720/QPLEX) with the hyperparameters in Table 9. Since its implementation is mostly based on the open-source implementation: [PyMARL](https://%20github.com/oxwhirl/pymarl) [], its most hyperparameters setting remains the same as the one in QMIX in addition to its own special hyperparameters.

**Table 9: Hyperparameters for QPLEX in SMAC environments.**
| Hyperparameter | Value |
| --- | --- |
| Batch size | 32 episodes |
| Buffer size | 5000 episodes |
| Epsilon in epsilon-greedy | $1.0\rightarrow 0.05$ |
| Epsilon anneal time | 50000 timesteps |
| Train interval | 1 episode |
| Discount factor $\gamma$ | 0.99 |
| Optimizer | RMSProp |
| RMSProp $\alpha$ | 0.99 |
| RMSProp $\epsilon$ | $10^{-5}$ |
| Gradient clipping | 10 |
| Number of layers in HyperNetwork | 1 |
| Number of heads in the attention module | 4 |

MBVD [] is evaluated based on the implementation in its supplementary material from [https://openreview.net/forum?id=flBYpZkW6ST](https://openreview.net/forum?id=flBYpZkW6ST) with the hyperparameters in Table 10.
Akin to QPLEX, its implementation is based on the open-source implementation: [PyMARL](https://%20github.com/oxwhirl/pymarl), its most hyperparameters setting remains the same as the one in QMIX in addition to its own special hyperparameters.

**Table 10: Hyperparameters for MBVD in SMAC environments.**
| Hyperparameter | Value |
| --- | --- |
| Batch size | 32 episodes |
| Buffer size | 5000 episodes |
| Epsilon in epsilon-greedy | $1.0\rightarrow 0.05$ |
| Epsilon anneal time | 50000 timesteps |
| Train interval | 1 episode |
| Discount factor $\gamma$ | 0.99 |
| Optimizer | RMSProp |
| RMSProp $\alpha$ | 0.99 |
| RMSProp $\epsilon$ | $10^{-5}$ |
| Gradient clipping | 10 |
| Number of layers in HyperNetwork | 1 |
| Number of heads in the attention module | 4 |
| Horizon of the imagined rollout | 3 |
| KL balancing $\alpha$ | 0.3 |
| Dimension of the latent state $\hat{s}$ | num agents x 16 |

## Appendix F Parameters Setting in SMAC

All our experiments are run on a machine with a single NVIDIA RTX 3090 GPU, a 36-core CPU, and 128GB RAM.
We provide the hyperparameters of MARIE for experiments in SMAC, shown as Table 11.
To enable the running of experiments in all scenarios with a single NVIDIA RTX 3090 GPU, we set the imagination horizon $H$ as 8 for other scenarios involving the number of agents $n>5$, 15 for $n\leq 5$.
In *so_many_baneling* and *2s3z*, we set the imagination horizon $H$ as 5.
Correspondingly, the number of policy updates in imaginations varies with imagination horizon $H$.

Figure: Figure 10: Performance between different loss functions for reward prediction in *3s_vs_5z*.
Refer to caption: x10.png

Additionally, it should be noted that the result in *3s_vs_5z* was obtained with using L1 loss function for reward prediction, which was the same as the reward prediction in MAMBA.
We found the performance in this scenario is quite sensitive to the selection of reward loss function, illustrated by Fig. 10.
Considering the significantly large action space of *2c_vs_64zg*, we enable the observation of agent id and last action for each agent and disable stacking the last 5 observations as input to the policy.

**Table 11: Hyperparameters for MARIE in SMAC environments.**
| Hyperparameter | Value |
| --- | --- |
| Batch size for tokenizer training | 256 |
| Batch size for world model training | 30 |
| Optimizer for tokenizer | AdamW |
| Optimizer for world model | AdamW |
| Optimizer for actor & critic | Adam |
| Tokenizer learning rate | 0.0003 |
| World model learning rate | 0.0001 |
| Actor learning rate | 0.0005 |
| Critic learning rate | 0.0005 |
| Gradient clipping for actor & critic | 100 |
| Gradient clipping for tokenizer | 10 |
| Gradient clipping for world model | 10 |
| Weight decay for world model | 0.01 |
| $\lambda$ for $\lambda$-return computation | 0.95 |
| Discount factor $\gamma$ | 0.99 |
| Entropy coefficient | 0.001 |
| Buffer size (transitions) | $2.5\times 10^{5}$ |
| Number of tokenizer training epochs | 200 |
| Number of world model training epochs | 200 |
| Collected transitions between updates | 100 or 200 |
| Epochs per policy update (PPO epochs) | 5 |
| PPO Clipping parameter $\epsilon$ | 0.2 |
| Number of imagined rollouts | 600 or 400 |
| Imagination horizon $H$ | {15, 8, 5} |
| Number of policy updates | {4, 10, 30} |
| Number of stacking observations | 5 |
| Observe agent id | False |
| Observe last action of itself | False |

## Appendix G Overview of MARIE Algorithm

Pseudo-code is summarized as Algorithm 1.

Figure: Algorithm 1 MARIE

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
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?
soundness: 2
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
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?
soundness: 2
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.
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.
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.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: - The paper presentation would benefit from increasing the font size of figures in the main text.
soundness: 3
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.
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.
weaknesses: - The paper presentation could be improved with captioning figures of experimental results with short conclusions

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 2
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.
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.
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.

Meta-review: 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.

Decision: Reject