Theory of Mind Guided Strategy Adaptation for Zero-Shot Coordination

1. Problem Statement

The paper addresses the challenge of Zero-Shot Coordination (ZSC) in Multi-Agent Reinforcement Learning (MARL). ZSC requires an agent to collaborate effectively with a novel, unseen partner without any additional training or explicit communication at test time.

• The Core Limitation: Standard population-based approaches train a single Best-Response (BR) agent to maximize expected returns across a diverse pool of partner agents. However, this often leads the BR agent to converge to a static, generalist policy. While this policy performs "reasonably well" on average, it fails to adapt specifically to the unique conventions or strategies of a new partner, resulting in suboptimal synergy and coordination failures.
• The Gap: There is a lack of mechanisms to leverage a diverse training pool to build an agent that can dynamically infer a partner's intent and switch to a specialized strategy, rather than relying on a single averaged policy.

2. Methodology: Theory of Mind-based Best Response Selection (TBS)

The authors propose TBS, a framework that bridges the adaptivity gap by combining population-based training with Theory of Mind (ToM) reasoning. Instead of training one generalist agent, TBS trains an ensemble of specialists and uses ToM to select the best one in real-time.

The framework consists of three main phases:

A. Partner Pool Construction & Clustering

1. Diverse Pool Generation: A diverse population of partner agents is generated using techniques like Maximum Entropy Population (MEP) or Random Reward Shaping.
2. Self-Tuning Spectral Clustering (STSC): Rather than treating every partner individually, the pool is clustered into behaviorally distinct groups.
   • A similarity matrix is constructed based on Cross-Play (XP) performance between agent pairs.
   • STSC is applied to this matrix to automatically determine the optimal number of clusters ($k$) and group agents with similar strategies.

B. Ensemble of Specialist BR Policies

• Instead of one global BR, the system trains a separate Best-Response policy ($\pi_{BR}^i$) for each identified cluster $C_i$.
• Each specialist policy is optimized to maximize returns specifically when paired with partners from its corresponding cluster, allowing for deeper, more specialized adaptation than a generalist policy.

C. ToM-Guided Adaptive Selection (Inference)

At test time, when paired with an unseen partner, the cooperator agent performs the following:

1. Intent Inference: The agent uses Theory of Mind (ToM) networks to infer the partner's high-level intentions (concepts).
   • Global ToM ($ToM_{global}$): Trained on all clusters to predict the partner's intent broadly.
   • Cluster-Specific ToM ($ToM_{C_i}$): Trained on specific clusters to predict intents within that behavioral context.
2. Cluster Identification: The system computes the Kullback-Leibler (KL) divergence between the concept distributions predicted by the Global ToM and each Cluster-Specific ToM.
3. Policy Selection: The cluster with the minimum divergence is identified as the partner's behavioral group. The corresponding specialist BR policy is selected for execution.
4. Initialization: At the start of an episode, a policy is selected randomly until sufficient interaction history allows for intent inference.

3. Key Contributions

1. Automatic Strategy Clustering: The introduction of self-tuning spectral clustering on cross-play performance matrices to automatically identify distinct strategy groups within a diverse partner population without manual hyperparameter tuning.
2. ToM-Based Adaptive Selection: A novel module that uses Theory of Mind to infer a partner's high-level intentions and dynamically select the most suitable specialist policy from an ensemble, moving beyond static generalist policies.
3. Robust ZSC Performance: Demonstration that combining strategy-specific BR agents with ToM-guided selection yields superior zero-shot coordination in the Overcooked environment, particularly as the diversity and size of the training pool increase.

4. Experimental Results

The method was evaluated in the Overcooked-AI environment (specifically the Onion Soup scenario) across seven distinct layouts under both fully observable and partially observable settings.

• Performance vs. Baselines: TBS consistently outperformed or matched the standard Best-Response (BR) baseline and significantly outperformed Random Selection and other baselines (CEM, CBPR).
• Scalability: A key finding is that while the standard BR baseline struggles to improve (or even degrades) as the training pool size increases (due to the difficulty of fitting a single policy to more diverse strategies), TBS performance continues to improve with larger pools. This confirms that TBS effectively leverages diversity.
• Ablation Studies:
  • Removing the clustering step or the ToM selection module significantly degraded performance, proving both components are essential.
  • The method is robust to the number of clusters (whether manually set or auto-determined).
  • High-level, intention-based concept sets performed slightly better than low-level action sets, though the method remained robust across different granularities.
• Toy Environment Validation: In a simplified communication game, TBS demonstrated the ability to adapt to increasing strategy diversity where the BR baseline converged to a static, suboptimal policy.

5. Significance

This work represents a significant shift in ZSC research from static generalization to dynamic adaptation.

• Theoretical Insight: It highlights that "one-size-fits-all" policies are a local optimum that hinders true adaptability. By explicitly modeling the partner's mental state (intentions), agents can navigate the "convention mismatch" problem inherent in multi-agent systems.
• Practical Impact: The framework provides a scalable solution for real-world multi-agent systems where partners may be unknown, heterogeneous, or dynamically changing. It suggests that explicit reasoning about partner intent is a viable and necessary mechanism for robust coordination in complex, partially observable environments.
• Future Direction: The paper identifies the need for theoretical analysis of generalization bounds and better handling of underspecified concept sets as future work.