R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Authors: Harsh Goel
Mohammad Omama
Behdad Chalaki
Vaishnav Tadiparthi
Ehsan Moradi Pari
Sandeep P. Chinchali
Venue: ICML 2025 poster (2025)
Keywords: Multi-agent Reinforcement Learning
Primary area: reinforcement_learning->multiagent
TL;DR: Improving multi-agent reinforcement learning by introducing a new information theoretic metric that enables diversity through world models

Abstract: Multi-agent reinforcement learning (MARL) has achieved significant progress in large-scale traffic control, autonomous vehicles, and robotics. Drawing inspiration from biological systems where roles naturally emerge to enable coordination, role-based MARL methods have been proposed to enhance cooperation learning for complex tasks. However, existing methods exclusively derive roles from an agent's past experience during training, neglecting their influence on its future trajectories. This paper introduces a key insight: an agent’s role should shape its future behavior to enable effective coordination. Hence, we propose Role Discovery and Diversity through Dynamics Models (R3DM), a novel role-based MARL framework that learns emergent roles by maximizing the mutual information between agents' roles, observed trajectories, and expected future behaviors. R3DM optimizes the proposed objective through contrastive learning on past trajectories to first derive intermediate roles that shape intrinsic rewards to promote diversity in future behaviors across different roles through a learned dynamics model. Benchmarking on SMAC and SMACv2 environments demonstrates that R3DM outperforms state-of-the-art MARL approaches, improving multi-agent coordination to increase win rates by up to 20%. The code is available at https://github.com/UTAustin-SwarmLab/R3DM.

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Harsh Goel
Mohammad Omama
Behdad Chalaki?
Vaishnav Tadiparthi?
Ehsan Moradi Pari
Sandeep Chinchali
Abstract
should
Multi-agent reinforcement learning (MARL) has achieved significant progress in large-scale traffic control, autonomous vehicles, and roboties. Draw- ing inspiration from biological systems where roles naturally emerge to enable coordination. role-based MARL methods have been proposed to enhance cooperation learning for complex tasks. However, existing methods exclusively derive roles from an agent's past experience during train- ing, neglecting their influence its future trajec- tories. This paper introduces key insight: an agent's role should shape its future behavior to enable effective coordination. Hence, we propose Role Discovery and Diversity through Dynam- ies Models (R3DM), a novel role-based MARL framework that learns emergent roles by maxi- mizing the mutual information between agents' roles, observed trajectories, and expected future behaviors. R3DM optimizes the proposed objec- tive through contrastive learning on past trajecto- ries to first derive intermediate roles that shape intrinsic rewards to promote diversity future behaviors across different roles through learned dynamics model. Benchmarking SMAC and SMACv2 environments demonstrates that R3DM outperforms state-of oftte-of--h-rrr MARL approaches, improving multi-agent coordination to increase win rates by up to 20%. The code is available at https///ittub..cmm/// Www.sttinSSwwaaaaaab/rrra
Prior Works
R3DM (Ours)
Role-1
Role-2
Figure I. In fire- ftrefighting scenario with two drones, standard role- based multi-agent RL methods fail distribute drones effectively, as roles are inferred from exhibited behavior. By linking roles to future expected bchavior via dynamics model, R3DM achieves better role diffcrentiation and coordination.
1. Introduction
Multi-agent Reinforcement Learning (MARL) has seen in- creasing progress in board games (Meta Fundamental AI Research Diplomacy Team (FAIR) et al., 2022; Perolat et al., 2022), traffic signal control (Chu et al., 2019; Chinchali et al., 2018; Goel al., 2023; Zhang et al., 2025), au- tonomous vehicles (Zhao et al., 2023; Han et al., 2023), stock markets (Bao & Liu, 2019), and collaborative robotics (Krnjaic et al., 2022; Zhang et al., 2022). A significant challenge is to learn policies for agents that enable effective coordination. Centralized Training Decentralized Execu- tion (CTDE) is a common paradigm for training collab- orative policies. CTDE methods such as QMIX (Rashid et al., 2020), QPlex (Wang et al., 2020a), MAPPO (Yu et al., 2022), MAAC (Iqbal & Sha, 2019) and VDN (Sunehag et al., 2018 8) commonly learn shared policy parameters from local observations across multiple agents. These approaches have been effective on benchmarks such as the Starcraft Multi-agent Challenge (SMAC) (Samvelyan et al., 2019; Ellis et al., 2024).
'Chandra Family Department of Electrical and Computer Engineering, The University of Texas Austin, USA "Honda Research Institute, USA. Correspondence Harsh Goel <harshg99@ @utexa..cuux.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
However, these methods, which rely primarily on learn- ing shared policy parameters for agents, often hinder the learning of diverse behaviors that individual agents need to exhibit to complete tasks reliably. To address this limita- tion, recent methods have explored various approaches: (i) encouraging individualized behaviors through diversity or skill-driven intrinsic rewards (Jiang & Lu, 2021a; Liu et al., 2023; Li et al., 2021; Liu et al., 2022), and (ii) explicitly training heterogeneous policies (Zhong al., 2024). While these methods promote the learning of diverse behaviors, the focus on developing individualized behaviors often compro- mises effective coordination (Hu et al., 2024). Furthermore, sacrificing parameter sharing in heterogeneous setups (Liu et al., 2022: Zhong et al., 2024) reduces sample efficiency.
2. We integrate this framework with an existing role- based MARL method, ACORM (Hu et al., 2024). By balancing task rewards with intrinsic rewards, our ap- proach enhances the diversity of role-specific behavior, thereby enabling exploration.
3.
We validate R3DM on SMAC (Samvelyan et al., 2019) and SMACv2 (Ellis et al., 2024) environments. Our method achieves superior coordination capabilities compared to existing role-based MARL approaches by showing up to 20% improvement in the final win rates in the challenging SMAC and SMACv2 environ- ments such as 3s52_vs.3s66 and protoss_10_vs.
2. Background
To balance diversity with efficient cooperation, role-based MARL methods like ACORM (Hu et al., 2024), CIA (Liu et al., 2023), RODE (Wang et al., 2021), GoMARI (Zang et al., 2023), and ROMA (Wang et al., 2020b) learn emer- gent roles from agents' past observations. These approaches encourage complementary and distinct behaviors within the parameter-sharing framework of CTDE training. However, deriving roles solely from past observations often fails to enforce effective coordination and cooperative behavior. For example, in fire-fighting task involving drones extinguish- ing two separate fires (Fig. 1), if both drones have similar initial observations, they are likely to adopt identical roles carly on. This leads to redundant behaviors, such as both drones moving toward the same fire, rather than distributing effectively to address both fires.
2.1. Preliminaries
Multi-agent tasks are modelled as a Decentralized Partially Observable Markov Decision Process (Dcc-OOMDP) G (1, O,S, A, P, R, Qt n, 7), where finite set of n agents, s' E S is the global state from the continuous state space S. Each agent i receives obser- vation of E O where O₂ is its observation space, via the observation model ((s", i) S x I + O,, This results in a joint observation space O- = =(1 x O₂ x... x O.). Each agent selects an action at E A, from its discrete action space A, by learning a policy T.: T x A, + [0, 1]. where of E (A, x O,)* = TV is the agent's 's Abservationnactionn history. This yields joint action a' [af...... ... ofl A, where A == (Aı x A2 x x... x An) is the global action space. The joint action policy is given by π. The team transitions to a next state s*+1 through a transition dynamics model P(s a²)) S A - [0, 1]. and receives reward r R(s",a") shared by all agents. Here, the discount factor y E 1). The joint policy induces the joint state value function QT, (s. a; d) at state s and action parameterized by neural network &
{0, I...., n 13 is
To address this issue, the central insight of this paper is that an agent's 's role should shape its future behavior, meaning an agent adopting different roles at any given moment will naturally follow distinct trajectories. For example, in a multi- drone firefighting scenario, a drone taking up a role to target any specific building would exhibit future observations and actions that naturally diverge as its targets change. This motivates approach, Role Discovery and Diversity through Dynamics Models (R3DM), which maximizes the mutual information (MI) between agents' roles at timestep, their observed trajectories up to that timestep, and expected future trajectories. By linking the current roles of agents to their future expected behaviors or trajectories, R3DM ensures agents learn more distinct yet complementary roles.
Qui (s, a; 3) = Egiioe 0000 NeT
Σ 2 =, 1=0
2.2. Value Function Factorization for Centralized Training with Decentralized Execution
Our key contributions are summarized as follows:
MARL through the CTDE paradigm, such as QMIX (Rashid et al., 2020), has been a major focus in recent research due to the balance between centralized learning and decentralized torization of the value function, which composes the global action value function (s', a); f(Q1 (7f af; Φ2), Qn, ai (Q)) s", a,""; Φ) through a function f. over individual utilities for each agent Q,(7/, al: oQ) with parameters Φq. By adhering to the Individual-
1. We propose novel information-theore objective between agents' roles, observed trajectories, and ex- pected future trajectories. This objective is optimized jointly by a contrastive learning framework (Hu et al., 2024) for learning intermediate roles, which shapes intrinsic rewards devised from observation dynamics models to promote role-specific behavior.
The key idea is the fac-
2

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Global-Maximum (IGM) principle, agents can inde- pendently maximize their utillty functions to obtain optimal global action, i.e., arg maxa: Q Qur(s', (s', a'; = a Qi af: (Q), arg af Qn (T oi [Q)]-
3.2. Intrinsic Rewards for Diversity in MARL
Diversity-based methods incentivize individualized or ex- ploratory behaviors; for instance, MAVEN (Mahajan et al., 2019) is the first work that fosters diverse exploratory be- haviors through latent space representations for hierarchical control. Additionally, EMC (Zheng et al., 2021) proposes to learn better behaviors through the means of an intrin- sic reward derived from the error in the predictions of the Q-values. CDS (Li et al., 2021) promotes individuality by designing intrinsic rewards to maximize the mutual infor- mation between agent identities and roles and EOI (Jiang & Lu, 202 encourages agents to explore by training an Doservation-to-ideniiy classifier, fostering individuality. In contrast to these methods, R3DM not only promotes in- dividuality when beneficial, but also learns emergent role embeddings that distinguish between agents' past and ex- pected future behaviors to enable better balance between specialization and coordination in complex environments with varying team dynamics.
2.3. Roles in MARL
Roles in MARL characterize distinct agent behaviors to fos- tr coordination in CTDE frameworks with shared parame- ters. Given a Dec-POMDP G = (1, 0, A, P. R, P, n, 7),, a role for agent i is a time-varying latent variable m E M determined from its observation-action
history of E (A, x O.)". Here, M is a fixed set of roles with cardinality [M]. Knowing these roles apriori is infeasible unless they are handcrafted for the task, therefore, each agent learns role-representation 2f E Z and role-conditioned policy (+|T,, 2/)= T x A, x Z + [0. 1] where Z c Rd is the role embedding space of dimension d,. These roles emerge dynamically and enable agents to exhibit specialized yet complementary behaviors to coordinate for the given task.
3. Related Work
4. Method
3.1. Role-Based Multi-Agent Reinforcement Learning
In this section, we present our approach, R3DM, which en- hances role learning in MARL through a novel information- theoretic objective that captures the relationship between agent's 's role, past behavior, and future trajectory. In Sec- tion 4.1, we derive a tractable lower bound of this objective, decomposing it into two key components: (i) deriving inter- mediate role embeddings from agents' observation- histories, and (ii) ensuring these embeddings guide diverse and distinct future behaviors to foster effective coordina- tion. We maximize the first component using the contrastive learning framework introduced in (Hu et al., 2024) (Sec- tion 4.2), therefore obtaining intermediate role embeddings from agents' histories. Next, we introduce intrinsic rewards to maximize the second component (Section 4.3), ensur- ing that these embeddings generate diverse yet role-specific future behaviors. This reward balances diversity across tra- jectories associated with different roles while preserving the alignment of each trajectory with its corresponding role.
ROMA (Wang et al., 2020b) introduces a role embed- ding space based on agents' current observations, and observation-action histories, marking it as the first work to achieve this. Expanding on ROMA, RODE (Wang et al., 2021) and SIRD (Zeng et al., 2023) associate each role with predefined subset of the joint action space. This approach simplifies the learning process by reducing action-space complexity. However, these methods maintain roles that access a static subset of actions during training, which pre- vents effective coordination. In contrast, LDSA (Yang et al., 2022), ALMA (Iqbal et al., 2022), and Xia et al. (2023) learn different sub-tasks within the MARL framework to dynami- cally assign roles to agents. Contrastive learning techniques have also been explored to learn role representations for improving multi-agent coordination. CIA (Liu et al., 2023) employs contrastive learning to learn roles by differentiating their credit or contributions to the team. Similarly, ACORM (Hu et al., 2024) leverages contrastive learning to learn roles that distinguish between different behavioral patterns. On the other hand, methods like SePS (Christianos et al., 2021) and GoMARL (Zang et al., 2023) promote heterogeneity and specialization by dynamically assigning agents into dif- ferent subgroups, either using a dynamics model (SePS) or through weights on the individual utilities in the mixing net- work (GoMARL). Different from role-based methods that derive emergent roles for agents based on their observation- action histories, R3DM introduces an intrinsic reward to encourage agents differentiate their expected future be- haviors across different roles.
4.1. Role-based Diversity Objective
Given an agent i, its role mi and concatenated observation- action trajectory 7 which comprises its observation- action history T. and future trajectory 7/ +++1++ of k steps, the MI objective is given as
p(T,+t ++ p(T(th)
mt)
] (2)
= Bie+t mm
log
However, maximizing this objective is intractable in practice. Therefore, we utilize the following theorem to obtain a lower bound of this objective.
3

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Theorem 4.1. Given a set of roles M with cardinality [M, role m² E M, and concatenated observation- action trajectory Tit which comprises its observation- action history गे and future trajectory T with k steps, if e = fo. (T:) denotes the embedding of the observation-action history obtained through a network O.. and ot ~ (x) | et) denotes role embeddinggs obtained from a network 8y. then
[o, ef) is a role-embeddinng from network 0,, then
[(({ let))
B Bef.st.mn log
2 log |M
p(2.)
g es)
Bet.e,mm log
g(2), c9) + Empress/m g( e(o))
where g(2 e.) is a function whose optimal value is pro- portional to p(z]le)), m is a role from M, and et is its p(2])
log p(2f ct))
}+
1(T! 1(7:th;m)) m,) 2 eeg mg
p(z))
corresponding observation-action history embedding.
I(T! - lictt 2f),
(3)
Contrastive Learning. Theorem 4.2 (proof provided in Appendix A.2) enables us to train the trajectory and role embedding networks, Og and 0,. First, agents are clustered into one of the |M roles based on their embeddings et. from which positive and negative pairs of role and agent embeddings are sampled. Following (Hu al., 2024), we employ contrastive learning using bilinear products (Laskin et al., 2020) to formulate a score function g(2 et) (from Theorem 4.2) that computes similarity between role and agent embeddings across clusters. We describe these in detail in the Appendix B.
where 1 [TT 1(,tt1:tt;;,) 22) is the MI between the future trajec- tory r/ baak and role embedding z at time t.
We provide the proof in Appendix A.1. This theorem in- mitively decomposes the objective into two parts i) learn- ing role embeddings through the observation action history
p(afle)) and ii) m(=))
by maximizing the term Bef,=t.mm log
maximizing the MI between the role embeddings and the future expected trajectory which results in intrinsic rewards as shown in Section 4.3. In the context of our firefighting example, drones would first derive intermediate role em- beddings based on their current history. These embeddings act as compact representations of their roles, enabling the drones to differentiate their future trajectories, which would ultimately guide them to address fires in distinct sectors.
4.3. Intrinsic Reward to optimize role based diversity
We introduce intrinsic rewards that maximizes the mutual information (MI) between role embeddings and future tra- jectories through Eq. in Theorem 4.1. These intrinsic rewards are shown to enable diverse future behaviors while simultaneously enforcing role-specific future trajectories, thereby balancing exploration with role-aligned future be- haviors. In our running example, these intrinsic rewards enforce the learning of role embeddings and subsequent policies that ensure that drones distribute effectively across the two fires. The intrinsic reward comprises: (i) Policy Intrinsic Reward, which encourages variability in an agent's policies through roles, and (ii) Dynamics Intrinsic Reward, which quantifies the predictive influence of role embeddings on an agent's 's future trajectory.
4.2. Role-embedding learning Objective
We maximize the term mei.m [log let) in Eq. 3 to learn intermediate role embeddings. However, maximiz- ing this term directly is computationally intractable. To address this, we use the framework developed in previous work (ACORM (Hu et al., 2024)), which leverages con- trastive learning to optimize the lower bound of this objec- tive (Radford et al., 2021; Oord et al., 2018). For the sake of completeness, we explain the contrastive learning frame- work that clusters similar behaviors while separating distinct ones, effectively distinguishing between unique behavioral patterns and moving beyond simple agent identifiers as seen in (Li et al., 2021). For instance, in our firefighting scenario, drones assigned to extinguish fires in different buildings would be grouped into distinct roles based on their trajecto- ries. Using this approach, we aim to derive intermediate role embeddings that capture meaningful distinctions in behavior. We utilize Theorem 4.2 to achieve this goal.
p(=t)
Theorem 4.3. Given a local trajectory -. {oo, au 01, ai....,0u oo} and a learned role represen- lation zi for an agent, the MI between the future trajectory ++1t+h and the role representation of can be expressed as T
-
p(at +f) p(a (T)
I( (r-+tatk, 2/) = Beg.et
Σ log
+
2+k-1 1
ppoott | T 25, at)
log
p(ott11
7 at)
(5)
Theorem 4.2. Let M denote a set of roles, and m E M denote a role. If e = fo. (T() is an embedding of the observation-actt actionnn-aciooooo history through a network O.. and 2!
where play I +(, is the probability of taking action of given the trajectory 7i and role representation 2/, and
4

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
p(o,t" | 2h, d is the probability of observing o given the trajectory Th action d'y, and role representation z/.
pendix B.5) through ELBO (Jordan al., 1999) as:
p(o T T..22.a)) (a))
)
The MI objective is decomposed into two terms: (i) the influence of the role on future action selection and (ii) the influence of the role on expected future observations. See Appendix A.3 for the proof.
Betrrfiat log
ploft a//
go [oolott 12 a/) p(oft2 |
2 B log
(8)
Policy Intrinsic Reward. The diversity induced by the role embeddings in the policy of the agents is charac- terized by the term Ex, -E log Ila in p(e
We leverage the Recurrent State Space Model (RSSM) from the Dreamer series (Hafner et al., 2019; 2023) for the observation dynamies qu (o.+1 | =f, a)) and p(o: -1 TH 2f, d)). The role conditioned variant of RSSM con- sists of the following components: 1) A sequence model Desca (h! ") with parameters dasqq encoding the tra- jectory 7| into a hidden representation h4 2) An ob- servation encoder qv. (d' h4g of) with parameters the mapping observations of into latent representation d', 3) A dynamics predictor Pakyn (dī h4. 2/) with parameters dayn. predicting the latent representation df from hi. and 4) An observation decoder qean (of I hi d', 2f) with pa- rameters Vdee. reconstructing observations of. As stated earlier, we learn two models: a role-agnostic RSSM model p(o 7/ a²) and role-conditioned RSSM model whose corresponding sequence model qu [o( 1 2/, at). These models defer in the implementation of the dynamics pre- dictor and observation decoder sub-networks, where for the role-agnostic model these are Poaya (d) | h1) with pa- rameters ddyn and Peate (of | hnf, df) with parameters @dec respectively. This model comprises a separate sequence model po.. (hg | Tj - with parameters @oog and an en- coder model Po... (h
Eq. 5. To encourage this diversity, we devise a pol- icy intrinsic reward. Since the policies in QMIX follow an e-greedy strategy, we obtain a lower bound for this ratio by leveraging the non-negativity of the KL diver- gence D kL (p (- | 25) || SoftMax (Qe (- T/ of:0)))))) Therefore, based on the equation
p(ad x4.22) ){ p(a: 73)
E //
log
SoftMax (Q: (at | | |/, (())))
E #//.//aaf
log
(6)
p(a; 7
we devise the intrinsic rewards as
ippll
++
DEL (SoftMax (Q:
11 (.1))))
1=t
| T. with parameters dene- Finally, the dynamics intrinsic reward that captures the influence of of on predicting the next observation is given by (see Appendix B.5 for derivation)
where T.) is the average action probabil- ity over all role embeddings z/. i.e., p ((-17)) = Eg [SoftMax (Q; (. T/, at: oo))] Note that E, is taken over role embeddings for all agents i E I at timestep t, and Dxz. denotes the KI. divergence. This intrinsic reward encourages diverse policies to be aligned with roles. We provide a detailed explanation in Appendix B.5.
z4-k-
7 K,yy =
Σ 1=t
[108 Veaane (o. /// d.+1, #t)
+ B2 log awayn (alt2 11/11, {
Dynamics Intrinsic Reward. We derive the dynam- ics intrinsic reward from the latter term in Eq. 5 Eegaat EE 1 log plo. a (r)aa) ]. which is equiva- lent to the mutual information between the future pre- dicted observations of (+I:++ and the embedding role z/. To compute this ratio we learn two models: (i) role- agnostic model that approximates observation dynamics p (o/+1 T/. af) (ii) role-conditioned dynamics model qu (ott1 r/, 2/, af) to approximate the posterior observa- tion dynamics p (o/tt1 7,,2% af)- We then obtain a lower bound due to the non-negativity of the KL Divergence Dg (o.+1 7/ 2f. at) ll (o) +1 T/ o(,))) (see Ap-
- [log Pesee (0/42 1/41 1
+ Ba log Poasa (dft1 (1) m/tt))):
(9)
where 31 and 33 are hyperparameters that balance the de- coder and latent-dynamics terms. The first bracket captures the likeliiid terms of the role-conditioned DreamerV3 model, and the second bracket captures the same terms un- der the role-agnostic model. Note that the hyperparameter 3 trades off between the diversity of trajectories across roles and the specificity of a trajectory to a given role. This incentivizes roles that yield distinct behaviors.
5

---

R3DM:
Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
4.4. Final Learning Objective
Environments. The experiments are conducted on two benchmarks: SMAC and SMACv2, which comprise several micro-managemeen scenarios to control each unit with lim- ited local observations, to defeat an enemy team controlled by a built-in game AI. SMAC includes various scenarios with differing terrain layouts and unit types with varying difficulty levels. We focus on the 6 hard and super-hard sce- narios 5m_vs_6m, MMM2, 3s5z.vs_ss6z, 27m_vs_30m 6h_vs_8r _r_ and Corridor that require more skillful coordina- tion. In constrast to SMAC, SMAC-v2 introduces stochas- ticity in the initialization conditions that further challenge agents to explore diverse behaviors that foster skillful coordi- nation, All experiments are conducted using SMAC version SC 2.4. 10, and we note that performance comparisons do not translate across different SMAC versions.
Given the task reward p² timestep t, we compute the total intrinsic reward rine Eier Pot summed over all agents. Note that 33 is hyperparameter to weigh the policy intrinsic reward relative to the dynamics intrinsic reward. Therefore, the final centralized Q learning objective is
LTD(O)= = Ir' + arint + max Qux(s (6)
e'+) +ymax
- Qiorrs" o))
(10)
Here, 7 is the discount factor, o are the parameters of the frozen target mixing Q network, and a is hyperparameter that weighs the intrinsic reward with respect to task reward. We detail the choice of the hyper-parameters in Appendix B and the entire algorithm in Appendix C.
Baselines. In addition to QMIX, compare the perfor- mance of R3DM against five state-of-the-art baselines, that learn roles through contrastive learning (ACORM & CIA), learn grouping mechanism for agents (GoMARL), and de- vise intrinsic rewards for diversity (CDS) and exploration (EMC), on the same random seeds.
0.0
5.1. Results on SMAC
0.
2.3
3.2
33 28
Staps
Enviforment Steps
vs, 30m
6h
Corridor
R3DM demonstrates significant improvements over base- line methods, particularly in challenging SMAC scenarios such as 6h_vs_8z, Corridor, and 3s5z_vs_3s66 where agents face stronger enemy teams. These results highlight the effec- tiveness of our approach in learning cooperative policies as well as achieving higher sample efficiency. For example, in Corridor and 3s5z.vs_3ss v3362, our method outperforms others by achieving faster convergence and higher test win rates compared to other baselines.
2.0
0.5
04
BM 1.
4.0
M0 18 24 6.6
OMX
Figure 2. Test Win Rate of R3DM compared to baselines on 6 maps in the SMAC. We observe that R3DM improves sample efficiency, and converges to higher win rates on super-hard envi- ronments such as 3x52. 1.55_vs_... vs_3,6,, Corridor, and v8.82.
We also note that R3DM outperforms ACORM, on which our method is based, in many environments. This is likely due to the limitation in ACORM, which restricts the ex- ploration of cooperative strategies due to the strong inter- dependence between agent identities and past trajectories. In contrast, R3DM not only promotes additional diversity through intrinsic rewards but also fosters role-specifie future behavior. Therefore, we exhibit improved sample efficiency through exploration and better performance due to the dis- covery of cootei-coodd behaviors amongst agents.
5. Results
In this section, we present the results of our experiments conducted on the challenging environments on both the Star- craft Multi-Agent Challenge (SMAC) benchmarks SMAC and SMACv2. Our primary goal is to evaluate the perfor- mance of R3DM by addressing the following questions.
1. Does R3DM facilitate the learning of winning coordi- nation strategies in multi-agent domains (Sections 5.1 and 5.2) through intrinsic rewards?
5.2. Results on SMACv2
The stochasticity introduced in SMACv2 such as variations in agent types, starting positions, and configurations, chal- lenges MARL algorithms to explore behaviors that gener- alize well. In this context, R3DM demonstrates superior performance over established baselines in terms of test win rates and test cumulative reward. We note that R3DM out- performs Vanilla QMIX due to the role learning mechanism
2. Does R3DM learn roles that qualitatively show distinct future behavior to enable cooperation (Section 5.3)?
3. What design or hyperparameter choices are crucial to R3DM (Section 5.4)
6

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
protoss 5.v5
terran 5 v
200g 3.V
vs
0.5
24
Enstrnmment Stape
Dnvvormant Stacs
Envismment Steps
Eistronment 1E Stops
Elifronment Stops
Erifronmeett Shops
protows
>erg
10
terran 10 vs
11
0.6 1.2
3.0
LS5
3.0 Eriel DENTINTE StOps
155
3.0
1.5
3.0
Ervoromment Staps
Divvormant StAOS
Elvtronment H. Stops
Enstronmentt Staps
civtoomment Steps
4CORM
130M 10urs)
GOMAAL
CI
4D5
OMIX
ENC
ACORM
R3DM (Ours)
6OMARL
CD5
CMRX
BYC
Figure 3. Comparison of Test Win Rate and Test Cumulative Reward on the SMACv2 suite of environments. We observe that R3DM showcases better returns highlighting better strategies learned in environments such as protoss_S.vs.. and terran where its test win rate is equivalent to the best- persfrmefoninn baseline ACORM. In zerg 5 vs 5, protoss 10 vs 11 and zerg 10 vs 11 environments, R3DM outperforms the baselines in terms of test win rates. Note that we report the means (solid line) and standard deviation (shaded regions) across 5 seeds.
that results in diverse behaviors. Furthermore, R3DM out- performs intrinsic reward methods such as CDS and EMC. CDS incentivizes diversity by ensuring distinct trajectories across all agents, therefore, this leads the learning of frag- mented behaviors ill-suited for skillful coordination. This is evident from the mean test cumulative reward plots, which are comparatively poorer across all environments. Addi- tionally, we notice that EMC does not make good learning progress, especially on the asymmetric 10 vs 11 environ- ments, primarily due to its emphasis on exploration.
distinct role embedding clusters in both methods, indicating proper role initialization. However, by t = 20, ACORM's role-conditioned policies show reduced diversity, with all agent clusters converging attack enemy units en masse. In contrast, R3DM demonstrates strategic differentiation: a green-coded stalker agent (highlighted in red) diverts three enemy zealots from the main group, while the remaining agents split into two specialized subteams to eliminate the weakened team. This implies that R3DM learns role em- beddings that enable the future behavioral diversity needed for coordination. Subsequently, we observe that the agents controlled by policies learned by ACORM the enemy agents as the episode progresses to 35 and = 50. On the other hand, since R3DM learned roles that enabled more distinct behavior earlier in the episode, agents effectively defeat enemy agents. We believe that this is observed due to the introduction of intrinsic rewards that optimize for the specificity of the trajectory also enabling exploration by increasing the entropy of the future behaviors demonstrated across all identified roles.
Finally, we discuss the performance of R3DM against Go- MARL, and role-based MARL methods CLA and ACORM. In the 5_vs.5 suite of environments, while R3DM exhibits similar win rates compared to ACORM we observe that it outperforms most baselines in test cumulative rewards. This indicates that R3DM learns more efficient winning strategies. In the more challenging assymetric 10_vs. 11 en- viromments, we observe that all learning algorithms fail to make sufficient learning progress due to the inherent stochas- ticity that results in wider range of unit compositions. We observe that R3DM is marginally more sample efficient and exhibits improved win rates and rewards in the protoss and zerg-based environments.
losing to
role, while
5.4. Ablation
We conduct 3 ablations on the design choices 1) the im- pact of the horizon of the imagined trajectories for reward computation, 2) cardinality of the number of roles, and 3) impact of contrastive learning.
5.3. Qualitative Analysis
We conduct qualitative analysis of the strategics learned by our method in contrast to the best-performing base- line ACORM to depict the behavioral differences between our method and ACORM. We plot the visualization of the learned policy in the 3s52_vs_ vs vs_66 environment and the TSNE plots of learned role representations at timesteps 1, 20, 35, and 50 in Fig. 4. Initial observations at ( = 1 reveal
Impact of Imagination Horizon for Rewards. We ana- lyze the effect of varying the number of imagination steps used to generate role-conditioned future trajectories with timesteps for intrinsic reward computation (note that these are equivalent). In R3DM, single imagination step is employed, as increasing this to two steps yields no statisti-
7

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
R3DM (Ours)
Dead
Dead
Cluder
Clatter
Pote tmbbeditgg Step3
Role Embbeödings st Stepp
Pote Eribecdinga at Step
Role Emieeciiings Step
20
-40
-10
-40
-63
TSNE
T5UE
TSNE
ACORM
Dead
Dead
Dead
Outter
Cluser
Cärder
Gunter
Holt Embedddings at Stop
Embbediings Stop
Eritieddings at
-30
-40
0
-00
-40
15ME-3
Figure 4. We show qualitative results the 3s_vs. 5z environment with the corresponding role embeddings and the clusters. R3DM learns better strategy compared to the baseline ACORM, where one stalker agent, as shown in timestep 20, successfully learns distinct role that lures enemy zcalots for the main team beat weakened enemy differentiated roles based on past observations, the resulting policies are inadequate to win against the enemy team.
the subsequent timesteps. While ACORM learns
cally significant performance improvement. Extending the horizon further (e.g., 5 or 10 steps) results in performance degradation, with the most pronounced decline occurring at 10 steps. We hypothesize that this degradation stems from compounding prediction errors in the model, which conditioned solely agent's local observation-action history. These errors propagate successive imagination steps, producing intrinsic rewards with increased variance and bias that destabilize policy optimization.
heterogeneity, whereas moderate number (3 clusters) opti- mally balances coordination and specialization.
Impact of Contrastive Learning, isolate the contribu- tions of contrastive learning (CL.) and intrinsic rewards, we compare three variants: 1) R3DM without intrinsic rewards (equivalent ACORM), 2) R3DM with intrinsic rewards but without CI., and 3) Full R3DM (intrinsic rewards CL.). The variant without CL underperforms R3DM but still sur- passes ACORM, demonstrating that intrinsic rewards, which are designed to maximize trajectory entropy conditioned on roles, drive significant performance gains. CL further en- hances results by enforcing distinct role representations, reducing conditional entropy between trajectories and role assignments. This highlights the complementary impact of intrinsic rewards (exploration) and CL (representation disentanglement) in the framework.
Impact of Role Cardinality. We evaluate R3DM with the number of role clusters |M] ranging from to 8 in the $s52_vs_3s62 environment. While the final performance remains statistically indistinguishable across configurations, training with 3 role clusters exhibits superior sample effi- ciency compared to higher cardinalities. This suggests that excessive role specialization (e.g., 8 clusters) introduces
8

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
3s5z 3366
3552 vs 3562
3s5z vs__v___aa 3562
0.8
Test Win Rate
0.0
0.0
0.6-
0.6
0.0
Test Win
Win
0.4
0.4
04
0.2
0.9
0.5
12
1.5
Environment Steps
Environment Steps
Environment Steps
rotioitoon 2) (Ouns)
mag notion
RSOM Hitkes
ACORM R32H (Quns)
ROOM (uo Contrntt Learnigs
Figure 5. We conduct an ablation study on the 3s52.vs.3666 environment to evaluate the impact of: 1) Imagination Horizon for Reward, 2) Number of Roles, and 3) Role optimization without Contrastive Learning. We observe that a R3DM with (a) shorter imagination horizons (l: = 1, outperform longer ones (k- 5, 10) due to reduced compounding errors, (b) moderate role cardinality (N, = 3) achieves faster convergence despite similar final performance across configurations, and (c) full R3DM with both contrastive learning and intrinsic rewards demonstrates superior performance compared to the partial implementations.
6. Limitations and Future Work
Impact Statement
While R3DM presents a key insight to improve role-based MARL, it has a few limitations. First, the number of roles in the environment needs to be set apriori, therefore, making a hyperparameter. Future work can explore removing this as a requirement and instead dynamically derive roles from the replay buffer. Additionally, R3DM computes intrinsic rewards through world model the local observation dynamics that do not take into account the influence of other agents' actions or roles on the observations of the ego agent. Future work can be extended to incorporate more sophisticated dynamics models that would more accurately compute the influence an ego-agent's role embedding on its future trajectory to learn better intrinsic rewards that would boost sample efficiency.
Our work presents new algorithm to advance the field of Multi-agent Reinforcement Learning, Beyond its potential impact on multi-player games and environments, we be- lieve that the implications of our work would have minimal societal consequences.
A Acknowledgement
We would like to thank Po-han Li and Jiaxun Cui for their helpful feedback on improving the paper. This work was supported in part by the Honda Research Institute, National Science Foundation Grants No. 2133481 and No. 2148186, and Cisco Systems, Inc., under MRA MAG00000005, through the 6G@ UT center within the Wireless Networking and Communications Group (WNCG) the University of Texas at Austin. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.
7. Conclusion and Future Work
To overcome the limitations of existing role-based MARL methods in fostering effective cooperation, we propose Role Discovery and Diversity through Dynamics Models (R3DM), novel role-based MARL framework. R3DM introduces mutual information-bases objective that estab- lishes direct connection between agents' roles, their ob- served trajectories, and their expected future behaviors. This enables the emergence of specialized and diverse behaviors while balancing exploration and role-specific specialization. R3DM optimizes the proposed objective through contrastive learning on past trajectories to first derive intermediate roles that subsequently shape intrinsic rewards promote di- versity in future behaviors across different roles through learned dynamics model. Experimental evaluations on chal- lenging benchmarks such as SMAC and SMACv2, show improvements in coordination capabilities, therefore improv- ing test win rates and cumulative rewards. These results highlight the potential of our framework and pave the way to incorporate model-based RL into MARL algorithms.
References
Bao, W. and Liu. X.-Y. Multi-agent deep reinforcement learning for liquidation strategy analysis. In Workshops at the Thirty-Sixth International Conference on Machine Learning (ICML): Al in Finance, 2019.
Chinchali, S., Hu, P., Chu, T., Sharma, M., Bansal, M.. Misra, R., Pavone, M., and Katti, S. Cellular net- work traffic scheduling with deep reinforeement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
Christianos, F. Schäfer, L., and Albrecht, S. V. Scaling multi-agent reinforcement learning with selective parameter sharing. In Proceedings of the 38th International Conference on Machine Learning (ICML), pp. 1978-1989, 2021. URL

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
http://proceedinnssmll...reee christianos211 b
Jiang, J. and Lu, Z. The emergence of individuality,
In International Conference Machine Learning, pp. 4992-5001. PMLR, 2021a.
Chu, T., Wang. J., Codecà, L., and Li, Z. Multi-agent deep reinforcement learning for large-scale traffic signal control. IEEE transactions on intelligent transportation systems, 21(3):1086-1095 2019.
Jiang, Y. and Lu, H. Emergent individuality in multi-agent systems through Observation-identigy mapping. Journal of Artificial Intelligence Research (JAIR), 72: 202 lb.
Ellis, B., Cook, J., Moalla, S., Samvelyan, M., Sun, M.,
Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., and Saul, L. K. An introduction to variational methods for graphical models. Machine learning, 37:183-233, 1999.
Mahajan, A., Foerster, J., and Whiteson, S. Smacv2: improved benchmark for cooperative multi-agent re- inforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Krnjaic, A., Steleac, R. D., Thomas, J. D., Papoudakis, G., Schäfer, L., To, A. W. K., Lao, K.-H., Cubuktepe, M., Ha- ley, M., Börsting, P., et al. Scalable multi-agent reinforce- ment learning for warehouse logistics with robotic and hu- man co-workers. arXiv preprint arXiv:2212. 1148, 2022.
Goel, H., Zhang. Y., Damani, M., and Sartoretti, G. Social- light: Distributed cooperation learning towards network- wide traffic signal control. In Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, AAMAS '23, pp. 1551-1559, Rich- land, SC, 2023, International Foundation for Autonomous Agents and Multiagent Systems. ISBN 978 17844903943..
Laskin, M., Srinivas, A., and Abbeel, Curl: Contrastive unsupervised representations for reinforcement learn- ing. In International conference machine learning, PMLR, 2020.
Hafner, D., Lillicrap. T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. arXiv
pp.
preprint arXiv: 1tt.111112.11000 2019.
Li, C.,, Wang, T., Wu, C., Zhao, Q. Yang. J., and Zhang, C. Celebrating diversity in shared multi-agent reinforcement learning. Advances in Neural Information Processing Systems, 34:3991-40022 2021.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap. T. Mastering diverse domains through world models. arXiv preprint arXiv:2301 10X1:22301.0404, .01144 2023.
Han, S., Zhou, S., Wang, J., Pepin, L., Ding, C.,, Fu, J., and Miao, F. A multi-agent reinforcement learning approach for safe and efficient behavior planning of connected autonomous vehicles. IEEE Transactions on Intelligent Transportation Systems, 2023.
Liu, S., Zhou, Y., Song, J., Zheng, T., Chen, K., Zhu, T., Feng, Z., and Song, M. Contrastive identity-aware learn- ing for multi-agent value decomposition. Proceedings of the AAAI Conference Artificial Intelligence, vol- ume 37, pp. 11595 11603, 2023.
He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R. Mo- mentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9733 2020.
Liu, Y., Li, Y., Xu, X., Dou, Y., and Liu, D. Heteroge- neous skill learning for multi-agent tasks. Advances in Neural Information Processing Systems, 35:37011- 37023, 2022.
Mahajan, A., Rashid, T., Samvelyan, M., and Whiteson, S. Maven: Multi-agent variational exploration. Advances in Neural Information Processing Systems (NeurPPS), pp. 7611-7622, 2019. URL https://www. cs ox. uk wwtppople/eeeeeeeetttttwpeeeeeea sshmonn. whiteson/pubs/maaeee hta.:///mmmaaaanisss....11..
Hu, Z., Zhang, Z., Li, H., Chen, C., Ding, H., and Wang.
Z. Attention-guided contrastive role representations for multi-agent reinforcement learning- In International Conference on Learning Representations (ICLR), 2024. URL
497243 (wwo888eebbbe662111118 pdf.
Meta Fundamental Al Research Diplomacy Team (FAIR), Bakhtin, A., Brown, N., Dinan, E., Farina, G., Flaherty, C., Fried, D., Goff, A., Gray, J., Hu, H., et al. Human- level play in the game of diplomacy by combining lan- guage models with strategic reasoning. Science, 378 (6624):0077-007,, 2022.
Iqbal, S. and Sha, F. Actor-attention-crittc for multi-agent re- inforcement learning. In Proceedings of the International Conference Machine Learning (ICML), pp. 2961- 2970, 2019.
Igbal, S., Costales, R., and Sha, F. Alma: Hierarchical learning for composite multi-agent tasks. Advances in Neural Information Processing Systems, 35:7155-7166, 2022.
Oord, A. v. d., Li, Y., and Vinyals, O. Representation learn- ing with contrastive predictive coding. arXiv preprint arXiv: 1xXiv:187.000 2018.
10

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Perolat, J., De Vylder, B., Hennes, D., Tarassov, E., Strub, E, de Boer, V., Muller, P., Connor, J. T., Burch, N., Anthony, T., et al. Mastering the game of stratego with model-free multiagent reinforcement learning. Science, 378(6623): 990-996, 2022.
Yu, C., Velu, A., Vinitsky, E., Gao, J., Wang, Y., Bayen, A., and Wu, Y. The surprising effectiveness of ppo in co- operative multi-agent games. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (NeurIPS), 2022.
Radford, A., Kim, J. W., Hallacy. C., Ramesh. A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natu- ral language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Zang. Y., He, J., Li, K., Fu, H., Fu, Q., Xing, J., and Cheng, J. Automatic grouping for efficient coopera- tive multi-agent reinforcement learning, In Advances in Neural Information Processing Systems (NeurIPS), 2023. URL https: http////ppneeeeeee//e////. httpleewwweeefoorrrrtea id-CG)72TyGJy,
Rashid, T., Samvelyan, M., Schroeder de Witt, C., Far- quhar, G., Foerster, J., and Whiteson, S. Qmix: Mono- tonic value function factorisation for deep multi-agent reinforcement learning. Journal of Machine Learning Research, 21(112):1-45 2020.
Zeng, X., Peng, H.,, and Li, A. Effective and stable role- based multi-agent collaboration by structural informa- tion principles. In Proceedings of the AAAI conference on artificial intelligence, volume 37, pp. 11772-11780, 2023.
Samvelyan, M., Rashid, T., De Witt, C. S., Farquhar, G., Nardelli, N., Rudner, T. G., Hung, C.-M., Torr, P. H., Foerster, J., and Whiteson, S. The starcraft multi-agent challenge. arXiv preprint arXiv: 1902.04999000444333 2019.
Zhang, R., Chen, G., Hou, J., Li, Z., and Knoll, A. Pipo: Policy optimization with permutation-invariaa con- straint for distributed multi-robot navigation. In 2022 IEEE International Conference on Multisensor Fusion and Integration for Intelligent Systems (MFI), pp. 1-7, IEEE, 2022.
Sunchag, P., Lever, G., Gruslys, A., Czamnekki W. M., Zam- baldi, V., Jaderberg, M., Lanctot, M., Sonnerat, N., Leibo, Z., Tuyls, K., and Graepel, T. Value-decomposition networks for cooperative multi-agent learning based team reward. In Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems (AAMAS), pp. 2085-2087, 2018.
Zhang, Y., Goel, H., Li, P., Damani, M., Chinchali, S., and Sartoretti, G. Coordlight: Learning decentral- ized coordination for network-wide traffic signal con- trol. IEEE Transactions on Intelligent Transportation Systems, 2025.
Wang. J., Ren, Z., Liu, T. Yu, Y., and Zhang. C. Qplex: Duplex dueling multi-agent q q-learning. arXiv preprint arXiv:2008.0106 2020a.
Zhao, R., Li, Y., Gao, F., Gao, Z., and Zhang, T. Multi-agent constrained policy optimization for conflict-free man- agement of connected autonomous vehicles at unsignal- ized intersections. IEEE Transactions on Intelligent Transportation Systems, 2023.
Wang, T., Dong, H., Yang, Y., Zheng, Z. Zhang, W., and Wang. Z. Roma: Multi-agent reinforcement learn- ing with emergent roles. In Proceedings of the 37th International Conference on Machine Learning (ICML), pp. 9876-9886 2020b. URL MRRLtttppppprcceeiiii mlr. wrwsspssssse.......p / w/t///1111//// http://wggg/a////o/// wwtp:///www///0..../
Zheng, L., Chen, J., Wang. J., He, J., Hu, Y., Chen, Y.. Fan, C., Gao, Y., and Zhang, C. Episodic multi-agent reinforcement learning with curiosity-driven exploration. Advances in Neural Information Processing Systems, 34: 3757-3769, 2021.
Wang, T., Zhang, W., Dong, H., Yang, Y., and Wang, Z. Rode: Learning roles to decompose multi- aggent tasks. In International Conference on Learning Representations (ICLR), 2021. URL https::////. /atpa:///ariiie /ttp/a/wwvv///////e//eee/e.aaeee org / 2010. 01523.
Zhong, Y., Kuba, J. G., Feng, X., Hu, S., Ji, J., and Yang. Y. Heterogeneous--agen reinforcement learning. Journal of Machine Learning Research, 25(1-67):1, 2024. 2024.
Ji,
Xia, Y., Zhu, J., and Zhu, L. Dynamic role discovery and assignment in multi-agent task decomposition. Complex & Intelligent Systems, 9(6):621 6222, 2023.
Yang, M., Zhao, J., Hu, X., Zhou, W., Zhu, J., and Li, H. Ldsa: Learning dynamic subtask assignment in cooper- ative multi-agent reinforcement learning. Advances in Neural Information Processing Systems, 35:1698-177 11100, 2022.

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
A. Proofs
A.1. Theorem 4.1
Theorem 4.1 Given a set roles A with cardinality |M I a role m E M, and concatenated observation-action trajectory
T! r which comprises its observation-action history 73 and future trajectory of 7/+li** with steps. if et fo. (7f) denotes
the embedding of the observation-action history obtained through network 0e. and zt So. (2f emheddings obtained from network 0,, then
ee)) denotes role
[(r!+(;m)) m!) Befaet.mg log (p(2]) (elle))
)]
+ I( ( ,f),
(11)
p(2{)
where [(r zf) the MI between the future trajectory 7. c++1i++ and role embedding af at time
Proof. We begin with the mutual information objective
[ t m³) = B m log p( mmn)) (mt) p(T!+H)
(12)
where the trajectory r! is concatenation of the observation-action history and future trajectory of agent i. Maximizing this objective is intractable as the role mţ is unknown, therefore, we obtain lower bound of this objective by learning a role encoder 0, to obtain an intermediate role representation zi from embeddings of an agent's 's observation-action history ej through an encoder fe.
We then make the following simplifications the ratio p(r)) p(
p(7)+h/m))) p(T)+t))
lep p(T
2 p(Z et)p(et 7f, m()p(r; defddttt
Jrp p( (T)(()))
pl 21, T/ )p( lef)
E af
(We assume that p(et ((, p(e(lm)))
PTT;"" ++
[ p()p(({))
pl p +:
2 pp(p et)
E Effiefm
p( 1:t+ )p(2f)
where T:- 1-11+ denotes the l steps of the agent's future trajectory that comprises future observations and future actions.
Therefore, the mutual information introduced in Eq. 12 between (+b and role m is
p( +1:1+
=f)p(z et)
I(r ; xo E.ett ,m log Boiatim
p(T +
Leveraging Jensen's inequality, we obtain a lower bound for this mutual information as follows:
p(T; ((11++ = =f)n(= (
I
m{) (r(tt;m²)) E E, mm;
log
E Eeet,ff.m im!
p( (( )p()))
log p(T, +i 2p( (ef)
Ext+ mffiffff
p(T Mitttittth p(()))
log p( +litt 21))
log p(2f ef) p(zt)
E E mf
+ Eateemm
[[] (+:++ ;f) + Esgettm log p(2t ef)
p(=))
12

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
This completes the proof, where we obtain a tractable lower bound of the original MI objective conditioned the unknown role variable mi. By optimizing this lower bound, optimize the original objective.
A.2. Theorem 4.2
Theorem 4.2
Let M denote set of roles, and mf E M denote role. ef (r;) is an embedding of the observation-actoo action history through a network Oc, and z ~ So, (2f ef) is role-embedding from network 0,, then
p(zz et) p(2))
g e
E Betief mt log
2 log [M] +
Bet,et af
log g(2), el.) + em/m/ Em/m e()),
where (<2), ef) is function whose optimal value is proportional to observation-action history embedding.
p(a el) p(a))
m is aole from M, and of is its corresponding
Proof. To prove this theorem, need to first show that the optimal value of the score function g( et) is proportional to the ratio p(=fie)) which describes the ratio between the conditional likelihood of obtaining the role embedding from the p(*
agent embedding and the marginal likelihood of the role embeddings.
Therefore, begin the proof by proving the following lemma.
Lemma A.1. The value of the function g(= et) that minimizes the following term
g(2t,e))
-Eei log g(2),e + Emp-em/m EM/m g(2),e(")
(13)
is proportional to
where m is a role from M and et is its corresponding Observation-action history embedding. p(e))
Proof. Following the proof in (Oord et al., 2018), it can be seen that the function in Eq. 13 is akin to the categorical loss of a model or a function g( ef). Therefore, g( ef) the optimal probability p(i = m ef, Z) that the agent i takes up the
role m based on the agent embedding et and the set of all role embeddings Z { i I} for all agents in I
We write the optimal probability p(i m ct; Z) using Bayes Theorem
m et) p(Z Z
p(i m Z),
Emy p(Z m/, ef
HJEn (p(=f)) p( e
(Since p(2} m/ et) p(z} (ee))
Em, IlEr (p(z)) p(
p(afee) cf)
p(z()
Emy plaglet) p(2{)
p(2 lef)
p(2!)
Since the function g ef) is equivalent to the optimal probability p(i = mm | et, 2), it proportional to concludes the proof of the Lemma.
p(z
This
p(2.)
Substituting the optimal value of the score function Eq. 13, we get
13

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
pizf p(x))
- -Bo
log
p(a!le)) p(at)
+ E EmfrEM//m
pl=jeet" p(a))
p(2{)
p(ef left))
1
Σ
p(2 et
p(2t)
mi E//EM/
p(21)
Bint EM/mf p(2 e,t))
mf
log
+ (IM) 1) p( cgy)
E
p(zz)
1 (IM - 1) p ee)) 1)
p(2{)
log
[.M p(=f)
E Eet
log
p(z (tt)
log p(=f ec)
log
|M
E Bef.at.mm
p(=f)
Therefore, we get the expression
g(2
p(af et) log p(2{)
2 log//M/ - Besaat.m - Bofiatim log et) + E Em/+eM/mi c(-)
Bet log p(= let) p(={)
cf)
log |M Betietm
log
g(2 et) + g(25. elt))
This concludes the proof.
A.3. Theorem 4.3
Theorem 4.3 Given a local trajectory 7f {oo, do, 01, a1...., ,r) and a learned role representation zf for an agent, the MI between the future trajectory +++ 1ot-ft and the role representation 2f can be expressed as
r+ Σ log
p(a x4,22) pla! T!)
+
plo 2122% at) ) p(o [/,
I(7 talhh ;2/) = Exg.et
;2,)
+
Σ log
where p(a! T/. 2!) is the probability of taking action of given the trajectory of and role representation 2f, and p(o/+1 | TH, af, a is the probability of observing of (+1 given the trajectory T/, action a/, and role representation af-
Proof. We leverage the state-action dynamics of the MDP to formulate this proof. We use the fact that the probability of the future expected trajectory 7 ++1+++ is the probability of the concatenated trajectory 75 conditioned the observation- action history T/. i.e. P(7 (a+4)) = p(T + T!))
The conditional probability can be expanded through the Markov chain in the MDP to obtain the probability of the trajectory [+1
in terms of the future actions and observations. We write p( p( 7 I p(a T/)p(o,+
7/. of) Likewise,
the conditional probability of the future trajectory given the role embedding 2 is p( 7f, + p(a!
T/. =f)(((o 1+1 T/. a. 2f). It follows that the MI between the future trajectories and the role representations can be written as:
14

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
p( 1 22/)
T( ; 2/)) = -Ba-+
log p(T;+
T( + (
BBack
log
p( t+k (
li=t ++ -1 p(o) I T/. 2()p( o 24)
B o
log
I p(a; rt)p(ott + | a.)
p(a x2,24() 2))
plo,tt 2f. d()
)
7 &
log
p(a) T/)
+ log
p(o,+1 [T/,
This concludes the proof.
15

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
B. Implementation Details
B.1. Contrastive Learning
We utilize the implementation of ACORM (Hu et al., 2024) to obtain the intermediate role embeddings necessary for the computation of the intrinsic rewards. First, the embeddings of the agent's observation action history of are periodically clustered into |M] groups using K-means, where each eς is assigned to a group Ct = {0, 1]. where i E and j - 1). Using the cluster assignments C, positive pairs are generated from agent embeddings (ef) and role (O...., | representations (zf) within the same cluster, while negative pairs are formed from embeddings across different clusters. More formally, for an agent i in cluster q (Cig = 1), positive keys are k+ = {f)( j E }), where @ = [5 : Cjg = 1), while negative keys are k- = (Jo: (ef) j Φ @).
We use bilinear products (Laskin et al., 2020; Hu et al., 2024) to compute similarities between role embeddings and agent embeddings, respectively, to formulate a score function g(2), ef) exp(>;W (e))), where W is learnable parameter matrix. This score function measures the similarity between the query role representation of of agent and key is e k+ belonging to the same cluster. Finally, the overall loss function is given by
ENER exp(z(Wkk)
L = - log ΣKER- exp(2(Wk) + EEEk- exp(2[WA))
(14)
where the network parameter Pr is updated by MOCO (He et al., 2020) through network parameters 6,, i.,, 8 0)0,+ + C0', where Ç is the momentum hyperparameter.
B.2. Networks
Our implementation details for the critic and role networks are similar to those in ACORM, while we utilize the implementa- tion of DreamerV3 for the RSSM dynamics models compute intrinsic rewards. We employ simple network architectures for the trajectory encoder, role encoder, and attention mechanism. The trajectory encoder consists of fully connected multi-layer perceptron (MLP) and a GRU network with ReLU activation, which encodes an agent's 's trajectory into a 1 128-dimensional embedding vector. The role encoder is fully connected MLP that converts the 128-dimensional trajectory embedding into a 64-dimensional role representation. For the mixing network, we follow the same configuration as (Rashid et al., 2020), which includes two hidden layers with 32 dimensions each and ReLU activation. Additionally, we follow the implementation of the mixing network in ACORM, which computes the attention weights between the embeddings of the global state trajectory and the individual role embedding representations. These attention weights are used compute the mixing weights for all agents' utility functions in conjunction with the global state. For this mixing subnetwork, the dimension of the embedding of the global state trajectory is equal to the state embedding, which is 64. Further details of these network structures are summarized in Table 1.
The learning hyperparameters are consistent with those of ACORM, where we use the Adam optimizer with a learning rate of 6 x 10-4 Exploration is conducted using an e-greedy strategy, where e linearly reduced from 1.0 to 0.02 over 80,000 timesteps and remains constant afterward. Episodes collected through online interactions are stored in a replay buffer with a capacity of 5,000 state transitions, and Q-networks are updated using batches of 32 episodes sampled from this buffer. The target Q-newwrrk is updated using soft update strategy with momentum coefficient of 0.005. Additionally, the contrastive learning loss is optimized jointly for every 100 Q-network updates. All the learned decentralized policies are evaluated every 5,000 updates using 32 generated episodes. For most runs, we set the number of clusters to 3. We detail all hyperparameters in Table 2, and we outline any deviations from the standard used hyperparameters within the following table.
Table 1. The network configurations used for ACORM based QMIX
Network Configurations
Value
Network Configurations
Value
role representation dim
64
hypernetwork hidden dim
32
agent embedding dim
128
hypernetwork layers num type of optimizer activation function
2
state embedding dim attention output dim attention head num attention embedding dim
64
Adam
64
ReLU
4
add last action
True
128
16

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Table 2. Hyperparameters used for ACORM under SMAC and SMAC-V2.
Hyperparameter
SMAC
SMAC-V2
buffer size
5000
5000
batch size
32, 64 for 6h vs_8z
32, 64 for protoss_ 10_v__0vvv 11, terran tewraenel.vs...... 10.vs_01,. zerg_10_vv. vs_
learning rate
6 10-4
6 10-4
learning rate decay
True
True
contrastive learning rate
8 10-4
8 10
momentum coefficient 3
0.005
0.005
update contrastive loss interval Tet
100
100
start epsilon 6s
1.0
1.0
finish epsilon ef
0.02
0.02, 0.00 for zerg 5 vs_5, 0.001 for terran 10_vs_11
e decay steps evaluate interval
80000
80000, 100000 for zerg-5.vs.
5000
5000
evaluate times
32
32
target update interval
200
200
discount factor
0.99
0.99
cluster num
3,5 for 5m.vs.6n
3, 4 for terran.
10_vs_11 1
B.3. Intrinsic Reward
The complete intrinsic reward is given by the following equation:
De (SoftMax (Q: (- 1- =Q)) 1| (+))
=ß3
81 Ilog Joann (o,t) h 141 d+t1, 2/) + B2 log quaya (d lbtt1 21))
log Pracee (o.t) 11tt1, d + B2 log Paasa (d/t2 | 11tt)),
where 4wayn (d) hi, 2 with parameters thayn is the role-conditioned latent dynamics predictor of the RSSM dynamics model in DreamerV3, Pou.. (d) (x)) with parameters dayn is the role-agnostic latent dynamics predictor of the RSSM model, Tyaso (of h4. d3, 2f) is the role-conditioned dynamics decoder of the model with parameters d/decs and Paa. (of 14 d)) the role-agnostic dynamics decoder of the RSSM model with parameters @dec We detail the components and the training of both the models in the subsequent section. Additionally, Qs (- 7f. 2/3 ΦQ) are the agent's utility functions conditioned on the trajectory, and we set p (- T!) as
p ((1+:)) = E EV SoftMax (Q, (: T/ 2f: 64)),
(15)
JET
where I is the set N agents. To simplify the hyperparameter tuning, we set 3g to and we set the decision horizon length 2 minimize training time. Note that the parameter 3g controls the significance of the log-likelihood of the future latent state prediction given the current trajectory and role representation.
B.4. DreamerV3 Implementation Details
We describe each of the architectural components and their corresponding network structures. The cole-conditioned variant of the RSSM model utilized in Dreamer V3 (Hafner al., 2023) comprises of the following components:
1. A sequence model Te (hg) 7/ 1111 with parameters seeq encoding the trajectory T5 into hidden representation ht.
17

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Table 3. Intrinsic Reward Hyperparameters.
Environment
Map
B1
32
5m_vs_6m
0.05
2.0
0.0
1.0
MMM2
0.05
2.0
0.0
1.0
3s5z_vs33s6z
0.1
1.0
0.0
1.0
SMAC
27m 30m
0.1
1.0
0.0
1.0
Corridor
0.1
1.0
0.0
2.0
6h_vs_8z
0.10
0.9
0.0
2.0
protoss.5_vs.
0.05
0.5
0.0
1.0
terran_5_vs_5
0.05
1.0
0.0
1.0
zerg-5_vs.5
0.05
0.2
0.0
0.5
SMAC-V2
protoss 10 vs 11
0.05
0.5
0.0
0.5
teman._ Crran_10_vvs 10_vs_1111111
0.05
0.5
0.0
0.5
zerg.10_vs_11
0.05
0.2
0.0
1.0
Table 4. Network and Training Hyperparameters for Dreamer V3 Implementation
Parameter
Value
Parameter
Value
Hidden Size
128
16 default
Batch Size
32 for SMACv2 5_vs_5 suite
64 for SMACv2 10 vs 11 suite and SMAC 6h vs 8z.
Deterministic Hidden Size
128
Max Batch Length
Max Environment Steps
Dynamics Stochastic
16
Model Learning Rate
le-4
Dynamics Latent Discretization
16
Dataset Buffer Size
5000
Encoder MLP Units
128
Initial Latent
'Learned'
Encoder MLP Layers
2
Weight Decay
0.0
Decoder MLP Units
128
KL Nats
512
Decoder MI.P Layers
2
Reconstruction Loss Scale
1.0
Activation
SiLU
Dynamics Loss Scale
0.5
Optimizer
Adam
Representation Loss Scale
0.1
Gradient Clip
1000
Latent Unimix Ratio
0.01
The sequence model is implemented an MLP that first maps the previous hidden state hi-1, action, and observation of -f to hidden features of size 128. This is subsequently passed GRU with recurrent depth and hidden embedding size of 128 to yield the hidden state features.
2. An observation encoder 1o (d) M4 of) with parameters vhe mapping observations of into discrete latent representa- tion dţ. The observation encoder comprises a 2-layer MLP with intermediate dimensions of size 128, to 16-dimension latent vector where each dimension is discretized to 16 bins.
3. A dynamics predictor 4way.. (de 14, z!) with parameters dayn. predicting the latent representation df from ht. This a single layer neural network.
4. An observation decoder Teaco (o) hC. df,,))) with parameters thaee, reconstructing observations of. This is 2-layer MLP network that maps the latent representation and hidden representation back to the observations. The intermediate dimension size here is 128.
We utilize the DreamerV3 torch implementation' and refer readers to their repository for more details on the implementation. Here, we describe the hyperparameters used in our training setup for the DreamerV3 RSSM model. We encourage the reader refer to the DreamerV3 paper (Hafner al., 2023) to understand the definitions of the hyperparameters.
'https://gtt thub. com/NM5 2rreammrr.....
18

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
B.5. Derivation of the Intrinsic Rewards We obtain lower bound for this ratio Ei=t log
plaiiffed p(a I 2f)).
Eq. 5 by leveraging the non-negativity of the KL We outline how the lower bound is obtained:
divergence D (SoftMax (Qs (.1 , (4:22:00))) 11
T/
Betirial log (p(at Th, r/) p(at T/))
)]
SoftMax (Qi (o) I (1:2300))) ( p(a) T/)
)
2 E 2,rtt
log
log p(aflr/ play 75)
We begin with the ratio Egg
and write it as the following
p(a. 2/)
SoftMax (Qi (at | 25 00))
log
log
SoftMax (Qi (a') T 2/: 20))
p(at | T!))
SoftMax (Q: (a' T =l: (l:)0)))))
Befor!
aag
log
De (SoftMax (Q: (- T/ 2f:0)) 1| (1
p(a! 7/))
SoftMax (Q: (at | r/, 21: (Q))
E Beg.rtiad
log
p(a) T!)
Therefore, we derive the intrinsic policy reward from the lower bound Bst.r! log SonMax (a.Maa( and
p(a
compute the intrinsic reward as the KL Divergence Dr. (SoftMax (Qr (- T/. 2t: ΦQ)) |I 1 +()) between individual utilities Qi (- T/ 2/i dq) and role agnostic action probability p (: (.)T)).
Likewise, we derive the dynamics intrinsic reward to obtain the lower bound due to the non-negativity of the KL Divergence (o | af,aa) || p(of+1 1 a()) as:
p(oott1 T/. 2/, a!)) p(o T/ ab))
]
golott 1 (4,2((a))
]
E Besttti
log
E Ratiriaa
log
p(o,tt |T,
[oog
pio Ir aa))
We can write the ratio E.
plo,t"lrf (rf.o.) as
p(o,t1 1+1
T/ 21, at)
=Eggrr/a
log
p(o,+ 1+1
T/ at)
plo, 1
T/ 24, at)
coolott1 (2
]
log
+ log
(0,
1//, 2/ dt)
plo," 1+1 (T,
4o (o,
2/,2%,a0
= -Egjorla
log
|l p(ott2
T/ 24,4))))
p(o:
a
(o)+) p(o
2/ at)
2 Estor
log
T/. of)
This yields the above-mentionns lower bound. From this, we derive the intrinsic rewards. Note that we use separate models, role conditioned DreamerV3 model (o 7/. 2f, at) and role-agnostic DreamerV3 model p(o/+1
of af).
Here, the RSSM model (o | T/. =f, af) can be written (o [+1 T/ 2/ at) Qmna (11+11 7 To)oo (d,
h ht+ =o0 (of h 1+1 d++1, =f), and p(o,+1 TT,
Porna (n 7 r:)oo (d h Po (of h
h 1.11).
19

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
Therefore, the ratio Eeg [log (F
is given as:
plo)
Ootucn (hit1 | TrDaa (dst)
11 =2) (of h h+1 de zf)
log
f ferm (hz²
Ponyn (dg
h4
x+1 Ponm (0) h h(+1 )
Jorayn (dlt 1 nott 2 Doideo
(0
11
1+1 ,,tt1221
}]
log
(sequence models will yield similar hidden states)
Poayn (cst) 16:t Pdsee (o
h
++1
.dit))
log Auaee
(o/+1 11/11, dl+1, 24 + log qwayn (d++1 li/t1, 21)
- -(lg Polcco (o ++1
111 d + log Peayn (d +1 nn )n(tt).
We add the intrinsic reward hyper-parameters 31 to regulate the contributions of the log-likelihood of the observations under the role-conditioned model. Likewise, add the hyperparameter B2 to regulate the relative weight of the log- ligelioels between of the observations through the decoder model and the future latent state under the dynamics predictor. The final dynamics intrinsic reward rf 1,yy computed at a future timestep + from a given timestep is given by:
dyy = (lo Team (0,+1 h1-1 d,t1, 2)) + 32 log Telann (d,t1 11tt1 =t)
-log log Posxe (0,+1 | h,41, d + B2 log Poaan (dk 1 h 11t
-
20

---

R3DM: Enabling Role Discovery and Diversity Through Dynamics Models in Multi-agent Reinforcement Learning
C. Algorithm
outline the algorithm for R3DM based on ACORM and highlight the differences in Blue.
Algorithm R3DM Based on ACORM and QMIX
Input: 02: agent's trajectory encoder, 0,: role encoder, | number of clusters, &'sqq Ve. Ydya, and tiacc role cole-conditioned D3eamerV3 RSSM network parameters for sequence model, encoder model, dynamics model, and decoder model respec- tively, Gseq. de. Ødyn. and Oiee: role-agnostic DreamerV3 RSSM counterparts, Tar time interval for updating contrastive loss, n. number of agents, B: replay buffer, T: time horizon of a learning episode, Ç: MOCO momentum hyperparameter, 0. key role encoder, a, 31, 32g .3: Intrinsic Reward hyperparameters Output: Parameters of individual Q-network ΦQ and mixing network Φ
1: Initialize all network parameters and replay buffer B for storing agent trajectories.
2: for episode = 1, 2, do
3: Initialize history agent embedding of and action vector af for each agent.
4: for 2.......
do
5: Obtain each agent's partial observation {o and global state st.
6: for agent = 2, ..... do
7: Calculate the agent embedding e = fa, a and the role representation = fo, (c)).
8: Select the local action a according to the individual Q-function Qi(e,, af).
9: end for
10: Execute joint action a' = [aj ag.................... TT and obtain global reward r².
end for
Append the complete trajectory to Bffofl [o if episode mod T
sf "e 12
13:
then Sample batch of trajectories from B. Partition agent embeddings {er] into |M clusters to obtain cluster allocation matrix Ci₃ = {0, 1} where
14:
15:
j E ....... ..... M - } using K-means.
16:
Update key role encoder parameters 0/, = (1 -),, + CØ2 for agent = 1, 2, . do Obtain cluster q such that Cig = Construet positive keys k+ k- --=18 {fo (ef) j E ))
17:
19:
(So j ()), where @ [j: Cjo 1), while negative keys
20:
end for
21:
Update contrastive learning loss according to Eq. 14.
22:
end if
23:
Sample minibatch B' E B and update RSSM network parameters v'sequ vee Vayns and Pice.
Sample minibatch B" B and update role-agnostic RSSM network parameters Dseg Des Dayns and dite-
25:
Sample minibatch 8" B of size K.
26:
for k = 1,2, ... K do
27:
for time = 1, 1,2 do
28:
Set rim. = 0.
29:
for agent i = .,..,, do
30:
Compute i,,, from Eq 7 from individual Q-network for and role representation 2f
31:
Compute r i,,,t from Eq 9 from =f, through models with parameters v/dyns Vitecs Ohayn and thee.
end for
33:
Set rft
r' = =EEE Bor I
like + dyn ri,
34:
Set = an + rk-
35:
end for
36: end for
37:
Update the parameters of individual Q-network &e and the mixing network ó with modified minibatch B"
38:
end for
21


Reviews:
[{"review_id": "7xApZbuSY4", "reviewer": "Reviewer_STTJ", "summary": "The authors propose R3DM, a new role-based MARL framework that enhances coordination by learning roles that shape agents' future behavior through maximizing mutual information and using intrinsic rewards derived from dynamics models.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their detailed and thoughtful feedback, which has greatly contributed to improving the scientific rigour of this paper. We detail our responses to the questions below and add the required ablation studies.\n\n1) Intrinsic Reward: \n\nWe clarify the intuition behind our intrinsic rewards. The rewards encourage agents to diversify their future behaviors, conditioned explicitly on their assigned roles. Formally, it represents the difference between the entropy of future trajectories and the conditional entropy of these future trajectories given specific roles. Therefore, the intrinsic reward inherently promotes diversity by increasing the overall trajectory entropy, ensuring diverse future behaviors while simultaneously maintaining the coherence and alignment of behaviors with respect to roles. \n\n2)  Clarification on the MI Objective : \n\nThe MI objective in R3DM is designed to ensure that an agent’s current role influences its future behavior while being grounded in its prior observations. This objective quantifies the dependency between the concatenated observation-action trajectory (which includes both the agent’s observation-action history and its future trajectory) and its role. Intuitively, maximizing this MI objective ensures that the roles derived from observation-action histories can effectively predict future actions and observations. This, in turn, enables better specialization within the team, which facilitates more sample-efficient multi-agent learning, as demonstrated in our experiments.\n\n3) Ablations:\n\nWe conduct 3 ablations on the design choices. The plots are in the link.\nhttps://drive.google.com/file/d/1b4IkSXhNhXFedi2-d8K8tTEJjCVz4hey/view?usp=sharing\n\na) Impact of the reward imagination: We analyze the impact of the number of imagination steps in the role-conditioned future trajectory, which is used to predict intrinsic rewards. In R3DM, we use a single imagination step to compute intrinsic rewards. We observe no statistically significant improvements when increasing the number of imagination steps to 2. However, with a higher number of imagination steps (e.g., 5 and 10), we observe a degradation in performance (more pronounced for 10 steps). We believe this degradation is due to compounding errors in the predictions by the model, which is conditioned solely on an agent’s past observations and actions. These errors lead to intrinsic rewards with higher variance and bias, ultimately degrading learning.\n\nb) Impact of the number of roles\n\nWe vary the number of role clusters used in R3DM progressively from 2 to 8 for the 3s5z_vs_3s6z environment. We don’t observe a significant difference in final performances when varying the number of role clusters. We observe that the learning with 3 role clusters is more sample efficient compared to other role clusters. \n\nc) Impact of Contrastive Learning (CL)\n\nWe compare the influence of intrinsic rewards and CL on the performance of R3DM. The variant of our R3DM without intrinsic rewards is equivalent to ACORM, as our method builds upon it. Next, we evaluate a variant of our algorithm that includes intrinsic rewards but excludes CL, which is used to enable more distinct role clusters. We observe that this degrades performance CL. Interestingly, this variant still outperforms ACORM, underscoring the significant impact of our proposed intrinsic rewards. These rewards enhance the entropy of future trajectories based on roles while reducing their conditional entropy with respect to intermediate role representations.\n\n4) Balancing task rewards with exploration diversity\n\nThe reviewer raises an insightful point regarding scenarios where it would be optimal for multiple agents (e.g., drones) to collaboratively target the same objective due to task difficulty. Our proposed approach accounts for such task-specific complexities by integrating intrinsic rewards with task-specific rewards. The balance between these two reward types is controlled by the hyperparameter \\alpha, which can be tuned based on the requirements of the specific task. This design allows our method to flexibly adjust the trade-off between promoting diversity and fostering cooperation, depending on the inherent difficulty and coordination demands of the task.\n\n5) Using a single dynamics model:\n\nThe reviewer rightly points out that employing two separate dynamics models (one role-conditioned and one role-agnostic) increases training computational overhead. Importantly, this overhead is confined to the training phase, and dynamics models are not needed during the testing or inference phase. We have explored the possibility of using a single dynamics model with multiple forward passes using different sampled role embeddings to compute intrinsic rewards. However, empirically, we observed this approach to be computationally slower due to the necessity of multiple forward passes with different role embeddings when computing the intrinsic reward.", "claims_and_evidence": "The core claims regarding R3DM's ability to learn effective coordination and outperform baselines on SMAC and SMACv2 are strongly supported by clear and convincing evidence through quantitative results, qualitative analysis, and a theoretical framework. The acknowledged limitations suggest avenues for future improvement but do not undermine the validity of the findings.", "methods_and_evaluation_criteria": "Generally yes, but: \n1.The authors could have also shown the effect of varying the number of roles.\n2. There is a lack of ablation studies, so it is unclear which components of the method are effective.", "theoretical_claims": "I did not check the proofs in the appendix in detail.", "experimental_designs_or_analyses": "There is a lack of ablation studies, so it is unclear which components of the method are effective.", "relation_to_broader_scientific_literature": "The method leverages and extends existing ideas in role-based MARL and contrastive learning, while introducing a novel approach to intrinsic reward design based on predicting the impact of roles on future trajectories using a dynamics model. This forward-looking perspective on role learning seems to be the key contribution.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper is well motivated in an intuitive manner and presents a novel method grounded in an information theoretic objective. The results are impressive, but there is a lack of ablation studies and a lack of insight on how varying the number of roles affects performance.", "other_comments_or_suggestions": "1. An overall block diagram describing the method would have been useful.\n\n2. Regarding the intrinsic reward - While the mathematical derivation is provided in the Appendix, the core idea of how such a reward encourages role-aligned diverse policies should be highlighted in the main text.\n\n3. The qualitative results could be improved by visualising the learned policies.\n\n4. Ablation studies showing the effect of each component of the proposed methods is needed.", "questions_for_authors": "1. The motivating example is useful, but I was wondering - what if it is indeed better for both drones to target the same building (eg: 1 drone alone cannot put out the fire)? This points to a more general question - does the approach take into consideration the difficulty of the task under consideration? How would one balance task difficulty with diversity?\n\n2. The MI objective could be explained in further detail-  why does this objective lead to the desired role properties?\n\n3. Having a fixed number of roles is a limitation. However, while the paper shows good performance with a fixed number of roles, how would the performance vary with different numbers of predefined roles?\n\n4. Having two dynamic models, one for role-conditioning and one for role-agnostic could be expensive. Have the authors thought about how a single dynamics model may be used?", "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": "", "extra_scores": {}}, {"review_id": "iBEmRtzfy0", "reviewer": "Reviewer_J2zb", "summary": "The paper proposes a method for improving the ability of agents to learn to effectively coordinate. The method is based on an existing idea that clusters learning agents into roles based on their observation history using contrastive learning but extends this idea by encouraging diversity among the roles to aid exploration and role selection. The paper evaluates the method in SMAC and SMACv2 and compares against several baselines, finding that their method outperforms the baselines in sample complexity in some scenarios (but rarely underperforms).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their time in providing detailed feedback on the paper to better highlight the key contributions and novelties of the paper. We detail our responses to the questions below.\n\n1) Clarification of the contributions\n\nOur primary contribution lies in introducing a novel information-theoretic objective specifically tailored for role-based multi-agent reinforcement learning (MARL). We theoretically demonstrate that this objective can be naturally decomposed into two complementary sub-objectives. One that derives intermediate role embeddings from observation-action histories, which can be effectively optimized using contrastive learning, and \nanother that refines these embeddings by aligning them more closely with future behaviors through intrinsic rewards. \nThis decomposition ensures that roles derived from past interactions are not only compact representations of historical data but also predictive of future actions, thereby enabling role-specific coordination and specialization. \n\nWhile our framework incorporates contrastive learning and clustering techniques borrowed from ACORM, our approach subsumes this to optimize the broader optimized objective. Hence, they are integral to derive these intermediate role embeddings to optimize our proposed objective. For the sake of completeness and transparency, we have included them in the methodology section of our paper. To avoid any potential confusion, we will make this distinction clearer in the final camera ready version of the paper by explicitly identifying these elements as preliminary or existing work.\n\n2)  Source Code \n\nWe have uploaded and anonymized source code with the necessary data used for the plots in the paper and rebuttal (ablation studies for reviewer STTJ) to the best of our ability here: https://anonymous.4open.science/r/R3DM-F1A0/README.md\n\n3) Clarification with respect to Entropy-based Methods\n\n[1] demonstrates that training a population of policies through adversarial diversity improves robustness by ensuring cross-policy compatibility. However, their approach requires maintaining multiple policies per agent to achieve competitive performance across permutations of partner policies, which increases complexity and computational overhead. We believe that this approach solves a related problem in the field of adhoc-teaming [2].  In contrast, our single-policy learning paradigm focuses exclusively on rapid task-reward convergence by leveraging role-driven trajectory diversity rather than optimizing for robustness across diverse partner policies. Entropy-based exploration methods such as CDS (which we benchmark against) are limited to static role assignments based on the identities of agents, which result in fixed behavior distributions that fail to adapt to evolving team dynamics. R3DM addresses this limitation by employing a mutual information objective (Theorem 4.1) to induce role emergence dynamically from observation-action histories. This allows agents to learn complementary roles based on the similarity or divergence of their trajectory histories. Additionally, R3DM’s intrinsic rewards are designed to capture diversity, specifically in future expected behaviors conditioned on the agent’s role, ensuring that roles evolve naturally and adaptively during interactions with the environment.\n\n\n4) Test Returns\n\nThe \"test returns\" mentioned in line 353 refer to cumulative rewards accumulated during testing episodes. We will clarify and correct this terminology in the final version of the paper to avoid ambiguity.\n\n\n5) Changing of role-clusters \n\nRoles dynamically change throughout the episode because they depend on each individual agent’s evolving observation-action history. Thus, as agents interact with the environment and their experiences diverge over time, their corresponding roles naturally adjust, causing fluctuations in role cluster assignments, even for agents that become inactive or \"dead.\n\n6) Error Bars\n\nThe error bars depicted in the learning curve figures represent the standard deviation computed across the 5 different random seeds used in our experiments. We will explicitly state this detail in the final manuscript to clarify the interpretation.\n\n[1] Cui, Brandon, et al. \"Adversarial diversity in hanabi.\" The Eleventh International Conference on Learning Representations. 2023.\n\n[2] Mirsky, Reuth, et al. \"A survey of ad hoc teamwork research.\" European conference on multi-agent systems. Cham: Springer International Publishing, 2022.", "claims_and_evidence": "Overall, the paper does a poor job of communicating what their contributions are, especially in relation to ACORN. Unless I've misunderstood, the contrastive learning objective and clustering strategy (sections 4.1 and 4.2) is identical to ACORN, but this is not made clear in the paper. If this is the case, using the terminology \"building on\" and \"inspired by\" is misleading on terms of the contribution of the paper and 4.1 and 4.2 should be phrased as preliminary or existing work. If the details in 4.1 and 4.2 do indeed differ from ACORN, these differences should be made clear.\n\nThe claims that R3DM improves in sample complexity over baselines is relatively convincing. The experiments are limited to the SMAC and closely related SMACv2 environments so the generalizability of the claim is not clear.", "methods_and_evaluation_criteria": "SMAC and SMACv2 are good domains for testing algorithm coordination abilities. A broader set of environments would substantially help strengthen the claims.", "theoretical_claims": "I did not check the correctness of any proofs.", "experimental_designs_or_analyses": "All of the experimental claims are somewhat limited by the fact that only SMAC and SMACv2 are used as evaluation environments.\n\nThe experimental claims that R3DM is more sample efficient than existing algorithms is decently well supported although an improvement over ACORN is limited to only a few of the scenarios. However, it's not clear that R3DM converges to better solutions (rather than just more quickly) than other methods because training curves are cut off before convergence in most scenarios. It would be valuable to see experiments run to a larger number of environment steps.", "relation_to_broader_scientific_literature": "The key contribution of the paper is the extension of ACORN to include a new intrinsic reward to optimize role based diversity. The paper's evaluations show that this helps in the SMAC and SMACv2 environments.", "essential_references_not_discussed": "One of the key points made in the paper is that the diversity of trajectories across roles can help improve learning. Existing literature [1] has explored the problems with entropy-based diversity in coordination settings. It would be interesting to see a discussion with respect to this line of work.\n\n[1] Cui, Brandon, et al. \"Adversarial diversity in hanabi.\" The Eleventh International Conference on Learning Representations. 2023.", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper proposes a new intrinsic reward that builds on top of the SOTA method for clustering-based coordination methods, ACORN. The idea is theoretically grounded and the experiments demonstrate the method improves the SOTA.\n\nWeaknesses:\n\nThe experiment settings are limited. The improvements over baselines are quite small.", "other_comments_or_suggestions": "130 col. 2: \"through [an] information-theoretic\"\n355 col. 2: spacing typo", "questions_for_authors": "What are the \"test returns\" recorded in SMACv2? (mentioned on line 353)\n\nI notice that the clustering of roles changes throughout and the episode visualized in Figure 4. For example, the dead agents switch between cluster 1 and cluster 2 from step 35 to 50. Also, the number of agents in the other clusters changes over time. Why does this happen?\n\nWhat are the error bars in the learning curve figures? Please specify this in the paper.", "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": "", "extra_scores": {}}, {"review_id": "HubiZqg2TF", "reviewer": "Reviewer_3pvG", "summary": "This paper tackles the problem of high quality multi agent reinforcement learning. Specifically, allowing individual agents to learn unique policies in order to better collaborate to achieve goals. The authors introduce a new approach to training, R3DM, which utilizes contrast learning to encourage individual agents to learn different policies based on their history and other agents in their group. The method is tested on several problems and appears to outperform standard methods.\n\n## Update after rebuttal\nThe initial review was positive and I did not feel the need to change it after the rebuttal.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback on the paper and for raising interesting and relevant questions. We aim to highlight this clearly in the camera-ready version of the paper.\n\n1) Impact and Details on Contrastive Learning\n\n Contrastive learning plays a crucial role in R3DM by deriving intermediate role embeddings, essential for enabling role-specific intrinsic rewards that drive diverse future behaviors. Contrastive learning is conducted throughout the training process. Specifically, following the training protocol outlined in ACORM (Hu et al., 2024), contrastive learning updates occur after every 100 episodes of data collection. We refer the reviewer to the ablation study under the rebuttal for reviewer STTJ for more details.\n\n2) Compute Cost Increase \n\nThe proposed R3DM algorithm does not incur additional computational overhead during the testing or deployment phase, as only the learned policy is executed. The overhead introduced by R3DM arises solely during training, specifically due to the learning of an additional dynamics model required for computing intrinsic rewards. This extra step significantly enhances coordination through role-specific behavior. This extra step results in the increases moderately increases the training time of R3DM by 50-75% compared to the baselines. With more efficient implementations of the dynamics model on frameworks such as Jax, we believe that the training run-time can be optimized further. \n\n3.  Diversity of Strategy\n\nOur approach inherently encourages strategic diversity as a mechanism to optimize coordination objectives, which in turn improves sample complexity and efficiency. Qualitative analyses provided in the paper illustrate that R3DM successfully facilitates the emergence of diverse and sophisticated strategies. For instance, agents have learned distinct tactics such as deliberately distracting subsets of enemy agents to weaken their overall strength, thereby enabling other teammates to effectively neutralize remaining threats.\n\n4. Have the authors compared their approach against using a shared base policy with small single layers for individual agents?\n\nYes. In our experiments, the baseline QMIX algorithm inherently implements a shared local critic structure that agents use to formulate policies for choosing the desired actions. This implementation persists in R3DM, where we use a shared local critic and the shared role embedding network across multiple agents. We provide the anonymized code here\nhttps://anonymous.4open.science/r/R3DM-F1A0/README.md.", "claims_and_evidence": "The authors provide adequate evidence for each of their claims.", "methods_and_evaluation_criteria": "The methods appear to be sufficient and make sense for the problem being studied.", "theoretical_claims": "There are some theoretical claims made but they do not form the basis of the results.", "experimental_designs_or_analyses": "The experiments are standard in the field and therefore sound.", "relation_to_broader_scientific_literature": "The methods introduced are of great interest to the scientific community, specifically, encouraging diverse behavior in multi agent systems without introducing the huge overhead of individual models for each of them.", "essential_references_not_discussed": "I do not believe any essential references were not discussed.", "other_strengths_and_weaknesses": "The study appears quite complete. It is well written and I find the study to be innovative.", "other_comments_or_suggestions": "Outside of the questions posed, I do not have suggestions.", "questions_for_authors": "1. Can you quantify the role the contrastive learning is playing? Does this need to be performed throughout all of the training or can it be done later once shared policies have been learned?\n2. What is the cost increase of using the algorithm?\n3. How diverse of a strategy can emerge from this algorithm?\n4. Have the authors compared their approach against using a shared base policy with small single layers for individual agents.", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents the idea that an agent's role should influence its future actions for effective coordination and then proposes R3DM, a new role-based MARL framework. R3DM learns emergent roles by maximizing mutual information among agents' roles, observed paths, and expected future behaviors, optimizing the objective via contrastive learning on past trajectories to get intermediate roles which shape rewards to boost future behavioral diversity across roles using a learned dynamics model. Overall, this paper is interesting and with solid technical contributions to MARL. Thus I tend to accept this paper.
Decision: Accept (poster)