解决扩散策略训练中梯度漂移问题,提升稳定性和性能
DiPOD: Diffusion Policy Optimization without Drifting Apart

- 通过自蒸馏与梯度更新交替,保持策略优化的紧密约束
- 在语言模型和连续控制任务中均实现更高奖励,训练更稳定
- 适合需要可靠强化学习后训练的扩散模型应用
强化学习后训练在提升扩散策略方面日益重要,但现有扩散策略梯度方法常不稳定,难以实现可靠策略改进。我们识别出根本原因是双重漂移现象:优化变分代理目标会导致证据下界(ELBO)与真实对数似然分离,进而使代理策略梯度与期望回报的真实策略梯度错位。为此,我们提出DiPOD,一种通过交替进行自蒸馏与策略改进梯度更新,全程维持紧约束的扩散策略优化框架。该方法简化为在每次扩散策略梯度更新中加入在线策略ELBO正则项。在扩散语言模型后训练和连续控制扩散策略任务中,DiPOD显著提升了训练稳定性,并获得比先前方法更高的奖励。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。