PAC-DP: PAC-Bayesian Diffusion Policy Learning

2026-07-27

AI Summary

This work addresses the limited generalization of diffusion policies in robotic manipulation due to data scarcity by introducing PAC-Bayes theory into diffusion-based policy learning for the first time. The authors propose a novel approach based on Bayesian neural networks that integrates PAC-Bayes generalization bounds with a KL-divergence regularization term, yielding a theoretically grounded training objective. This formulation enhances generalization performance without substantially increasing computational overhead. Experimental results demonstrate that the proposed method significantly improves denoising efficacy, reduces variational negative log-likelihood, and achieves higher task success rates across multiple robotic manipulation benchmarks, particularly in low-data regimes.

Abstract

Diffusion Policies (DPs) are able to perform complex manipulation tasks. However, DPs are typically trained by minimizing a denoising objective, which provides limited control over generalization in the finite-data regimes common in robotics. In this letter, we propose PAC-DP, an approach that increases the performance of DPs in robotic manipulation tasks. By modeling the DP as a Bayesian neural network, and defining a PAC-Bayes generalization bound, we derive a novel training objective that augments the standard denoising loss with a Kullback-Leibler divergence regularizer between the posterior and prior parameter distributions. From the theoretical perspective, our approach provides a principled approach to regularize the training of DPs without significantly increasing the training time. From the practical point of view, experimental results demonstrate improved denoising performance, lower variational negative log-likelihood, and higher success rates across multiple robotic manipulation benchmarks. Crucially, the largest improvements are observed in low-data training regimes and complex tasks, establishing PAC-DP as a theoretically grounded framework for robot policy learning.

Problem

• Diffusion Policies
• generalization
• finite-data regimes
• robotic manipulation
• PAC-Bayes

Innovation

• PAC-Bayes
• Diffusion Policy
• Bayesian Neural Network
• Generalization Bound
• KL Regularization