Theory of Mind Guided Strategy Adaptation for Zero-Shot Coordination

TL;DR

This work tackles zero-shot coordination by addressing the adaptivity gap seen when training a single BR against a diverse partner pool. It introduces TBS, which builds a BR ensemble specialized to behavior clusters and uses Theory-of-Mind reasoning to adaptively select the most appropriate BR at test time. By clustering partners via self-tuning spectral clustering on cross-play similarity and training both cluster-specific BRs and ToM models, TBS achieves robust ZSC in Overcooked across multiple layouts and observability conditions, outperforming a standard BR baseline and several ablations. The approach has practical implications for deploying cooperative agents in dynamic, multi-agent environments where partners vary across tasks and time, enabling more reliable and scalable zero-shot collaboration.

Abstract

A central challenge in multi-agent reinforcement learning is enabling agents to adapt to previously unseen teammates in a zero-shot fashion. Prior work in zero-shot coordination often follows a two-stage process, first generating a diverse training pool of partner agents, and then training a best-response agent to collaborate effectively with the entire training pool. While many previous works have achieved strong performance by devising better ways to diversify the partner agent pool, there has been less emphasis on how to leverage this pool to build an adaptive agent. One limitation is that the best-response agent may converge to a static, generalist policy that performs reasonably well across diverse teammates, rather than learning a more adaptive, specialist policy that can better adapt to teammates and achieve higher synergy. To address this, we propose an adaptive ensemble agent that uses Theory-of-Mind-based best-response selection to first infer its teammate's intentions and then select the most suitable policy from a policy ensemble. We conduct experiments in the Overcooked environment to evaluate zero-shot coordination performance under both fully and partially observable settings. The empirical results demonstrate the superiority of our method over a single best-response baseline.

Figures (11)

• Figure 1: Left: Toy communication environment. Alice (in red) is shown a random number and can take actions A through D or bail out. Bob (in blue) sees Alice's action and can guess the random number or bail out. After the round is over, the random number is revealed. Right:$J_{\text{inter-XP}}$ of different PBT ZSC algorithms as a function of partner agent pool size in a toy environment. Shaded regions represent bootstrapped 95% confidence intervals capturing variability over evaluation seeds.
• Figure 2: Overview of TBS: During training, (1) a diverse partner pool is constructed, (2) partners are clustered by behavioral similarity and specialist BR policies are trained per cluster, and (3) global and cluster-specific ToM networks are learned. During testing, the cooperator infers an unseen partner’s cluster and selects the corresponding BR policy for adaptive coordination.
• Figure 3: Onion soup environment layouts. From left to right: Large Room, Counter Circuit, Bothway Coordination, Asymmetric Advantages, Forced Coordination, Coordination Ring, Cramped Room.
• Figure 4: ZSC Performance of different methods in 7 Onion Soup layouts in fully observable (left) and partially observable (right) settings. The final “Overall” bar in each plot denotes the average performance across all layouts. Error bars indicate bootstrapped 95% confidence intervals capturing variability over different evaluation seeds. Our TBS method equals or improves on the baseline BR method in almost all layouts.
• Figure 5: Experiments ablating the main components of TBS (left), the number of agents in the training pool (middle) and the number of strategy clusters when using spectral clustering with a fixed number of clusters (right). Error bars and shaded regions indicate bootstrapped 95% confidence intervals capturing variability over different evaluation seeds Left: Average ZSC performance of different methods over all 7 onion soup layouts. Both clustering and ToM-based strategy selection improve ZSC performance. Middle: ZSC performance on the difficult Counter Circuit layout as a function of the number of agents in the training pool. Larger training pools result in better ZSC performance for both TBS and BR, but TBS improves more. Right: Average performance of TBS with different numbers of clusters over all 7 onion soup layouts. Average scaled rewards are all around 0.7 to 0.8, indicating that our method is robust to the number of clusters.