KDPE: A Kernel Density Estimation Strategy for Diffusion Policy Trajectory Selection

2025-08-14

AI Summary

In behavior cloning, diffusion-based policies achieve state-of-the-art performance but suffer from two key limitations: (1) stochastic denoising degrades trajectory quality, and (2) the supervised learning paradigm is sensitive to dataset outliers, causing policy deviation from the true action distribution. To address these issues, we propose a manifold-aware kernel density estimation (KDE) trajectory filtering method. Specifically, we construct an action density function in the diffusion-generated action space by incorporating geometric priors of the underlying action manifold, enabling efficient identification and rejection of low-probability, high-risk trajectories. Our approach incurs negligible computational overhead, requires no additional training, and preserves the original diffusion model architecture. Evaluated on simulated single-arm manipulation and real-robot grasping tasks, it significantly improves trajectory stability and task success rate over the baseline diffusion policy—demonstrating strong robustness and generalization capability.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Humans and AI: Human-Aware Planning and Behavior Prediction
• Machine Learning: Learning with Manifolds

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

Abstract

Learning robot policies that capture multimodality in the training data has been a long-standing open challenge for behavior cloning. Recent approaches tackle the problem by modeling the conditional action distribution with generative models. One of these approaches is Diffusion Policy, which relies on a diffusion model to denoise random points into robot action trajectories. While achieving state-of-the-art performance, it has two main drawbacks that may lead the robot out of the data distribution during policy execution. First, the stochasticity of the denoising process can highly impact on the quality of generated trajectory of actions. Second, being a supervised learning approach, it can learn data outliers from the dataset used for training. Recent work focuses on mitigating these limitations by combining Diffusion Policy either with large-scale training or with classical behavior cloning algorithms. Instead, we propose KDPE, a Kernel Density Estimation-based strategy that filters out potentially harmful trajectories output of Diffusion Policy while keeping a low test-time computational overhead. For Kernel Density Estimation, we propose a manifold-aware kernel to model a probability density function for actions composed of end-effector Cartesian position, orientation, and gripper state. KDPE overall achieves better performance than Diffusion Policy on simulated single-arm tasks and real robot experiments. Additional material and code are available on our project page https://hsp-iit.github.io/KDPE/.

Problem

• Mitigating stochasticity impact on action trajectory quality
• Filtering harmful trajectories from Diffusion Policy outputs
• Improving robot policy performance with low computational overhead

Innovation

• KDPE filters harmful trajectories using Kernel Density Estimation
• Manifold-aware kernel models action probability density
• Achieves better performance with low computational overhead