Technical Summary: Expressivity and Statistical Trade-offs in Diffusion Policy Learning

1. Problem Statement

Reinforcement learning (RL) in continuous action spaces often relies on stochastic policy parameterizations. While classical families (e.g., Gaussian, Beta, mixtures) offer tractable sampling and likelihood evaluation, they frequently lack the expressivity required to model complex, multimodal, or highly non-Gaussian action distributions found in modern tasks like multi-stage manipulation and locomotion.

Diffusion-based policies have emerged as a powerful alternative, representing state-conditioned action distributions as the terminal laws of diffusion processes with parameterized drifts. Despite their empirical success, the theoretical understanding of these policies remains limited. Specifically, it is unclear:

1. What mathematical quantity governs the expressivity of diffusion policies in approximating optimal deterministic policies?
2. How does this increased expressivity impact statistical learning when policies are learned from finite data?

Existing theory for diffusion models focuses primarily on generative modeling (sampling and score estimation) rather than the decision-making objectives inherent to RL.

2. Methodology and Framework

The authors propose a theoretical framework analyzing diffusion policies defined by state-conditioned stochastic differential equations (SDEs):
$$d\bar{a}^x_t = f(t, \bar{a}^x_t, x) dt + \sigma(t, \bar{a}^x_t, x) dB_t, \quad \bar{a}^x_0 \sim \mu_0$$
where the policy $\pi(\cdot|x)$ is the law of the terminal state $\bar{a}^x_T$.

Key Analytical Tools

• Drift Lipschitz Budget ($K$): The central quantity identified is the Lipschitz constant $K$ of the drift function $f$ with respect to the action variable $a$. The authors measure the complexity of the policy class through this budget.
• Approximation Theory: The paper establishes upper and lower bounds on the ability of $K$-Lipschitz drifts to concentrate near deterministic optimal actions.
• Finite-Sample Learning: The drift $f$ is parameterized by ReLU neural networks. The analysis decomposes the total error into three components:
  1. Diffusion Approximation Error: The error due to the diffusion process not perfectly concentrating on the optimal action.
  2. Neural Network Realization Error: The error in approximating the optimal drift function with a finite-width network.
  3. Statistical Estimation Error: The error arising from learning from a finite sample size $n$.
• Policy Gradient via Girsanov's Theorem: Since the terminal density of the diffusion process is generally intractable, the authors derive a tractable policy-gradient formula. By applying Girsanov's theorem, they replace the intractable score function with a pathwise likelihood-ratio term involving the drift sensitivity $\nabla_\theta f$.

3. Key Contributions and Results

A. Expressivity of Diffusion Policies

The paper quantifies the expressivity of diffusion policies through the drift Lipschitz budget $K$:

• Upper Bound: Diffusion policies with $K$-Lipschitz drifts can approximate deterministic optimal policies with a mean-squared localization error of order $O(1/K)$. Consequently, the value function approximation error is also of order $O(1/K)$ under mild continuity assumptions on the reward and transition kernel.
• Lower Bound: The authors prove a matching lower bound, showing that under nondegenerate diffusion noise, no $K$-Lipschitz drift can localize the terminal action near a finite set of target actions faster than $O(1/K)$. This establishes $1/K$ as the intrinsic limit of expressivity for this class.
• Implication: Increasing $K$ improves the ability of the policy to approximate deterministic optima, but it also increases the complexity of the function class.

B. Finite-Sample Statistical Trade-offs

When the drift is parameterized by ReLU neural networks, the authors derive explicit finite-sample oracle inequalities. The total value gap is bounded by the sum of approximation and estimation terms. Balancing these terms yields two distinct regimes:

1. Generic Neural Drift Regime:
   For a generic $K$-Lipschitz neural drift class, the statistical estimation error scales with $\sqrt{sK/n}$ (where $s$ is network size). Balancing the approximation error ($1/K$) and statistical error yields a finite-sample performance gap of order:
   $$\tilde{O}\left(n^{-2/(m+6)}\right)$$
   where $m$ is the state dimension. This rate is achieved by choosing the diffusion budget $K \asymp n^{2/(m+6)}$.

2. Dissipative Drift Regime:
   If the drift class satisfies a one-sided dissipativity condition (e.g., mean-reverting drifts), the stability of the policy class improves. The statistical term reduces to $\sqrt{s/n}$, removing the dependence on $K$. This yields a sharper rate:
   $$\tilde{O}\left(n^{-2/(m+4)}\right)$$
   achieved by choosing $K \asymp (n/\log n)^{2/(m+4)}$.

C. Algorithmic Contribution

The paper provides a policy-gradient theorem for continuous-time diffusion policies. It demonstrates that the gradient of the value function can be expressed as an expectation involving the drift sensitivity and a stochastic integral derived from Girsanov's theorem. This provides a tractable training identity that does not require the explicit density of the terminal action distribution.

D. Empirical Validation

Numerical experiments validate the theoretical regimes:

• Generic Regime: In a nonlinear multi-step MDP, training a generic neural drift shows a U-shaped trade-off in $K$. Increasing $K$ improves expressivity but destabilizes optimization under fixed sample sizes. The empirically optimal $K$ scales as $n^{2/(m+6)}$.
• Dissipative Regime: In a contextual bandit setting with a mean-reverting (dissipative) drift, the policy regret decreases monotonically with $K$ until it hits the center-estimation floor. The results confirm the sharper scaling $K \asymp (n/\log n)^{2/(m+4)}$.

4. Significance and Claims

The paper claims to provide the first theoretical foundations for diffusion-based policies in RL that explicitly link the structural quantity of the diffusion process (the drift Lipschitz budget $K$) to both approximation capabilities and finite-sample statistical error.

• Theoretical Insight: It identifies $K$ as the "expressivity bottleneck" for diffusion policies near deterministic optima.
• Practical Principle: The framework suggests a concrete implementation principle: choose the diffusion budget $K$ according to the available sample size $n$, and then select a neural network architecture with the corresponding fixed Lipschitz coefficient.
• Algorithmic Utility: The derived policy-gradient formula enables the training of SDE-based diffusion policies without requiring intractable density evaluations, bridging the gap between theoretical analysis and practical implementation.

The authors maintain a modest stance, noting that their bounds are for specific classes of drifts and that future work could explore sharper concentration inequalities or incorporate optimization errors from practical training.