Multi-agent In-context Coordination via Decentralized Memory Retrieval

1. Problem Statement

The paper addresses the challenge of fast adaptation in cooperative Multi-Agent Reinforcement Learning (MARL) within Decentralized Partially Observable Markov Decision Processes (Dec-POMDPs).

• Context: While In-Context Learning (ICL) has shown success in single-agent Reinforcement Learning (RL) by allowing agents to adapt to new tasks using few-shot examples without parameter updates, extending this to cooperative multi-agent settings is difficult.
• Core Challenges:
  1. Partial Observability: In decentralized execution, agents only see local observations, leading to a biased or incomplete understanding of the global task.
  2. Credit Assignment: Agents typically receive a shared team-level reward, making it difficult to assess individual contributions. This often leads to the "lazy agent" problem, where agents fail to learn effective policies because they cannot distinguish their own impact on the team's success.
  3. Task Misalignment: Decentralized policies may struggle to align with shared goals when adapting to unseen tasks using only local context.

2. Methodology: MAICC Framework

The authors propose MAICC, a framework that enables agent teams to coordinate rapidly on unseen tasks by retrieving relevant historical trajectories as context. The architecture consists of three main phases:

A. Multi-Agent Trajectory Embedding Models (Training Phase)

To facilitate efficient retrieval, the authors design a Centralized Training with Decentralized Execution (CTDE) embedding architecture:

• Centralized Embedding Model (CEM): Trained with access to global team observations and actions. It captures fine-grained team-level information and serves as a "teacher."
• Decentralized Embedding Models (DEMs): Trained using only local observations. They are optimized to approximate the CEM's embeddings via knowledge distillation (minimizing KL divergence). This allows agents to infer team-level task characteristics from local data during execution.
• Token Design: The models use observation, action, and post-step information tokens (global reward, done signal, task completion flag). Notably, the Return-to-Go (RTG) token is omitted during embedding training to prevent retrieving trajectories from irrelevant tasks that happen to have similar return values.

B. Retrieval-Based In-Context Decision Training

• Retrieval Mechanism: During testing, an agent's current sub-trajectory is encoded by the DEM. Using Maximum Inner Product Search (MIPS), the system retrieves the top-$k$ most similar trajectories from a constructed memory.
• Decision Model: A causal Transformer (Decision Model) is trained to predict actions conditioned on the concatenation of the retrieved in-context trajectories and the current sub-trajectory. Unlike the embedding model, this model does utilize RTG tokens from retrieved trajectories to guide action selection toward desired returns.

C. Decentralized In-Context Fast Coordination (Testing Phase)

• Hybrid Memory Mechanism: To balance exploration and exploitation, agents maintain a memory buffer $B'$ that combines:
  • Offline Dataset ($D$): Multi-task data collected during pre-training.
  • Online Replay Buffer ($B$): Data collected during the current test episodes.
  • Exponential Decay: A coefficient $\beta_t = \exp(-\lambda \frac{t}{T})$ controls the sampling probability. Early episodes prioritize offline data (exploration), while later episodes increasingly leverage high-value online data (exploitation).
• Hybrid Utility Score: To solve the credit assignment problem, the retrieval score is augmented with a utility score $S_{util}(\tau) = \alpha \cdot \text{norm}(R) + (1-\alpha) \cdot \text{norm}(\tilde{R})$.
  • $R$: Global team return.
  • $\tilde{R}$: Predicted individual return (estimated by the DEM from action embeddings).
  • This ensures retrieved trajectories are beneficial for both the team and the individual agent, mitigating the "lazy agent" problem.

3. Key Contributions

1. First ICRL for Dec-POMDPs: MAICC is the first approach to apply In-Context Reinforcement Learning specifically to decentralized cooperative multi-agent settings.
2. Distillation-based Embedding: The introduction of a CEM-DEM distillation framework allows decentralized agents to access team-level task information despite partial observability.
3. Hybrid Utility Retrieval: The novel memory mechanism combines offline/online data with a hybrid utility score (individual + team returns) to ensure effective credit assignment and trajectory selection.
4. Theoretical Guarantee: The paper provides a theoretical bound on the online cumulative regret of MAICC, showing it scales efficiently with the horizon and state space size.

4. Experimental Results

The method was evaluated on standard cooperative MARL benchmarks:

• Environments: Level-Based Foraging (LBF) and StarCraft Multi-Agent Challenge (SMAC v1 and v2).
• Baselines: Compared against MADT (Multi-Agent Decision Transformer), AT (Agentic Transformer), RADT (Retrieval-Augmented DT), and HiSSD.
• Performance:
  • MAICC consistently outperformed all baselines across six test scenarios, demonstrating significantly faster adaptation to unseen tasks.
  • In complex SMACv2 scenarios with high task diversity, MAICC was the only method to exhibit clear in-context adaptation.
  • Ablation Studies: Confirmed that removing the CEM distillation, the hybrid utility score, or the exponential decay memory mechanism significantly degraded performance.
  • Visualization: t-SNE visualizations showed that the proposed embedding design successfully clusters trajectories by task, whereas using RTG tokens or incomplete loss functions led to overlapping clusters and poor retrieval.

5. Significance

This work bridges the gap between the adaptability of large language models (via in-context learning) and the complexities of cooperative multi-agent systems. By solving the issues of partial observability and credit assignment through decentralized memory retrieval and hybrid utility scoring, MAICC offers a scalable solution for deploying AI agents in dynamic, real-world multi-agent environments where retraining is impractical. It paves the way for agents that can "learn on the fly" by recalling relevant past experiences without needing to update their neural network weights.