PADiff: Predictive and Adaptive Diffusion Policies for Ad Hoc Teamwork

AI summary

Overview

Research area: Multi-agent reinforcement learning, specifically ad hoc teamwork (AHT), combined with diffusion-based generative models for policy learning.

Technical level: Advanced. The paper assumes familiarity with reinforcement learning, conditional diffusion models (including discrete diffusion such as D3PM), FiLM-style feature modulation, and Dec-MDP formalisms.

Scope: The paper introduces PADiff, a diffusion-based ego-agent policy for cooperating with previously unseen teammates, built from two new components (a Predictive Guidance Block and an Adaptive Feature Modulation Net) and evaluated in three cooperative benchmarks.

What This Paper Is About

In ad hoc teamwork, an agent must cooperate with teammates it has never seen before and whose strategies it cannot know in advance. The authors argue that standard reinforcement learning collapses to one dominant behavior and therefore cannot represent the many different ways cooperation might succeed (for example, passing to one teammate, passing to another, or shooting directly). PADiff's goal is a policy that keeps multiple cooperation modes available while also predicting what teammates are likely to do and adapting to their changes in real time.

Key Contributions

1. The first diffusion-based approach to ad hoc teamwork. PADiff models the ego agent's policy as a conditional diffusion process so that the policy can represent multimodal action distributions and support diverse cooperation patterns, rather than collapsing to a single dominant behavior as the authors say RL-based methods do.
2. Adaptive Feature Modulation Net (AFM-Net). A denoising network that applies two FiLM-like feature-wise modulation layers to scale and shift intermediate features using the current state context, combined with Layer Normalization, residual connections, and dropout regularization, to adapt in real time to non-stationary teammates without needing attention mechanisms.
3. Predictive Guidance Block (PGB). A module integrated into the denoising process that uses intermediate features to predict two auxiliary targets: Collaborative Return (expected cumulative future team reward) and Collaborative Goal (a predicted future state standing in for teammate intention). Gradients from these predictions shape the intermediate features so that action generation aligns with long-term team objectives.
4. Empirical validation across three cooperative environments. The authors report that PADiff consistently outperforms strong baselines with an average performance gain of 35.25% across all tested environments.

Main Findings

• PADiff beats all reported baselines in every environment and pool size. Averaged return comparison (mean ± standard error, 50 trials) from Table 1:

• Strongest competitor varies by environment. TAGET reports the best baseline scores in Predator-Prey (61.2 on PP-4, 60.1 on PP-8), while CQL reports the best baseline score on Overcooked-4 (0.64) and Overcooked-8 (0.46).

• The policy is genuinely multimodal. When a single state from the Predator-Prey environment was fed to the model 1,000 times, the resulting action distribution showed multiple distinct peaks under kernel probability density estimation (six sampled states are shown). In a visualized Predator-Prey scenario starting from the same initial state, the ego agent takes different cooperative paths.

• AFM-Net matters. Replacing AFM-Net with either a simple MLP or a U-Net lowered performance in all environments, which the authors attribute to AFM-Net being better suited to AHT than image-oriented U-Net or plain MLP denoisers.

• Both predictive tasks matter. Removing either the CoGoal predictor or the CoReturn predictor from the PGB produced substantial degradation, and the full model outperformed all ablated variants.

• Teammate pool diversity was verified with cross-play. Using Soft-Value Diversity (SVD) from the CSP framework, the authors instantiated four independent multi-agent populations per environment; cross-population cooperation scores were lower than within-population scores, with diagonal entries of the cross-play matrices performing best.

Methodology in Plain English

The researchers frame ad hoc teamwork as a Decentralized MDP extended with a space of possible teammate policies, and designate one agent as the "ego agent." Training data is gathered by sampling teammates from a deliberately diverse pool built with Soft-Value Diversity, then letting the ego agent interact and storing the resulting trajectories in an offline dataset. Three training populations (each with three unique policy checkpoints) supply training data, while a fourth population of twelve checkpoints is held out purely for testing, split into a group of 4 and a group of 8 unseen strategies.

The ego agent's policy is a conditional diffusion model: it starts from noise and repeatedly denoises into an action, with the state, the diffusion step, and a compressed representation of recent team dynamics supplied as conditions. For discrete action spaces the paper uses D3PM-style categorical diffusion with Uniform transition matrices. Because plain diffusion denoisers are not designed for changing teammates, the authors replace the standard denoiser with AFM-Net, which uses MLP-generated scale and shift parameters to modulate its intermediate features, plus layer normalization, residual connections, and dropout. Separately, the PGB reads those intermediate features during training and predicts both the expected cumulative team reward (CoReturn) and a future state representing teammate intention (CoGoal); these predictions add auxiliary loss terms whose gradients reshape the features the denoiser produces. At inference, the PGB is dropped and only the denoising network is used, since it has internalized team awareness. Models were trained for 20 epochs on the collected offline datasets, evaluated every two epochs against the held-out pool, with returns averaged over 50 trials and 95% confidence intervals computed as ±1.96·σ/√n.

Why This Matters

Impact on research. The paper positions diffusion policies as a way to escape the single-mode collapse the authors attribute to expected-return RL in ad hoc teamwork, and argues that maximum-entropy regularization only disperses behavior rather than structurally modeling multimodal distributions. It also contributes a design pattern, injecting auxiliary prediction tasks into the denoising process rather than only at inference time, which is relevant to any setting where a generative policy must anticipate other agents.

Real-world applications (as motivated in the paper):

• Robotic soccer, where an autonomous agent plays with unfamiliar teammates whose styles and strategies are unknown.
• Disaster response, where agents must coordinate with human partners operating under unknown procedures.
• Autonomous driving, where vehicles interact with other road users without established communication protocols.
• Multi-agent systems generally being deployed in the real world, where predefined coordination cannot be assumed.

Industry relevance. The approach targets settings where an AI system must be dropped into a team it did not train with, which is common in robotics fleets, human-AI collaboration tools, and simulation-based training. Because the PGB is discarded at inference, the deployed model is a denoising network, which the authors present as a route to real-time decision-making without extra inference-time overhead.

Future Directions

• Characterizing the accuracy/compute trade-off. The appendix describes a grid over the number of diffusion steps and dropout rate and states that larger numbers of diffusion steps showed diminishing returns, but the provided content is truncated before the exact findings; how many denoising steps are needed for real-time control remains an open question.
• Continuous and higher-dimensional action spaces. The discrete diffusion formulation used here follows D3PM with Uniform transition matrices; extending the same predictive-guidance idea to continuous control is not demonstrated.
• Scaling team size and environment complexity. The reported benchmarks use grid-world and kitchen-style tasks with small teams; behavior with larger teams or richer observation spaces is not reported.
• Testing with human or genuinely out-of-distribution teammates. Evaluation uses held-out checkpoints from the same SVD-based population generation process; collaboration with humans or with strategies produced by entirely different algorithms is not evaluated.

Target Audience

Researchers and graduate students working on multi-agent reinforcement learning, ad hoc teamwork, and generative models for decision-making; practitioners building agents that must cooperate with unfamiliar partners; and readers interested in how auxiliary prediction objectives can be embedded inside a diffusion denoising loop rather than applied only at sampling time.

Authors’ abstract

Ad hoc teamwork (AHT) requires agents to collaborate with previously unseen teammates, which is crucial for many real-world applications. The core challenge of AHT is to develop an ego agent that can predict and adapt to unknown teammates on the fly. Conventional RL-based approaches optimize a single expected return, which often causes policies to collapse into a single dominant behavior, thus failing to capture the multimodal cooperation patterns inherent in AHT. In this work, we introduce PADiff, a diffusion-based approach that captures agent's multimodal behaviors, unlocking its diverse cooperation modes with teammates. However, standard diffusion models lack the ability to predict and adapt in highly non-stationary AHT scenarios. To address this limitation, we propose a novel diffusion-based policy that integrates critical predictive information about teammates into the denoising process. Extensive experiments across three cooperation environments demonstrate that PADiff outperforms existing AHT methods significantly.