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
Venue: Accepted by TMLR ()

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 across different number of agents in a centralized architecture, 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. Extensive results on Starcraft Multi-Agent Challenge (SMAC) and MAMujoco demonstrate superior sample efficiency and overall performance compared to strong model-free approaches and existing model-based methods.

Published in Transactions on Machine Learning Research (04/2025)
Decentralized Transformers with Centralized Aggregation are Cample-Efficient Multi-Agent World Models
Yang Zhang', Chenjia Bai²t, Bin Zhao', Junchi Yan", Xiu Li", Xuelong Li
'Tsinghua University
*Institute of Artificial Intelligence (TeleAT), China Telecom $ SShanghai Raotong University
Shanggaai AI Laboratory
Reviewed on OpenReview: https: // openreview. net/ forun? id- #TBBEgXVVV
Abstract
Learning world model for model-free Reinforcement Learning (RL) agents 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 across different number of agents a centralized architecture, 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 centralized representation aggregation from all agents. We cast the dynamics learning as 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 pioncering Transformer-based world model for multi-agent systems, we introduce Perceiver Transformer as an effective solution to enable centralized representation aggregation within this context. Extensive results on Starcraft Multi-Agent Challenge (SMAC) and MAMujoco demonstrate superior sample efficiency and overall performance compared to strong model-free approaches and existing model-based methods. Our code is available at https //gt thub. com com/bbreezzyoon/MMRRI
1 Introduction
Multi-Agent Reinforcement Learning (MARL) has made remarkable progress, driven largely by model- free algorithms (Nguyen et al., 2020). However, due to the complexity of multi-agent systems arising from partial observability and multiple agents' inter-dependeucces that induce rapidly growing state-action space complexities, such algorithms usually demand extensive interactions to learn coordinative behaviors (Hernandez-Leal et al., 2020). 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 (Hafner et al., 2020; Lukasz Kaiser et al., 2020; Hafner et al., 2021; 2023; Hansen et al., 2022; 2024). However, extending the design of world model in the single-agent domain to the multi-agent context encounters unique challenges due to the inter-agent interactions inherent to multi-agent environments.
The challenges primarily stem from two different means for multi-agent dynamics learning: centralized and decentralized. Centralized world models can naturally encapsulate agent inter-dependency through joint dynamics modeling, but suffer from the combinatorial explosion of the joint action space. This leads exponential growth in sample complexity the number of agents incrcases (Hernandoz-Leal et al., 2020; Nguyen et al., 2020), making them impractical for robustly scaling across diverse leam configurations (see Figure in our ablation studies). Conversely, applying decentralized world model to approximating the local dynamics of each agent mitigates the scalability issue yet incurs non-stationarity, as unexpected interventions
"Work doue during Yang Zhang's interuship both Shanghai Al Lab and TeleAL. i "Correspondigg Authors.

---

Published in Transactions on Machine Learning Research (04/2025)
from other agents may occur in each agent's individual environment (Olichock et al., 2016). Furthermore, beyond these unique challenges inherent in modeling multi-agent dynamics, existing model-based MARL approaches (Willemsen et al., 2021; Egorov & Shpilman, 2022; Xu et al., 2022) excessively neglect the fact that the policy Icarned in imaginations of the world model heavily relies on the quality of imagined trajectories (Micheli et al., 2023). thereby necessitates accurate long-lerm predictions, especially with respect to the non-stationary local dynamics.
Inspired by the capability of Transformer (Vaswani et al., 2017) in modeling complex discrete sequences and long-term dependency (Brown et al., 2020; Devlin et al., 2019; Micheli et al., 2023), seek to construct a Transformer-basee world model within the multi-agent context for decentralized local dynamics together with centralized feature aggregation, combining the benefits of two distinct designs.
In this paper, we introduce MARIE (Multi-Agent uuto-Regressive Imagination for Efficient learning), the first Transformer-basee multi-agent world model for sample-cfiicient policy learning. Specifically, the highlights of this paper are:
To tackle the inherent challenges within the multi-agent context, we build an effective world model via scalable decentrulized dynamics modeling and essential centralized representation aggregaling, which mirrors the principle of Centralized Training and Decentralized Execution (CTDE).
2. To enable accurate and consistent long-term imaginations from the non-stationary local dynamics, we cast the decentralized dynamics learning sequence modeling over discrete tokens by leveraging highly expressive Transformer architecture the backbone. In particular, successfully present the first Transformer-basee world model for multi-agent systems.
3. While it remains open for how to effectively enable centralized representation with the Transformer as the backbone, achieve it by innovatively introducing Perceiver Transformer (Jaegle et al., 2021) for efficient global information aggregation across all agents.
4. Experiments on the Starcraft Multi-Agent Challenge (SMAC) benchmark in low data regime and additional experiments on MAMujoco 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 exccution (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 prriiiionmakiing Model-free MARL methods with this paradigm can be divided into 2 categories: value-based (Sunchag et al., 2018; Rashid et al., 2018; Son et al., 2019; Wang et al., 2021) and policy-based (Lowe et al., 2017; Foerster et al., 2018; Iqbal & Sha, 2019; Ryu et al., 2020; Liu et al., 2020; Kuba et al., 2021; Peng et al., 2021; Yu et al., 2022; Zhang et al., 2024b;a). In contrast to model-free approaches, model-based MARL algorithma remain fairly understudied. MAMBPO (Willemsen et al., 2021) incorporates MBPO-style (Janner et al., 2019) techniques into multi-agent policy learning under the CTDE framework. Tesseract (Mahajan et al., 2021) 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 (Egorov & Shpilman, 2022) integrates the backbone proposed in DreamerV2 (Hafner et al., 2021) with attention mechanism across agents to sustain an effective world model in environments with an arbitrary number of agents, which leads to notably superior sample elficiency to existing model-free approaches. In terms of model-based algorithm coupled with planning, MAZero (Liu et al., 2024) expands the MCTS planning- -aasnd Muzero (Schrittwieser et al., 2020) framework to the model-based MARL settings. However, learning- basod or planning- 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. In contrast, the world
policies these
2

---

Published in Transactions on Machine Learning Research (04/2025)
model in our MARIE is solely utilized to accelerate policy learning during the training phase, rather than enhancing the policy with additional informative features for decision-making. Consequently, MARIE can be seamlessly integrated with any model-free algorithm, enabling the rapid deployment and inference of a well-trained policy without requiring an accompanying world model.
Besides, a recent. method CoDreamer (Toledo Prorok, 2024) extends DreamerV3 (Hafner et al., 2023) to the multi-agent setting with a graph neural network as a communication module. A related in-depth discussion is deferred to $F.
Learning behaviors within the imagination of world models. The Dyna architecture (Sutton, 1991) 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 (Ha & Schmidhuber, 2018) is initially introduced by composing a variational and a recurrent network to mimic the complete environmental dynamics, then an artificial agent is trained entirely inside the hallueinated imagination generated by the world model. Sim (Lukasz Kaiser al., 2020) shows that PPO policy (Schulman et al., 2017) learned in a predictive model delivered super- humann performance in Atari domains. Dreamer (Hafner et al., 2020) built the world model upon a Recurrent State Space Model (RSSM) (Hafner et al., 2019) 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 (Hafner et al., 2021) further demonstrates the advantage of discrete latent states over Gaussian states. For MARL, MAMBA (Egorov & Shpilman, 2022) 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 (Vaswani et al., 2017), TransDreamer (Chen et al., 2022) and TWM (Robine et al., 2023) explored variants of Dreamer wherein the backbones of the world model were substituted with Transformers. Instead of incorporating deterministic and stochastic latent states, IRIS (Micheli et al., 2023) 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 proposod MARIE concentrates on est establishing effective Transformer-basel world models in multi-agent. contexts with shared dynamics and global representations.
(VAE) (Kingma & Welling, 2014)
Preliminaries. We focus on fully cooperative multi-agent systems where all agents share team reward signal. We formulate the system as a decentralized partially observable Markov decision process (Dec-POMDP) (Oliehoek et al., 2016), which can be described by a tuple S,, P. R, S2, 7) N = =N n denotes a set of agents, S is the finite global state space, AA = II A is the product of finite actions spaces of all agents, i.e., the joint action space, P: PPS: S x A x S > [0, 1] is the global transition probability function, R: S A - R is the shared reward function, Ω I Q is the product of finite observation spaces of all agents, i.e, the joint observation space, O = 10', e } is the set of observing functions of all agents. Oi S + Ω2 maps global states to the observations for agent i, and y is the discount factor. Given global state st at timestep agent i is restricted to obtaining solely its local observation of O'(s,, takes an action of drawn from its policy x (-102,) based on the history of its local observations d'z which together with other agents' actions gives joint action a, = (a)...., a.) E A, equivalently drawn from a joint policy (-(l=))) = II (-)o)))). Then the agents receive shared reward r; = R(st, at), and the environment moves to next state St+1 with probability P(se+1/5t, a.). The aim of all agents is to learn joint policy π that maximizes the expected discounted return J [(T) = Esp.ag.... E 7 R(xt, az)].
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 hyperparameters are provided in SA.
3

---

Published in Transactions on Machine Learning Research (04/2025)
Joint obs-action sequence (n (K + 1) x D)
Agents 1:72
Pi
of gf+1
Latent array (n x D)
MLP Encoder E
Cross Attention
Codebook Z
Shared Transformer
Quantize()
Agntts 1:n
x
OO0
x
Perceiver Transformer
MLP Decoder D
Agent-wise Aggregation
e
Observation lokens
Aggregated features
01-10 g
Action tokens
Invalid outputs
Figure Overview of the proposed world model architecture in MARIE. VQ-VAE (left) maps local observations o' of each agent i into discrete latent codes (ri,..., zk); where (E, is shared across all agents. Together with discrete actions, this process forms local discrete sequences (..,, 2 : " di...) of each agent. Then the Perceiver (right) performs aggregation of joint discrete sequences of all agents (rt,,,...,)) x ......, ",,,,,,, F,, a.)) independently at each timestep t, and inserts the aggregated global representations (et, ef....,e.) et) into original local discrete sequences respectively. The resulting sequences (..., æ) Pi.1..... ot,K' d/, e!...) contain rich information between transitions in local dynamics and are fed into the shared Transformer (middle), which Iearns observation token predictions in an autoregrossive mannor. Predictions of individual reward rt and discount. 7f at limestep are computed based on all historical sequences (a (1' T' azg e..).
3.1 Discretizing Observation
We consider that trajectory 7' of agent i consists of T local observations and actions, as
= = (oj,ah..... of of....,,o o7, a+).
To utilize the expressive Transformer architecture, we need to express the trajectory into a discrete token se- quence for modeling. Accounting for continuous observations, a prevalent but naive approach for discretization involves discretizing the scalar into one of m fixtd-width bins in each dimension independently (Janner et al., 2021). Such discretization would encode the observation with more tokens when faced with a higher dimension of the observation, resulting in longer sequence the Transformer's input. Thus it causes higher computa- tional complexily of sequence modeling via the Transformer because of the quadratic complexity with respect to the sequence length. In contrast, a discrete codebook of learned compact representations can concisely encode a observation with less tokens. To this end, we employ the idea from neural discrete representation learning (van den Oord et al., 2017), and learn Vector Quantised- Variational AutoEncoder (VQ- (AE) to play role that resembles the tokenizer Natural Language Processing (Devlin et al., 2019; Brown et al., 2020). The VQ-VAE composed of an encoder E, decoder D, and codebook Z. We define the discrete codebook Z {2 C R"= where N is the size of the codebook and n, is the dimension of codes. The encoder E takes an observation of E Ro as input and outputs K n.-dimeeniooall latents 2 R Kxx, reshaped from the dircct outputs of encoder. Subsequently, the tokens { E (0, 1,.., 1} for representing of is obtained by nearest neighbour look-up using the codebook Z where ot arg min, || - =zyll- Then the decoder D (0, 1, .... N 1]K Re converts tokens back into an reconstructed observation of. By learning this discrete codebook, we compress the redundant information via a succinct sequence of tokens, which helps improve sequence modeling. See $4.3 for discussion.
4

---

Published in Transactions on Machine Learning Research (04/2025)
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 (Janner et al., 2021; Brohan et al., 2023). Therefore, a trajectory 7² of agent can be treated as a sequence of tokens,
,' (-. o6 -) (..., (..,,,,,1,1
KK ay....)
(1)
where si. the j-th loken of the observation of agent at timestep Given arbitrary sequences of observation and action tokens in Eq. 1, we try to learn over discrete multimodal tokens.
The world model consists of 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 & predicts the future local observation {d1+1,)ン= ) the future individual reward f and discount Ne based on the agent's individual historical observation-action history (a'se.. and aggregated global features et of the agent. The modules are shown in Eqs. 2-5.
Transition:
4111. Po( Iest. (g with of t111 pollttin a
(2)
Reward:
P/~ ~ po (f d (())
(3)
Discount:
1 ~ Po (fill' d'y. d'si)
(4)
Aggregation:
(et, er, fo(2),1 1 " ... 21, o....., F4,1 F7 Ft,K a.')
(5)
Transition Prediction. In the transition prodiction in Eq. 2, the k-th observation token is additionally conditioned on the tokens that were already predicted #1+1,<k A (11-110=+11. "...,,.)....()1u_--) ensuring the autoregressive token prediction to facilitate modeling over the trajectory sequence. Inter-step auto regression is as intuitive as predicting the future based on all information in the past while intra-step auto regression can be interpreted as learning how to compose the language provided by VQ-VAE to correctly express the observation within a certain timestep, since the tokens for encoding observations can be viewed as a special inner language like the human's.
Discount and Reward Prediction. The discount predictor outputs Bernoulli likelihood and lets estimate the probability of an individual agent's episode ending when learning behaviors from model predictions. And we simply adopt smooth L1 loss for training the prediction of reward.
Agent-wise Aggregation. Due to the partial environment, the non-stationarity issue stems from the sophisticated agent-wise wise inter-dependency on local observations generation. To address it, we introduce a Perceiver (Jacgle al., 2021) perform agent-wiso representation aggregation which plays a similar role communication. To sustain the decentralized manner in transition prediction, we hope every agent can possess its own inner perception of the whole situation. Nonetheless, with discrete representation for local observation, the observation-action pair of agent i at timestep & is projected into sequence (1,, "i. "..,,,,,, K at) of length K + 1. It leads to a joint observation-actin sequence of length n(K + 1) at timestep, which linearly scales with the number of agents.
A naive approach for extracting aggregated features for each agent is using self-attention (Egorov Shpilman, 2022; Liu et al., 2024) which takes as input this sequence of length n(K + 1) and outputs a sequence of the same length containing aggregated features of all agents, described
(#t,1..., T[,K
K. at .... ".,,..,, TLKK a))
Sell-Attontion, (e),1, "..........eee....... C..Kt er, K ((,41))
Aggrogating
where et,, is the j-th aggregated feature for agent at timestep t. However, when composing the informative sequence of local trajectories by inserting these aggregated features into the sequence of length H(K + 1) in Eq. 1, the length of the local sequence involving aggregated features would be twice as long, i.e., 2H + 1). Due to the quadratic computational complexity of Transformer, it may hinder the efficient sequence modeling over this sequence.
To this end, we choose the Perceiver as the agent-wise representation aggregation module, which excels dealing with the case that the size of inputs scales linearly and then generates compact output sequence.
5

---

Published in Transactions on Machine Learning Research (04/2025)
mm
MLP Encoder
Agent-wise Aggregation
Agent-wise Aggregation
x
x
aj
ei
2,2
x2.
x2.
az
e
Agents 1:12
MLP Decoder D
MLP Decoder D
Figure 2: Imagination procedure in MARIE. We unroll the imagination of all agents {1, n) in parallel. Initially, each agent's observation derived from a joint observation sampled from a replay buffer. A policy, depicted in red arrows, gencrates actions based on reconstructod 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 autoregressive manner, shown by blue arrows. This parallel imagination iterates for H steps. The policies π are exclusively trained using imagined trajectories.
Equipped with a flexible querying mechanism and self-attention mechanism, the Perceiver aggregates the joint representation sequence ( x ag,... L7,1 24 Ti, az) of length n(K + 1) into a sequence of n features (e),
(#t, LL,K K d/ ..., F1,1, Fi,K at)
Perceiver > (e) Gt ct,.,,eee))) ct:)
Aggregating
where each feature of serves an intrinsic global abstraction of the environmental contexts perceived from agent i's viewpoint. By introducing Perceiver, we provide feasible solution for reducing the modeling complexity when using transformer-basee local dynamics.
Overall Learning Objective. The world model é is trained with trajectory segments of a fixed horizon H sampled from the replay buffer 'D in self-supervised manner. The transition predictor, discount predictor, and reward predictor are optimized to maximize the log- likelihood of their corresponding targets:
log pol (rs ei log Pogpo(ni (r d'z ett LDyale. 0) EAWNE En1
H
reward loss
discount loss
K logpolat+1, se azeer e xi+1
(6)
transition loss
where (et, e. ...,,,, et² fo(x), 22.2 Ti a. ...,, Th11 22.2 TZKK af), Vt.
We jointly minimize this loss function in Eq. 6 with respect to the model parameters of local dynamics (i.e., )) and global representation (i.e., 0) using the Adam optimizer (Kingma Ba, 2015).
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 &b and respectively. In the following, we take agent i as an exemplar case for clarity and omit the superscript for demoting the index of the agent to avoid potential confusion. By benefiting 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 ô, as input, and samples an action a, ( jo.). The world model then predicts the individual reward f., individual discount f and next local observation ōzt1- Starting from initial observations sampled from the replay buffer, this imagination procedure is rolled

---

Published in Transactions on Machine Learning Research (04/2025)
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 (ε( = ET /Z127/-"1)
In our approach, we train the actor and critic in a MAPPO-like (Yu et al., 2022) 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 Thansformer-based world model, 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. X-target in Dreamer (Hafner et al., 2020) is used to updated the value function. The details of behavior learning objectives and algorithmic description of MARIE are presented in $B and ŞI, respectively.
allow fast inference in the
4 Experiments
We consider the most common benchmark StarCraftIT Multi-Agent Challenge (SMAC) (Samvelyan et al., 2019) 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 singlo-agent Atari domain (Lukasz Kaiser et al., 2020). Additional experiment results on MAMujoco (Peng al., 2021) (i.e., continuous action space case) is provided in $F.1.
4.1 Experiment Setup and Evaluations
StarCraftII Multi-Agent Challenge. SMAC (Samvelyan et al., 2019), 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 AT and the other by our algorithm. The initial position, number, and type of units in each army varies from scenario to scenario, as docs the presence or absence of elevated or impassable terrain. And the goal 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 assigument for each agent's 's action. Another notable property of SMAC is that not all actions are accessible during decision- mocision-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 cstimation within the imagination horizon. Thus, in this benchmark, we additionally add one more head for the prediction of available action mask. During the imagination of MARIE, the available action mask is estimated by this head, instead of being generated manually according to the meaning of each element in the reconstructed observation. The latter introduces too much prior knowledge about StarCraft and can be considered as benchmark hacking.
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 (Egorov & Shpilman, 2022) and restrict the number of samples from the real environment to 100k for Easy scenarios, 200k for Hard scenarios and 400k for Super/ard scenarios, to establish a low data regime in SMAC. We compare MARIE with three strong model-free baselines MAPPO (Yu et al., 2022), QMIX (Rashid et al., 2018) and QPLEX (Wang et al., 2021), and two strong model-based baselines with the same policy learning paradigm as ours MBVD (Xu et al., 2022) and MAMBA (Egorov & Shpilman, 2022) on SMAC benchmark. Specially, as multi-agent variant of Dreamer V2 (Hafner el al., 2021), MAMBA achieves powerful 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. The hyperparameters of MARIE and other basclines are listed in SD and SH. Particularly, the hyperparameters of model-free basclines in low data regime are directly referred to Egorov & Shpilman (2022) and Liu et al. (2024).
Experimental Setup.
7

---

Published in Transactions on Machine Learning Research (04/2025)
Median
IQM
Mean
Optimality Gap
MARIE (Ours) MAMBA MBVD QPLEX
QМIX
MAPPO
0.25 0.50 0.75
0.25
0.50
0.75
0.4
0.6
0.8
0.2
0.4
0.6
0.8
SMAC Win Rate
Figure 3: MARIE exhibits statistically significant improvements on SMAC compared to the baselines. We report mean, median, and inter-quartile mean win rate, computed with stratified bootstrap confidence intervals. By default, all algorithms are evaluated using runs.
MAPPO
OMIX
OPLEX
MAMBA
MBVD
MARIE
25
8m
1c355z
3s vs 3z
0.8
1.0
1.0
0.7
0.6
Win Rate
0.5
Win
0.4
Win
0.4
0.2
0.2
0.1
0.0
0.0
0.0
0.0
00.0
0.4
0.6
0.8
1.0
-0.0
0.4
0.6
1.0
-0.0
0.2
0.4
0.6
1.0
-0.0
0.2
0.4
0.6
0.8
1.0
Environment Steps
Envirenment Steps
Environment Steps
Environment Steps
3s, vs. 4z
3s vs. 5z
2c vs. 642g
carridor
0.7
0.8
0.25
0.5
0.20
Rabe
0.5
0.15
0.3
0.4
0.10
0.3
0.2
0.0
0.0
0.00
0.0
-0.0
0.2
0.4
0.6 0.8
1.0
-0.0
0.4
0.8
1.2
1.6
2.0
-0.0
0.4
0.8
1.2
1.6
2.0
-0.0
0.8
1.6
2.4
3.2
4.0
Environment Steps
Envirenment Steps
Environment Steps
Environment Steps
Figure 4: The gap in sample efficiency and final performance between MARIE and the chosen baselines becomes more pronounced with increasing difficulty of scenarios. Curves of evaluation win rate for methods in 8 chosen SMAC maps. See Table for win rates. Y axis: win rate; X axis: number of steps taken in the real environment.
Table 1: MARIE outperforms every considered baselines on all task with more than 3 agents. We report mean evaluation win rate and standard deviation on 13 SMAC maps for different methods over random seeds. And we bold the values of the maximum and highlight them with blue color.
Methods
Maps
Difficuity
Stepe
QMIX
QPLEX
[Ours)
(Egaroy & Shpplman, 2022)
(Yu es al., 2022)
(Radiid et al,, 2018)
(Wang et al., 2021)
(Xa et al., 2022)
85.0g9
77-7p5a:
18:40114)
43.6g99-
68-3(7-)
60.8m11
95.5:7.9)
95.5c2.31
86.7
70.31144)
81.81130)
36.7(24.)
06.9
96.007-3
0.0gum)
15.7(um-)
8.7(14
2x3t
80.5(9-3)
31.2(124))
37.7(15.5)
50.2
53.4(--1
Eosy
9v.54)
87.7
80.5(124)
51.472177
88.7
73.9
Je__
98.90 9x
89.3010.1
1.221.50
0.0g00
0.0gg.
0.0gg
3o. dz
73.0k6-2)
29 3:13)
0.01a.m)
0.0go.t)
0.0kp.o)
0.000
Sun
88. 0cs.2)
66.002.75
70.3rsaa)
69. 69.5m12e)
83.4
74.7
MMM
#7. #7.60.%)
50.2 at
5.5c4
31.10749
60.305-)
20.5
50 many baneling
94.8833.3)
91.604.33
43.8115.6)
20.0gs.ov
32:2/6:1
15.0c1
Se. Se
Hont
200K
78.4:11.2)
13.4m14m
0.0ruu
0.0gum
0.0gum
0.0ma
2c Ve_ [o__g
25.0
0.8/4.7)
7.8
0.5qu.s
0.1go
0.2 e)
corridor
Sopd
26.511n.)
0.4
0.0guu)
0.0gum)
Model Param Size
35.35 MB
23.72 MB
0:21 MB
0.19 MB
0.23 MB
0.81 MB
4.2 Main Results on SMAC
MARIE achieves superior sample efficiency and final win rate model-based and model-free multi-agent. baselines. Agarwal et al. (2021) discuss the limitations of mean and median scores, and
8

---

Published in Transactions on Machine Learning Research (04/2025)
show that substantial discrepancies arise between standard point estimates and interval estimates in RL benchmarks. To deliver a rigorous statistical evaluation, we summarize in Figure 3 the win rate with stratified bootstrap confidence intervals for mean, median, and inter-quartile mean (IQM). Overall, we find MARIE achieves significantly better sample efficiency and a higher win rate compared with other strong baselines. Furthermore, we report. the detailed averaged win rales over four seeds in every chosen scenario in Table and provide additional learning curves of several chosen scenarios in Fig. 4. As presented in Table and Fig. 4, 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. 38 us 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. Moreover, special scenario 2c us 642g deserves attention, which features only 2 agents but with a considerably large action space of up to 70 discrete actions for each agent. Although the performance of MARIE in 2r. us 6429 suffers a relative large variance due to the overly large action space, MARIE achieves a remarkably non-trivial mean win rate just via learning in the imagination. Note that 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 _c_us__ 21__622 proves that our Transformer-bascl world model has higher prediction accuracy and a deeper understanding of the underlying mechanics.
On the other hand, MARIE contains more than twice the number of parameters as MAMBA, which may partially account for its more pronounced improvement. This observation further implics that the learning of world models may benefit from scaling laws, albeit at the cost of increasing computational overhead and longer training time.
4.3 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 T= of of at, a?....,. of,..... Given that r already contains the joint observations and actions, we disable the aggregation module in this centralized variant. As illustrated in Figure 5, our comparisons span scenarios involving 2 to 7 agents. When the nunber of agents small enough, reducing the multi-agent system to single-agent one over the joint observation and action space would not cause prominent scalability issue, as indicated by the result in 2s_ us_ 1sc. However, the scalability issuc is exacerbated by 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-actio 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 fcatures, MARIE delivers stable and superior sample efficiency.
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. 6, in the 3-agents scenario (e.g., 3s vs 3z), the correlation among each agent's 's local observation tends to be negligible. Therefore, the nearly independent generation of each agent's local obscrvation 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 giset-wise aggregation, the shared Transformer struggles to infer accurate future
9

---

Published in Transactions on Machine Learning Research (04/2025)
Decentralized Marner
Centralized Manner
25 vs. 1sc agents)
35 vs. 32 (3 agants)
2s32 (5 agents)
50 many baneling (7 agents)
1.0
1.0
1.0
1.0
Vin Rate
Wib Rate
0.6
0.6
0.6
0.6
0.2
0.2
0.2
0.2
0.0
-0.0
0.2
0.4
0.6
0.8
1.0
-0.0
0.2
0.4
0.6
0.8
1.0
-0.0
0.2
0.4
0.5
0.8
1.0
-0.0
0.2
0.4
0.5
0.8
1.0
Environment Steps
Environment Steps
Envirenment Steps
Environment Steps
Figure 5: Decentralized dynamics modeling allows for better scaling to environments with growing number of agents. We conduct ablation study what manner to integrate into the design of the world model. Decentralized Manner denotes the standard implementation of MARIE, while Centralised Manner denotes that the world model is designed for learning the joint dynamics of all agents over the joint trajectory. Sample efficiency of the centralized variant encounters a significant drop due to the scalability issue while MARIE robust to scenarios with various number of agents.
MARIE
MARIE wio aggregation
35, Vs. (3 agents)
s0 mony baneling (7 agents)
1c3s5z (9 agents)
1.0
0.8
0.8
0.8
Win Rate
Win Rate
0.6
0.6
0.6
Win
0.2
0.2
0.2
0.0
0.0
-0.0
0.2
0.4
0.6
0.8
1.0
-0.0
0.2
0.4
0.6
0.8
1.0
-0.0
0.2
0.4
0.6
0.6
1.0
Environment Steps
145
Environment Steps
105
Environment Steps
205
Figure 6: Local dynamics struggles to infer accurate future local observations without agent- aggregation. We plot the comparison results between MARIE with and without the usage of the aggregation module.
Tokenizer (16 takens)
Bins Discretization (Noau tokens)
2m 17 (Nee 16)
25 vs. (Mea - 17)
(Dawe - 42)
2533 (Maze
1.0
0.8
Win Rate
Winn Rate
Winn Ratee
0.6
0.0
0.6
0.86
0.2
0.2
0.0
0.0
5.0
4
10
10
2
16 Timerost(hours)
Timecost(hours)
Timecosthaarr)
Timecost(hours
Figure 7: Vector quantization is more effective than Bins Discretization for encoding observations with discrete tokens. We conduct an ablation study on the type of discretization for local observations. Tokenizer denotes the standard implementation of MARIF; Bins Discrelization denotes the variant of MARIE where the Rose-dimensions observation discretization is performed by projecting the value into one of m fixed-width bins in each dimension independently. X-axis: cunulative run time of algorithms in the same platform. VQ-VAB encapsulates local observations within a succinet sequence of tokens, computationally efficiently promoting the learning of the -ransformer-basedd world model.
local observations, thus hindering policy learning in the imaginations of the world model and resulting in notable degradation in the win rate evaluation.
VQ-VAE encapsulates local obscrvations within a succinet 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 tokens from Z, perhaps a more naive tokenizer is projecting the value each dimension into one of m fixcd-width bins (Janner et al., 2021), resulting in a Robs-long token sequence for each observation, which we term Bins Discretization. We set the number of bins equal to the size of codebook |2] and compare these two types of tokenizers in different environments with
10

---

Published in Transactions on Machine Learning Research (04/2025)
Agent 1 imagination
Agent local imagination
Agent local imagination
2.50
2.50
2.40
MAREF
MARE MAVIBA
HARE
2.75
NAMBA
1.75
Accunulated Erros
Accumulated Errors
Accumulated Errors
1.30
1.30
2.30
3.25
1.25
1.25
3.50
2.20
0.00
0.75
0.75
0.75
0.50
0.25
0.25
0.25
10
15
20
25
13
15
20
5
12
13
20
Horizons
Horizons
Horizons
Figure 8: MARIE exhibits significantly lower compounding prediction error than MAMBA across all agents with respect to prediction horizon. We compare the imagination accuracy of MARIE to that of MAMBA over the course of planning horizon in 3s 5z scenario.
various Robs. As shown in Fig. 7, the performance of the two tokenizers are comparable only in 2s us Isc where Robs is close to 16. Even worse, Bins Discretization experiences a pronounced decline as fobs increases in more complex environments (e.g., 35 us 42) under identical training durations. We hypothesize that for a single local observation, a Mobe--oken- verbose sequence yielded by Bins Discrelization 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 notably slow convergence. Meanwhile, the result in 2m us Iz indicates Bins Discrelization may ignore the correlation of different dimensions, which would be helpful in sequence modeling.
4.4 Model Analysis
Error Accumulation. A quantitative evaluation of the model's accumulated error versus prediction horizon is provided in Fig. 8. Since learning the world model is tied to a progressively improving policy both in MARIE and MAMBA, we scparately use their final policies to sample 10 episodes for fairness. We then compule Ln errors per observation dimension wetween 1000 trajectory segments randomly sampled from these 20 episodes and their imagined counterpart. The result Fig. 8 suggests architecture differences play large role in the world model's '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 us 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 (Janner el al., 2019).
Attention Patterns. During model prediction, we delve into the attention maps inside the shared Transformer and the cross attention maps 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 the previous transition, while the other is regularly striated wherein the model attends to specific tokens in multiple prior transitions. During the agent-wisc aggregation, we also identify two distinct patterns - individualily 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 SC for further details and visualization results.
5 Conclusion and Limitation
We have introduced a model-based multi-agent algorithm MARIE, which utilizes a shared Transformer as a 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 lcarning, it significantly boosts the sample efficiency and improves final performance compared state-of tate--ffthe-art model-free methods and existing model-based methods with the same learning paradigm, in the low data regime. But it should
11

---

Published in Transactions on Machine Learning Research (04/2025)
be also noted that there are potential limitations on the current evaluation on the main experiment with 4 limited seeds, e.g,, the limitations of mean and median scores (Agarwal et al., 2021). Thus, we also provide a standardized performance evaluation following the protocol provided by Agarwal et al. (2021) in $4.2 and SG. To further deliver a rigorous statistical validation, evaluation with more scods is definitely necessary. As the first Transformer-based multi-agent world model for sample-efficien 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 safety. While it has the great potential to illuminate 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 very long prediction horizons, due to the aulo-regressive property.
References
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 1-22304-293200 2021.
Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? In International Conference on Learning Representations, 2022. URL
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Tomas Jackson, Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Kuang-Huei Lee, Sergey Levine, Yao Lu, Utsav Malla, Deeksha Manjunath, Igor Mordatch, Ofir Nachum, Carolina Parada, Jodilyn Peralta, Emily Perez, Karl Pertsch, Jornell Quiambao, Kanishka Rao, Michael Ryoo, Grecia Salazar, Pannag Sanketi, Kevin Sayed, Jaspiar Singh, Sumedh Sontakke, Austin Stone, Clayton Tan, Huong Tran, Vincent Vanhoucke, Steve Vega, Quan Vuong, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich. RI-1: Robotics transformer for real-world control at scale. In Robotics: Science and Systems (RSS), 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020. URL Paper pdf.. neurips. wtt.ppp.ccppppp...cp sateor_ffl//pppp/e///ee /000
Vosse,
Chang Chen, Yi-Fu Wu, Jacsik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprinl ar Xiw:2202.00981 2022.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North Americon Chapler of the Associalion for Compulational Linguislies: Human Language Technologies, NAACI-HITT, 2019. doi: 10.18653/V1 //////1//NDD 11/N19-1423. URL https: //doi. www..o..gg..o.. 10. 18653/v1/n19111111
Vladimir Egorov and Alexei Shpilman. Scalable multi-agent madel-based reinforcement learning. In Proceedings of the 21st International Conference on Autonomous Agents and Multiagent Systems, 2022.
Benjamin Ellis, Jonathan Cook, Skander Moalla, Mikayel Samvelyan, Mingfei Sun, Amuj Mahajan, Jakob Nico- laus Foerster, and Shimon Whiteson. SMACv2: An improved benchmark for cooperative multi-agent reinforcement learning. In Thirty-seventh Conference on Newral Information Processing Systems Datasets and Benchmarks Track, 2023. URL https://openreriex http:///aarrrreeee..eeeerrrrraaa wtt.://ooorrrrre--11111......
Jakob N. Foerster, Gregory Farquhar, Triantafyllos Afouras, Nantas Nardelli, and Shimon Whiteson. Coun- terfactual multi-agent policy gradients. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
12

---

Published in Transactions on Machine Learning Research (04/2025)
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, 2018.
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixcls. In Proceedings of the 36th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Represenlations, 2020. URL https://opeereviii net/ wwwoorodddddiooneee
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021. URL https://openrevier net/forum iid=0oibyy2000n
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. aixiv preprint xrXiv22301.0410. 2023.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2022. URL https:// http:p/////oocedngg..oo mlr wpppsss/vv16/hhaane2222a ppppppeeera...tt..pppee//////aaa
html.
Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for continuous control. In The Twelfth International Conference on Learning Represenlations, 2024. URL https: wwwparrrvee.nettooorrreeeeeeeeee
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08 15, 2016.
Pablo Hernandez-Leal, Bilal Kartal, and Matthew E. Taylor. A very condensed survey and critique of multiagent deep reinforcement learning. In Proceedings of the 19th International Conference Aulonomous Agents and MultiAgent Systems, 2020.
Shariq lqbal and Fei Sha. Actor-attention-criiic for multi-agent reinforcement learning. In Proceedings of the 36th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2019. URL
.ml...sss........mmm...........m wallrrpress//eee v97 007//qqal10 19a. html.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, 2021.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Newral Information Processing Systems, 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems, 2021.
Andrej Karpathy. mingpt: A minimal pytorch re-implementation of the openai gpt (generative pretrained transformer) training, https://git thub. /otn//aappah///inGGr, 2020.
Diedcrik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio
and Yann LeCun (eds.), 3rd International Conference on Learning Representations, 2015. URL http:
//arxiv.org www./bs////eeeeeeeeeeeeeeeeeeeee httg///////1..eeeee 11/.1111.66999 6980.
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. In 2nd International Conference Learning Representations, 2014.
Jakub Grudzien Kuba, Ruiqing Chen, Muning Wen, Ying Wen, Fanglei Sun, Jun Wang, and Yaodong Yang. Trust region policy optimisation in multi-agent reinforcement learning. In International Conference on Learning Representations, 2021.
13

---

Published in Transactions on Machine Learning Research (04/2025)
Qihan Liu, Jianing Ye, Xiaoteng Ma, Jun Yang. Bin Liang, and Chongjie Zhang. Efficient multi-agent reinforcement learning by planning. In The Twelfth International Conference on Learning Representalions, 2024. URL https::/openrrviier wttpiiiwwwwweeeeee wwt./erumn
Yong Liu, Weixun Wang, Yujing Hu, Jianye Hao, Xingguo Chen, and Yang Gao. Multi-agent game abstraction via graph attention neural network. In The Thirty-Fourrts AAAI Conference on Artificial Intelligence, 2020.
doi: 10. 10.1609/
URL https: //op::://ddi.. org/ 10. 1001 aa4.i....i............ 122105.6211.
Ryan Lowe, Yi Wu, Aviv Tamar, Jean Harb, Pieter Abbeel, and Igor Mordatch. Multi-agent actor-critic for mixed cooperatiee-comptttivv environments. In Proceedings of the 31st International Conference on Neural Information Processing Systems, 2017.
Anuj Mahajan, Mikayel Samvelyan, Lei Mao, Viktor Makoviychuk, Animesh Garg, Jean Kossaifi, Shimon Whiteson, Yuke Zhu, and Animashree Anandkumar. Tesseract: Tensorised actors for multi-agent reinforce- ment learning. In Proceedings of the 38th International Conference Machine Learning, Proceedings of Ma- chine Learning Research. PMLR, 2021. URL https://proceeddnggg mmlr wttp://w......pppesssesse https///vvmmmmaaaaaaaaaa html.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL forum? derrm?dddvhFulAcbbb..
Thanh Thi Nguyen, Ngoc Duy Nguyen, and Saeid Nahavandi. Deep reinforcement learning for multiagent. systems: A review of challenges, solutions, and applications. IEEE Transactions Cybernetics, 50: 3826-3839, 2020. doi:
Frans A Oliehoek, Christopher Amato, et al. concise introduction to decentralized POMDPs, volume 1. Springer, 2016.
Bei Peng, Tabish Rashid, Christian Schroeder de Witt, Pierre-Alexandre Kamienny, Philip Torr, Wendelin Boehmer, and Shimon Whiteson. FACMAC: Factored multi-agent centralised policy gradients. In Advances in Newral Information Processing Systems, 2021. URL https: https: //opprrve///eeeieew http://pperrrrreeeeeeeeooooeeeee htte//frrmm/ddkee74/0000
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Tabish Rashid, Mikayel Samvelyan, Christian Schroeder, Gregory Farquhar, Jakob Foerster, and Shimon Whitcson. QMIX: Monotonic value function factorisation for deep multi-agent reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2018. URL https://proceeding mml. ppess//80/rashiiei htt.ahhhhi.aalll.........m..m html.
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-base world models are happy with 100k intoractions. In The Eleventh International Conference on Learning Representations, 2023. URL https://openrevier http:://pppaeeeeeeeeeeeeerrrrara het//torumm:icTTBBDDGCCipp
Heechang Ryu, Hayong Shin, and Jinkyoo Park. Multi-agent actor-critic with hierarchical graph attention network. In Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
Mikayel Samvelyan, Tabish Rashid, Christian Schroeder de Witt, Gregory Farquhar, Nantas Nardelli, Tim G. J. Rudner, Chia-Man Hung. Philip H. S. Torr, Jakob Focrster, and Shimon Whiteson. The starcraft. multi-agent challenge. In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, 2019.
Julian Schrittwieser, Ioannis Antonoglon, Thomas Hubert, Karen Simonyan, L. Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588:604 - 609, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alee Radford, and Oleg Klimov. Proximal policy optimization algorithms. or Xiu preprint ar arXio:1 2017.
14

---

Published in Transactions on Machine Learning Research (04/2025)
Kyunghwan Son, Daewoo Kim, Wan Ju Kang, David Earl Hostallero, and Yung Yi. QTRAN: Learning to factorize with transformation for cooperative multi-agent reinforcement learning. In Proceedings of the 36th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2019. URL httpp://proceeddngs..oot mlr. wtwpp/////sosnee//eee httppp/ceecaaiaa.a.pppp///////o/
html.
Peter Sunchag, Guy Lever, Audrunas Gruslys, Wojeiech Marian Czarnecki, Vinicius Zambaldi, Max Jaderberg, Mare Lanciol, Nicolas Sonnerat, Joel Z. Leibo, Karl Tuyls, and Thore Graepel. Value-decomposition networks for cooperative multi-agent learning based on team reward. In Proceedings of the I 17th International Conference on Autonomous Agents and MultiAgent Systems, 2018.
Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigort Bulletin, 2(4):160 163, 1991.
Edan Toledo and Amanda Prorok. Codrcamer: Communication-based decentralised world models. arXiv preprinl Xrv::22406. 13600, 2024.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In Advances in Newral Information Processing Systems, 2017. ISBN 9781510860964.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Usækoreit, Llion Jones, Aidan N Gomez, Lukass Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017. URL https: ppprocediing..eurrte 3t5e24445//ce11110033..coo444444
Paper. pdf.
Jianhao Wang, Zhizhou Ren, Terry Liu, Yang Yu, and Chongjie Zhang. QPLEX: Duplex dueling multi-agent g-learning. In International Conference on Learning Representations, 2021. URL https://openreviee
net//oorumidddRRaOOxxx00
Daniël Willemsen, Mario Coppola, and Guido CHE de Croon. Mambpo: Sample-efficient multi-robot. reinforcement learning using learned world models. In 2021 IEEE/RSJ International Conference Intelligent Robots and Systems (IROS). IEEE, 2021.
Zhiwei Xu, Dapeng Li, Bin Zhang, Yuan Zhan, Yunpeng Baiia, and Guoliang Fan. Mingling foresight with imagination: Model-based cooperative multi-agent reinforcement learning. In Advances in Newral Information Processing Systems, 2022. URL https //tp://vee/eeeeee/e nww/teeuutiiieeoocmm (o-.orunnidd11yyy2i------ 11BYpZKW6TT.
Chao Yu, Akash Velu, Eugene Vinitsky, Jiaxuan Gao, Yu Wang, Alexandre Bayen, and Yi Wu. The surprising effectiveness of PPO in cooperative multi-agent games. In Thirty-sizth Conference on Neural Information Processing Systems Dalasets and Benchmarks Track, 2022.
Qiaosheng Zhang, Chonjia Bai, Shuyue Hu, Zhen Wang, and Xuclong Li. Provably efficient information- directed sampling algorithms for multi-agent reinforcement learning. arXio preprint arXiv:2404.1999 19292, 2024a.
Yang Zhang, Shixin Yang, Chenjia Bai, Fei WI, Xiu Li, Xuelong Li, and Zhen Wang. Towards efficient Ilm grounding for embodied multi-agent collaboration. arXiv preprint ar Xrx: 25.. 14814, 2024b.
Eukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020. URI. https://openreview wwwpweewwwwwweeeeeeeeeeeeeeeeeee met/tfruuu 1iddiddxxJPHDDD
15

---

Published in Transactions on Machine Learning Research (04/2025)
A World Models Details and Hyperparameters
A.1 Observation Tokenizer
Our tokenizer for local observation discrctization is based on the implementation' of a vanilla VQ- VQ-VAE (van den Oord el al., 2017). 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 hyperparameterss 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-VAL as input instead to avoid the distribution shift between policy learning and policy execution.
For training this vanilla. VQ-VAE, we use straight-through estimator to enable gradient backpropagation through the non- blediiffernntabble quantization operation in the quantization of VQ-VAE. The loss function for learning the autoencoder is as follows:
EvQ.vAR(E,D.. = E. [llo' - 0012 | |sg[E((o)]] 2 B1|sg[= E[
(7)
where N = (1,2,..,, denotes the set of agents, sg[-] denotes the stop-gradient operation and 3 is the coefficient of the commitment loss ||s][=]] Fll In practice, we found the codebook 2 can suffer from codebook collapse when learning from scratch. Thus, we adopt the Exponential Moving Averages (EMA)
(van den Oord et al., 2017) technique to alleviate this problem.
Table 2: VQVAE hyperparameters.
Hyperparameter
Value
Encoder& Decoder
Layers
3
Hidden size
512
Activation
GFenddesssyykks & Gimpel, 2016)
Codebook
Codebook size (N)
512
Tokens per observation (K)
16
Code dimension
128
Coef. of commitment loss (3)
10.0
A.2 Transformer
The shared Transformer serving as the local dynamics model is based on the implementation of minGPT (Karpathy, 2020). Given fixed imagination horizon H, it first takes token sequence of length H(K + +1) composed of observation tokens and action tokens, and embeds it into H(K + 1) 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) x D. This tensor is forwarded through fixed Transformer blocks. Here, we adopt GPT2-like blocks (Radford et al., 2019) as the basic blocks. The hyperparameters are listed as Table To enable training across all envirouments on a single NVIDIA RTX 3090 GPU, we adapt imagination horizon based on the number of agents.
A.3 Perceiver
The Perceiver (Jaegle et al., 2021) is based on the open-source implementation?. 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
Code can be found httpp://gttabb..mm/mmsssttttttte wwt.e///wiiaaattteeetttoaaaaaaaa htw.://ieeei.ttttoooooooaatetrre pytorc
2Code can be found in http://ggtttttt..eee http:://gtthha..mmmm/caaaaaaaaan http://.ww/eeeciiaiaaacccceereee
16

---

Published in Transactions on Machine Learning Research (04/2025)
Table 3: Transformer hyperparameters.
Hyperparameter
Value
Imagination horizon (4)
{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
Table 4: Perceiver hyperparameters.
Hyperparamcter
Value
Length of latent querying
72 (number of agents)
Cross attention heads
8
Inner Transformer layers
2
Transformer attention heads
8
Dimension per attention head Embedding dropout
64
0.1
Attention dropout Residual dropout
0.1
0.1
aggregation: (i) the embedding tensor of shape (K + 1) x D at each timestep, mentioned in Appendix A.2, is concatenated with others from all agents, thereby getting n(K + 1) x 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 compressed from length n(K+1) n; (iii) the compressed sequence is then forwarded through standard transformer with bidirectional attention inside the Perceiver. The hyperparameters are listed as Table 4.
B Behaviour Learning Details
In MARIE, we use MAPPO-like (Yu et al., 2022) 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 all reconstructed observations ô, of all agents. Therefore, the actor v remains a 3-layer MLP with ReLU activation, while the critic È is enhanced with an extra layer of self-attention, built on top of the original 3-layer MLP, i.e., we overwrite the critic Vi( ~ EE3 (Z for agent Similar to off-the-sher CTDE model-free approaches, we adopt parameter sharing across agents.
Critic loss function We utilize A-return in Dreamer (Hafner et al., 2020), which employs an exponentially- weighted average of different k-steps TD targets to balance bias and variance the regression target for the critic. Given imagined trajectory (ō,, ot pt H for agent i, X-return is calculated recursively as,
Ai + (1 (o))( +V
t <<
Vi(6,) =
(8)
(6,)
17

---

Published in Transactions on Machine Learning Research (04/2025)
Table 5: Behaviour learning hyperparameters.
Hyperparameter
Value
Imagination Horizon (H)
{15, 8, 5}
Predicted discount label y
0.99
À
0.95
? Clipping parameter
0.001
0.2
The objective of the critic & is to minimize the mean squared difference C3 with A-returns over imagined trajectories for each agent as
C3 = B
Σ t=1 (V) sg(()(()))))
(9)
where sg(+) denotes the crop-gradient operation.
We optimize the critic loss with respect to the critic
parameters & using the Adam optimizer.
Actor loss function The objective for the action model (-(-)))) is to output actions that maximize the prediction of long-term future rewards made by the critic. To incorporate intermediate rewards more directly, Irain the actor to maximize the same X-return that was computed for training the critic. In terms of the non-stationarity issue in multi-agent scenarios, adopt PPO updates, which introduce important sampling for actor lcarning. The actor loss function for agent is:
C -E Epp Vold
Σ min (li))))) e, c,11+)44 + (7((T(((1)))
(10)
where r((b) / Pord is the policy ratio and A = (0) (.))) is the advantage. We optimize the actor loss with respect to the actor parameters v using the Adam optimizer. In the discount prediction of MARIE, we set its learning target y to be 0.99. Overall hyperparameters are shown in Table 5.
C Extended Analysis Attention Patterns
To provide qualitative analysis of our world model, we select typical scenarios Is vs 52 where our method achieves the most signifieant 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. 9 and Fig. 10.
The prediction of local dynamics entails two distinct attention patterns. The left one in Fig. 9 can be interpreted as a Markovian pattern, 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, individuality and commonality. The top one in Fig. 10 illustrates that each agent flexibly attends to different tokens according to their specific nceds. In contrast, the bottom one exhibits consistent attention allocation across all agents, with atlention highlighted in nearly identical positions. The diverse patterns in the Transformer and Perceiver may be the key to accurate and consistent imagination.
also identify two distinct patterns:
D Baseline Implementation Details
MAMBA (Egorov & Shpilman, 2022) is evaluated based on the open-source implementation: https:
//github. com/ aaamm/br-iiiiie aabrraa-labaaaaaaeeeeeeeaaaaaaea -amabbbammaa with the hyperparameters in Table 6.
18

---

Published in Transactions on Machine Learning Research (04/2025)
Table 6: Hyperparameters for MAMBA in SMAC environments.
Hyperparameter
Value
Batch size
256
À for λ-return computation
0.95
Entropy coefficient
0.001
Entropy annealing
0.99998
Number of policy updates Epochs per policy update Clipping parameter e
4
5
0.2
Actor Learning rate
0.0005
Critic Learning rate Discount factor y
0.0005
0.99
Model Learning rate Number of model training epochs
0,0002
60
Number of imagined rollouts
800
Sequence length
20
Imagination horizon H Buffer size
15
2.5 x 10a
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
Table 7: Common hyperparameters for MAPPO in SMAC environments.
Hyperparameter
Value
Batch size
mum X buffer length mum agents
Mini batch size
batch size / mini-batch
Recurrent data chunk length
10
GAE À
0.95
Discount factor 7
0.99
Value loss
huber loss
Huber delta
10.0
Optimizer
Adam
Optimizer learning rate
0.0005
Optimizer epsilon
1 x 10-5
Weight decay
0.0
Gradient clipping
10
Network initialization
orthogonal
Use reward normalization
True
Use feature normalization
True
19

---

Published in Transactions on Machine Learning Research (04/2025)
0
0
a245 ests
efts'
Figure 9: 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 DocDPO 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 u8 scenario.
Head
agent
4gent
agent (o).ok tokens
agent (oF, tokens
agent (or.n)) tokens
agent 4's (o), tokens
agent 5's af) toloon
Cros Atlastion Head
agent
agent
agent I's ( tokens
agent 2's (o tokens
agent 3's (02.9)) tokens
agent 4's (o), tokens
agentt 5's (o). (f) tokens
Figure 10: Cross attention patterns in the Perceiver. We observe the individuality and commonality in the agent-wisc 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- aggregation the 2s3z scenario.
MAPPO (Yu et al., 2022) is evaluated based on the open-source implementation: https://github.coo com/ marlbenchmark/on on- htth//ioppppppp with the common hyperparameters Table 7.
QMIX (Rashid al., 2018) is evaluated based on the open-source implementation: https://github.com oxwhir1/pymarl with the hyperparameters in Table 8.
QPLEX (Wang et al., 2021) is evaluated based on the open-source implementation: attps://giis thu. com/wjhtn hom/wjh720/QPLEX with the hyperparameters in Table 9. Since its implementation is mostly based on the open-source implementation: PyMARL (Samvelyan et al., 2019), its most hyperparameters setting remains the same as the one in QMIX in addition to its own special hyperparameters.
MBVD (Xu et al., 2022) is evaluated based on the implementation its supplementary material from with the hyperparameters in Table 10. Akin to QPLEX,
https://openrevieu
20

---

Published in Transactions on Machine Learning Research (04/2025)
Table 8: Hyperparameters for QMIX in SMAC environments.
Hyperparameter
Value
Batch size
32
Buffer size
5000
Epsilon in
1.0 + 0.05
Epsilon anneal time
50000
Train interval
episode
Discount factor y
0.99
Optimizer
RMSProp
RMSProp
0.99
RMSProp
10-5
Gradient clipping
10
Table 9: Hyperparameters for QPLEX in SMAC environments.
Hyperparameter
Value
Batch size
32
Buffer size
5000
Epsilon in epsilon-greedy
1.0 - 0.05
Epsilon anneal time
50000
Train interval
1 episode
Discount factor 7
0.99
Optimizer
RMSProp
RMSProp a
0.99
RMSProp
10-5
Gradient clipping
10
Number of layers in HyperNetwork
1
Number of heads in the atlention module
4
21

---

Published in Transactions on Machine Learning Research (04/2025)
Table 10: Hyperparameters for MBVD in SMAC environments.
Hyperparameter
Value
Batch size
32
Buffer size
5000
Epsilon in epsilon-gren -geedyy
1.0 + 0.05
Epsilon anneal time
50000
Train interval
1 episode
Discount factor Y
0.99
Optimizer
RMSProp
RMSProp a
0.99
RMSProp e
10-5
Gradient clipping
10
Number of layers in HyperNetwork
Number of heads in the attention module
4
Horizon of the imagined rollout
3
KL balancing a Dimension of the latent state
0.3
num agents x 16
its implementation is based the open-source implementation: PyMARL, its most hyperparameters setting remains the same as the one in QMI in addition to its own special hyperparameters.
E
Additional Experiments
E.1
Evaluations on MAMujoco
The Multi-Agent MuJoCo (MAMuJoCo) (Peng el al., 2021) environment is a multi-agent extension of MuJoCo. While the MuJoCo tasks challenge a robot to learn an optimal way of motion, MAMuJoCo models each part robot as an independent agent for example, a leg for a spider or an arm for a swimmer and requires the agents to collectively perform efficient motion. With the increasing variety of the body parts, MAMujoco can be also considered testbed for evaluating the coordination among heterogeneous agents, which poses big challenge for learning the multi-agent dynamics inside it, especially in decentralized manner.
While MAMBA (Egorov& & Shpilman, 2022) originally does not take the continuous action space case into considcration, which is a obvious limitation in it, we would like to evaluate MARIE in such case, c.g., MAMujoco, to beiter demonstrate that our method can be also effective other multi-agent domains. MAMujoco, we discretize the scalar in each dimension of continuous actions into one of 256 fixed-width bins independently to obtain discrete action tokens for local dynamics lcarning. As the behaviour learning in MARIE adopts MAPPO-like and on-policy manner, choose two strong on-policy PPO-based baselines - MAPPO (Yu et al., 2022) and HAPPO (Kuba et al., 2021). Additionally, we also include MAMBA as a model-based baseline for comparison. Since MAMBA (Egorov & Shpilman, 2022) was originally originally designed for domains with discrete action space, significant effort was required to adapt and evaluate on MAMujoco, which features continuous act space. The experiments are conducted in HalfCheetah-oee 222,, HulfCheetah-u2-B3 and Walker2d-un (2---12222 The learning curves of the return averagod over seeds are presented Figure 11. Notably, MAMBA fails to enhance policy learning in the Walker2d-u2-2rr -22-33 scenario and remains exceptionally time-consuumngg Consequently, report its results only for 1 million environment steps in this scenario.
As illustrated in Figure 11, our MARIE consistently shows superior sample elficiency and achieves the best performance in of 3 scenarios with limited 2M environment steps. For the performance difference between MAMBA and MARIE HalfCheetah-n2--3 significantly from using the internal recurrent features of the world model as inputs in this scenario, while the policy in method only takes the reconstructed observation input in order to support fast deployment.
we hypothesize that MAMBA's policy learning benefits
22

---

Published in Transactions on Machine Learning Research (04/2025)
MAPPO
НАPPO
MAMBA
MARIE
HalfCheetah-v2 -tttheeeetth--2233333
HalfCheetah-v2-B33
Walker2d Dalker2d-v2-22 -2---2222 v2-2x3
5000
5000
4000
4000
4000
Returns
3000
Returns
3000
Returns
3000
2000
2000
2000
1000
1000
1000
0
-0.0
0.4
1.2 1.6
0.4
1.2
2.0
0.8 1.2
2.0 1et
Environment Steps
Environment Steps
Environment Steps
Curves of the performance for MARIE, MAMBA, HAPPO and MAPPO in 3 chosen MAMujoco scenarios. Y axis: return; X axis: number of steps taken in the real environment.
Figure 11:
MARIE
QMIX
MAPPO
HAPPO
MAMBA
zerg_5_vs_
protoss vs_5
terran_5_vss
0,40
0.30
0.35
0.25
0.4
0.30
Win Rate
0.20
Win Rate
0.25
Win Rate
0.3
0.15
0.20
0.15
0.2
0.10
0.10
0.1
0.05
0.05
0.00
0.00
0.0
-0.0
0.4
0.8
1.2
1.6
2.0
-0.0
0.4
0.8
1.2
1.6
2.0
-0.0
0.4
1.2
1.6
2.0
Environment Steps
Environment Steps
Environment Steps
Curves of the performance for MARIE, MAMBA, QMIX, HAPPO and MAPPO in 3 chosen SMACv2 scenarios. Y axis: Win Rate; X axis: number of steps taken in the real environment. While MARIE shows competitive performance to MAMBA on zerg 5 __v____ MARIE is superior to all other baselines in terms of sample efficiency and final performance in the rest scenarios.
Figure 12:
in the environment without the participation of the world model. We attribute the performance gap between MARIE and other two modcl-free baselines in HalfCheetah-u2-91t to the access to global oracle state in the chosen baselines. The policy in our algorithm is purely learned from the inner imaginations of the world model where there is only reconstructed local observation. Considering MAMujoco a multt-hetterogenooo--ggoe benchmark which necessitates a more precise credit assignment during training, it would be much more helpful for policy learning to have access to the true global oracle state than in other benchmarks. But overall, MARIE presents faster convergence rate, implying that Transformer-based world model generate accurate imaginations and bring remarkable sample efficiency.
E.2 Evaluations on SMACv2
Given known scrious flaws in SMACvI (c.g., the tricky open-loop policy issue), we extend our evaluation of MARIE to SMACv2 (Ellis al., 2023), which introduces more stochasticity and partial observability. In this comparison, we benchmark MARIE against four baselines: MAPPO, HАPPО, QMIX and MAMBA. For each random seed, we adopt the same evaluation protocol the main experiment on SMACv1. Importantly, the hyperparameters of MARIE remain unchanged, detailed in şH. Here, directly use the results of MAPPO and QMIX provided in the official SMACv2 repository Illustrated in Figure 12, while MARIE shows competitive performance to MAMBA on zerg 5_ vs 5, MARIE is superior to all other baselines in terms of sample efficiency and final performance in the rest 2 scenarios.
Results of QMIX and MAPPO available tttps://github.com http::////tttttttccccccc///////a http://///aaaaaa///ec////a/a//aa
23

---

Published in Transactions on Machine Learning Research (04/2025)
E.3
Comparison with QMIX and QPLEX with different epsilon annealing time
Considering the potentially inappropriate influence of a large annealing time used in the epsilon-greedy algorithm when evaluated in the low data regime evaluation, we run QMIX and QPLEX with smaller annealing time, and compare the performance of them with ours. The result is reported in Table 11. The reported result shows that the original hyperparameters used in the main experiment, which are also directly referred to Egorov & Shpilman (2022) and Liu et al. (2024), are reasonable since the performance of QMTX and QPLEX under the original hyperparameters is superior to the ones with a smaller annealing time most scenarios. Besides, our MARIE still consistently oulperforms QMIX and QPLEX with smaller annealing time.
Table 11: Mean evaluation win rate and standard deviation for QMIX and QPLEX with different epsilon anneal time over random seeds. We bold the values of the maximum.
Maps
Steps
MARIE
QMIX 50000)
QMIX (1, 10000)
QPLEX (f, 50000)
QPLEX (t, 10000)
1c3652
85.0:9.4)
43.6/29.2)
33.3(15.0)
68.3(7.4)
44.8(11.0)
2m va 1z
95.5(7.9)
70.3(14.8)
36.1(28.2)
84.8(10.8)
93.2(4.7)
2s Ise
96.9:7.1)
0.0(o.0)
3.9
15.7(19.5)
43.2(32
2432
80.5g9.3)
(7.7(15.5)
(2.1(20.3)
50.2(8.4)
28.3(11.5)
3m
100K
99.5(0.4)
541.((22.7)
63.8(14.6)
88.7(6.9)
85 (51.()))))
3s vs 3z
98.9:1.5)
0.0ço.0)
0.0ço.0)
0.0ro.n)
0.0(o.0)
38 vs. 42
73.0
0.0çc.o)
0.0ç0.0)
0.0(o.0)
0.0(o.0)
8m
88.0 88.0(3.9)
69.5(12.8)
68.6(13.6)
83. 83.4(6.4)
79.7(9.8)
MMM
87. 87.6:5.0)
31.1
18.9(4.3)
69.3(35.1)
20.2(7.7)
many baneling
94.8:5.9)
20.0(8.9)
30.7(18.5)
32.2 1)
37.7(9.2)
2c vs 642g
25.9(14.3)
0.5(c.5)
0.0çg.0)
(.1
0.0(o.0)
200K
3s vs
78.4(11.2)
0.0(c.0)
0.0(o.0)
0.0(o.0)
0.0(o.0)
corridor
400K
71.0(13.8)
0.0ro.0)
0.0ra.0)
0.0(0.0)
0.0ga.0)
E.4 Comparison with existing Transformer-based world models
Existing Transformer--ased world models are primarily designed for single-agent scenarios, but they can be naturally adapted to multi-agent settings, modeling either independently local dynamics joint dynamics. Fortunately, we have included IRIS as a Transformer-based world model baseline in our ablation experiments. Specifically, the Centralized Manner and MARIE w/o aggregation variants from our ablation experiments correspond to IRIS baseline variants under different deployment strategies. But different from their original implementation, these IRIS baseline variants also uses the same actor-critic method as MARIE during learning in imaginations phase (i.c., using PPO instcad of REINFORCE for behaviour Icarning). With a shared behaviour learning phase, we can analyze the reason why existing single-agent Transformer- model cannot be directly adapted to MARL. As shown in Figure 13, without incorporating CTDE principle, the learning of single-agent world model would be disrupted by the scalability and non-stationarity issues.
world
E.5 Comparison against MAMBA with Different Imagination Horizon
We report the performance of MARIE and MAMBA on $s__s_5 52 with different imagination horizon in Table 12. And the result shows that the performance gap between the two is not related to the choice of imagination horizon. Interestingly, larger imagination horizon may help policy learning in imagination. We hypothesize that longer imagined trajectories help alleviating shortsighted behaviours in policy learning.
E.6 Comparison between different agent-wiss wise aggregation modules
To convincingly demonstrate the effectiveness of our designed Perceiver aggregation module, we report the FLOPs for both Perceiver and Self-Attention aggregation with varying number of agents, shown in Table
24

---

Published in Transactions on Machine Learning Research (04/2025)
50 many baneling
MARIE
1.0
IRIS (centralized)
(decentralized)
0.8
Win Rate
0.4
0.2
0.0
-0.0
0.2
1.0
Environment Steps
105
Figure 13: Comparison with two direct extension of IRIS in the so
so_ many baneling scenario.
Table 12:
The performance of MARIE and MAMBA on Is us 52 with different imaginat ion horizon.
Maps
Method
Horizon H = 8 Horizon H = 15
Horizon 25
MARIE
0.40 + 0.34
0.75 + 0.09
0.78 + 0.11
38
5%
MAMBA
0.00 + 0.00
0.13 + 0.14
0.16 + 0.13
13. Note that the Perceiver aggregation in our algorithm consists of one transformer layer with cross-attention and two transformer layers with self-attention. The details can be found in Table 4. To make fair comparisons, the Sele-Attention aggregation also consists of three transformer layers with sell-attention.
Table The FLOPs Perceiver aggregation and Self-Attention aggregation with varying number of agents.
Aggregation Module
agents
agents
5 agents
9 agents
Perceiver 0.016G FLOPs 0.024G FLOPs 0.041G FLOPs
0.073G FLOPS
Self- Attention
0.33G FLOPs 0.201G FLOPs 0.335G FLOPs
0.603G FLOPs
Additionally, compare the performance of these aggregation methods in the 2m 1z scenario after 50000 steps. These results show that Perceiver aggregation offers a more computationally efficient solution compared to self-attention aggregation.
Table 14: The performance of different aggregation methods the 2m us 12 sconario after 50000 steps.
Map
Perceiver Aggregation (3 layers) Self-Attention Aggregation (3 layers)
2m 1z
0.96 + 0.07
0.54 + 0.46
F Additional Discussion between CoDreamer and MARIE
Additionally, a recent method CoDreamer (Toledo & Prorok, 2024) extends DreamerV3 (Hafner el. al., 2023) to the multi-agent setting, using GAT V2 (Brody et al., 2022) for communication among agents' world models and policics. Though the aggregation modules in CoDreamer and ours are both built upon the Transformer architecture, our focus lies in computational efficiency of aggregation while focuses on the underlying topological graph structure among agents. However, fundamental difference is the backbone
25

---

Published in Transactions on Machine Learning Research (04/2025)
used for modeling the local dynamics. While we cast the local dynamics learning the sequence modeling over discrete tokens, which can be achieved by using auto-regressive Transformers with causal attention mechanism, CoDreamer directly adopts the RSSM framework in DreamerV3.
G
Additional Results with Standardized Performance Evaluation Protocol
For finer comparisons, we also provide probabilities of improvement in Figure 14.
Algorithm X
Algorithm Y
MARIE
MBVD
MARIE
QPLEX
MARIE
QMIX
MARIE
MAPPO
MARIE
MAMBA
0.80
0.85
0.90 P(X Y)
0.95
1.00
Figure 14: Probabilities of improvement (Agarwal et al., 2021).
26

---

Published in Transactions on Machine Learning Research (04/2025)
Table 15: 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 Tokenizer learning rate World model learning rate Actor learning rate Critic learning rate
Adam
0.0003
0.0001
0.0005
0.0005
Gradient clipping for actor & critic
100
Gradient clipping for tokenizer Gradient clipping for world model
10
10
Weight decay for world model À for A-return computation Discount factor y Entropy coelficient
0.01
0.95
0.99
0.001
Buffer size (transitions)
2.5 10a
Number of tokenizer training epochs
200
Number of world model training epochs
200
Collected transitions bet ween updates
{100, 200)
Epochs per policy update (PPO epochs)
5
PPO Clipping parameter Number of imagined rollouts Imagination horizon H Number of policy updates
0.2
600 or 400
[15, 8, 5]
[, 10, 30)
Number of stacking observations
5 False
Observe agent
Observe last action of itself
False
Table 16: Computational time consumption of MARIE in SMAC.
Environment Steps
100000
200000
400000
Training Time day
2-3 days
4 days
Usage of GPU Mem
22GB
22GB
22GB
H Parameters Setting and Computational Consumption of MARIE
All our experiments are rum on machine with a single NVIDIA RTX 3090 GPU, 36-core CPU, and 128GB RAM. We provide the hyperparameters of MARIE for experiments in SMAC, shown Table 15. To enable the running of experiments in all SMAC scenarios with single NVIDIA RTX 3090 GPU, we set the imagination horizon H as 8 for other scenarios involving the number of agents > 5, 15 for n 5. In ___many_ia ___man__beelling baneling and 2s3z, we set the imagination horizon H as 5. Correspondingly, the number of policy updates in imaginations varies with imagination horizon H. As for the scenario 2c vs 642g, considering the significantly large action space in it, we enable the observation of agent id and last action for each agent and disable stacking the last observations as input the policy.
Based the above reported setting, we present rough computational consumption in Table 16.
27

---

Published in Transactions on Machine Learning Research (04/2025)
Overview of MARIE Algorithm
Pscudo-code is summarized as Algorithm 1.
Algorithm MARIE
// |main loop of training
for epochs do
collect cecrect__xxereeceeeeeee numm_ transitions) do for learning_wool firis_mode_____
Irain world model() end for for learning behavioor steps per epoch do train signns(gets) end for end for
function collect__esee (xppreeceee)rii)ee()):::: Ou f env.reset()
for / = O,...,1- do / processed by VAE 6, < D(E(o,))
Sample of Th (ajlot) Vi O Ft, done= + env env. env.step(a.)
done True then o++1 env.resct()
Nt 0.
else
it + 0.99
end if end for + (o, a Tes 10
function train world model():
Sample for, at, rt, T=r
M-r+#H1
Update (E, D, Z) via LVQ-VAE observations (=r
for agent: i = 1,...,m do
Update o, 0 via (oy 0) over local trajectories (oj, 71 1 1=++
end for
function train agents():
Sample an initial observation Og D
{ K E(og), of +- D(E(o))), Vi
for t= 0..... H 1 do
Sample of ~ (ajlo)) Vi
Aggregate (s]_...,, a. ... x7 r....... Ti af) into (el,...,e?) Sample 11,1,,, fi, Ai
II
.f) via the Perceiver 0 ps ( Hi, A/a.) ag .h, Sy af, Vi Vi
6 + -D( -( end for
for agent i 1, 1...., do
Update actor m and critic Va via LDya 0) over imagined trajectories (0), at, ef 4/)==111 end for
28


Reviews:
[{"review_id": "lxDb53ehzp", "reviewer": "Reviewer_Cex6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\nThe experiments in this paper are quite convincing, although more trials would be helpful to assess the variance. The results on MAMuJoco and SMAC2 seem to suggest that this method robustly outperforms its competitors.\n\nThe ablation experiments are comprehensive and provide supporting evidence for the claims made by the paper regarding the efficacy of aggregation, etc.\n\nThe architecture of the paper is well-justified and innovative. \n\n**Weaknesses:**\nMany components of the model require some centralized aggregation during training. This is a disadvantage in many practical settings which really require each agent to be fully autonomous from the others.\n\nThe experiments are limited to SMAC and MAMuJoco, which may not capture the full spectrum of performance of the algorithm.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new framework for multi-agent model-based reinforcement learning known as MARIE. The paper utilizes a local dynamics model, which aggregates a transformer-based representation seq. The work also contains a suite of experiments on the StarCraft benchmark SMAC with several additional ablation experiments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The aggregation model (perceiver) seems fairly arbitrarily chosen based on existing open-source methods. Do the authors suspect that improvements could be made to this method, for instance with other learned representations of the history?\n\nHave the authors tested the method in other (preferably inexpensive) tasks? If not, there is no need to do so, but any supplemental evidence to support the paper’s claims would be appreciated.\n\nThe clarity of the paper could be much improved. The paper needs to be proofread for typos; I was able to spot some below, but I believe the flow in many sections (introduction especially) could be improved.\n\nThe following are a list of typos, which are by no means comprehensive:\n\nPage 3: model architecture and hyperparameter ->  model architecture and hyperparameters;\n\tAgent i Consists -> agent i consisting\n\nPage 4: historical sequence -> historical sequences\n\nPage 5: global feature -> global features;\n\tAggregated feature -> aggregated features;\n\tBy insert these -> by inserting these\n\nPage 6: Benefited from the shared local dynamics -> By benefiting from …\n\nEqref is not used correctly, so there are phrases like “Eqs. equation 2- equation 5”\n\nConclusion: \"bright the future\" -> \"illuminate the future\"", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "UWAld2md0U", "reviewer": "Reviewer_FQP3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- The overall contribution of this paper is strong and relevant to the community. While model-based RL has demonstrated effectiveness in single-agent RL, its application to multi-agent RL is under-explored. Learning a joint dynamics model for all agents becomes challenging when the number of agents is large, whereas learning individual dynamics models ignores the shared information. This paper reconciles these aspects using a decentralized dynamics model with centralized feature aggregation, providing a scalable instantiation of multi-agent model-based RL.\n- The effectiveness of the proposed method is validated across a suite of community standard benchmarks such as SMAC, MAMujoco, and SMAC v2. The individual design choices are justified via detailed ablation studies.\n- The paper is well-written and logically coherent.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes MARIE, a decentralized multi-agent dynamics model to improve the sample efficiency of multi-agent RL algorithms. The dynamics model features decentralized dynamics prediction and centralized feature aggregation -- instead of jointly predicting all agent's future observations conditioned on current observations and actions, the model predicts each individual's future observations conditioned on its own current observations, actions, and the aggregated features from all agents. This design choice makes the model scalable to a large number of agents. \n\nSpecifically, the proposed MARIE model uses a Transformer backbone to predict each agent's future observations, rewards, and discounts conditioned on its current observations, actions, and aggregated features from all agents. The image observations are quantized into discrete latents using a VQ-VAE. The aggregated features are obtained using a Perceiver Transformer -- learnable query vectors cross-attending to all agents' observations and actions. This model is used in a standard model-based RL pipeline, consisting of on-policy data collection, off-policy dynamics model learning, and policy learning in the imagination of the dynamics models (via reconstructed images). The policies are trained with an actor-critic algorithm. Each actor has access to an agent's individual observation and the critic is augmented with access to all agents' observations. \n\nThe main results are on the StarCraftII Multi-Agent Challenge (SMAC) benchmark, where MARIE outperforms prior model-based (MAMBA) and model-free (MAPPO, QMIX, etc.) MARL algorithms in terms of sample efficiency and asymptotic win rate. The supplemental material contains more evaluations on standard benchmarks such as MAMujoco and SMAC v2, where similar results hold. The authors further conduct ablation studies on design choices such as decentralized prediction, centralized feature aggregation, and VQ-VAE quantization.  These results demonstrate the effectiveness of the model-based MARL paradigm and the specific design choices made in MARIE.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Given the improvements of SMAC v2 over v1, I suggest rerunning all experiment on SMAC v2 and moving it to the main paper. The SMAC v1 results can be kept in the appendix.  \n- In Figure 1, the predicted observations have the same notation as the ground truth observations. I suggest adding a hat to the predicted observations to clarify.", "broader_impact_concerns": "The paper's technical contributions are evaluated in simulated benchmarks and do not yet require adding a Broader Impact Statement. If such a system is deployed in a human-centered environment in the future, then potential safety concerns need to be addressed.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "mAM2Jx4UoI", "reviewer": "Reviewer_29Y4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is clearly written, the figures are neat. The method is well motivated, and well evaluated. The evaluation is broad (many baselines are compared). The design choices are supported by ablations studies. Overall, this paper is a good contribution to the field of MARL. Let us list some opportunities for improvement.\n\nIntroduction motivation is a bit off. The authors insist a lot on the environments containing a large amount of interacting agents (already in abstract « arising from a large number of agents »). Readers thus expect a very large amount of agents (at least 10s). I don’t think that the experimental evaluation reflects that, and thus is a bit off. Maybe adjust the claim or instead write that already with a few number of agents, the accumulated or long term prediction error might explode.\n\nI am not convinced by the better sample efficiency to MAMBA (only). \nI still find MARIE superior in win rate, but maybe this claim needs to be better supported or adjusted. Or is it that MAMBA would catch up one MARIA’s win rate in e.g. the bottom row of figure 2, but has not yet seen enough sample ?\n\nI would advise listing the precise scientific questions that your evaluation are answering in form of an enumeration, before or after experimental setup:\nQ1/ Does MARIE achieve superior win rate performances compared to MB and MF multi agents baselines?\nQ2/ Is MARIE more sample efficient than the existing baselines?\nQ3/ What is the impact of the Perceiver on MARIE’s performances?  … etc (list as many relevant questions as you want)\nYou actually are already nicely presenting bold sentences (that could correspond to answers to these questions), e.g.  on page 9: « Agent-wise aggregation helps MARIE capture the sophisticated interdependency on the generation of each agent’s local observation. »\n\nWhile the figures are clear and neat, the captions do not allow to directly understand what claim it supports, but the readers is required to infer the claim.\nI would advise starting each caption with a bold sentence that summarize the supported claim.\nFor example:\n* Table 1. **MARIE outperforms every considered baselines on all task with more than 3 agents.**\n* Figure 4: **Decentralized execution allows for better scaling to environments with growing number of agents.** (why not show that on the 1c3s5z version, with 9 agents ?)\n\nThis is already done for the Figure 1 and 2 (you can make the first sentence bold).\n\n\nI would advise bringing Figure 13 to the main text (replacing maybe Table 1, that can be placed in Appendix), as first evaluation figure, to first convince the reader of the superior performances of MARIE, before providing more detailed evaluations and ablation studies.\n\nTypos and remarks:\nPage 1:\n* long term prediction*S*\nPage 4:\n* break the sentence starting with « However, when faced with ». It’s too long, difficult to follow.\nPage 5:\n* of local trajectories by insert*ING* these aggregated\n\nReferences need to be cleaned: e.g. How attentive are graph attention networks? Has been accepted to ICLR 2022, yet the arXiv version is cited.\n\nQuestion:\nSample efficiency: transformers are not much sample efficient, what about architecture that incorporate more inductive bias?\nI do not understand (p4):\n« 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. »\nWhy would you want to learn a non-static gamma, different for each agent?\nIs there any way to extract interpretations from the perceiver extracted features?\n* Why does your method and MAMBA achieve better win rates on the super hard corridor task than on the hard 2c_vs_64zg one ?\n* What is the total number of parameters/weights of MARIE in comparison to the other baselines ? It would be a nice add-on to Table 1.\n* Please write the 4 random seeds used also at in the experimental setup paragraph. Do you provide the mean (and stds) over the averaged 10 evaluations of each seeded agents ?\n* Do you retrain the other methods’ agents yourself or take reported results ? \n* Do you intend to provide access to the code after the paper gets accepted ?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces MARIE, a model-based Multi Agent RL framework. MARIE learns a centralized world model, incorporated in each agent (in a decentralized manner) at execution time. They used VQ-VAE to discretize the sequence of past observations of each agent. They then provide these together with the selected actions to a perceiver, in order to extract global representations. These are then added to each agent's state/action sequence. For each agent, the next state, as well as the reward and the discount factor, are then predicted by a transformer (shared among the agents). The authors show that such algorithms achieve more accurate future state predictions, and lead to better performances in a set of tasks from the Starcraft Multi Agent Challenges, as well as from MAMujoco. \nThey finally conduct different ablation studies to support their design choices.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Here are a list of simple changes that could be brought to the manuscript. After these, this paper will be valuable for the readers of TMLR. \n\n* Provide more convincing evidence of the better sample efficiency of MARIE over MAMBA or tune down the claim.\n* Adjust the introduction to reduce the bias towards massive MA environments\n* Restructure the experimental section to provide an explicit list of accurate scientific questions that support the claims and will be shown in the experimental evaluation.\n* Improve the captions by explicitly providing the supported claim.\n* Clean the references", "broader_impact_concerns": "No broader impact statement is necessary.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers are in agreement that the manuscript may be published in its current form.
Decision: Accept (poster)