Skip to content

Multi-agent In-context Coordination via Decentralized Memory Retrieval

🕒 Published (v1): 2025-11-13 07:08 UTC · Source: Arxiv · Venue: AAAI · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

MAICC extends In-Context Reinforcement Learning (ICRL) to cooperative multi-agent settings (Dec-POMDPs) by training a centralized embedding model that distills team-level trajectory representations into decentralized per-agent models, enabling parameter-free fast adaptation to unseen tasks. A selective memory mechanism balances offline and online trajectories via exponential decay, and a hybrid utility score combining global and predicted individual returns mitigates the lazy-agent problem. The method outperforms all baselines on LBF and SMAC v1/v2 benchmarks in few-episode adaptation speed.

Problem

Prior In-Context RL (ICRL) methods (AT, RADT, DPT) succeed only in single-agent settings; directly applying them to cooperative MARL fails because (1) decentralized execution gives each agent only partial observations, producing biased task embeddings, and (2) agents receive only shared team rewards, making credit assignment ambiguous and causing the "lazy agent" problem. No prior ICRL method addresses Dec-POMDPs.

Method

Centralized Embedding Model (CEM): A causal transformer with intra-team token visibility trained on multi-task offline trajectories using three auxiliary losses — behavior cloning (Lµ), reward prediction (LR), and next-observation prediction (LT). RTG tokens are deliberately excluded to prevent retrieval of superficially similar but task-irrelevant trajectories.

Decentralized Embedding Models (DEMs): Per-agent transformers operating on local observations/actions only, trained via KL-divergence distillation from CEM embeddings so they approximate team-level representations at execution time.

Retrieval: At inference, each agent embeds its current sub-trajectory with its DEM, then retrieves top-k trajectories via Maximum Inner Product Search (cosine similarity) from a constructed memory buffer.

Selective Memory (βt): Memory buffer B′ mixes offline dataset D and online buffer B with coefficient βt = exp(−λ·t/T); early episodes weight offline data (exploration), later episodes weight online data (exploitation).

Hybrid Utility Score: Retrieval score S = cossim(z_c, z_q) + α·norm(R) + (1−α)·norm(R̃), where R is global return and R̃ is individual return predicted by DEM's MLP head. This combines task-relevance, team performance, and individual credit.

Decision Model: A shared-parameter causal transformer (GPT-2 size) trained via behavior cloning conditioned on concatenated retrieved context + current sub-trajectory; uses RTG from retrieved trajectories to guide return-conditioned action selection.

Key Contributions

  • First ICRL framework designed for Dec-POMDPs, explicitly handling partial observability in cooperative MARL.
  • CEM→DEM distillation pipeline that gives decentralized agents access to approximate team-level trajectory semantics at execution.
  • Selective memory mechanism with exponential time decay, theoretically grounded with an online cumulative regret bound of Õ(CH^(3/2) √(ωAT)).
  • Hybrid utility score that integrates global and predicted individual returns for credit-aware trajectory retrieval.
  • Ablation evidence that each component (CEM, exponential decay, all three loss functions, hybrid α) is necessary.

Results

  • MAICC consistently outperforms all six baselines across all six evaluation scenarios (LBF 7×7-15s, LBF 9×9-20s, SMACv1 Protoss/Terran/Zerg, SMACv2 all) over 50 test runs (5 seeds × 10 tasks each).
  • On SMACv2: all (highest task diversity), default MAICC achieves mean return 14.51 ± 0.46; all ablations and baselines score lower (best ablation variant: 13.61 ± 0.40 without hybrid score; AT and RADT perform near zero on complex SMAC scenarios).
  • MAICC-S (no CEM distillation) substantially underperforms full MAICC, especially on complex SMAC scenarios.
  • MADT and HiSSD (no online adaptation) degrade severely on LBF tasks with limited observability, confirming in-context adaptation is critical.
  • Removing exponential decay (β=0 or β=1) drops SMACv2 return to 12.16 or 11.17 vs. 14.51 default.
  • Removing RTG from embedding (row A) drops to 13.52; omitting LR loss (row C) drops to 10.55.

Limitations

  • Memory construction relies solely on exponential time decay; the authors acknowledge uncertainty-based scheduling could improve generalization.
  • Theoretical regret bound carries the distribution-shift factor C = sup P(M)/P_D(M), which can be large when test tasks diverge significantly from training distribution.
  • Scalability with very large agent counts is untested; transformer inference cost scales quadratically with context length, constraining retrievable context size k.
  • Offline dataset D is collected by QMIX, so trajectory quality and diversity are bounded by the pretraining policy quality.
  • No evaluation in continuous-action or real-world robotics settings; benchmarks are discrete-action game environments.

Relevance to Agentic AI / LLM Agents

MAICC directly instantiates a core agentic AI capability — few-shot, parameter-free task adaptation — in a multi-agent setting, paralleling how LLM agents use retrieved demonstrations as context to solve novel problems. The CEM→DEM distillation mirrors how a centralized "world model" or supervisor can inject global state awareness into autonomous decentralized executors, a pattern increasingly relevant to multi-agent LLM frameworks where agents have partial information but must coordinate. The hybrid utility score and selective memory are practical solutions to problems that multi-agent LLM systems also face: credit assignment across agents and balancing stale long-term memory against fresh episodic context. The work demonstrates that retrieval-augmented, in-context decision-making generalizes beyond language to embodied cooperative settings, strengthening the principled foundations of agentic retrieval-augmented generation (RAG) architectures.