DiRecT: 基于滚动时域去噪的安全扩散规划

# 基于滚动时域去噪的安全扩散规划
来源：https://arxiv.org/html/2606.15359
Paolo Giaretta MIT pgiarett@mit\.edu &Zeyang Li MIT zeyang@mit\.edu &Navid Azizan MIT azizan@mit\.edu

###### 摘要

扩散模型通过学习动作和轨迹上的多模态分布，已成为规划和控制的有力工具。然而，在推理时可靠地强制执行安全性仍是其在安全关键任务中部署的主要障碍。现有方法通常将每个去噪迭代投影到可行集上，尽管约束仅定义在最终的干净轨迹上。因此，对含噪中间样本强制可行性可能会过度约束采样动力学，从而显著降低样本质量。为解决这一局限性，我们引入了DiRecT（基于终端约束滚动时域去噪的扩散规划），这是一种无需训练的算法，通过随机最优控制（SOC）实现扩散模型的约束采样。DiRecT仅在最终的干净样本上施加约束，避免了对中间去噪动力学的不必要限制。受模型预测控制的启发，我们推导出一个有原则的滚动时域替代方案，用于解决原本难以处理的约束SOC公式，从而产生一种高效的算法，将随机去噪与约束满足清晰分离，逐步将样本引导至可行的最终轨迹，同时不扭曲学习的扩散动力学。此外，DiRecT具有高度灵活性：它可以利用现成或特定领域的优化器，融入环境动态的先验知识，并优化额外的软奖励。在安全规划基准上的大量实验表明，DiRecT在部署安全性和任务性能上显著优于现有的基于扩散的规划基线。我们的代码可在https://github.com/azizanlab/DiRecT获取。††脚注文本：通讯作者：Zeyang Li (zeyang@mit\.edu)。

## 1 引言

安全可靠的规划仍然是一个核心挑战，需要算法能够在满足约束的同时适应多样化的环境、目标和约束类型。经典规划方法[25 (https://arxiv.org/html/2606.15359#bib.bib1),33 (https://arxiv.org/html/2606.15359#bib.bib2),36 (https://arxiv.org/html/2606.15359#bib.bib3),32 (https://arxiv.org/html/2606.15359#bib.bib4),51 (https://arxiv.org/html/2606.15359#bib.bib5),57 (https://arxiv.org/html/2606.15359#bib.bib7)]在许多结构化设置中取得了显著成果，但也面临重要局限。基于搜索的方法在大规模多智能体问题中可能变得计算上难以承受[70 (https://arxiv.org/html/2606.15359#bib.bib9)]，而基于优化的方法通常是局部的且对初始化敏感[51 (https://arxiv.org/html/2606.15359#bib.bib5),57 (https://arxiv.org/html/2606.15359#bib.bib7)]。这些限制推动了对数据驱动生成式规划器的兴趣增长，这类规划器从离线数据中学习可重用的先验知识，并能建模多样化、高维的行为。继其在图像生成中的成功[17 (https://arxiv.org/html/2606.15359#bib.bib20)]，扩散模型[60 (https://arxiv.org/html/2606.15359#bib.bib12),27 (https://arxiv.org/html/2606.15359#bib.bib11),62 (https://arxiv.org/html/2606.15359#bib.bib13)]已成为规划和控制的强大框架，能够捕捉复杂的多模态状态-动作分布[29 (https://arxiv.org/html/2606.15359#bib.bib24)]，同时支持灵活的推理时引导[1 (https://arxiv.org/html/2606.15359#bib.bib26)]。尽管取得了这些进展，将基于扩散的规划器部署到现实世界的安全关键任务中仍然具有挑战性。例如，生成的轨迹必须避免与障碍物碰撞，因为即使是单次违规也可能导致灾难性失败。由于扩散模型依赖于高维随机去噪动力学（通常由大型神经网络参数化），即使在可行轨迹上训练，它们也可能生成不安全的计划。当轨迹必须满足数据集中未捕获的新约束时，这一挑战进一步加剧。这些局限性激发了对约束扩散规划器的测试时机制日益增长的兴趣。受图像生成中的引导技术[17 (https://arxiv.org/html/2606.15359#bib.bib20),28 (https://arxiv.org/html/2606.15359#bib.bib21),15 (https://arxiv.org/html/2606.15359#bib.bib22)]启发，早期方法通过在去噪过程中将约束满足编码为软引导信号[45 (https://arxiv.org/html/2606.15359#bib.bib29)]来引导样本。然而，软引导只能鼓励可行性，并不能保证硬约束的满足。因此，近期工作更关注直接在采样轨迹上强制执行硬约束，旨在为安全关键规划提供更强保证。特别是，大多数硬约束扩散采样器通过沿去噪轨迹进行投影或约束驱动更新来强制可行性[14 (https://arxiv.org/html/2606.15359#bib.bib28),43 (https://arxiv.org/html/2606.15359#bib.bib33),39 (https://arxiv.org/html/2606.15359#bib.bib60),71 (https://arxiv.org/html/2606.15359#bib.bib30),69 (https://arxiv.org/html/2606.15359#bib.bib31)]。这造成了一个不匹配：规划约束定义在最终的干净轨迹上，而中间去噪迭代是含噪的，它们本身不必可行。因此，在整个采样过程中施加约束可能会过度约束学习的反向动力学，显著降低样本质量。此外，基于投影的公式天然无法提供统一的机制来处理硬约束与额外的软奖励或成本。为了解决这些局限性，我们引入了DiRecT（基于终端约束滚动时域去噪的扩散规划），这是一种用于基于扩散规划的约束采样的无需训练的算法。我们通过随机最优控制（SOC）[31 (https://arxiv.org/html/2606.15359#bib.bib35)]的视角来制定推理时约束强制执行，其中学习到的反向扩散过程作为名义随机动力学，控制输入将采样器引导至最终干净轨迹的可行性。关键是，约束仅施加在这个终端干净样本上，避免了对含噪中间去噪迭代的不必要限制。求解由此产生的带终端约束的SOC问题在计算上是难以处理的。因此，我们利用扩散模型的结构，结合模型预测控制（MPC）[52 (https://arxiv.org/html/2606.15359#bib.bib40)]的思想，推导出一个有原则且可扩展的滚动时域替代方案。在每个去噪步骤中，DiRecT预测当前含噪迭代所隐含的最终干净轨迹，在此预测上求解一个约束优化问题，并将所得优化转化为对当前含噪迭代的受控更新。本文的贡献如下：

- •我们识别了先前约束扩散采样器的一个关键局限性：对含噪中间去噪迭代强制可行性可能会过度约束采样动力学并降低样本质量。相反，我们将无训练的约束扩散规划建模为一个带终端约束的SOC问题，仅在最终干净轨迹上强制可行性，同时保持接近学习到的扩散动力学。
- •我们提出DiRecT，一种无训练的约束扩散采样器，将难以处理的约束SOC问题简化为一系列可处理的干净轨迹优化子问题。通过优化预测的干净轨迹并将这些优化转化为对含噪迭代的受控更新，DiRecT在不扭曲学习到的去噪过程的情况下，将样本引导至满足约束。
- •我们展示了DiRecT具有高度的灵活性，支持现成和特定领域的优化器、等式和不等式约束、环境特定的动态先验，以及推理时的额外软奖励或成本。
- •我们在多种机器人规划应用中评估了DiRecT，包括Maze2D中的迷宫导航、D3IL中的机器人操作、多机器人运动规划（MRMP）以及PushT中的多样接触丰富操作。在这些任务中，DiRecT在约束满足和任务成功方面始终优于现有的基于扩散的规划基线。

## 2 相关工作

本节将我们的工作置于更广泛的文献背景中。首先，我们回顾了基于扩散的规划器的约束采样技术，强调了激发我们方法的关键局限性。其次，我们讨论了先前探索扩散模型与随机最优控制之间联系的工作，阐明了我们的约束采样视角与这些方法的不同之处。最后，我们提及了一个平行但互补的研究方向：基于训练的扩散模型规划方法。由于空间限制，详细内容见附录A (https://arxiv.org/html/2606.15359#A1)。

## 3 背景

##### 扩散模型。
我们采用连续时间公式，并将扩散模型定义为一对伊藤过程：一个正向过程，通过随机微分方程（SDE）逐步破坏从数据分布p0p_{0}中抽取的样本：
dXt=ft(Xt)dt+gt(Xt)dWt,X0∼p0,t∈[0,1]dX_{t}=f_{t}(X_{t})dt+g_{t}(X_{t})dW_{t},\;X_{0}\sim p_{0},\;t\in\left[0,1\right],
以及一个反向过程，通过正向SDE的时间反转[3 (https://arxiv.org/html/2606.15359#bib.bib42)]，从简单先验分布p1≈N(0,Id)p_{1}\approx\mathcal{N}\left(0,I_{d}\right)开始生成样本：
dXt=[ft(Xt)−gt2(Xt)∇Xtlogpt(Xt)]dt+gt(Xt)dW̄t,X1∼p1,dX_{t}=\left[f_{t}(X_{t})-g_{t}^{2}(X_{t})\nabla_{X_{t}}\log p_{t}(X_{t})\right]dt+g_{t}(X_{t})d\bar{W}_{t},\quad X_{1}\sim p_{1},(1)
其中(1 (https://arxiv.org/html/2606.15359#S3.E1))沿时间反向积分，WtW_{t}和W̄t\bar{W}_{t}分别表示正向和反向维纳过程，st(xt)=∇xtlogpt(xt)s_{t}(x_{t})=\nabla_{x_{t}}\log p_{t}(x_{t})是边际分布ptp_{t}的得分函数。此外，我们将关注点限制在常见的高斯仿射调度上，其正向加噪过程的条件分布具有闭式解qt(xt|x0)=N(αtx0,σt2Id)q_{t}(x_{t}|x_{0})=\mathcal{N}\left(\alpha_{t}x_{0},\sigma_{t}^{2}I_{d}\right)，其中噪声调度系数{αt,σt}t∈[0,1]\{\alpha_{t},\sigma_{t}\}_{t\in\left[0,1\right]}与SDE的漂移和扩散函数相关：f(x,t)=α̇tαtx,g2(t)=ddtσt2−2α̇tαtσt2f(x,t)=\frac{\dot{\alpha}_{t}}{\alpha_{t}}x,\quad g^{2}(t)=\frac{d}{dt}\sigma_{t}^{2}-2\frac{\dot{\alpha}_{t}}{\alpha_{t}}\sigma_{t}^{2}。遵循去噪得分匹配[67 (https://arxiv.org/html/2606.15359#bib.bib16),62 (https://arxiv.org/html/2606.15359#bib.bib13)]，得分函数st(xt)s_{t}(x_{t})由神经网络stθ(xt)s^{\theta}_{t}(x_{t})近似，该网络在训练中最小化条件得分匹配（CSM）损失：
LCSM(θ)=Et,x0,ε[λ(t)∥stθ(αtx0+σtε)+εσt∥2],\mathcal{L}_{\mathrm{CSM}}(\theta)=\mathbb{E}_{t,x_{0},\varepsilon}\left[\lambda(t)\left\lVert s^{\theta}_{t}(\alpha_{t}x_{0}+\sigma_{t}\varepsilon)+\frac{\varepsilon}{\sigma_{t}}\right\rVert^{2}\right],(2)
其中t∼U[0,1]t\sim\mathcal{U}[0,1]，x0∼p0x_{0}\sim p_{0}，ε∼N(0,Id)\varepsilon\sim\mathcal{N}(0,I_{d})，λ(t)\lambda(t)是时间相关权重。

##### 采样与Tweedie公式。
利用学习到的得分函数对扩散模型进行采样，通过数值模拟反向动力学(1 (https://arxiv.org/html/2606.15359#S3.E1))实现。常见的随机采样器包括DDPM[27 (https://arxiv.org/html/2606.15359#bib.bib11)]、Euler–Maruyama[35 (https://arxiv.org/html/2606.15359#bib.bib61)]和高阶求解器[42 (https://arxiv.org/html/2606.15359#bib.bib18)]。尽管我们关注的是随机动力学，但少步确定性采样器通常通过对概率流ODE[62 (https://arxiv.org/html/2606.15359#bib.bib13)]进行数值积分得到，该ODE与(1 (https://arxiv.org/html/2606.15359#S3.E1))具有相同的边际分布。例子包括DDIM[61 (https://arxiv.org/html/2606.15359#bib.bib19)]和基于ODE的求解器[41 (https://arxiv.org/html/2606.15359#bib.bib17)]。这些采样器的一步实现提供了正向过程后验条件均值的近似，通过Tweedie公式[54 (https://arxiv.org/html/2606.15359#bib.bib23)]与得分函数相关：
E[X0|Xt]≈x^0θ(xt,t)=xt+σt2stθ(xt)αt.\mathbb{E}[X_{0}\,|\,X_{t}]\approx\hat{x}^{\theta}_{0}(x_{t},t)=\frac{x_{t}+\sigma_{t}^{2}s^{\theta}_{t}(x_{t})}{\alpha_{t}}.(3)
因此，我们根据方便互换使用stθs^{\theta}_{t}和x^0θ\hat{x}^{\theta}_{0}，假设它们由(3 (https://arxiv.org/html/2606.15359#S3.E3))关联。

##### 扩散规划器。
遵循基于扩散的规划公式[29 (https://arxiv.org/html/2606.15359#bib.bib24),13 (https://arxiv.org/html/2606.15359#bib.bib25)]，我们将干净样本视为一个有限时域计划，而非单一配置。令H∈NH\in\mathbb{N}为预测时域，H={0,...,H}\mathcal{H}=\{0,\ldots,H\}。我们将计划表示为τ=(τ0,...,τH)∈DH+1\bm{\tau}=(\bm{\tau}_{0},\ldots,\bm{\tau}_{H})\in\mathcal{D}^{H+1}，其中每个元素τk∈D⊆RD\bm{\tau}_{k}\in\mathcal{D}\subseteq\mathbb{R}^{D}根据规划器参数化可能编码一个状态、一个动作或一个状态-动作对。扩散规划器学习这些计划的分布，并通过从高斯先验样本去噪生成τ\bm{\tau}。在闭环部署中，规划器仅执行前M≤HM\leq H个元素，然后重新规划，执行的前缀拼接形成滚动轨迹。

## 4 方法

我们现在形式化扩散模型的硬约束采样并推导DiRecT。推导在精神上类似于[38 (https://arxiv.org/html/2606.15359#bib.bib86)]中为流匹配模型开发的测试时引导框架，但在扩散模型设置和最终算法形式上有所不同。由于空间限制，我们将算法的完整推导推迟到附录C (https://arxiv.org/html/2606.15359#A3)。给定一个预训练的得分模型，一个约束集S⊆Rd\mathcal{S}\subseteq\mathbb{R}^{d}，以及一个成本函数C:Rd→RC:\mathbb{R}^{d}\rightarrow\mathbb{R}，我们旨在从模型中采样去噪计划τ\bm{\tau}，这些计划需满足：（i）安全（τ∈S\bm{\tau}\in\mathcal{S}），（ii）最小化成本CC，（iii）保持接近学习到的数据分布。我们现在扩展问题设置，展示安全规划如何自然地框架为求解一个带终端约束的随机最优控制问题。然后我们推导出一个滚动时域替代公式，进而得到一个可处理的算法。在本节中，我们使用一般的扩散中心符号。