🤖 machine learning

CooT: Learning to Coordinate In-Context with Coordination Transformers

CooT is a novel framework that leverages in-context learning to enable multi-agent systems to rapidly and robustly adapt to unfamiliar partners without parameter updates, outperforming existing population-based, fine-tuning, and Meta-RL approaches on benchmarks like Overcooked and Google Research Football.

Original authors: Huai-Chih Wang, Hsiang-Chun Chuang, Hsi-Chun Cheng, Dai-Jie Wu, Shao-Hua Sun

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Huai-Chih Wang, Hsiang-Chun Chuang, Hsi-Chun Cheng, Dai-Jie Wu, Shao-Hua Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "New Partner" Dilemma

Imagine you are a professional dancer. You have spent years practicing with your regular partner. You know exactly how they move, when they spin, and how they react. You are perfect together.

Now, imagine you are suddenly paired with a stranger for a dance-off. You don't know their style. They might be slow, fast, or clumsy.

Old AI methods are like dancers who only practiced with one specific partner. When they meet a stranger, they try to force the stranger to dance their way, or they just freeze because they don't know the steps.
Other methods try to learn the new partner by dancing with them for hours (thousands of times) to figure out their style. This is too slow and expensive, especially if you only get a few minutes to dance together.

The goal of this paper is to teach an AI how to instantly "read" a new partner and dance in sync with them after just a few steps, without needing to relearn how to dance from scratch.

The Solution: CooT (The "Context" Dancer)

The authors created a new AI framework called CooT (Coordination Transformer).

Think of CooT not as a dancer who memorizes steps, but as a super-observant improviser.

How it Learns (The Training Phase):
Before the big dance-off, CooT watches thousands of videos of different dance pairs. It doesn't just watch the moves; it watches the relationship.
- It sees Partner A (who likes to spin left) and watches how Partner B (the "best response") adjusts to spin right to match them.
- It sees Partner C (who moves slowly) and watches how Partner D waits patiently.
- The Analogy: CooT is like a student who has watched every type of dance couple in the world. They haven't danced with you yet, but they have seen how people react to your style of dancing.
How it Works (The "In-Context" Magic):
When CooT meets a new human or AI partner, it doesn't change its brain (it doesn't "retrain"). Instead, it uses In-Context Learning.
- The Analogy: Imagine CooT has a "scratchpad" (a context window). As you dance together, it writes down what you just did in the last few seconds on this scratchpad.
- When it needs to decide its next move, it looks at the scratchpad: "Oh, my partner just moved left. In the videos I watched, when a partner moved left, the best move was to step right."
- It uses those recent observations to instantly guess your style and adapt.

Why is this better than the old ways?

Vs. "Self-Play" (Practicing with yourself):
Old AIs practice by dancing with a mirror image of themselves. They get really good at dancing with themselves, but terrible with strangers. CooT learns to dance with everyone.
Vs. "Fine-Tuning" (Learning on the fly):
Some AIs try to learn by dancing with you for hours, making mistakes, and correcting them. This is slow. CooT is like a genius who figures out your style after one or two songs just by watching you. It adapts instantly without needing to "re-train" its brain.
Vs. "Population" Methods (Trying to be good at everything):
Other AIs try to be a "jack of all trades" by training with a huge group of random partners. But if you meet someone truly weird, they might still fail. CooT looks at your specific recent actions and adapts to you specifically.

The Results: Did it work?

The researchers tested CooT in two very different "dance floors":

Overcooked: A chaotic video game where two chefs must cook soup together in a tiny kitchen. (Like a high-stress kitchen shift).
Google Research Football: A soccer game where players must pass and move together. (Like a team sport).

The Findings:

Speed: CooT started coordinating well almost immediately, getting better as the game went on.
Human Test: When real humans played with CooT, they rated it as the best partner. They felt CooT understood them, adapted to their style, and didn't get in their way.
Resilience: If the partner suddenly changed their strategy (e.g., stopped passing and started running alone), CooT noticed the change in the next few seconds and adjusted its plan immediately.

The Bottom Line

CooT is a new way for AI to work with humans or other AIs. Instead of forcing everyone to learn the same rigid rules, CooT acts like a chameleon. It observes the person it is working with, remembers their recent behavior, and instantly shifts its own behavior to match them perfectly.

It's the difference between a robot that says, "I will do my pre-programmed dance," and a robot that says, "I see you're moving left, so I'll move right to help you."

Technical Summary: CooT: Learning to Coordinate In-Context with Coordination Transformers

1. Problem Statement

The paper addresses the ad-hoc teamwork (AHT) problem in multi-agent systems, where an agent must coordinate effectively with previously unseen partners. Existing approaches face significant limitations:

Self-Play (SP): Often converges to specific conventions that fail when interacting with unfamiliar collaborators.
Population-Based Methods: While they improve robustness through diversity (e.g., HSP, MEP), they lack mechanisms for efficient online adaptation beyond the training distribution.
Fine-Tuning: Gradient-based adaptation is impractical in few-shot settings due to high interaction costs and the need for thousands of trajectories to learn effective strategies.
Context-Based Meta-RL: Prior methods (e.g., PACE, PECAN) often rely on compressed latent representations of partner behavior. These can struggle to capture rich, temporally structured behaviors and may fail when test-time behaviors deviate from training data.

The core challenge is to enable an agent to adapt to diverse partner behaviors in real-time with minimal interaction (few-shot), without requiring parameter updates or explicit communication.

2. Methodology: Coordination Transformer (CooT)

The authors propose CooT, a framework that leverages In-Context Learning (ICL) to treat partner adaptation as a learning task. Instead of generalizing across tasks, CooT generalizes across diverse partner behaviors.

2.1 Formalization: Hidden-Utility Markov Game (HU-MG)

The problem is formulated as a HU-MG. Agents operate in a shared environment but are driven by hidden reward functions ( $r_w$ ) observable only to the partner. These hidden utilities induce diverse behavioral preferences (conventions) even within identical environments. The goal is to infer these hidden preferences through observation and align actions accordingly.

2.2 Training Data Generation

CooT is trained on a dataset of interaction trajectories between pairs of agents:

Behavior-Preferring Agents ( $\pi^p_i$ ): Agents trained with specific hidden reward functions (linear combinations of environmental events) to exhibit diverse behavioral styles.
Best-Response Agents ( $\pi^{br}_i$ ): Agents trained via Multi-Agent Proximal Policy Optimization (MAPPO) to maximize returns against their corresponding behavior-preferring partners.
Dataset Construction: The model learns from tuples $(s_h, C, \hat{a})$ $(s_{h}, C, \overset{a}{^})$ , where:
- $C$ is a context window of past interaction trajectories between a behavior-preferring agent and its best response.
- $s_h$ is a query state sampled independently from the context (preventing simple trajectory continuation).
- $\hat{a}$ is the ground-truth best-response action taken by $\pi^{br}_i$ at state $s_h$ .

2.3 Model Architecture and Training

Architecture: CooT utilizes a Decision-Pretrained Transformer (DPT) backbone (specifically GPT-2).
Objective: The model is trained via supervised learning to predict the best-response action distribution $\hat{p}_h(\cdot) = M_\theta(\cdot | s_h, C)$ . The loss function is the cross-entropy between the predicted action and the ground-truth best-response action.
Key Distinction: Unlike standard Decision Transformers that condition on task objectives or rewards, CooT conditions on partner interaction history to infer the partner's latent style and predict the optimal collaborative response.

2.4 Online Deployment

At test time, CooT operates with a strict no-prior protocol:

It initializes with an empty context buffer.
It interacts with an unseen partner for an episode.
The resulting trajectory is appended to the context buffer (FIFO), pushing out the oldest data.
The model predicts actions based on the current state and the accumulated context $C$ , allowing for continuous, gradient-free adaptation to the partner's behavior.

3. Key Contributions

Novel Framework: Introduction of CooT, which reframes partner adaptation as an in-context learning problem, enabling rapid few-shot coordination without parameter updates.
Training Paradigm: A method to train transformers on "best-response" trajectories conditioned on partner contexts, effectively teaching the model to map observed partner styles to optimal collaborative strategies.
Empirical Validation: Comprehensive evaluation on two benchmarks:
- Overcooked: A discrete grid-world requiring tight coordination.
- Google Research Football (GRF): A continuous, high-dimensional environment with larger teams (3 vs 1).
Human Evaluation: A user study demonstrating that human partners perceive CooT as more adaptive and reliable compared to state-of-the-art baselines.

4. Experimental Results

CooT was evaluated against a broad set of baselines, including Behavior Cloning (BC), Population-based methods (HSP, MEP), gradient-based fine-tuning (HSP-ft), and context-based Meta-RL (HSP-meta, PACE).

Performance: CooT consistently outperformed all baselines across diverse Overcooked layouts and the GRF scenario. It achieved the highest Best Response Proximity (BR-prox) and average returns, particularly in coordination-heavy layouts (e.g., Coord. Ring Multi-recipe) where partner-specific adaptation is critical.
Adaptation Speed: CooT demonstrated rapid few-shot adaptation, showing significant performance improvements within the first 15 episodes and stabilizing quickly. In contrast, gradient-based fine-tuning (HSP-ft) showed instability and slow convergence, while other context-based methods (PACE, HSP-meta) often underperformed or failed to improve over time.
Robustness: In non-stationary settings where partners switched strategies mid-task, CooT recalibrated its strategy within an average of 3.67 episodes, whereas other methods struggled to recover.
Human Study: In a study with 36 participants, CooT achieved the highest mean return and was selected as the preferred collaborator by 50% of participants. Qualitative analysis confirmed that humans perceived CooT as explicitly adaptive, noting its ability to improve coordination over time.

5. Significance and Claims

The paper claims that CooT offers a practical solution for real-world human-AI collaboration where online fine-tuning is infeasible. By treating partner adaptation as a learning task via in-context learning, CooT achieves:

Efficiency: Effective alignment with new partners using only a handful of interactions.
Stability: Robustness against sudden partner changes without the risk of catastrophic forgetting or instability associated with gradient updates.
Generalization: The ability to generalize across diverse partner behaviors that were not explicitly seen during training, provided they fall within the distribution of the hidden-utility space.

The authors position this work as a shift from "task generalization" to "partner generalization," demonstrating that conditioning on full interaction histories can support effective few-shot adaptation in complex multi-agent environments.

6. Limitations

The paper acknowledges several limitations:

No Explicit Communication: The method relies solely on action histories and does not utilize explicit communication channels.
Benchmark Scope: Experiments are conducted in simulated benchmarks (Overcooked, GRF) rather than embodied settings with rich perception and control.
Cold-Start: The strict no-prior protocol means the agent starts with no context; future work could explore using pre-filled contexts from similar partners to reduce early interaction costs.
Safety: The authors note that adaptive mechanisms could potentially be misused to exploit predictable partners or reinforce unsafe habits, emphasizing the need for safety constraints in future deployments.