Reidentify: Context-Aware Identity Generation for Contextual Multi-Agent Reinforcement Learning

We sincerely thank the reviewer for the thoughtful and constructive feedback. We respond below to the key concerns raised:

Q1: Why is the context vector treated as a latent variable? Why is it allowed to vary within an episode?

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.

Our 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.

Q2: Does not using an identity predictor at test time create a train-test mismatch?

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.

Q3: Visualization and interpretability of the learned identities and context vectors.

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.

Q4: Are the discrete contextual IDs used as one-hot vectors in the Action Regulator?

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.

Q5: Equation (1) seems notationally incorrect.

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.

We 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.

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:

Q1: Lack of ablation or alternative identity modeling comparisons.

A1: Thank you for highlighting this. We have added the following ablation experiments:

• CAID w/o AI: Removes the identity modeling entirely—all agents share the same input embedding.
• CAID w/ Fixed ID: Assigns static identities based on agent indices (e.g., [1, 2, ..., n]).
• CAID w/ Random ID: Assigns random but fixed identities per episode.

Results (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.

Q2: No analysis of computational cost due to causal Transformer.

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.

Q3: No evaluation of method stability under varying conditions.

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

Q4: Identity interpretability and task similarity are unexplored.

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.

We 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.

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:

Q1: How is policy generalization defined? How does it differ from definitions in [1] [2]?

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].

Q2: The motivation is unclear and the limitations of prior work are not systematically discussed.

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.

Q3: The meaning of the “draw” operation in Equation (5) is unclear.

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.

Q4: Missing comparison with ROMA/RODE; no ablation on Action Regulator.

A4: Thank you for these insightful comments. We address them as follows:

① 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.

② 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.

Q5: Zerg performs worse than other races in Figure 4.

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.

Thank 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.

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:

Q1: The paper does not compare with dynamic role assignment methods.

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.

Q2: The relevance of Meta-RL references is unclear.

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.

Q3: Does CAID support multi-task adaptation within the same MARL domain?

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.

We really appreciate your comments and they really help us improve our paper! And we also appreciate it if you have any further comments.