Reidentify: Context-Aware Identity Generation for Contextual Multi-Agent Reinforcement Learning
Authors: Zhiwei Xu
Kun Hu
Xin Xin
Weiliang Meng
Yiwei Shi
Hangyu Mao
Bin Zhang
Dapeng Li
Jiangjin Yin
Venue: ICML 2025 poster (2025)
Keywords: Multi-Agent System
Reinforcement Learning
Primary area: reinforcement_learning->multiagent
TL;DR: We introduce Context-Aware Identity Generation (CAID), a framework for improving generalization in multi-agent reinforcement learning (MARL) under varying problem setups.

Abstract: Generalizing multi-agent reinforcement learning (MARL) to accommodate variations in problem configurations remains a critical challenge in real-world applications, where even subtle differences in task setups can cause pre-trained policies to fail. To address this, we propose Context-Aware Identity Generation (CAID), a novel framework to enhance MARL performance under the Contextual MARL (CMARL) setting. CAID dynamically generates unique agent identities through the agent identity decoder built on a causal Transformer architecture. These identities provide contextualized representations that align corresponding agents across similar problem variants, facilitating policy reuse and improving sample efficiency. Furthermore, the action regulator in CAID incorporates these agent identities into the action-value space, enabling seamless adaptation to varying contexts. Extensive experiments on CMARL benchmarks demonstrate that CAID significantly outperforms existing approaches by enhancing both sample efficiency and generalization across diverse context variants.

Reidentify: Context-Aware Identity Generation for Contextual Multi-Agent Reinforcement Learning
Zhiwei Xu
Kun Hu
Xin Xin
Weiliang Meng?
Yiwei Shi
Hangyu Mao 5 Bin Zhang? Dapeng Li
Jiangjin Yin
Abstract
Predator
Prodater
Predator
Prodator
Predoter
Generalizing multi-agent reinforcement learning (MARL) accommodate variations in problem configurations remains critical challenge in real- world applications, where even subtle differences in task setups can cause pre-trate treitra policies to fail. To address this, we propose Context-Awer Identity Generation (CAID), a novel framework to enhance MARL performance under the Con- textual MARL (CMARL) setting. CAID dynam- ically generates unique agent identities through the agent identity decoder built causal Trans- former architecture. These identities provide con- textualized representations that align correspond- ing agents across similar problem variants, fa- cilitating policy reuse and improving sample effi- ciency. Furthermore, the action regulator in CAID incorporates these agent identities into the action- value space, enabling seamless adaptation to vary- ing contexts. Extensive experiments on CMARL. benchmarks demonstrate that CAID significantly outperforms existing approaches by enhancing both sample efficiency and generalization across diverse context variants.
Fko
Pray
Fredsnor 3
State w/ Context 1
Same Contextual State
State w/ Context
Figure 1. An example of CMARI. The central part depicts a fun- damental state from the classic predator-prey task. By introducing different contexts to this state, the resulting environmental states on the left and right sides exhibit substantial differences.
and node optimization in communication networks (Mao et al., 2020), yielding promising results. Nevertheless, significant challenge in reinforcement learning, which is equally prominent in MARL, lies in generalizing across task variants (Kirk et al., 2023). Many existing approaches strug- gle to maintain comparable performance when faced with even minor task modifications, as the policies trained on original tasks often fail adapt effectively. Since dynamic problems common in real-world scenarios, training dis- tinct policies for every variation is infeasible, particularly in multi-agent systems.
Several single-agent studies have focused on generalization across distinct problem variants within the same domain. One approach integrates meta-learning with reinforcement learning to optimize gradient directions that balance multiple objectives (Finn al., 2017a; Yu al., 2019; Nagabandi et al., 2019), thereby improving the robustness of reinforcement learning algorithms in multi-task settings. Another line of research leverages transfer learning to fine- tune pre-trained policies (Taylor & Stone, 2009), enabling adaptation the distributional shifts between original and target tasks. This approach particularly prominent in robotics, where sim-to-real methods (Zhao al., 2020) address the discrepancies between simulated environments and the real world. Additionally, inspired by the success of large models (OpenAI, 2023) in natural language processing (NLP), the intelligent decision-making field has begun scaling up policy networks to improve model performance (Reed al., 2022: Lee et al., 2022). However,
1. Introduction
In recent years, the field of multi-agent reinforcement learn- ing (MARL) has advanced rapidly, shifting its focus from theoretical exploration to practical applications. For in- stance, MARL algorithms have been successfully applied to energy scheduling in power systems (Yan & Xu, 2020), coordinated control in drone formations (Ge et al., 2018),
'Shandong University National University of Defense Technology "Institute of Automation, Chinese Academy of Sci- ences "University of Bristol Kuaishou Technology HHazzhong Agricultural University. Correspondence to: Jiangjin Yin <jiangjinyin@ @ mail.hzau.co www......u...........
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
these efforts often face limitations: they either achieve generalization only in simple tasks, narrowly adapt to specific problem variants through fine-tuning, or use model scaling to combine unrelated task strategies. As result, achieving robust generalization across problem variants remains a significant and unresolved challenge in the field.
2. Related Work
The generalization of RL algorithms to task variations has garnered significant attention, particularly in the single- agent RL domain. The MAML framework (Finn et al.. 2017b), for instance, optimizes initialization parameters to enable models to achieve strong performance on new tasks with minimal gradient updates. First-order optimization algorithms (Nichol et al., 2018), such as Reptile (Nichol & Schulman, 2018), further reduce the computational complexity of MAMI,, making more feasible in real- world scenarios. Many meta RL approaches also focus training history-dependen policies, often implemented recurrent neural networks (RNNs) (Chung et al., 2014) like RL (Duan et al., 2016), which dynamically adapt to interaction histories. MetaGenRL (Kirsch et al., 2020) introduces a meta-learned, low-complexity neural objective funetion based on diverse agent experiences, leveraging off off-ppolcy second-order gradients for improved sample efficiency. Beyond meta RL., transfer RL represents another critical avenue for addressing generalization. For example, a component-base based transfer learning framework (Sodhani et al., 2021) leveraging abstract representations has been pro- posed to facilitate the mastery of complex tasks using prior models. REPAINT (Tao et al., 2021) further enhances deep RL efficiency by combining representation learning and instance transfer techniques. Recent research on in-context learning (Laskin et al., 2023) has explored training models across episodes with task-agnostic processes, enabling generalization to diverse tasks without explicit task-specific adjustments. However, many of these approaches still require fine-tuning to adapt to new tasks, and their applicability is sometimes limited to very simple situations.
Recently, Contextual Reinforcement Learning (CRL) (Ben- jamins al., 2023; Hallak et al., 2015) is introduced to formalize the generalization problem across related problem variants. CRL seeks to find policy capable of solving group of similar Markov Decision Processes (MDPs), where each MDP presents a variation of the problem, defined by a context variable that reflects the specific characteristics of each variation (Cobbe al., 2020; Rajan et al., 2023). Sim- ilarly, we introduce Contextual Multi-Agent Reinforce- ment Learning (CMARL), which focuses on learning poli- cies for multiple agents across various tasks. These tasks may vary in aspects such as initial states and agent types. For example, while the two initial scenarios may seem dis- tinct, they can be transformed into identical configurations by applying a specific way, as illustrated in Figure 1. Conse- quently, in CMARL tasks, some related environment states can often be unified into a single contextual state, thereby simplifying the representation of the task. Furthermore, this transformation disrupts the alignment of agent numbering schemes in the two contexts, which underscores the inherent difficulty in ensuring consistent agent identification across different contexts. Current MARL methods cannot automat- ically assign the appropriate identity to each agent based on the context of the task. As a result, agents must learn separate policies for each situation, significantly reducing sample efficiency in CMARL tasks.
To improve the performance of existing algorithms in the CMARL setting, we propose a novel framework called Context-A ware IDentity Generation (CAID). The main idea behind CAID is to dynamically assign each agent a unique identity. This is achieved through context-aware approach integrating global state information with individual agent attributes. By assigning each agent distinctive identifica- tion within the system, the framework effectively reduces the complexity of the contextual decision space. These iden- tities not only serve as unique labels for agents but also provide foundation for interaction and cooperation among them. Furthermore, the context-aware design ensures that external information about agents is required only during the training phase, allowing for compatibility with most algorithms following the centralized training decentralized execution (CTDE) paradigm. Extensive experiments on benchmarks such as the Vectorized Multi-Agent Simula- tor (VMAS) (Bettini et al., 2022), Traffic Signal Control (PyTSC) (Bokade & Jin, 2024) and StarCraft Multi-Agent Challenge (SMACv2) (Ellis al., 2023) demonstrate the effectiveness of CAID,
Some MARL algorithms enable agents to adaptively define roles form dynamic groups, allowing them to learn opti- mal strategies under varying task conditions. REFIL (Iqbal et al., 2021), for example, improves agent robustness to vari- ations in initial states by randomly partitioning environmen- tal entities. ROMA (Wang et al., 2020) and RODE (Wang et al., 2021b) dynamically assign roles to each agent, where each role corresponds to a specific sub-action space, thereby effectively reducing the complexity of the action space. COPA (Liu et al., 2021) employs a global "coach" agent to coordinate partially observable "player" agents through limited information exchange. This design enables COPA to demonstrate zero-shot policy generalization even in teams with dynamic compositions. COLA (Xu et al., 2023) ap- plies the principle of viewpoint invariance from computer vision to map the state space into discrete representation via contrastive learning. E2GN2 (McClellan et al., 2024) is framework that improves sample efficiency and generaliza- tion through equivariant graph neural networks (GNNs) (Wu et al., 2021). By integrating equivariant and invariant fea- tures, E2GN2 overcomes scalability challenges common
2

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
in traditional methods. However, these approaches are not specifically tailored for Contextual MARL tasks. In contrast, our proposed CAID assigns an identity to each agent, focus- ing on maintaining consistent identifiers for corresponding agents across different states. This design minimizes the need for retraining when tasks undergo minor changes.
framework can be integrated into these methods to enhance their overall performance.
3.3. Contextual Multi-Agent Reinforcement Learning
Contextual Reinforcement Learning (CRL) (Benjamins et al., 2023) extends traditional reinforcement learning by addressing scenarios where collection of related tasks needs to be solved, with each task influenced by a specific context. The context provides additional information that may not be directly observable but can significantly affect decision-making. CRL formalizes such problems using Con- textual Markov Decision Processes (CMDPs) (Hallak et al., 2015), where the state space is augmented to include the en- vironment state and context variable that captures the con- ditions under which the agent operates. Formally, the state is represented as 3 = (s, e), where is the environment state and e denotes the contextual information relevant to the task.
3. Preliminaries
3.1. Decentralized Partially Observable Markov Decision Process
In this paper, we investigate a fully cooperative multi- agent task, which can be modeled as Decentralized Partially Observable Markov Decision Process (Dec- POMDP) (Oliehoek & Amato, 2016). A Dec-POMDP framework commonly used to formalize cooperative drtision-making problems in partially observable and decen- tralized environments. Formally, it is defined by the tuple G = (S, U, A, P,r, Z, O, n1,7), 7), where S represents the set of possible states of the environment, and A = {l.....,) } denotes the set of n agents involved. At each time step, each agent a E A selects an action u E U based on its local ob- servation =" E Z, obtained through the observation function O(s, a): : S x A + Z. Here, s S denotes the true state of the environment. The joint action of all agents is denoted as U. The state transition dynamics are determined by the function P(se+1|$t /st, u²): S x U x S + [0. 1]. All agents share a common reward function r(s, u) S x U + R, and E [0. 1) is the discount factor. The primary objective in the Dec-POMDP framework is to maximize the discounted return Σj=o T'
In CMDPs, the context is assumed to remain fixed through- out an episode, ensuring consistency during the agent's s interaction with the environment. However, the context can vary across episodes, leading to distribution of possible contexts p(c). Consequently, the initial state distribution in a CMDP is given by p(3) p(c)p(s|e), where p(s le) represents the distribution of the environment state s condi- tioned on the context c. For fixed context e, the CMDP behaves traditional Markov Decision Process (MIP). CRL aims to determine an optimal policy x² that maximizes the expected return across all CMDPs. The expected return for policy π within a specific CMDP Me is denoted as Mo), The overall objective is to maximize this return across the distribution of contexts:
3.2. Value Decomposition Methods
max EeNlle) (T. Mc)]]
Value decomposition methods (Sunehag et al., 2018; Rashid et al., 2018; Son et al., 2019; Yang et al., 2020) are among the most widely adopted techniques in MARL, particularly for addressing challenges related to coordination. A fun- damental concept in value decomposition methods is de- composability, which ensures alignment between global and individual agent objectives. This is formalized through the Sadividuall-lobbll- 2019), which stipulates that the optimal action for each agent arg max,, Q2 (7 ).),, must be consistent with the optimal joint action of all agents arg max, Qio (T, u). Mathemati- cally, this condition is expressed as:
In multi-agent settings, CRL introduces additional complex- ity, as each CMDP now represents a multi-agent control problem where multiple agents interact both with the en- vironment and with one another. In this paper, we extend CRL to Contextual Multi-Agent Reinforcement Learning (CMARL) by modeling each CMDP as Dec-POMDP, a framework well-suited for real-world scenarios. The objec- tive of CMARL is to find an optimal joint policy:
(IGM) assumption (Son et al.,
mmxx (Eo-o(o) [R (TT, Mc)]] Mo~GG
where MM. represents the Dec-POMDP induced by the given context e. Our motivation stems from real-world CMARL tasks where the semantics of an environment can shift sig- nificantly even within a single episode. For instance, in a traffic control task, the agent behavior required during morn- ing rush hours may differ substantially from that during the evening, despite both being within one episode. We allow the context vector to evolve over time, effectively treating the episode as a sequence of sub-episodes, each governed by a different latent context.
arg max Viot (T, u) arg max (2(F", Va E A,
(1) where T E T represents the joint action-observation histo- ries of all agents. Qtor is the global action-value function, and Qn denotes the individual action-value function for agent a. Several value decomposition methods have been developed based on this principle, and our proposed CAID
3

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
Encoder
Decoder
nism (Vaswani et al., 2017). This mechanism captures the relationships between input tokens while modeling the data comprehensively from multiple perspectives.
Context
id'
id'
id
id
Softmax & Sample
Building on this insight, we propose the Contextual State Encoder. It produces a context variable è by integrating the environment state s with the local observations z of all agents. This context serves as a substitute for the unavailable ground truth and concurrently assigns unique identities to each agent.
Contextual Identity
Casual Transformer
2
id
id
The structure of the contextual state encoder is illustrated in Figure 2. Inspired by the Seq2Seq architecture (Sutskever et al., 2014), which is widely used in NLP, we employ the encoder component as the contextual state encoder. In CMARL scenarios, the environment state s is highly dy- namic, making challenging to capture contextual informa- tion assign consistent identities based merely on To address this, we represent the input as a sequence of length n + 1, comprising the environment state and the local observations of all agents 2 = {z (=',=".... ,?,....,, ""). Formally, the encoder f Ennoddr generates the contextual information as follows:
Figure 2. The overall architecture of the contextual state encoder and the agent identity decader. The input includes the global state and observations from all agents, while the output comprises the inferred context and the identity assigned to each agent.
4. Methodology
This section presents a detailed overview of the CAID frame- work, which comprises three core modules. The first two modules are the Contextual State Encoder and the Agent Identity Decoder, both built upon a Transformer network architecture (Vaswani et al., 2017). These modules work collaboratively to infer contextual information from the en- vironment and generate unique identities for each agent. The third module, the Action Regulator, transforms the actions selected by each agent based on its assigned iden- tity, ensuring context-aware decision maaiiin.. A key feature of CAID is its ability to map the environment state space into compact, low-dimensional space. This mapping fa- cilitates the identification of similar contextual scenarios, allowing agents to infer analogous contexts and assign iden- tical identities even when operating in significantly different environment states. By leveraging these identities, agents can effectively reuse previously learned policies, improv- ing adaptability and efficiency in complicated multi-agent environments.
è e= Bncodeer (s, 2).
(2)
While the contextual state encoder does not directly recon- struct the ground truth c, it achieves two key objectives: (I) mapping the complex state space into a compact, low- dimensional representation, and (2) extracting relationships between the environment state and agents' local observa- tions. These relationships capture critical subset of the contextual information, allowing for more efficient policy learning and identity assignment in multi-agent systems.
4.2. Agent Identity Decoder
After generating the contextual information, the next step focuses on assigning unique identities to agents. In most ex- isting MARI. frameworks (Lowe et al., 2017; 17 Kurach et al., 2020), agent identities are determined using heuristic rules. However, in scenarios with similar contextual settings, such methods often fail to ensure consistent alignment of agent identities across different environments. To overcome this limitation, we propose a context-aware identity assignment module capable of automatically assigning each agent a unique identity id based on the generated contextual information.
4.1. Contextual State Encoder
In practical CMARL tasks, the context and its distri- bution p(c) are frequently inaccessible, rendering the context-augmented state 8 unavailable during training. As result, agents must learn their policies based solely on the environment state 8, which lacks explicit contextual information. When changes in c cause variations s, agents must relearn their policies, which dramatically increases training difficulty.
As shown in Figure 2, the Agent Identity Decoder is designed using a causal Transformer decoder. Inspired by natural language generation tasks and the Pointer Net- work (Vinyals et al., 2015), the identities are determined in an autoregressive manner. The process begins with special token (BOS), commonly employed in NLP tasks to indi- cate the start of decoding. Representing the agent identity decoder as Decoder, the probability of identity assignment
Inspired by the capabilities of large language models (LLMS) (Touvron et al., 2023), we observe that these mod- els often produce consistent responses even when the or- der of input words changes. Research indicates that this robustness is attributed to the underlying Transformer archi- tecture, particularly its multi-head attention (MHA) mecha-
4

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
Que
same identity may have identical absolute semantics in their action spaces, but the relative semantics of their actions vary significantly across different states, It is necessary to transform their action spaces accordingly.
Mixing Network
Q
(+), ni. idf)
(r,, uf, id?)
During the centralized training phase, the contextual state encoder and agent identity decoder generate and id". How- ever, these identities cannot be used as inputs the agent networks during the decentralized execution phase. We draw inspiration from hypernetworks (Ha et al., 2017) to overcome this limitation and propose an Action Regulator to align agents' action spaces based on their identities. The action regulator operates on the Q-values output by each agent and consists of two linear layers. Both layers take the agent's identity id as input, generating the weights w and bias b for hypernetwork. The transformed Q-value output of the agent network is then expressed as:
Regulator 1
Regulator n
id)
(T m!)
idF (r m²))
MLP
MLP
Agent1
Agent
kdp"
(x2, u
(al", u.-:)
Figure 3. Illustration of value decomposition methods with the
action regulators.
for each agent can be expressed
P(a)= Softmax DDDecddrr (s, z, id! ...... id id"
(3)
(6) By introducing an additional hypernetwork layer to perform an affine transformation on the raw Q-values, the action regulator effectively aligns the action spaces of agents, ensuring that their decisions are consistent with their assigned identities.
Qu ut, (df))d)) (i(id)) (Tf, ui) + b(idg).
Each identity idx E (1, n is discrete and must be unique to distinguish agents. Enforcing identity uniqueness is a critical constraint, formally expressed as follows:
Vi * id' + id..
To satisfy this constraint, the agent identity decoder em- ploys dynamic masking mechanism M, which ensures that newly generated identities do not duplicate previously assigned ones. The mask M is defined as:
To enable seamless integration of the action regulator into existing value decomposition methods for MARI, must satisfy the Individual-Global-e scribed in Equation (1). Precisely, the global optimal ac- tion should align with the individual optimal actions of the agents. We adopt simple yet effective approach to meet this requirement: ensuring that remains positive by tak- ing its absolute value. This guarantees the transformation preserves the consistency between global and individual optimal actions.
(IGM) paradigm, as de-
Mii]- = 0 identity id has already been assigned, otherwise.
During each decoding step, the probability distribution is dynamically adjusted based on Equation (3) through the application of the specified mask, resulting in:
In summary, the action regulator serves as critical com- ponent of the CAID framework, enabling identity-based decision-making and ensuring compatibility with existing multi-agent reinforcement learning methods while maintain- ing theoretical guarantees of the IGM condition.
P(a)= = Softmax Decoder (s, 2, id M).
(4)
To prevent premature convergence to suboptimal identity assignments, identities are sampled rather than greedily selected. Then the identity assignment for each agent a can be formulated as:
4.4. End-to-End Training
The three main components of CAID, along with their in- puts and outputs. have been detailed earlier. The causal Transformer architecture includes a contextual state encoder and an agent identity decoder. The encoder generates a contextual state è based on the environment state and the local observations of all agents z. This contextual state $ (s. é) replaces the ground truth and serves input to the mixing network in value decomposition methods, pro- viding a more comprehensive representation of CMARL tasks compared to the original environment state s. The agent identity decoder sequentially predicts each agent's 's identity in an autoregressive manner. These identities are
id = Draw (P(a)), Va E A.
(5)
This stochastic approach reduces the likelihood of a partic- ular agent consistently dominating a specific identity, pro- moting more balanced identity assignments among agents.
4.3. Action Regulator
After assigning unique identities to agents, the next step to incorporate this identity information into decision- making processes. In multi-agent systems, agents with the
5

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
subsequently used as inputs the action regulator, align- ing agents" action spaces. The entire CAID framework is trained end-to-end, where the three modules are optimized jointly with the reinforcement learning module. The loss function for this process is as follows:
except those introduced specifically by CAID, are kept con- sistent with the original methods. Unless explicitly stated, CAID refers to the variant implemented QMIX. Details on algorithm hyperparameters are provided in Appendix A.
5.1. StarCraft Multi-Agent Challenge
L= (Yhot - Qroc (T, u, 8))²
Based on the renowned real-time strategy game StarCraft II, SMAC (Samvelyan et al., 2019) is among the most widely used platforms for multi-agent micro-management exper- iments. It contains a variety of suni-scenarios, including homogeneous, heterogeneous, symmetric, and asymmet- rie problem types. Each task requires controlling a team of agents to defeat Al team controlled by heuristic strategies. Agents can only access information within their local obser- vation range and share global reward function. However, recent studies (Ellis et al., 2023) have shown that agents in SMAC can achieve near-optimal policies by relying solely on the current timestep and fixed agent IDs, completely disregarding their local observations. This finding high- lights critical limitation of SMAC: a lack of stochasticity, which is essential in real-world problems. To address this limitation, SMACv2 was introduced, where agents' posi- tions and types can change dynamically in each episode, posing new challenges for existing MARL algorithms. We conducted extensive experiments in this classic CMARL environment to evaluate various algorithms, and scenario details are provided in Appendix B.1.
where ytor is the target joint value function, computed as Ytot = r + y maxu Qtor uu', 8'). All parameters within CAID are optimized by minimizing the temporal-differenee (TD) error. A significant challenge arises because each agent's identity is sampled, preventing direct gradient back- propagation to the agent identity decoder. To address this, we utilize Straight-Thne Throoghh Gradients (Bengio et al., 2013), a technique easily implemented using automatic dif- ferentiation. The agent identity in Equation (5) is modified as follows:
id = Draw (P(a))+P(a)) (a((aa)SStoppie StopGraa (P(a)), Ya E A,
(8) where the StopGrad function prevents gradient flow through its input during backpropagation. This approach ensures that agent identities are sampled to avoid premature convergence to suboptimal solutions while still allowing gradients to propagate back to optimize the agent identity decoder.
The overall CAII) framework is illustrated in Figure 3. By integrating the contextual state á and agent identities into value decomposition methods, CAID provides a more ex- pressive representation of CMARL tasks. This design en- ables CAID to be incorporated into existing value decompo- sition methods, improving its capability to handle complex multi-agent scenarios.
The experimental results in SMACv2 are presented in Fig- ure 4. Across 12 tested scenarios, the inherent stochastic- ity of the environment occasionally led to scenarios where achieving victory was entirely infeasible, preventing conver- gence to near sear-perfect win rates, Our primary observation is that CAID outperforms the baseline QMIX algorithm across all tasks. Furthermore, CAID outperforms other al- gorithms in most scenarios, particularly in the Terrain and Protoss series scenarios. In the Zerg series scenarios, the results derived from integrating CAID into QMIX were constrained by QMIX's suboptimal credit assignment capa- bilities. Nonetheless, CAID consistently demonstrates the fastest initial convergence rate among all tested algorithms.
5. Experiments
In this section, we evaluate CAID in three well-known CMARL environments: StarCraft Multi-Agent Challenge (SMACv2) (Ellis et al., 2023), Vectorized Multi-Agent Sim- ulator (VMAS) (Bettini et al., 2022) and Traffic Signal Con- trol (PyTSC) (Bokade & Jin, 2024). The tasks in these domains exhibit considerable variability across episodes, primarily in the agents' positions, agent types, and target lo- cations. First, the performance of CAID is evaluated through a comparison with several classical algorithms, including Weighted QMIX (Rashid et al.. 2020), QPLEX (Wang et al., 2021a), and the baseline QMIX (Rashid et al., 2018), along with recently proposed methods such as RIIT (Hu et al., 2021), COLA (Xu et al., 2023), and VMIX (Su et al., 2021). Then the contributions of individual modules within the framework are discussed. To ensure the reliability of the results, each experiment is repeated five times with different random seeds. For a fair comparison, all hyperparameters,
5.2. Vectorized Multi-Agent Simulator
In addition to the StarCraft II game environment, we seek additional effective CMARL environments. 'he Vector- ized Multi-Agent Simulator (VMAS) is vectorized and fully differentiable simulator designed for efficient MARL benchmarking. It features a high-performance,, vectorized 2D physics engine and diverse suite of challenging multi- robot scenarios. The modular and user-friendly design of the simulator facilitates the creation of custom scenarios, encouraging contributions from the research community. Some scenarios in VMAS represent tasks such intelligent
6

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
QMIX
CW-QMIX
OW-QMIX
QPLEX
RIIT
VMIX
COLA
CAID (Ours)
terran 5 _
protoss 5 v
zerg 5 _s_
00
Median Win
30
16
Median Test Wii (%5)
40
terran
20
protoss
28
zerg
15
16
10
6
2
Medaan Test Win (95)
terran. 10 vs_ 10
protoss 10
10 vs_10
80
48
to
0
terran 10__s_11 vs_ 11
protoss 10_vs_1 vs
zerg 10 vS__
Medan Tess Win
(ml)
T [mll)
T [mll)
Figure 4. Performance comparison with baselines in different SMACv2 scenarios.
formations and obstacle avoidance for drones uncrewed vehicles, which are of significant practical relevance. Most notably, nearly all tasks can be configured to follow the CMARL setting, where both the initial and goal positions of agents are randomly assigned in each episode. Detailed descriptions of the scenarios employed VMAS can be found in Appendix B.2.
5.3. Ablation Study
An ablation study was carried out on CAID to evaluate its performance in Traffic Signal Control (TSC) scenar- ios. PyTSC (Bokade & Jin, 2024) is a development en- vironment tailored for research in traffic signal control, de- signed to facilitate the rapid development of reinforcement learning- based solutions. In PyTSC, the tasks emulate ur- ban traffic management by controlling traffic lights to opti- mize vehicle movement through intersections. The project currently integrates with SUMO (López et al., 2018) and CityFlow (Zhang al., 2019) simulation backends, offer- ing researchers practical tools to utilize open-source traffic signal control datasets. Notably, it supports the CMARL setting, where the contexts (c.g., traffic flow) in each episode are dynamically randomized. Details regarding the environ- ment are provided in Appendix B.3. Three variants of CAID were designed, each modifying single component of the original algorithm. To investigate the importance of the contextual state encoder, we proposed CAID w/o CS, where the contextual state & is excluded from the mixing network in the base algorithm. For the agent identity decoder, we introduced CAID w/o AI, in which the agent network does not utilize the newly assigned identity information. Lastly, in CAID w/o ST, new identities for agents are generated using a greedy algorithm instead of
We used VMAS compare the performance of three clas- sic value decomposition MARL algorithms, VDN (Sunehag et al., 2018), QMIX, and QPLEX, along with their CAID variants. The learning curves for the CAID variants and their original counterparts are illustrated in Figure 5. It is evident that, in most cases, the CAID variants outperform the original baseline algorithms. This comparison among the three distinct value decomposition methods and their respective variants demonstrates that CAID is applicable to existing value decomposition approaches and can improve their performance. Furthermore, as a distinct value decom- position method, VDN lacks a mixing network module, implying that the contextual state inferred within CAID- VDN is not used. Consequently, the performance improve- ments observed in CAID-VDN are entirely attributable to the reidentification of agents. This finding underscores the contribution of the agent identity decoder. In the following sections, we systematically evaluate the role and importance of each component in the CAID framework.
7

---

Reidentify; Context-Aware Identity Generation for Contextual MARL
CAID (Ours)
Origin
VON
QMIX
QPLEX
Mecian Retur
25
25
Balance
0.00
0.25
0.50
0.78
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.76
1.00
Median Retum
0.5
Dispersion
-0.5
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Median Return
28
25
25
Passage
0
-28
-25
0.00
0.25
0.50
0.76
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Median Ratum
10
10
Reverse
Transport
0.00
0.25
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.25
0.50
0.75
1.00
Mediann Return
-4.25
-3.5
-4.0
Wheel
+4.0
-4.5
-4.75
0.25
1.00
0.75
0.25
T (1106)
(1e6)
T (166)
Figure Comparison of our approach against baseline algorithms on Vectorized Multi-Agent Simulator.
Table 1. Mean of metrics for various algorithms across scenarios in PyTSC.
Jinan
Hangzhou
New York
Algos
Queue
Delay
Travel Time
Queue
Delay
Travel Time
Queue
Delay
Travel Time
CAID
463.39
0.471
329.28
146.50
0.651
336.52
479.48
0.901
369.05
w/o CS
476.52
0.473
326.86
155.79
0,653
340.00
488.39
0.908
372.13
w/o AI
663.59
0.498
364.05
171.37
0.659
346.77
480.41
0.902
368.38
w/o ST
490.70
0.482
332.34
150.11
0.662
330.65
487.28
0.903
371.38
QMIX
563.13
0.507
338.15
170.58
0.660
349.25
496.08
0.902
381.57
sampling, potentially leading to suboptimal solutions. The ablation experiment results are illustrated in Table The best performances of the above algorithms are bold.
5.4. Visualization
We use a dataset consisting of 20,000 timesteps. The left side the figure shows the 2D t-SNE embedding of the contextual states produced by the Contextual State Encoder, while the right side presents the 2D t-SNE embedding of the raw states. Each point corresponds to state, and the color of each point encodes specific permutation of agent identities in that state. Specifically, we enumerate all permutations of the agent indices (0, 1, 2, 3, 4), resulting in 120 possible arrangements. Each permutation mapped to unique integer in [0, 119) using bijective func- tion based on Cantor expansion (Lehmer code), enabling to assign distinct color to each identity configuration. For comparison, traditional methods typically assign agent identities using heuristic rules, leading to only single
We selected three key performance metries that provide more direct assessment of performance: queue length, delay, and travel time. The results demonstrate that all three ablated variants degrade the performance of CAID. Among them, CAID w/o CS shows the least performance drop, followed by CAID w/o ST. These findings highlight the importance of the identities produced by the agent identity decoder in enabling agents adapt effectively dynamic environments. In summary, the designs of the contextual state encoder and agent identity decoder are critical to the effectiveness of CAID.
8

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
fixed permutation- by a single color in the figure. Notably, the middle of the figure highlights pair of mirror-symmetrii and semantically similar states. In the raw state embedding (right), these states appear distant from each other, whereas the contextual state embedding (left), they are close together and share similar identity patterns across agents.
consequences of our work, none which we feel must be specifically highlighted here.
References
Bengio, Y., Léonard, N., and Courville, A. C. Estimating propagating gradients through stochastic neurons for conditional computation. CoRR, abs/1308.3433 2013.
Benjamins, C., Eimer, T., Schubert, F., Mohan, A., Döhler, S., Biedenkapp, A., Rosenhahn, B., Hutter, F,, and Lin- dauer, M. Contextualize me- the case for context in reinforcement learning. Trans. Mach. Learn. Res., 2023, 2023.
Bettini, M., Kortvelesy, R., Blumenkamp, J., and Prorok, A. VMAS: A vectorized multi-agent simulator for collective robot learning. In Bourgeois, J., Paik. J., Piranda, B., Werfel, J., Hauert, S., Pierson, A., Hamann, H., Lam, T. L., Matsuno, F., Mehr, N., and Makhoul, A. (eds.), Distributed Autonomous Robotic Systems 16th Inter- national Symposium, DARS 2022, Montbéliard, France, 28-30 November 2022, volume 28 of Springer Proceed- ings in Advanced Robotics, pp. 42-56. Springer, 2022. doi: 00.1007/978-3-03--544999
Contectual State
SMACv2 Snapshots
Original State
Figure 6. The analysis and visualization of the agent identitics in the terrain 5_5 scenario.
6. Conclusion
In this paper, we introduce the Context-Aware Identity Gen- eration (CAID) framework, which leverages the global state and local observations from all agents to construct contex- tual states and dynamically assign agent identities. This de- sign enables agents adapt effectively different task vari- ations. The CAID framework can be fully trained in an end- to-end manner using the reinforcement learning paradigm. It is compatible with existing MARL value decomposition al- gorithms, significantly improving sample efficiency. The ex- tensive experiments in diverse contextual MARL scenarios demonstrate compelling performance, providing strong evi- dence for the practicality of the proposed framework. We be- lieve this work represents meaningful step toward bridging the gap between multi-agen reinforcement learning and real- world applications. Promising directions for future research include enhancing the robustness of CAID in contextual MARL tasks with dynamic agent populations and extending its applicability to zero-shot generalization scenarios. Fur- thermore, identifying more informative representation of contextual MARL, similar to the agent identities generated by CAID, would be an intriguing avenue for exploration.
4.
Bokade, R. and Jin, X. Pytse: A unified platform for multi- agent reinforcement learning in traffic signal control. CoRR, abs/2410.18202 2024. doi: 10.44850/ARRIV. 2410.18202.
Chung. J., Gülçehree Ç., Cho, K., and Bengio, Y. Empirical evaluation of gated recurrent neural networks sequence modeling. CoRR, abs/14 abs/1442.3355,, 2014.
Cobbe, K., Hesse, C., Hilton, J., and Schulman, J. Lever- aging procedural generation to benchmark reinforcement learning. Proceedings of the 37th International Confer- ence on Machine Leurning, ICML. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp.
PMLR, 2020.
Duan, Y., Schulman, J., Chen, X., Bartlett, P. L., Sutskever, I., and Abbeel, P. RI$-2S: Fast reinforcement learning via slow reinforcement learning. CoRR, abs/ abs/161,,02279, 2016.
Acknowledgements
This work was supported in part by the National Natural Science Foundation of China under Grant 62302189 and in part by the Fundamental Research Funds for the Central Universities under Grant 2662022XXQD002.
Ellis, B., Cook, J., Moalla, S., Samvelyan, M., Sun, M., Mahajan, A., Foerster, J. N., and Whiteson, S. Smacv2: An improved benchmark for cooperative multi-agent rein- forcement learning. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023. NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
Finn, C., Abbeel, P., and Levine. S. Model-agnostic meta- learning for fast adaptation of deep networks. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th Inter- national Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1126- 1135. PMLR, 2017a.
M., Bousquet, O., and Gelly, S. Google research football:
A novel reinforcement learning environment. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12. 2020, pp. 4501-4510 AAAI Press, 2020. doi: 10.1609/AAAI 3311044587..
Finn, C., Abbeel, P., and Levine. S. Model-agnostic meta- learning for fast adaptation of deep networks. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th Inter- national Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1126- 1135. PMLR, 2017b.
Laskin, M., Wang, L.,, Oh, J., Parisotto, E., Spencer, S., Steigerwald, R., Strouse, D., Hansen, S. S., Filos, A.. Brooks, E. A., Gazeau, M., Sahni, H., Singh. S., and Mnih, In-context reinforcement learning with algo- rithm distillation. In The Eleventh International Confer- ence on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Ge, X., Han, Q., and Zhang, X. Achieving cluster formation of multi-agent systems under aperiodic sampling and communication delays. IEEE Trans. Ind. Electron., 65 (4):3417-3426, 2018. doi: 10. 10.1109/1E..2017.7551111
Lee, K. Nachum, O., Yang, M., Lee, I.., Freeman, D., Guadarrama, S., Fischer, I., Xu, W., Jang, E., Michalewski, H., and Mordatch, I. Multi-game deci- sion transformers. In Koyejo, S.. Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Ha, D., Dai, A. M., and Le. Q. V. Hypernetworss. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Confer- ence Track Proceedings. OpenReview.net, 2017.
Hallak, A., Castro, D. D., and Mannor, S. Contextual markov decision processes. CoRR, abs/1502.02259,, 2015.
Liu, B., Liu, Q.. Stone, P., Garg, A., Zhu, Y., and Anandku- mar, A. Coach- player multi multi-agent reinforcement learning for dynamic team composition. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Confer- ence on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 6860-6870. PMLR, 2021.
Hu, J., Jiang, S., Harding, S. A., Wu, H., and Liao, S.-w Rethinking the implementation tricks and monotonicity constraint in cooperative multi-agent reinforcement learn- ing. arXiv preprint arXiv:2102.03479 2021.
Iqbal, S., de Witt, C. A. S., Peng, B., Boehmer, W., White- son, S., and Sha, F. Randomized entity-wise factorization for multi-agent reinforcement learning. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 1596-4606 PMLR, 2021.
López, P. Á., Behrisch, M., Bieker- Walz, L.,, Erdmann, J., Flötteröd, Y., Hilbrich, R., Lücken, L.,, Rummel, J., Wag- ner, P., and WieBner, E. Microscopic traffic simulation using SUMO. In Zhang, W., Bayen, A. M., Medina, J. J. S., and Barth, M. J. (eds.), 21st International Confer- ence on Intelligent Trunsportation Systems, ITSC 2018, Maui, HI, USA, November 4-7, 2018, pp. 2575-2582 IEEE, 2018. doi: N0.1109/TSC.22088855600
Kirk, R., Zhang, A., Grefenstette, E., and Rocktäschel, T. A survey of zero-shot generalisation in deep reinforcement learning. J. Artif. Intell. Res., 76:201-264, 2023, doi: 00.1613/JAIR.1.41117
Lowe, R., Wu, Y., Tamar, A., Harb, J., Abbeel, P., and Mor- datch, I. Multi-agent actor-critic for mixed cooperative- competitive environments. In Guyon, I., von Luxburg, U. Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural In- formation Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 6379-6390 2017.
Kirsch, L., van Steenkiste, S., and Schmidhuber, Im- proving generalization in meta reinforcement learning using learned objectives. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Kurach, K., Raichuk, A., Stanczyk, P., Zajac, M., Bachem, O., Espeholt, I.., Riquelme, C., Vincent, D., Michalski,
Mao, H., Liu, W., Hao, J., Luo, J., Li, D., Zhang, Z., Wang, J., and Xiao, Z. Neighborhood cognition con-
10

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
sistent multi-agent reinforcement learning. In The Thirty- Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Ar- rificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 7219-7226. AAAI Press, 2020. doi: W00.://AAA1334100.0.0
Information Processing Systems 2020, NeurIPS 2020, De- cember 6-12, 2020, virtual, 2020.
Reed, S. E., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A., Barth-Maron, G., Gimenez, M., Sulsky, Y., Kay, J., Springenberg, J. T., Eccles, T., Bruce, J., Razavi, A., Edwards, A., Heess, N., Chen, Y., Hadsell, R., Vinyals, O., Bordbar, M., and de Freitas, N. A generalist agent. Trans. Mach. Learn. Res., 2022, 2022.
McClellan, J., Haghani, N., Winder, J., Huang, F., and Tokekar, P. Boosting sample efficiency and generalization in multi-agent reinforcement learning via equivariance. CoRR, abs/2410.02581, 2024. doi: 10.48550/ARXIV. 2410.02581.
Samvelyan, M., Rashid, T., de Witt, C. S., Farquhar, G., Nardelli, N., Rudner, T. G. J., Hung, C., Torr, P. H. S., Foerster, J. N., and Whiteson, S. The starcraft multi-agent challenge. In Elkind, E., Veloso, M., Agmon, N., and Taylor, M. E. (eds.), Proceedings of the 18th Interna- tional Conference on Autonomous Agents and MultiAgent Systems, AAMAS '19, Monireal, QC, Canada, May 13- 17, 2019, pp. 2186-2188 International Foundation for Autonomous Agents and Multiagent Systems, 2019.
Nagabandi, A., Clavera, I., Liu, S., Fearing, R. S., Abbeel, P., Levine, S., and Finn, C. Learning to adapt in dynamic, real-world environments through meta-reinforcement learning. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Sodhani, S., Zhang, A., and Pineau, J. Multi-task rein- forcement learning with context-based representations. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 202 I, 18-24 July 2021, Virtual Event, volume 139 of Pro- ceedings of Machine Learning Research, pp. 9767-9779. PMLR, 2021.
Nichol, A. and Schulman, J. Reptile: a scalable metalearn- ing algorithm. arXiv preprint arXiv: 1803009999 2(3):4, 2018.
Nichol, A., Achiam, J., and Schulman, J. On first-order meta-learning algorithms. CoRR, abs/1803.02999 0.099, 2018.
Son, K., Kim, D., Kang. W. J., Hostallero, D., and Yi, Y. QTRAN: learning to factorize with transformation for cooperative multi-agent reinforcement learning. In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceedings of the 36th International Conference on Machine Learn- ing, ICM1,2019, 9-15 June 2019, Long Beach, California, USA volume 97 of Proceedings of Machine Learning Research, pp. 5887-5896. PMLR, 2019.
Oliehoek, F. and Amato, C. A concise introduction to decen- tralized pomdps. In SpringerBriefs in Intelligent Systems, 2016.
OpenAI. GPT-4 technical report. CoRR, abs/2303.08774, 2023. doi: 10.48550/ARXIV.23035
Rajan, R., Diaz, L. B., Guttikonda, S., Ferreira, F., Biedenkapp, A., von Hartz, J. O., and Hutter, F. MDP playground: An analysis and debug testbed for reinforce- ment learning. J. Artif. Intell. Res., 77:821-890, 2023. doi: 10.1613/JAIR.II 1.1314.
Su, J., Adams, S. C., and Beling, P. A. cole--dccoppositio multi- agent actor-critics. In Thirty-Fifth AAAI Confer- ence on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial In- telligence, IAAI 2021, The Eleventh Symposium on Edu- cutional Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, pp. 1 11352-11360. AAAI Press, 2021. doi: 10..6/9/AAAL..VSS..........
Rashid, T., Samvelyan, M., de Witt, C. S., Farquhar, G., Foerster, J. N., and Whiteson, S. QMIX: monotonic value function factorisation for deep multi-agent reinforcement learning. In Dy, J. G. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learn- ing. ICML 2018, Stockholmsmüssan. Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 4292-4301. PMLR, 2018
Sunchag, P., Lever, G., Gruslys, A., Czarnecki, W. M., Zambaldi, V. F., Jaderberg, M., Lanctot, M., Sonnerat, N., Leibo, J. Z., Tuyls, K., and Graepel, T. Value- decomposition networks for cooperative multi-agent learning based on team reward. In André, E., Koenig, S., Dastani, M., and Sukthankar, G. (eds.), Proceedings of the 17th International Conference on Autonomous Agents and MultiAgena Systems, AAMAS 2018, Stockholm, Swe- den, July 10-15, 2018, pp. 2085-2088 International Foun- dation for Autonomous Agents and Multiagent Systems Richland, SC, USA/ / ACM, 2018.
Rashid, T. Farquhar, G., Peng, B., and Whiteson, S. Weighted QMIX: expanding monotonic value function factorisation for deep multi-agent reinforcement learn- ing. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M.. and Lin, H. (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
Sutskever, I., Vinyals, O., and Le, Q. V. Sequence se- quence learning with neural networks. In Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N. D., and Wein- berger, K. Q. (eds.), Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada, pp. 3104-3112, 2014.
Wang, T., Gupta, T., Mahaaann, A., Peng, B., Whiteson, S., and Zhang, C. RODE: learning roles decompose multi- agent tasks. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021b.
Wu, Z.,, Pan, S., Chen, F., Long, G., Zhang, C.,, and Yu, P. S. A comprehensive survey on graph neural networks. IEEE Trans. Neural Networks Learn. Syst., 32(1):4-244 2021. doi: 10.1 (0000//NNNS.2000022777
Tao, Y., Genc, S., Chung. J., Sun, T., and Mallya, S. RE- PAINT: knowledge transfer in deep reinforcement learn- ing. In Meila, M. and Zhang. T. (eds.), Proceedings of the 38th International Conference on Machine Learn- ing, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 10151-10152. PMLR, 2021.
Xu, Z., Zhang, B., Li, D., Zhang, Z., Zhou, G., Chen, H., and Fan, G. Consensus learning for cooperative multi- agent reinforcement learning. In Williams, B., Chen, Y., and Neville, J. (eds.), Thirty-Seventh AAAI Conference on Artificial Intelligence, AAAI 2023. Thirty-Fifth Confer- ence on Innovative Applications of Artificial Intelligence, IAAI 2023, Thirteenth Symposium on Educational Ad- vances in Artificial Intelligence, EAAI 2023, Washington, DC, USA, February 7-14, 2023, pp. 11726-11734, AAAI Press, 2023. doi: 10. 10.1609/AAAII /TT/:/AAA..333110266....
Taylor, M. E. and Stone, P. Transfer learning for reinforce- ment learning domains: A survey. J. Mach. Learn. Res., 10:1633-1685, 2009. doi: 10.5555/157700911553333
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.. Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E,, and Lample, G. Llama: Open and efficient foundation language models. CoRR, abs/2302.13971, 2023. doi: 10.48500ARXXIV.222.13377
Yan, Z. and Xu, Y. A multi-agent deep reinforcement learn- ing method for cooperative load frequency control of a multi-area power system. IEEE Transactions on Power Systems, 45(6):459944600) 2020.
Yang, Y., Hao, J., Liao, B., Shao, K., Chen, G., Liu, W., and Tang, H. Qatten: A general framework for cooperative multiagent reinforcement learning. CoRR, abs/2002.03939, 2020.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Guyon, I., von I Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-6008, 2017.
Yu, T., Quillen, D.,, He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. Meta-world: A benchmark and eval- uation for multi-task and meta reinforcement learning. In Kaelbling, L. P., Kragic, D., and Sugiura, K. (eds.), 3rd Annual Conference on Robot Learning, 2011.2019, Osaka, Japan, October 30 - November 1, 2019, Proceed- ings, volume 100 of Proceedings of Machine Learning Research, pp. 1094-1100. PMIR, 2019.
Vinyals, O., Fortunato, M., and Jaitly, N. Pointer networks. In Cortes, C., Lawrence, N. D., Lee, D. D., Sugiyama, M., and Garneut R. (eds.), Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp. 2692-2790000 2015.
Zhang, H., Feng, S., Liu, C., Ding, Y., Zhu, Y., Zhou, Z., Zhang, W., Yu, Y., Jin, H., and Li, Z. Cityflow: A multi- agent reinforcement learning environment for large scale city traffic scenario. In Liu, L.., White, R. W., Mantrach, A., Silvestri, F., McAuley, J. J., Yates, R., and Zia, L. (eds.), The World Wide Web Conference, WWW 2019, San Francisco, CA, USA. May 13-17, 2019, pp. 3620- 3624. ACM, 2019. doi: 10.
Wang, J., Ren, Z., Liu, T., Yu, Y., and Zhang, C. QPLEX: duplex dueling multi-agent q-learning In 9th Interna- tional Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenRe- view.net, 2021a
Zhao, W., Queralta, J. P.. and Westerlund, T. Sim-to-real transfer in deep reinforcement learning for robotics: a survey. In 2020 IEEE Symposium Series Compu- tational Intelligence, SSCI 2020, Canberra, Australia, December 1-4, 2020, pp. 737-744. IEEE, 2020. doi: 1LL1099/SSC44888.22000.0088
Wang, T., Dong, H., Lesser, V. R., and Zhang, C. ROMA: multi-agent reinforcement learning with emergent roles. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 77766686..111 PMLR, 2020.
12

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
A. Implementation Details
A.I. Algorithmic Description
The pseudo-code of CAID shown in Algorithm
Algorithm 1 Context-Aware Identity Generation (CAID)
1: for each episode do
2: Get the global state 81 and the local observations 21 {z] =f, ,., =ft of agents for 1 to T do
4: for a 1 do
5:
Select action uy according to the agent network
6: end for
7: Carry out the joint action u, = {ui...., "
8: Get the global reward 7t+1, the next local observations 2(+1, and the next state $1+1
9:
end for
Store the trajectory in the replay buffer D. Sample batch of episodes B Uniform(D).
12:
Compute the context variable è using Equation (2).
13:
Generate agent identities based on Equation (8).
14:
Evaluate the transformed Q-values for all agents using Equation (6).
15:
Update the parameters of the CAID model using Equation (7).
16:
Periodically update the parameters of the target network.
end for
A.2. Hyperparameters
Unless specified otherwise, the hyperparameter configurations across different environments are presented in Table 2. These settings are identical to those provided in PyMARL2 (Hu et al., 2021). All experiments in this study were conducted using NVIDIA GeForce RTX 2080 Ti graphics cards and Intel(R) Xeon(R) Silver 4114 CPUs. For all methods, exploration during training is achieved via independent e-greedy action selection, with linearly annealed from 1.0 to 0.05 over 50,000 steps. In SMACv2, training ends after 5 million timesteps, whereas concludes after million timesteps in VMAS and million timesteps in PyTSC
Table Hyperparameter settings.
Description
Value
Learning rate
0.001
Type of optimizer
Adam
How many episodes to update target networks
200
Reduce global norm of gradients
10
Batch size
128
Capacity of replay buffer
5000
Batch size for parallel execution
8
Discount factor
0.99
B. Introduction for Environments
In this paper, we employed three experimental environments: SMACv2, VMAS, and PyTSC. Given the limited research on Contextual MARL and the recent introduction some new environments, is essential to provide detailed overview of these environments tailored to the CMARL setting. The following sections aim familiarize readers with the objectives of these tasks and the corresponding evaluation metrics, offering foundation for understanding this emerging field.
13

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
(a) SMACv2
(h) VMAS
(c) PyTSC
Figure 7. Screenshots of the two experimental platforms used in this paper.
B.1. SMACv2
SMACv2 (StarCraft Multi-Agen Challenge version 2) an enhanced benchmark environment designed to assess the effectiveness of cooperative MARL algorithms. Compared to its predecessor, SMACv2 incorporates procedurally generated scenarios. requiring agents to adapt their collaborative strategies to dynamically changing environments, including diverse enemy configurations and map layouts. This design intensifies the challenges posed by partial observability, compelling agents to base their decisions real-time sensory input rather than relying scripted behaviors. In each episode, the positions and types of both ally and enemy agents are subject to change.
The primary focus SMACv2 is on evaluating the agents' generalization capabilities in previously unscen scenarios. This is achieved through metrics such as win rates, collaboration efficiency. and adaptability in novel environments. By introducing more demanding scenarios and imposing stricter evaluation standards, SMACv2 establishes rigorous platform for advancing research multi-agent reinforcement learning algorithms. It fosters algorithmic development in complex, dynamic contexts and ensures robustness and reliability in real-world applications. Descriptions of the scenarios are provided in Table 3.
Table 3. Maps in different scenarios.
Race
Unit Types
Probability of Generation
Scenarios
terran_5_vs.
Marine
0.45
terran 5_v.___....
Terran
Marauder
0.45
terran_ 1at__10v_____l _v_
Medivac
0.1
terran. 10_vs_111
protoss_5 5 v.vs..
Stalker
0.45
protoss_5 5 5_v
Protoss
Zealot
0.45
motooss_O_vs...
Colossus
0.1
protoss 10.vs. v0.vv... 11
zerg 5
Zergling
0.45
zerg_5_vs_6
Zerg
Baneling
0.1
zerg. 10.vv.1vll. vs_10
Hydralisk
0.45
zerg. 10.vo....... vs_11
B.2. VMAS
VMAS (Vectorized Multi-Agent Simulator) is efficient and scalable simulation framework for MARL, featuring differentiable and vectorized 2D physics engine. VMAS includes comprehensive suite of complex scenarios that encompass both cooperative and competitive tasks. It supports highly customizable features, including sensors, elastic collisions, rotations, joints, and communication mechanisms. Each scenario integrates tailored reward mechanisms and termination conditions to rigorously evaluate the agents' collaborative strategies and division of labor under diverse task objectives.
VMAS offers a range
of bre-b multi-agent reinforcement learning scenarios, addressing both cooperative and adversarial
14

---

Reidentify; Context-Aware Identity Generation for Contextual MARL
Balance
(b) Dispersion
(c) Passage
(d) Reverse Transport
(e) Wheel
Figure 8. Illustration of benchmark tasks in VMAS.
tasks. Descriptions of representative scenarios are as follows:
Balance. Agents collaborate balance spherical object on a linear platform while transporting a designated target. Any fall of the object platform results substantial penalty, Successful completion requires coordinated task allocation to maintain equilibrium while steadily advancing toward the goal.
Dispersion. Starting from common origin, agents must efficiently disperse to distinct target points. This scenario highlights the importance of decentralized coordination, minimizing clustering, and optimizing task completion.
Passage. A group of robots must traverse a wall with a narrow passage to reach their assigned destinations. The challenge involves avoiding collisions and orchestrating the use of the passage in sequential and cooperative manner, testing path-planning capabilities and localized teamwork.
Reverse Transport. Agents cooperate to transport high-mass package to target location. Given the package's significant weight, individual agents cannot accomplish the task alone. Success demands precise synchronization of force direction and magnitude among all agents.
Wheel. Agents encircle a rotating beam anchored at the origin, coordinating their efforts to achieve target angular velocity. The task requires an optimal distribution of forces to precisely control the beam's speed without overexertion, which could cause deviations from the target.
These scenarios designed to rigorously evaluate agents' collaborative efficiency, adaptability, and task performance across diverse objectives and constraints. Notably, all scenarios can be configured to the CMARL settings.
B.3. PYTSC
PyTSC is simulation environment tailored to multi-agent reinforcement learning (MARL) in the domain of traffic signal control (TSC). Its primary goal is to reduce global congestion across the traffic network, quantified by metrics such total delay, average waiting time, or vehicle queue length, through the coordinated control of multiple traffic signal agents. In fully cooperative MARL framework, agents adapt to dynamic traffic patterns and develop efficient control strategies by interacting with both the environment and other agents, ultimately achieving global optimization.
The evaluation of PyTSC focuses on two key dimensions: traffic efficiency and algorithmic performance. Traffic efficiency
15

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
assessed using metrics such as average vehicle waiting time, average travel time, and total network delay, while algorithmic performance measured by factors like convergence speed and policy stability. Additionally, PyTSC allows for testing the adaptability of algorithms under dynamic traffic flow scenarios, such as morning and evening peak periods, providing comprehensive assessment of MARL approaches in TSC applications.
Table 4. Scenarios in PyTSC environments.
Scenarios
Simulators
Network Types
Agent Types
Total Agents
Jinan
CityFlow
Real-world
Homogeneous
12
Hangzhou
CityFlow
Real-world
Homogeneous
16
New York
CityFlow
Real-wer Real-world
Homogeneous
16
C. Additional Results
C.1. Results of Dynamic Role Assignment Methods
ROMA
RODE
CAID (Ours)
5_vs_5
protoss __
zerg_5 vs_
80
Medion vin
30
15
Ierran_
protoss
4.40
28
Median test Wh
30
15
10
terran_ 10, vs_
protoss 10, v0
2erg. 10, _S_10
Median TeatWin
60
46
50
30
20
terran_10 10__v__111
protoss 10_
zerg_ 10_ _0_vs_11
Median Teat Win
(mi])
(mil)
(mll
Figure 9. Performance comparison with dynamic role assignment methods in different SMACv2 scenarios. We selected ROMA and RODE two representative methods role assignment. methods were configured with identical hyperparameters, including the optimizer and techniques such as eligibility traces, to ensure consistency with other algorithms implemented in PYMAR... It is worth noting that RODE fails perform properly in SMACv2 environment. This limitation arises because the original RODE implementation modifies the SMAC environment by sorting enemy units in the environment file. Such changes simplify the original SMAC environment, which undermines its complexity and makes the implementation incompatible with the more challenging and standardized SMACx2, As result, RODE is not applicable to the Contextual MARI setting. CAID achieves superior performance and faster convergence.
16

---

Reidentify: Context-Aware Identity Generation for Contextual MARL
C.2. Additional Ablation Study
CAID
w/o CS
w/o ST
w/o Al
CAID
w/o Al
w/ Fixed ID
w/ Random ID
Median test Win (%) Modian Tegt Win (%)
protoss vs 5
Median Text Win (%) Median Test Win (%)
protoss 5_vs_5
80
8
80
20
20
0
2
3
protoss_ 6_vs_6 6
protoss __vs_6
15
(1e6)
(186)
Figure 10. Left: The learning curves of the three variants discussed in Section 5.3 the SMACv2 benchmark. CAID w/o Al refers the variant where only the Action Regulator component of CAID is removed, while all other components remain unchanged. The results demonstrate that all three variants underperform compared the full CAID model, highlighting the essential role each module plays in the overall performance of CAID. Right: The learning curves of different variants for identity modeling, including ablation and alternative strategies. CAID with Fixed ID denotes configuration where each agent is assigned a static identity based on its index in the agent list. CAID with Random ID refers to setting where agents are assigned randomly generated but episode-consisten identities.
17


Reviews:
[{"review_id": "2h3Haph8sw", "reviewer": "Reviewer_v3VM", "summary": "The authors introduce Context-Aware Identity Generation framework which is able to generalize between tasks in one Contextual MARL domain. CAID integrates dynamically assigned identity information into action decoding for each agent, which is claimed to provide smooth adaptation to varying contexts. Combined with Action Regulator, which uses identities to produce actions from agents, and Contextual State Encoder, which encode MARL interaction history to context sequence, it was shown that CAID outperforms various baselines on classic MARL environments, such as StarCraft SMAC, Vectorized Multi-Agent Simulator and Traffic Signal Control environments. Moreover, authors compare CAID without components and show that each component is important for better performance.", "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 sincerely thank the reviewer for the detailed and constructive comments! We are grateful for your careful reading of both the main paper and the supplementary material. Below, we respond to your concerns point by point:\n\n**Q1**: The paper does not compare with dynamic role assignment methods.\n\n**A1**: Thank you for raising this important point. In the initial submission, we focused on comparisons with RIIT, COLA, and VMIX, as these methods are widely adopted and perform strongly on the newly CMARL benchmarks such as SMACv2. However, we fully acknowledge the significance of role-based methods like ROMA and RODE in learning strategies. To address your concern, we have included experimental comparisons with ROMA and RODE (See the figure in https://anonymous.4open.science/r/CAID-7A6C/Role_methods.jpg). Our preliminary results indicate that CAID achieves superior performance and faster convergence, especially in environments with contextual variations such as agent permutations, type switches, and initial state shifts.\n\n**Q2**: The relevance of Meta-RL references is unclear.\n\n**A2**: We appreciate your thoughtful observation. We originally referenced Meta-RL to provide context for generalization strategies in reinforcement learning. However, we agree that CAID does not follow the standard meta-learning paradigm. Rather, it addresses the generalization challenge by generating context-aware agent identities that unify behaviors across task variants. To improve clarity and focus, we will revise the Related Work section to remove or rephrase the Meta-RL discussion, better highlighting CAID's distinct contribution in dynamic identity generation.\n\n**Q3**: Does CAID support multi-task adaptation within the same MARL domain?\n\n**A3**: Thank you for this question. As introduced in Section 3.3 of our paper, CAID is specifically designed for the *Contextual Multi-Agent Reinforcement Learning (CMARL)* setting. While CMARL may involve variations across episodes—such as changes in agent initial positions, or types—it differs fundamentally from traditional multi-task reinforcement learning. In CMARL, the core assumption is that tasks share a common structure and can be semantically aligned through a latent contextual representation. The objective is not to learn separate policies for separate tasks, but rather to enable policy reuse by reasoning over contextual variations. Thus, CAID is built to tackle *context shift* rather than *task switch*. We appreciate the reviewer for raising this important point and will clarify the distinction between CMARL and conventional multi-task RL more explicitly in the revised paper.\n\n\nWe really appreciate your comments and they really help us improve our paper! And we also appreciate it if you have any further comments.", "claims_and_evidence": "The main claims are supported with experiments: it is clear that CAID outperforms various methods, and that all three components of CAID are important for CAID performance. But it wasn’t shown that CAID roles dynamic assignment is more optimal than role assignment from methods introduced in Related Work (ROMA, RODE, COPA). Moreover, authors compare CAID with only RIIT, COLA and VMIX methods, which don’t use role assignments for more effectiveness: COLA uses consensus builder, which utilizes DINO to learn labels, which does not imply any identification learning.", "methods_and_evaluation_criteria": "Authors compare CAID with lots of baselines on a large number of datasets, making comparison fair with calculating experiments with 5 seeds, and setting up hyperparameters according to original methods. Evaluations are fair and make sense for the problem.\n\nBut I also wrote my concern about evaluation methods in Claims and Evidence, it seems better to compare CAID with ROMA, RODE or COPA, to make effectiveness of dynamic role assessment more clear", "theoretical_claims": "Paper does not provide complex theoretical claims, and those that are provided are correct. Authors provide detailed description of their method, so I don’t have questions about CAID methodology.", "experimental_designs_or_analyses": "As noted in Methods and Evaluation criteria paragraph, all experiments designs sounds valid. Authors provide detailed description of hyperparameters analysis in appendix A. On the other hand, there is some misunderstandings on analysis of experiments, because ablation study shows only importance of identity decoder compared to CAID without identity decoder, so it does not show in full effectiveness of dynamical identification learning in comparison with previous methods.", "relation_to_broader_scientific_literature": "The paper is related to MARL literature and propose novel and intriguing architecture for managing MARL agents roles compared to previous architectures.", "essential_references_not_discussed": "The paper covers similar literature well, however, I did not quite understand the part about the meta-RL reference. In my opinion, paper in general discusses agent’s dynamical identity management, and it is the main feature of CAID algorithm, so I think the part about meta-RL is redundant in context of CAID algorithm.", "other_strengths_and_weaknesses": "**Strengths:**\n\n- This paper is well-written, and easy to follow\n- This paper provides good experiment setup and different baselines co compare CAID with\n- This paper provides intriguing CAID architecture with interesting insights of identity use in MARL agents\n- This paper provides detailed experiments about importance of each part of CAID\n\n**Weaknesses:**\n\n- This paper does not provide experiments with baselines uses different strategies of identity management (ROMA, RODE, COPA), and does not show effectiveness of proposed identity management in comparison to previously proposed.", "other_comments_or_suggestions": "None.", "questions_for_authors": "In Related Work, authors also have a paragraph of discussing Meta Reinforcement Learning. I suggest that authors claim that CAID can adapt to similar tasks in one MARL domain, but authors did not highlight it clearly in abstract and experiments. Do the authors think that CAID is able to adapt to tasks in one MARL domain, and maybe authors can confirm it with experiments? I thought it was wrong to deduct points for this claim, as it was not highlighted clearly in abstract, but if authors also propose it, it would be better to provide experiments for multitask adaptation of CAID.", "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": "uWFlgdIf78", "reviewer": "Reviewer_PyYp", "summary": "In multi-agent reinforcement learning (MARL), generalization poses a significant challenge. Existing MARL methods exhibit vulnerability when confronted with even slight variations in task settings, requiring the retraining of policies for each task variant. This paper introduces a Context-Aware Identity Generation (CAID) framework, which utilizes global states and local observations from all agents to construct contextual states and dynamically assign agent identities.", "questions": "", "limitations": "", "rating": 1, "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 the thoughtful and detailed comments. Below we respond to the key points raised in the \"Other Strengths and Weaknesses\" section:\n\n**Q1**: How is policy generalization defined? How does it differ from definitions in [1] [2]?\n\n**A1**: Thank you for this important question. In Section 3.3, we formalize a new generalization framework called *Contextual Multi-Agent Reinforcement Learning (CMARL)*. Unlike traditional multi-task or multi-agent generalization settings (as in [1] [2]), CMARL assumes that different environmental states across episodes can be aligned through latent contextual variables (e.g., rotation, mirroring, permutation). Our goal is to unify semantically equivalent contexts into a compact representation and dynamically assign identities to corresponding agents. This enables policy reuse under *context shift* rather than *task shift*. We will revise the Related Work section to clearly distinguish CMARL from the generalization definitions in [1] [2].\n\n**Q2**: The motivation is unclear and the limitations of prior work are not systematically discussed.\n\n**A2**: We appreciate your constructive feedback. We will clarify our motivation by systematically identifying limitations in existing MARL methods: (1) They typically fail to capture semantic transformations in contextual environments (e.g., agent position mirroring), and (2) they lack dynamic identity alignment, making it difficult to reuse policies across variants. We will make explicit that CAID addresses these issues by generating contextual representations and dynamically consistent agent identities for improved generalization.\n\n**Q3**: The meaning of the “draw” operation in Equation (5) is unclear.\n\n**A3**: Thank you for pointing this out. “Draw” denotes sampling from the agent identity distribution computed via the decoder. We apply Straight-Through Gradients to enable gradient backpropagation. We will revise the equation and add detailed explanations in the main text.\n\n**Q4**: Missing comparison with ROMA/RODE; no ablation on Action Regulator.\n\n**A4**: Thank you for these insightful comments. We address them as follows:\n\n① We have added comparisons with ROMA and RODE (See the figure in https://anonymous.4open.science/r/CAID-7A6C/Role_methods.jpg). CAID still outperforms these methods on SMACv2.\n\n② We carried out the ablation variant “CAID w/o AI” (without Action Regulator) on SMACv2. Results (See the left figure in https://anonymous.4open.science/r/CAID-7A6C/Ablation_study.jpg) show that removing this module significantly degrades performance, confirming its necessity for consistent action semantics under identity shift.\n\n**Q5**: Zerg performs worse than other races in Figure 4.\n\n**A5**: Regarding the Zerg performance drop, we believe that Zerg units have higher variance in attack patterns and durability, which makes the value decomposition more sensitive to alignment. As the foundational algorithm of CAID, QMIX exhibits greater performance degradation under these conditions due to its suboptimal credit assignment capability.\n\nThank you again for your valuable suggestions. They have greatly helped us improve the clarity and rigor of our work. And we also appreciate it if you have any further comments.", "claims_and_evidence": "See the section “Other Strengths And Weaknesses”.", "methods_and_evaluation_criteria": "See the section “Other Strengths And Weaknesses”.", "theoretical_claims": "I have checked the correctness of any proofs for theoretical claims.", "experimental_designs_or_analyses": "I have checked the soundness/validity of any experimental designs or analyses. For a more detailed description, please refer to the section \"Other Strengths and Weaknesses,\".", "relation_to_broader_scientific_literature": "No.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths\n1.  The writing is clear and easy to understand.\n2.  The experiments selected a variety of experimental scenarios (SMACV2, VMAS) and diverse algorithms (including zero-shot generalization). The diversity of experimental scenarios helps demonstrate the effectiveness of the algorithms.\n\nWeaknesses\n1. How is the generalization of multi-agent policies defined? How does the generalization in this paper differ from the definitions of generalization in related works [1][2]? I recommend that the paper clearly highlight these differences in the Related Works section.\n2. The motivation of the paper is somewhat unclear and requires further elaboration. In particular, the statement in Section 2 (Related Works) that \"However, these approaches are not specifically tailored for Contextual MARL tasks.\" needs more clarification. From my understanding, the aim of the paper is to address the generalization problem by proposing a context-based learning approach for MARL. However, the rationale behind choosing this approach, as well as a detailed discussion on the shortcomings of existing MARL methods in handling generalization, are lacking and need to be presented in a more systematic manner.\n3. What does the 'draw' operation in Equation (5) mean? Additional clarification should be provided here.\n4. The experimental setting of the paper needs to add relevant experiments. ①The proposed method seems to adopt a role-based assignment approach, and it would be useful to compare it with works like ROMA and RODE. ②Furthermore, the Action Regularization method discussed in the paper does not impose any constraints on the action space. I am curious about how this ablation method performs on the SMACV2 benchmark.\n5. Why does Zerg perform worse than Terran and Protoss in Figure 4? I don't understand why this phenomenon occurs, as there shouldn't be a significant difference at the algorithmic level. Simply stating in the paper that 'the results derived from integrating CAID into QMIX were constrained by QMIX's suboptimal credit assignment capabilities' is not sufficient. It would be more convincing to include additional experiments here.\n\nReferences\n\n[1] Mahajan, Anuj, et al. \"Generalization in cooperative multi-agent systems.\" arXiv preprint arXiv:2202.00104 (2022).\n\n[2] Qiu, Wei, et al. \"Rpm: Generalizable multi-agent policies for multi-agent reinforcement learning.\" In ICLR. 2023.", "other_comments_or_suggestions": "No.", "questions_for_authors": "See the section “Other Strengths And Weaknesses”.", "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": "KsTrcZtK5m", "reviewer": "Reviewer_139G", "summary": "This paper introduces a novel approach called Context-Aware Identity Generation (CAID) to improve the generalization ability and sample efficiency of Multi-Agent Reinforcement Learning in contextual environments. CAID leverages a causal Transformer structure to generate dynamic agent identities, while incorporating an action regulation module that embeds identity information into the action-value space.", "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 the constructive and thoughtful feedback. We are encouraged that you found the CAID framework innovative and recognized its potential for improving MARL generalization. Below, we address your concerns in detail:\n\n**Q1**: Lack of ablation or alternative identity modeling comparisons.\n\n**A1**: Thank you for highlighting this. We have added the following ablation experiments:\n\n* CAID w/o AI: Removes the identity modeling entirely—all agents share the same input embedding.\n* CAID w/ Fixed ID: Assigns static identities based on agent indices (e.g., [1, 2, ..., n]).\n* CAID w/ Random ID: Assigns random but fixed identities per episode.\n\nResults (See the right figure in https://anonymous.4open.science/r/CAID-7A6C/Ablation_study.jpg) show significant performance drops across all variants, especially under strong context perturbations. The original CAID consistently maintains performance, validating the importance of contextual identity learning.\n\n**Q2**: No analysis of computational cost due to causal Transformer.\n\n**A2**: Thank you for the suggestion. The causal Transformer employed in our method is lightweight, comprising only **a single Transformer block**. We conducted a comparative analysis of the training times for QMIX and CAID based on the recorded logs. Across all scenarios, the additional forward pass introduced by CAID accounts for **less than 15%**, thereby ensuring real-time performance in real-world deployments. Furthermore, the modules introduced in CAID are exclusively involved during centralized training. In the decentralized execution stage, CAID maintains an execution efficiency comparable to that of conventional algorithms.\n\n**Q3**: No evaluation of method stability under varying conditions.\n\n**A3**: We included robustness analysis across different random seeds, agent initial condition perturbations. Results in our paper show CAID yields acceptable level of performance variance compared to baselines\n\n**Q4**: Identity interpretability and task similarity are unexplored.\n\n**A4**: We agree this is an important direction. We have added t-SNE visualizations of learned identity vectors (See the figure in https://anonymous.4open.science/r/CAID-7A6C/Visualization.jpg), which show clustering of agents with similar roles across varied environments. Additionally, we believe the KL-divergence of contextual encodings can be as a first step toward quantifying task similarity. We plan to explore more sophisticated similarity metrics and cross-task identity propagation in future work.\n\nWe deeply appreciate your constructive feedback and have incorporated your suggestions into our revised submission. Your insights significantly enhance the completeness of our work, and we hope these improvements will help raise your evaluation.", "claims_and_evidence": "Most of the conclusions in the paper are supported by experimental results, but there is a lack of comparative experiments to illustrate how performance would degrade without identity modeling or how different identity generation methods would impact the results.", "methods_and_evaluation_criteria": "Overall, the proposed methods and evaluation criteria are reasonable.", "theoretical_claims": "yes.", "experimental_designs_or_analyses": "The experimental design of the paper is generally reasonable. Benchmark tasks such as SMACv2, VMAS, and PyTSC are selected to evaluate the generalization ability and sample efficiency of CAID.", "relation_to_broader_scientific_literature": "The core contribution of this paper is closely related to the research direction of improving the generalization ability of multi-agent reinforcement learning (MARL), especially in terms of identity modeling and context adaptability.", "essential_references_not_discussed": "1.The UPDeT model enhances performance by utilizing policy decoupling and Transformer architecture, making it advantageous for deployment in tasks with varying numbers of agents.\n2.Multi-Agent Transformer (MAT) treats the MARL problem as a sequential modeling task, demonstrating the potential of Transformer-based architectures in MARL.", "other_strengths_and_weaknesses": "Strengths:\n1. Utilize a causal Transformer to generate dynamic agent identities, rather than relying on fixed or predefined identity representations.\n2. Embed identity information directly into the action-value space through the action regulator module.\n3. Achieve competitive results across multiple multi-agent benchmark tasks.\n\nWeaknesses:\n1.The paper proposes that dynamic agent identity modeling enhances generalization, but it lacks comparative experiments showing how performance would change if identity modeling were removed or if different identity generation methods were used.\n2.The introduction of the causal Transformer structure may lead to higher computational costs, yet the paper does not provide a complexity analysis.\n3. There is no analysis of the stability of the proposed method under different conditions.\n4. The paper does not explore whether the generated agent identities are interpretable, nor does it analyze whether identity representations can be intuitively understood across different environments.\n5. While the paper states that agents in similar tasks can share useful information, it does not clearly explain how task similarity is measured or how identity information is propagated between agents.", "other_comments_or_suggestions": "1. Add ablation experiments to test the impact of removing identity modeling or using different identity modeling approaches (e.g., random identity, fixed identity) on the results.\n2. Add a complexity comparison analysis.\n3. Provide convergence curves or related analysis.\n4. Visualize the dynamic changes in identity encoding as tasks vary.\n5. Offer a method for measuring task similarity.", "questions_for_authors": "If you resolve my concerns, I will raise my rating.", "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": "2Q643Kucyd", "reviewer": "Reviewer_Tgnb", "summary": "The paper proposes an algorithm (CAID) for contextual/multi-scenario (where each scenario is defined by a different MDP) multi-agent reinforcement learning (MARL) in the centralized training decentralized execution (CTDE) paradigm. Each scenario is characterized by a context vector which is unobservable (even during training). Furthermore, in each scenario, different agents may have different identities (also unobservable). Focusing on cooperative MARL, the proposed solution uses existing value decomposition based techniques in conjunction with a new transformer architecture to learn the context vector and agent identities. The architecture includes (i) an encoder that estimates the context vector from the current state and observations of all agents, (ii) a decoder that outputs identities of all agents, and (iii) an action regulator that modifies the Q-values of individual agents based on their identities. The agent identity decoder and the context encoder depend on the full state of the environment and hence, are only used during training; only the individual agents' Q-functions are used during test time. The proposed approach is evaluated empirically and compared with baselines to show reasonable improvement.", "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 the thoughtful and constructive feedback. We respond below to the key concerns raised:\n\n**Q1**: Why is the context vector treated as a latent variable? Why is it allowed to vary within an episode?\n\n**A1**: Thank you for raising this point. In realistic Contextual MARL (CMARL) tasks, the contextual information (e.g., layout changes, perspective shifts) is often not explicitly observable or encoded as scenario IDs. Therefore, we model the context vector as a latent representation derived from the agent-environment interaction history rather than assuming it is accessible.\n\nOur motivation stems from real-world CMARL tasks where the semantics of an environment can shift significantly even within a single episode. For instance, in a traffic control task, the agent behavior required during morning rush hours may differ substantially from that during the evening, despite both being within one episode. We allow the context vector to evolve over time, effectively treating the episode as a sequence of sub-episodes, each governed by a different latent context. This dynamic modeling improves CAID's ability to adapt to diverse temporal patterns and enhances its learning efficiency under non-stationary environments. We will clarify this motivation in the revised paper.\n\n**Q2**: Does not using an identity predictor at test time create a train-test mismatch?\n\n**A2**: This is an excellent question. To avoid requiring any identity information during execution, we designed the Action Regulator to integrate identity into the training objective without modifying the test-time inference pipeline. Specifically, during training, the Q-values are adjusted via the Action Regulator using identity embeddings. This transforms the target Q-values in a context-aware manner, effectively shifting the action value space. However, both during training and testing, the actual actions taken are always based on the raw Q-values directly output from the agent’s policy network. This ensures that identity modeling benefits learning, while the execution remains identity-independent, avoiding any test-time mismatch.\n\n**Q3**: Visualization and interpretability of the learned identities and context vectors.\n\n**A3**: Thank you for the suggestion. We now include t-SNE visualizations of both the learned context vectors and the raw states (See the figure in https://anonymous.4open.science/r/CAID-7A6C/Visualization.jpg). The results show meaningful clustering across semantically similar scenarios and agents with shared roles. We believe this helps support the interpretability and necessity of these components.\n\n**Q4**: Are the discrete contextual IDs used as one-hot vectors in the Action Regulator?\n\n**A4**: Yes, the identity decoder outputs a categorical distribution, from which we sample a discrete identity and convert it into a one-hot representation to feed into the Action Regulator MLP. We will clarify this implementation detail in the revised paper.\n\n**Q5**: Equation (1) seems notationally incorrect.\n\n**A5**: Thank you for catching this inconsistency. In the revised version, we will correct this by ensuring both sides of the equation have consistent dimensionality. We appreciate your attention to this detail.\n\nWe thank you again for your thoughtful questions. Your comments have helped us clarify key design choices and highlight future directions for improving test-time identity modeling. We hope our revisions address your concerns and improve your evaluation of the work.", "claims_and_evidence": "Overall, the proposed approach seems to perform better than baselines in many environments. However, the improvement appears to be small in many cases (e.g., CAID learns with fewer samples but baseline learns to perform better eventually).  Therefore, the claim regarding significant improvement seems to have relatively weak evidence (but is still decent).", "methods_and_evaluation_criteria": "The chosen environments and the experimental setup makes sense. Method-wise there are some choices that don't seem completely natural to me.\n- It is unclear if assuming that the context vector is hidden is necessary. Since the context vector is only needed during training time, it would be more efficient to generate/learn the context vector using the knowledge of the scenario. Also, the context vector is allowed to change from one time-step to another (within the same episode); this does align well with the intuition of what the context vector is trying to capture.\n- The contextual agent identities are not used during test time. This seems to create a mismatch between training and testing because, during test time, agents are assumed to have fixed identities that don't change across scenarios. It is unclear why there is no agent-level identity predictor that can be used during test-time.", "theoretical_claims": "There are no theoretical claims in the paper.", "experimental_designs_or_analyses": "The experimental setup and results appear reasonable and sound from my reading of the paper. I did not check the supplementary material or any code.", "relation_to_broader_scientific_literature": "- The idea of learning contextual agent identities appears broadly relevant and can enable better skill transfer in MARL. This aspect of skill transfer is often not studied and seems novel to this paper.\n- As authors mentioned, the proposed architecture is compatible with any value decomposition based algorithm for cooperative MARL.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "- The paper is reasonable well written.\n- Experiments are performed on wide range of environments. The authors also include an ablation study to understand the impact of various components of the architecture on the performance.\n- As mentioned earlier, the intuition behind the various components of the proposed architecture is not completely clear. It would be good to visualize the learned context vectors and/or contextual identities (to make a stronger case for the different choices made by the authors).", "other_comments_or_suggestions": "- Equation (1) is not quite correct mathematically, left side is a vector and right side is a scalar. Just a notational issue.", "questions_for_authors": "- The contextual ids are discrete. In the Action Regulator, are these ids treated as one-hot vectors (to feed to the MLP)?", "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 introduces Context-Aware Identity Generation (CAID), a novel framework designed to improve generalization in multi-agent reinforcement learning within the Contextual MARL (CMARL) setting, where task configurations vary across episodes. CAID employs a causal Transformer to dynamically generate agent identities based on context and integrates these identities into the value function estimation process.

Initial reviews were mixed, generally appreciating the novelty and empirical results on benchmarks like SMACv2 but raising valid concerns. These included the need for comparisons against established role-assignment methods (like ROMA/RODE), more extensive ablation studies to validate the contribution of different components (especially the dynamic identity generation), clarity on the specific CMARL problem setting versus other generalization frameworks, and requests for analysis of computational overhead and interpretability of the learned identities.

The authors provided a comprehensive and convincing rebuttal. They furnished additional experimental results comparing CAID against ROMA and RODE, performed the requested ablation studies (demonstrating the value of dynamic identities over fixed or no identities), offered t-SNE visualizations suggesting meaningful identity/context representations, and clarified the CMARL definition, the Action Regulator mechanism, and computational costs. This rebuttal successfully addressed the primary concerns of three reviewers (v3VM, 139G, Tgnb), who subsequently raised their scores or indicated their intention to do so, moving towards acceptance. One reviewer (PyYp) remained unconvinced, maintaining concerns about the precise definition of generalization used and the method's applicability, although their engagement in the discussion period after acknowledging the rebuttal was limited compared to the others.

Considering the substantial positive response to the authors' thorough rebuttal, the additional evidence provided, and the resulting majority support from the reviewers who found their initial concerns resolved, the paper makes a solid contribution to addressing generalization in MARL.
Decision: Accept (poster)