ReCollab: Retrieval-Augmented LLMs for Cooperative Ad-hoc Teammate Modeling

TL;DR

Ad-hoc teamwork requires rapid inference of unseen teammate policies under partial observability. The authors propose CoLLAB, an LLM-based rubric-driven world model for classifying teammate types, and ReCoLLAB, a retrieval-augmented variant that grounds predictions with exemplar trajectories. They construct a discriminative behavior rubric via feature selection using mutual information $I(f_j;\\tau)$ and type prototypes $(\\mu_{j,\\tau},\\sigma_{j,\\tau})$, then route to the corresponding best-response policy $\\pi^{1,\\hat{\\tau}}$. In cooperative Overcooked across layouts Cramped Room, Asymmetric Advantage, and Coordination Ring, ReCoLLAB achieves near Pareto-optimal trade-offs between classification accuracy and cumulative return, outperforming CoLLAB and several baselines. The results underscore retrieval grounding as a key factor for robustness in ad-hoc teamwork and point to future work on richer fingerprints and online policy adaptation.

Abstract

Ad-hoc teamwork (AHT) requires agents to infer the behavior of previously unseen teammates and adapt their policy accordingly. Conventional approaches often rely on fixed probabilistic models or classifiers, which can be brittle under partial observability and limited interaction. Large language models (LLMs) offer a flexible alternative: by mapping short behavioral traces into high-level hypotheses, they can serve as world models over teammate behavior. We introduce \Collab, a language-based framework that classifies partner types using a behavior rubric derived from trajectory features, and extend it to \ReCollab, which incorporates retrieval-augmented generation (RAG) to stabilize inference with exemplar trajectories. In the cooperative Overcooked environment, \Collab effectively distinguishes teammate types, while \ReCollab consistently improves adaptation across layouts, achieving Pareto-optimal trade-offs between classification accuracy and episodic return. These findings demonstrate the potential of LLMs as behavioral world models for AHT and highlight the importance of retrieval grounding in challenging coordination settings.

Figures (4)

• Figure 1: System diagram of CoLLAB and ReCoLLAB. The Overcooked environment produces observed trajectories from a controlled agent interacting with a teammate. These trajectories are transformed into prototype feature vectors (e.g., action histograms, dwell times, cumulative reward), which are matched against behavior rubrics to model teammate types. In CoLLAB, an LLM classifies the teammate type directly from the behavior rubric. In ReCoLLAB, the LLM additionally conditions on retrieved exemplar trajectories from a database, grounding its classification. Both approaches output a predicted teammate type with associated confidence and rationale, which is used to select the best-response policy from a library of trained policies.
• Figure 2: Pareto Frontier Study. We plot the teammate-type classification accuracy vs. the obtained cumulative reward and estimate the Pareto optimal frontier. ReCoLLAB consistently lies near or directly on the Pareto frontier.
• Figure 3: Probe-length ablations. Teammate-type classification accuracy and episodic returns as a function of probe length $P$.
• Figure 4: Number of retrievals ablations. Teammate-type classification accuracy and episodic returns as a function of the number of retrieved exemplars $k$.