Technical Summary: Partner-Aware Hierarchical Skill Discovery for Robust Human-AI Collaboration

Problem Statement

Effective multi-agent collaboration, particularly in human-AI teaming, requires agents capable of adapting to novel partners with diverse and dynamic behaviors. Conventional Deep Hierarchical Reinforcement Learning (DHRL) methods typically focus on agent-centric rewards, optimizing for diversity without explicitly accounting for partner influence. This limitation leads to "shortcut learning," where agents exploit spurious correlations in state features rather than capturing information relevant to partner interactions. Consequently, learned skills often fail to generalize across novel partners, resulting in poor coordination and behavioral inconsistency. Existing approaches that maximize mutual information (MI) between skills and states often suffer from this issue, as they are bounded by the entropy of the skills and do not ensure sensitivity to partner behavior, leading to static skills with limited adaptability.

Methodology

The authors propose Partner-Aware Skill Discovery (PASD), a DHRL framework designed to learn skills conditioned on partner behavior. The core innovation is a contrastive intrinsic reward that shapes skill representations to be consistent across partners with similar behaviors while maintaining discriminability across diverse strategies.

Framework Overview

PASD operates within a hierarchical structure consisting of a high-level manager ($\pi_{hi}$) that selects latent skills ($z$) and a low-level controller ($\pi_{lo}$) that executes primitive actions conditioned on the current state and the active skill. The framework utilizes a population of partners ($\mathcal{D}_p$) sampled from a distribution during training.

Key Mechanisms

1. Contrastive Intrinsic Reward:
   Instead of maximizing standard MI between skills and states, PASD maximizes a variational lower bound on the MI between skills and sub-trajectories. This is achieved via an InfoNCE-style contrastive objective.

   • Positive Pairs: Sub-trajectories generated by the same skill interacting with different partners.
   • Negative Pairs: Sub-trajectories generated by different skills.
   • Objective: The intrinsic reward ($r_{int}$) encourages the policy to select skills that produce consistent state embeddings across partners with similar behaviors (intra-skill consistency) while ensuring distinct embeddings for different skills (inter-skill discriminability).

2. Training Objective:
   The total reward for both high-level and low-level policies is a weighted combination of the extrinsic team reward ($r$) and the contrastive intrinsic reward ($r_{int}$):
   $$\tilde{r} = (1 - \lambda)r + \lambda r_{int}$$
   During early training, $\lambda$ is high to promote exploration of diverse partner-conditioned patterns. As training progresses, $\lambda$ is annealed to prioritize task performance while retaining the learned behavioral consistency.

3. Optimization:
   The high-level manager, low-level controller, and termination function are optimized end-to-end using the Proximal Policy Optimization (PPO) algorithm.

Key Contributions

1. PASD Framework: Introduction of a DHRL framework that conditions skill discovery on partner behavior, enabling robust coordination with novel partners.
2. Novel Intrinsic Reward: Proposal of a contrastive intrinsic reward derived from InfoNCE loss. This reward captures shared patterns from parallel rollouts with diverse partners, mitigating shortcut learning caused by spurious state information and directly leveraging partner-relevant information.
3. Comprehensive Evaluation: Extensive evaluation in the Overcooked-AI benchmark across diverse partner populations (varying skill levels and play styles) and human proxy models. The study demonstrates that PASD learns transferable skills that generalize effectively across a wide range of partner behaviors.

Experimental Results

The authors evaluated PASD against baselines including Fictitious Co-Play (FCP), DIAYN, and Hierarchical Population Training (HiPT) in the Overcooked-AI environment.

• Performance with AI Partners: PASD consistently outperformed all baselines across five standard layouts (Cramped Room, Asymmetric Advantages, Coordination Ring, Counter Circuit, Forced Coordination). For instance, in the Cramped Room layout, PASD achieved a mean reward of 165.8, significantly higher than FCP (137.7) and HiPT (117.9).
• Performance with Human Proxy Models: When paired with behavior cloning (BC) models trained on human-human trajectories, PASD maintained superior performance, achieving the highest returns in all layouts (e.g., 150.00 in Cramped Room vs. 118.75 for FCP).
• Human Subject Study: In a controlled study with 19 human participants, PASD achieved higher joint rewards than HiPT across all layouts, improving human-AI collaboration by 22% to 47%.
• Qualitative Analysis: Analysis of learned skill embeddings (via cosine similarity heatmaps and UMAP projections) revealed that PASD successfully disentangles skills. Embeddings for the same skill formed compact clusters regardless of the partner, while different skills remained distinct, confirming the model's ability to capture partner-conditioned dynamics rather than agent-centric noise.

Significance and Claims

The paper claims that PASD addresses a critical gap in multi-agent reinforcement learning: the inability of standard hierarchical methods to adapt to the dynamic behaviors of novel partners. By shifting the focus from agent-centric diversity to partner-aware consistency, the method mitigates shortcut learning and promotes behavioral robustness.

The authors assert that PASD provides a foundation for more adaptive and efficient collaborative agents by learning skills that are not only diverse but also meaningful in the context of human-AI interaction. The results suggest that conditioning skill discovery on partner behavior allows agents to generalize effectively to unseen partners, including those with human-like play styles, without requiring extensive human data for training. The work highlights that robust coordination requires reasoning over temporally extended behaviors that are sensitive to the specific strategies of collaborators.