ReGuide：从测试时引导到自改进扩散策略

# 从测试时引导到自我改进的扩散策略
来源：https://arxiv.org/html/2606.28939
###### 摘要

行为克隆扩散策略虽然表达能力强大，但仍然容易受到协变量偏移的影响：与演示状态的小偏差可能累积成任务失败。现有方法要么通过专家修正或合成增强来扩展训练分布，要么在测试时利用学习模型的引导来引导冻结策略。前者可能成本高昂或依赖于假设，而后者在执行后丢弃了修正后的轨迹。我们引入了ReGuide，这是一个自我改进的框架，将引导式 rollout 视为可复用的在策略恢复数据。ReGuide 首先使用**阶段条件引导（PCG）**来生成矫正性 rollout：它构建阶段特定的潜在目标，仅在发生漂移但可恢复的区域内应用引导，并通过估计的清洁动作进行引导，以匹配动力学模型的训练分布。成功的引导式 rollout 随后通过**ReGuide-FT**（微调当前检查点）或**ReGuide-FS**（在增强数据集上从头训练）吸收回策略中；两者也可以组合并迭代。在 Robomimic Can、Square、Transport 和 Tool Hang 任务上，ReGuide 将基础策略成功率提升了1.3–7.7倍，在仅测试时设置中优于 LPB，并且匹配数据量的消融实验表明，性能提升来自引导恢复数据，而非仅仅是额外的 rollout。

> 关键词：扩散策略，测试时引导，模仿学习

## 1 引言

基于模仿学习的策略容易受到协变量偏移的影响：小误差使机器人偏离演示状态，之后策略在训练分布之外的观测上被查询，误差随时间累积[21 (https://arxiv.org/html/2606.28939#bib.bib2)]。这个问题对于长周期操作尤为严重，早期偏差可能使后续子任务无法恢复。扩散策略通过建模多模态动作分布[3 (https://arxiv.org/html/2606.28939#bib.bib4)]提高了表达能力，但并未消除这种分布偏移问题。常见的解决方案是扩展训练分布。DAgger 类方法在学习者访问过的状态下收集专家修正[21 (https://arxiv.org/html/2606.28939#bib.bib2),11 (https://arxiv.org/html/2606.28939#bib.bib10),26 (https://arxiv.org/html/2606.28939#bib.bib3)]，而数据增强方法利用任务结构、模拟或局部动力学假设来合成额外示例[13 (https://arxiv.org/html/2606.28939#bib.bib11),1 (https://arxiv.org/html/2606.28939#bib.bib12),12 (https://arxiv.org/html/2606.28939#bib.bib14)]。这些方法解决了正确的失败模式，但需要持续的专家访问或对扰动是否保持任务正确性的假设。相比之下，最近的测试时引导方法如 DynaGuide[5 (https://arxiv.org/html/2606.28939#bib.bib5)]和潜在策略屏障（LPB）[23 (https://arxiv.org/html/2606.28939#bib.bib1)]利用来自学习动力学模型的梯度来引导冻结的扩散策略，无需重新训练或新的演示即可提高恢复能力。然而，这些方法在执行后丢弃了引导轨迹。

在本文中，我们提出一个简单但非常有效的解决方案：**将引导式 rollout 用作迭代自我改进的训练数据**。成功的引导式 rollout 精确捕捉了原始演示中缺失的信息：学习策略容易在哪里漂移，以及哪些矫正动作能使系统回到任务完成状态。回收此类 rollout 仅利用轨迹级别的成功信号，而非像 Dagger 类算法中使用的昂贵的逐状态专家重新标注，从而提供在策略恢复数据。

虽然这个想法说起来简单，但将引导式 rollout 转化为有用的训练数据需要几个技术选择，以防止自生成数据降低策略性能。引导式 rollout 仅在引导信号可靠时才有用。长周期任务具有阶段结构，单一全局目标集可能将 rollout 拉向错误阶段或压制有效行为模式。动力学梯度也仅在数据分布附近可靠：在接近专家流形时应用引导可能扰动正确动作，而在远离时应用引导可能依赖于外推预测。最后，标准扩散引导通过带噪声的去噪迭代进行微分，而动力学模型通常是在清洁动作块上训练的。

我们引入了 ReGuide，这是一个自我改进的扩散策略框架，通过控制引导式 rollout 的生成、过滤和吸收回策略的方式，将上述想法转化为实用算法。ReGuide 首先使用**阶段条件引导（PCG）**来生成可靠的矫正性 rollout：PCG 从演示中构建阶段特定的潜在目标，通过每个阶段多个目标保持多模态，将引导门控至漂移但可恢复的区域，并按照 MPGD[9 (https://arxiv.org/html/2606.28939#bib.bib7)] 通过估计的清洁动作应用引导梯度。成功的引导式 rollout 随后通过两种互补的更新机制回收用于策略训练：**ReGuide-FT**，使用演示和引导式 rollout 的排练式混合对当前检查点进行微调，以及 **ReGuide-FS**，在增强数据集上重新训练新策略。关键的是，更新后的策略可以在相同引导机制下再次 rollout，以生成新一批更高质量的恢复数据，从而形成迭代的 rollout-收集-训练循环，无需额外专家演示即可改进策略，直到自生成数据达到其性能上限。

在 Robomimic 任务[17 (https://arxiv.org/html/2606.28939#bib.bib23)]上，ReGuide 将基础策略成功率提升了 1.3–7.7 倍。PCG 在仅测试时设置中优于 LPB，匹配数据量的消融实验表明，在相同数据量下，引导式 rollout 比非引导式 rollout 更有价值，并且第二次 ReGuide-FT 迭代进一步提高了性能，表明引导式 rollout 生成能够支持迭代自我改进。

## 2 相关工作

**扩散策略与测试时引导。** 我们以 Diffusion Policy[3 (https://arxiv.org/html/2606.28939#bib.bib4)] 为基础视觉运动策略。最近的工作在推理时引导预训练的扩散策略而不修改其权重：DynaGuide[5 (https://arxiv.org/html/2606.28939#bib.bib5)] 使用外部动力学模型的梯度，潜在策略屏障（LPB）[23 (https://arxiv.org/html/2606.28939#bib.bib1)] 将专家潜在嵌入视为隐式分布内屏障，PPGuide[24 (https://arxiv.org/html/2606.28939#bib.bib9)] 学习性能预测器进行引导，生成预测控制（GPC）[19 (https://arxiv.org/html/2606.28939#bib.bib6)] 执行基于模型的超前规划。这些方法改善了执行时行为，但引导轨迹在每集之后被丢弃。ReGuide 则利用引导作为数据生成机制：成功的引导式 rollout 被添加回训练集以改进策略。ReGuide 在引导设计上也有所不同：我们将流形保持引导扩散[9 (https://arxiv.org/html/2606.28939#bib.bib7)] 调整为通过估计的清洁动作进行微分，并用阶段条件目标集和两阈值门控取代全局目标。LPB 是最接近的架构基线，因为它也在潜在空间中对扩散策略应用动力学梯度引导。

**模仿学习的自我改进与数据增强。** 行为克隆中的协变量偏移通常通过扩大学习者的训练分布来解决。交互式方法如 DAgger[21 (https://arxiv.org/html/2606.28939#bib.bib2)] 和 RaC[11 (https://arxiv.org/html/2606.28939#bib.bib10)] 在策略访问过的状态下收集新的监督，但需要持续的专家访问。其他方法无需进一步专家干预即可增强数据[13 (https://arxiv.org/html/2606.28939#bib.bib11),1 (https://arxiv.org/html/2606.28939#bib.bib12),12 (https://arxiv.org/html/2606.28939#bib.bib14),18 (https://arxiv.org/html/2606.28939#bib.bib13)]。最接近的是 CCIL[13 (https://arxiv.org/html/2606.28939#bib.bib11)]，它从演示中学习 Lipschitz 正则化动力学模型，并在演示状态附近生成矫正动作标签。ReGuide 在机制和循环结构上都有所不同：它使用测试时动力学梯度来恢复已经偏离演示的 rollout，过滤成功轨迹，并在得到的在策略恢复数据上迭代重新训练或微调策略。

**其他相关工作。** 更多相关工作讨论见附录 A (https://arxiv.org/html/2606.28939#A1)。

## 3 预备知识

**扩散策略。** 给定专家演示 $\mathcal{D}_{\text{expert}}=\{(o_t,a_t)_{t=1}^T\}$，其中 $o_t\in\mathcal{O}$ 和 $a_t\in\mathcal{A}$ 分别表示时间 $t$ 的观测和动作，模仿学习算法通过最小化监督损失来学习策略 $\pi_\theta:\mathcal{O}\rightarrow\mathcal{A}$。一个策略被称为 **扩散策略** [3 (https://arxiv.org/html/2606.28939#bib.bib4)] 当 $\pi_\theta$ 被参数化为一个条件去噪扩散模型，以过去观测为条件。遵循 Chi 等人 [3 (https://arxiv.org/html/2606.28939#bib.bib4)]，该策略以过去 $T_o$ 个观测的窗口 $O_t=(o_{t-T_o+1},...,o_t)$ 为条件，并生成一个长度为 $T_p$ 的未来动作块 $A_t=(a_t,a_{t+1},...,a_{t+T_p-1})$，执行前几个动作后重新规划。实际上，去噪器以低维潜在嵌入 $z_t=h(O_t)$ 为条件，而非原始观测窗口。

为了采样动作块 $A_t\sim\pi_\theta(\cdot\mid O_t)$，扩散策略从高斯噪声 $A_t^K\sim\mathcal{N}(0,I)$ 开始，并通过 $K$ 个逆向扩散步骤迭代去噪。在每一步 $k$，学习的噪声预测器 $\epsilon_\theta(A_t^k,k,z_t)$ 产生逆向更新：

$$A_t^{k-1} = \frac{1}{\sqrt{\alpha_k}}\left(A_t^k - \frac{1-\alpha_k}{\sqrt{1-\bar{\alpha}_k}} \epsilon_\theta(A_t^k,k,z_t)\right) + \sigma_k \xi, \quad \xi\sim\mathcal{N}(0,I),$$ (1)

其中 $\alpha_k$、$\bar{\alpha}_k$ 和 $\sigma_k$ 由噪声调度决定，$\xi$ 在最后一步（$k=1$）设为零。从 $k=K$ 向下迭代到 $k=1$ 得到执行的动作块 $A_t=A_t^0$。

**扩散策略的引导。** 扩散模型允许在推理时通过可微奖励函数的梯度来偏置生成过程的自然机制。具体来说，当通过逆向扩散从 $x_K\sim\mathcal{N}(0,I)$ 到 $x_0$ 生成样本 $x$ 时，分类器引导 [4 (https://arxiv.org/html/2606.28939#bib.bib28)] 及其无分类器变体 [10 (https://arxiv.org/html/2606.28939#bib.bib29)] 在步骤 $k$ 修改预测的噪声为：

$$\hat{\epsilon}(x_k) = \epsilon_\theta(x_k) - \eta\sqrt{1-\bar{\alpha}_k} \nabla_{x_k} r(x_k),$$ (2)

其中 $r$ 是奖励（例如，目标类的分类器对数概率，生成图像的文本-图像保真度），$\eta$ 控制引导强度。梯度 $\nabla_{x_k} r(x_k)$ 将每个去噪步骤推向在 $r$ 下得分更高的样本，因此最终 $x_0$ 倾向于满足奖励准则，同时仍然位于由 $\epsilon_\theta$ 学习的数据流形上。这一原理已被用于在无需重新训练底层扩散模型的情况下，将图像生成引导至目标类、文本描述和美学特性。

最近的工作将这个想法从图像生成转移到表示为扩散模型的控制策略。在离线强化学习中，几种方法使用学习价值函数的梯度来引导扩散策略 [6 (https://arxiv.org/html/2606.28939#bib.bib30),14 (https://arxiv.org/html/2606.28939#bib.bib32),8 (https://arxiv.org/html/2606.28939#bib.bib31)]，将动作生成推向高价值行为。在模仿学习中，DynaGuide [5 (https://arxiv.org/html/2606.28939#bib.bib5)] 和 LPB [23 (https://arxiv.org/html/2606.28939#bib.bib1)] 则使用学习的动力学模型来保持 rollout 接近专家分布。一般形式类似于上述图像生成情况：在去噪步骤 $k$，噪声预测通过可微代价 $\delta$ 的梯度增强：

$$\hat{\epsilon}(A_t^k) = \epsilon_\theta(A_t^k) - \eta\sqrt{1-\bar{\alpha}_k} \nabla_{A_t^k} \delta(A_t^k),$$ (3)

其中 $x_k \leftrightarrow A_t^k$，奖励 $r$ 被代之以惩罚偏离专家流形的代价 $\delta$。直觉与之前相同：每个去噪步骤被推向预计后果在 $\delta$ 下得分良好的动作。

**动力学模型。** 动力学模型预测在给定动作下环境如何演变：给定潜在状态 $z_t$ 和候选动作块 $A_t$，它返回环境预期达到的潜在状态 $\hat{z}_{t+T_p}=d_\phi(z_t, A_t)$，其中 $z_t=h(o_t)$，$h$ 是与扩散策略共享的视觉编码器。遵循 [23 (https://arxiv.org/html/2606.28939#bib.bib1)]，我们使用代价函数 $\delta(A_t^k)=\|d_\phi(z_t, A_t^k)-z^*\|^2$，其中 $z^*$ 是目标潜在状态。梯度 $\nabla_{A_t^k}\delta$ 然后将每个去噪步骤推向预计后果接近 $z^*$ 的动作块，将引导信号从“这个动作看起来像专家的动作”转变为“这个动作会导致类似专家的未来状态”。

我们采用 DINO-WM 架构 [27 (https://arxiv.org/html/2606.28939#bib.bib8)] 用于 $d_\phi$，遵循 [23 (https://arxiv.org/html/2606.28939#bib.bib1)]。动力学模型与扩散策略共享其视觉编码器 $h$，以使潜在空间距离作为引导信号有意义，并在扩散策略生成的 rollout（包括成功和失败）上进行训练。

## 4 我们的方法：ReGuide

**算法 1** ReGuide
1: 演示数据 $\mathcal{D}_{\text{demo}}$，迭代次数 $N$
2: $\mathcal{D}_1 \leftarrow \mathcal{D}_{\text{demo}}$
3: 在 $\mathcal{D}_1$ 上训练 $\pi_1$
4: 通过算法 3 (https://arxiv.org/html/2606.28939#alg3) 构建阶段目标和距离分布
5: for $i=1$ to $N$ do
6:     收集成功 rollout $\mathcal{D}_i^g$