Diffusion Policy Optimization without Drifting Apart

2026-06-11

AI Summary

This work addresses the instability in diffusion-based policy gradient methods during reinforcement learning post-training, which stems from “double drift” caused by variational proxy optimization, leading to inaccurate policy gradients. To resolve this issue, the authors propose the DiPOD framework, which alternates between self-distillation and policy-gradient updates regularized by an evidence lower bound (ELBO) on the policy distribution. This design tightly aligns the variational lower bound with the true log-likelihood, thereby stabilizing training dynamics. Empirical results demonstrate that DiPOD significantly enhances training stability and achieves superior return performance over existing approaches in both diffusion language model post-training and continuous control tasks.

Abstract

RL post-training has become increasingly pivotal for improving diffusion policies, but existing diffusion policy-gradient methods are often unstable and cannot achieve reliable policy improvement. We identify the cause as the double-drift phenomenon: optimizing a variational surrogate can let the ELBO separate from the true log-likelihood, which then makes the resulting proxy policy gradient misaligned with the true policy gradient of expected return. We propose \textbf{DiPOD}, a diffusion policy optimization framework that maintains tight-bound behavior throughout training by interleaving self-distillation with policy-improving gradient updates. This leads to a simple and practical algorithm: augmenting each diffusion policy-gradient update with an on-policy ELBO regularizer. Across diffusion language model post-training and continuous-control diffusion policies, DiPOD substantially stabilizes training and reaches higher rewards than previous methods.

Problem

• diffusion policy
• policy gradient
• RL post-training
• training instability
• policy improvement

Innovation

• Diffusion Policy Optimization
• Double-Drift Phenomenon
• Self-Distillation
• ELBO Regularization
• Policy Gradient Alignment