Expressivity and Statistical Trade-offs in Diffusion Policy Learning
Viet Vu, Renyuan Xu, Jiacheng Zhang, Yufei Zhang
stat.ML, cs.LG, math.OC
Submitted: 2026-07-08
License: http://creativecommons.org/licenses/by/4.0/
The gist: Diffusion-based policies have recently emerged as powerful policy parameterizations for reinforcement learning, representing state-conditioned action distributions as terminal laws of diffusion
Terminology
Abstract
Diffusion-based policies have recently emerged as powerful policy parameterizations for reinforcement learning, representing state-conditioned action distributions as terminal laws of diffusion processes with parameterized drifts. This terminal-law representation has shown substantial expressive flexibility in practice, enabling diffusion policies to model complex, multimodal, and highly non-Gaussian action distributions; however, it remains unclear what mathematically drives this expressivity and how to fully exploit it when the policy is learned from finite data. In this paper, we identify the drift Lipschitz budget K as a central quantity governing the expressivity and statistical behavior of diffusion policies. We quantify expressivity through approximation: diffusion policies with K-Lipschitz drifts can concentrate near optimal deterministic policies and achieve value approximation error of order 1/K; moreover, we prove a matching lower bound under nondegenerate diffusion noise. This increased expressivity comes with a statistical cost. When the drift is parameterized by neural networks, increasing K improves approximation but increases statistical complexity. Balancing these two terms yields a finite-sample performance gap of order (n-2/(m+6)) for generic neural-network drifts, and a sharper rate (n-2/(m+4)) for one-sided dissipative drift classes, where n is the sample size and m is the dimension of the state space. Numerical experiments provide empirical evidence for the sample-dependent trade-off in K, supporting both theoretical regimes. Our framework also suggests a practical implementation principle: choose the diffusion budget K according to the available sample size, and then select a neural-network architecture with the corresponding fixed Lipschitz coefficient.
Sources
- Model-Based Diffusion Sampling for Predictive Control in Offline Decision Making
- IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies
- GenPO++: Generative Policy Optimization with Jacobian-free Likelihood Ratios
- DiffuseLoco: Real-Time Legged Locomotion Control with Diffusion from Offline Datasets
- 3D Diffuser Actor: Policy Diffusion with 3D Scene Representations
- Score-Based One-step MeanFlow Policy Optimization
- On diffusion processes with drift in $L_{d+1}$
- Efficient Online Reinforcement Learning for Diffusion Policy
- Covering Numbers for Deep ReLU Networks with Applications to Function Approximation and Nonparametric Regression
- Proximal Policy Optimization Algorithms
- Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching
- Policy Representation via Diffusion Probability Model for Reinforcement Learning
- Reparameterization Flow Policy Optimization
- Truncated Rectified Flow Policy for Reinforcement Learning with One-Step Sampling
Related papers
- Behavior of prediction performance metrics with rare events
- Optimal Estimation of Generic Dynamics by Path-Dependent Neural Jump ODEs
- A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors
- Scalable Posterior Uncertainty for Flexible Density-Based Clustering
- A Transport Geometry of Benign Overfitting
- Efficient quadratic entropy with distance sketches