Expressivity and Statistical Trade-offs in Diffusion Policy Learning

2026-07-08

AI Summary

This work investigates the source of diffusion policies’ capacity to model complex action distributions under limited data. Through theoretical analysis, it establishes—for the first time—the Lipschitz constant \( K \) of the drift term as the central parameter governing the trade-off between expressivity and generalization error, revealing its role in balancing approximation error against statistical complexity and providing matching upper and lower bounds. Leveraging nonparametric approximation and statistical learning theory, the study derives a general performance bound of \( \widetilde{O}(n^{-2/(m+6)}) \) for neural network-based drift functions, along with an improved rate of \( \widetilde{O}(n^{-2/(m+4)}) \) under one-sided dissipativity conditions. Empirical results validate the practical principle that \( K \) should be adaptively tuned according to sample size.

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 $\tilde{O}(n^{-2/(m+6)})$ for generic neural-network drifts, and a sharper rate $\tilde{O}(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.

Problem

• diffusion policy
• expressivity
• statistical trade-off
• Lipschitz budget
• finite-sample learning

Innovation

• diffusion policy
• Lipschitz budget
• expressivity-statistical trade-off
• approximation error
• finite-sample analysis