无漂移扩散策略优化

# 扩散策略优化：避免漂移分离 来源：https://arxiv.org/html/2606.13795 Haozhe Jiang¹,² Haiwen Feng¹,² Pieter Abbeel¹ Jiantao Jiao¹,³ Angjoo Kanazawa¹‡ Nika Haghtab¹‡ ###### 摘要 RL后训练对于改进扩散策略日益重要，但现有的扩散策略梯度方法往往不稳定，无法实现可靠的策略改进。我们将其原因归结为双漂移现象：优化变分代理可能导致ELBO与真实对数似然分离，进而使得生成的代理策略梯度与期望回报的真实策略梯度方向不一致。我们提出DiPOD，一种扩散策略优化框架，通过在策略改进梯度更新之间穿插自蒸馏步骤，在整个训练过程中保持紧界行为。这引出一个简单实用的算法：为每个扩散策略梯度更新添加一个在线策略ELBO正则化项。在扩散语言模型后训练和连续控制扩散策略中，DiPOD显著稳定了训练，并达到了比先前方法更高的奖励。 00脚注：¹加州大学伯克利分校；²Impossible, Inc.；³NVIDIA。 33脚注：共同指导。

## 1 引言

参考图注

图1：DiPOD示意图，使用期望回报景观在策略参数空间上。我们希望策略既能保证策略改进，又具有紧致的ELBO，如图中彩色曲线所示。先前基于代理的算法（图中奖励曲线来自实际FPO（McAllister等人，2025 (https://arxiv.org/html/2606.13795#bib.bib4)）在GSM8K（Cobbe等人，2021 (https://arxiv.org/html/2606.13795#bib.bib11)）上的实验；见图3 (https://arxiv.org/html/2606.13795#S4.F3)）初始阶段能够实现策略改进，但随后变得不稳定，因为ELBO不再紧致。相比之下，DiPOD交替进行自蒸馏步骤和适当的梯度更新，如蓝色箭头所示。梯度更新保证局部策略改进，自蒸馏则将参数带回ELBO紧致且期望奖励不变的区域。

扩散模型正成为语言、代码和数学推理中离散生成的有力范式，同时扩散和流模型更广泛地支持图像（Ho等人，2020 (https://arxiv.org/html/2606.13795#bib.bib15)；Song等人，2020 (https://arxiv.org/html/2606.13795#bib.bib16)；Song和Ermon，2019 (https://arxiv.org/html/2606.13795#bib.bib17)）、视频（Ho等人，2022b (https://arxiv.org/html/2606.13795#bib.bib18)，a (https://arxiv.org/html/2606.13795#bib.bib19)）和机器人控制（Black等人，(https://arxiv.org/html/2606.13795#bib.bib20)；Bjorck等人，2025 (https://arxiv.org/html/2606.13795#bib.bib21)）等连续域中的强大性能。这一前景对于使用扩散LLM（dLLM）（Nie等人，2025 (https://arxiv.org/html/2606.13795#bib.bib10)；Xie等人，2025 (https://arxiv.org/html/2606.13795#bib.bib2)；Song等人，2025 (https://arxiv.org/html/2606.13795#bib.bib22)）进行推理尤其引人注目，因为扩散LLM提供快速并行采样和灵活的非自回归解码（Jiang等人，2025 (https://arxiv.org/html/2606.13795#bib.bib50)；Ermon，2026 (https://arxiv.org/html/2606.13795#bib.bib51)；Chen等人，2026 (https://arxiv.org/html/2606.13795#bib.bib52)）。然而，用强化学习对扩散策略进行后训练仍然面临根本性困难：标准的策略梯度方法依赖于对数似然log π_θ(a|o)，而扩散模型的这个量通常难以处理。最近的方法（Zhao等人，2025 (https://arxiv.org/html/2606.13795#bib.bib1)；McAllister等人，2025 (https://arxiv.org/html/2606.13795#bib.bib4)；Wang等人，2025 (https://arxiv.org/html/2606.13795#bib.bib7)）因此引入了多种可处理的似然代理，但这些方法仍未系统理解基于代理的策略梯度更新在何时保持可靠。这一缺口对于dLLM尤为重要，因为不稳定的RL后训练导致推理性增益远落后于自回归对应模型。在这项工作中，我们提出DiPOD——扩散策略优化，避免漂移分离——一种针对扩散和流策略的自然的原理性框架，用于实现可靠的策略梯度更新。当应用于dLLM后训练时，DiPOD显著稳定了学习，并在GSM8K、MATH500、Countdown和数独等基准上提高了推理性能，包括成为首个在零样本设置下饱和数独的方法（第4.2节和图1 (https://arxiv.org/html/2606.13795#S4.T1)）。作为扩散策略专用的RL后训练方法，DiPOD旨在满足两个期望条件。首先，(A) 策略梯度更新应改进下游目标，因为这是后训练的主要目标。然而，仅奖励改进不足以改进扩散模型，尤其是当使策略成为连贯扩散模型的概率结构在此过程中被破坏时。因此，我们额外要求 (B) 微调应通过保持代理目标紧致来保留扩散模型结构。DiPOD满足这两个标准。如图1 (https://arxiv.org/html/2606.13795#S1.F1) 所示，DiPOD通过穿插 (i) 保留策略的自蒸馏步骤（确保似然与其代理之间差异最小）与 (ii) *适当的* 策略梯度更新¹¹这是一个技术术语（见定义2 (https://arxiv.org/html/2606.13795#S2.SS0.SSS0.Px3)），指在似然与其代理之间无大差异时准确的常见策略梯度更新类型。来防止代理更新偏离预期的策略改进方向。从这个框架中，我们推导出一个惊人简单的即插即用实现：在扩散策略梯度更新中添加每个更新的ELBO正则化项（算法2 (https://arxiv.org/html/2606.13795#alg2)）。

#### 变分推理方法用于RL后训练。为什么首先用策略梯度更新扩散是困难的？将去噪链视为MDP使似然变得可处理（Black等人，2023 (https://arxiv.org/html/2606.13795#bib.bib26)；Ren等人，2024 (https://arxiv.org/html/2606.13795#bib.bib29)），但它将训练与采样器绑定，并且不适合dLLM，因为dLLM的后训练应保留灵活的解码顺序和推理预算。因此，近期非MDP dLLM方法直接近似似然，使用平均场或单步估计器、部分依赖恢复或强加的自回归顺序（Zhao等人，2025 (https://arxiv.org/html/2606.13795#bib.bib1)；Xie等人，2025 (https://arxiv.org/html/2606.13795#bib.bib2)）；这些可能改变优化后的策略与执行扩散策略的一致性。变分推理（VI）方法则用证据界（如ELBO（McAllister等人，2025 (https://arxiv.org/html/2606.13795#bib.bib4)）或EUBO（Wang等人，2025 (https://arxiv.org/html/2606.13795#bib.bib7)））替代log π_θ，保留扩散原生采样灵活性的同时与预训练目标自然连接。关键问题是这种忠实性仅是局部的。变分扩散RL方法在初始阶段可以工作得很好：在预训练良好的初始化附近，证据界紧致，ELBO仍然是真实对数似然的良好局部代理。然而，RL更新并不保持这种紧界状态。随着优化进行，ELBO与对数似然之间的差异可能增大，我们的分析指出了阿喀琉斯之踵：一旦ELBO偏离对数似然，策略梯度保真度也随之偏离。我们将这种耦合效应称为双漂移现象。DiPOD通过将紧界行为作为显式设计约束来解决这个问题。我们的关键观察是，许多基于VI的估计量是*适当的*（定义2 (https://arxiv.org/html/2606.13795#S2.SS0.SSS0.Px3)）：如果扩散模型经过完美训练使证据界紧致，代理梯度匹配真实对数似然梯度，基于代理的策略梯度步骤与真实策略梯度一致。DiPOD利用这一点，通过自蒸馏反复将模型拉回紧界状态——这在不改变策略分布的情况下，在参考展开分布下收紧证据界——然后才进行策略梯度步骤。在图1 (https://arxiv.org/html/2606.13795#S1.F1) 所示的理想交替形式中，梯度更新保证局部策略改进，而自蒸馏在不改变策略输出分布的情况下恢复紧致性。在实践中，我们实现了一个简单的近似交替方法：对于每一批展开数据，我们使用通常的扩散策略梯度估计器*加上*一个在同一批展开数据上的ELBO最大化项来更新参数。这个额外的ELBO项作为正则化项，减少在线策略的变分差异，改善代理梯度与真实策略梯度之间的一致性。实验表明，这个简单的DiPOD实现惊人地有效：它在离散dLLM后训练任务（例如GSM8K Cobbe等人 (2021 (https://arxiv.org/html/2606.13795#bib.bib11))、MATH500 Lightman等人 (2023 (https://arxiv.org/html/2606.13795#bib.bib12))、数独Arel (https://arxiv.org/html/2606.13795#bib.bib14)、Countdown Pan等人 (2025 (https://arxiv.org/html/2606.13795#bib.bib13))）和连续控制流策略（例如人形机器人运动跟踪；详见第4节 (https://arxiv.org/html/2606.13795#S4)）中显著稳定了训练并改进了奖励优化。

总之，我们做出三个贡献：
- • 我们识别出变分扩散RL中的一个基本失效模式：双漂移。随着RL更新放宽证据界，ELBO偏离对数似然，进而导致代理策略梯度更新偏离预期的策略梯度方向（第3.1节 (https://arxiv.org/html/2606.13795#S3.SS1)）。
- • 我们提出DiPOD，一个原理性的扩散策略优化框架，将保持紧界状态作为显式设计目标，通过交替进行策略梯度更新与保留策略的自蒸馏来实现（第3.2节 (https://arxiv.org/html/2606.13795#S3.SS2)）。
- • 我们推导出一个简单的、即插即用的实用算法，通过在每个更新中添加ELBO正则化项来实现这一原则，显著提高了扩散RL在实际中的稳定性和奖励优化（算法2 (https://arxiv.org/html/2606.13795#alg2) 和 4 (https://arxiv.org/html/2606.13795#S4)）。

### 1.1 相关工作

#### 扩散模型。扩散模型是一类强大的生成模型，应用于图像生成（Ho等人，2020 (https://arxiv.org/html/2606.13795#bib.bib15)；Song等人，2020 (https://arxiv.org/html/2606.13795#bib.bib16)；Song和Ermon，2019 (https://arxiv.org/html/2606.13795#bib.bib17)）、视频生成（Ho等人，2022b (https://arxiv.org/html/2606.13795#bib.bib18)，a (https://arxiv.org/html/2606.13795#bib.bib19)）和机器人学（Black等人，(https://arxiv.org/html/2606.13795#bib.bib20)；Bjorck等人，2025 (https://arxiv.org/html/2606.13795#bib.bib21)）。最近，扩散语言模型作为自回归模型的一个有前景的替代方案出现，用于快速灵活的解码（Nie等人，2025 (https://arxiv.org/html/2606.13795#bib.bib10)；Xie等人，2025 (https://arxiv.org/html/2606.13795#bib.bib2)；Song等人，2025 (https://arxiv.org/html/2606.13795#bib.bib22)）。我们通过变分推理视角（Kingma和Gao，2023 (https://arxiv.org/html/2606.13795#bib.bib23)；Sohl-Dickstein等人，2015 (https://arxiv.org/html/2606.13795#bib.bib24)）研究扩散策略，这也是DiPOD用于推理似然代理和证据界的视角。

#### 策略梯度。策略梯度方法直接优化参数化策略以最大化期望回报（Williams，1992 (https://arxiv.org/html/2606.13795#bib.bib25)；Schulman等人，2017 (https://arxiv.org/html/2606.13795#bib.bib6)），在高维控制（Schulman等人，2016 (https://arxiv.org/html/2606.13795#bib.bib30)）、运动（Rudin等人，2022 (https://arxiv.org/html/2606.13795#bib.bib27)）、操作（Schwarke等人，2023 (https://arxiv.org/html/2606.13795#bib.bib28)）和现代语言模型后训练（Shao等人，2024 (https://arxiv.org/html/2606.13795#bib.bib3)）中处于核心地位。DiPOD针对扩散策略的这种在线策略优化设置。

#### 扩散策略的强化学习。将策略梯度应用于扩散策略的主要障碍是精确对数似然难以处理。先前的工作要么将去噪过程本身视为MDP（Black等人，2023 (https://arxiv.org/html/2606.13795#bib.bib26)；Ren等人，2024 (https://arxiv.org/html/2606.13795#bib.bib29)），要么用变分证据界替代似然（McAllister等人，2025 (https://arxiv.org/html/2606.13795#bib.bib4)；Wang等人，2025 (https://arxiv.org/html/2606.13795#bib.bib7)），要么简化扩散语言模型中的依赖以获得可处理的似然代理（Zhao等人，2025 (https://arxiv.org/html/2606.13795#bib.bib1)；Yang等人，2025 (https://arxiv.org/html/2606.13795#bib.bib5)；Xie等人，2025 (https://arxiv.org/html/2606.13795#bib.bib2)；Tang等人，2025 (https://arxiv.org/html/2606.13795#bib.bib32)）。这些方法在可处理性、采样器灵活性以及优化代理与执行扩散策略之间的一致性方面进行权衡。DiPOD聚焦于变分路线，并通过在策略优化过程中保持证据界与真实似然对齐来稳定它。更全面的相关工作讨论见附录A (https://arxiv.org/html/2606.13795#A1)。

## 2 预备知识

#### 策略梯度方法。强化学习（RL）的目标是学习一个策略π_θ，最大化环境中的期望回报。这里θ表示策略的参数。在时间步t，策略从环境中获取观测值o_t，然后执行动作a_t ~ π_θ(·|o_t)，环境提供奖励r_t作为反馈。回报定义为累积奖励直到环境达到终止状态。策略梯度算法使用期望回报梯度的估计量更新策略参数θ。大多数策略梯度方法可以看作使用以下梯度估计量的变种：
E_θ [∇_θ log π_θ(a_t|o_t) Â_t(o_t, a_t)]，       (1)
其中Â_t是时间步t的估计优势函数。期望取自环境和策略的随机性。这里E_θ表示a_t根据π_θ(·|o_t)生成。知名的实用算法如PPO（Schulman等人，2017 (https://arxiv.org/html/2606.13795#bib.bib6)）和GRPO（Shao等人，2024 (https://arxiv.org/html/2606.13795#bib.bib3)）可以理解为对这个更新引入修改以提高稳定性和样本效率。估计方程1 (https://arxiv.org/html/2606.13795#S2.E1)