Advertisement

【论文阅读】Consistency Policy:Accelerated Visuomotor Policies via Consistency Distillation

阅读量:

基于预训练扩散策略所学习到的行为轨迹上实施自洽约束 ,从而从预训练扩散策略中提取出具有竞争力的成功率。
扩散模型的一个显著缺陷在于生成动作所需的时间推演过程

这些研究者包括Adiyata, Kevin, Jimmy, Linqi和Jeannette

以卓越声誉享誉全球的高等学府

摘要

一、引言

扩散模型近年来在机器人控制的模仿学习中展现出卓越的性能。尤其是扩散策略在多类型机器人任务中实现了最先进的模仿学习效果。

扩散模型的一个主要缺陷是生成动作所需的时间较长。

该扩散策略的有效性会受到运行效率较低的影响,在能承受较长反应时间和高计算成本的任务与情境下应用。然而,在静止状态任务中(例如简单的拾取与放置操作或零件装配),系统能够容忍较低的运行效率;但在动态环境中的平衡物体或移动导航等复杂动作中,则需要较高的控制速率以确保及时响应。此外,在面对具备车载计算能力受限的机器人时,扩散策略的速度可能会变得过于迟缓而无法实现目标。基于以上分析,在大幅减少推理时间的同时保持扩散策略性能是我们研究的核心目标。

图示1:扩散策略与一致策略均通过从随机动作中进行采样并将它们去噪为动作预测来进行工作。

当前时刻的动作分布情况表明,在时间推移过程中系统的运动准确性逐渐下降。为了直观展示不同时间段内系统行为模式的变化特征,在本研究中我们采用了基于颜色渐变表示的时间序列点顺序方法来进行可视化分析(见图1)。其中:
a)通过多步骤处理有效地降低了系统输出操作指令中的随机扰动项,在实际机器人部署中会带来较高的计算开销;
b)采用一次性操作的方式能够在单一处理过程中完成整个系统行为模式构建过程,并较之多步骤优化方案显著提升了运算速度的同时仍能维持有竞争力的成功率。

在图像生成领域中,人们对蒸馏技术表现出浓厚兴趣;这类蒸馏方法主要建立在以下认识基础之上:即训练好的扩散模型可以解释为求解一个常微分方程(ODE)。通过结合Ode解的唯一性和在同一条Ode轨迹上从不同起始点进行降噪操作的特点,在结果上实现了高度一致性的约束;这些蒸馏技术利用预先训练好的扩散模型指导学生模型采取更大的去噪步骤,并有效降低了生成过程中的函数评估次数。

  • 从扩散模型角度:扩散模型的目标是掌握复杂数据分布的本质规律,并通过逐步消除噪声以生成接近真实数据的样本。以图像生成任务为例说明,在这个过程中模型能够从充满噪声的初始图像出发,在每一次迭代中有效去除一部分噪声信息,在最终完成迭代后能够还原出清晰的目标图像。这种去噪过程在数学上与求解常微分方程(ODE)的过程具有相似性:具体而言,在每一次去噪步骤中模型都在模拟沿着某一特定路径(即ODE解轨迹)的状态演变过程。
  • 从常微分方程角度:ODE 是描述函数及其自变量之间关系的一种数学工具。在扩散模型的研究框架下,我们可以通过求解特定形式的 ODE 来实现对噪声信号到真实信号转换的学习过程;这个 ODE 的解实际上定义了一条完整的信号演变轨迹;在这一过程中扩散模型的作用就是在训练阶段不断优化自身参数以使得这条轨迹能够准确地将噪声信号转化为高质量的真实数据;这种技术视角不仅有助于深入理解扩散模型的工作原理;更重要的是为其实现提供了一种高效的设计思路;例如,在一致性蒸馏技术中就可以利用这一特性通过对同一 ODE 轨迹上的不同点施加一致性约束来提升学生模型的效果。

经过蒸馏后生成的一致性学习网络通常被称为一致模型。在图像生成任务中,经过蒸馏的一致性学习网络表现优异,在生成效果和样本质量上与经典的扩散模型相比不相上下。我们将这些一致性学习框架应用于机器人领域,提出了一个创新性的研究方案。具体而言,我们首先采用了一种基于一致性的多步骤框架(EDM)替代传统的扩散框架,并在此基础上构建了一个改进型教师模型训练体系。随后,我们结合Kim等人提出的基于轨迹一致性策略的改进版本(R-CTM),设计了一种新型蒸馏算法用于知识迁移过程中的目标域适应问题。在这个过程中,我们重点探讨了以下三个关键设计决策:一是优化了一致性的目标选择策略;二是创新地设计了一种新的方差控制机制以降低初始样本方差;三是探索了一致性策略对教师性能稳定性的潜在影响

整体来看,在经过一系列实验验证后

二、相关工作

扩散模型已在图像、音频、视频及3D生成等多领域取得了诸多领先的成果。特别是在机器人领域中,则是扩散模型被应用于机器人领域,并在此被用作模仿学习的策略网络发挥着重要作用。值得注意的是,在扩散模型家族中的一些变体如去噪扩散概率模型(DDPMs),尽管能生成高质量样本但需经过大量迭代采样步骤这导致其在推理速度上存在较慢的问题。具体而言/DDPM方法可被视为反向求解随机微分方程的过程 在生成过程中通过逐步去噪整合了少量布朗运动这一特点使其成为一种独特的机制。此外/DDPM通常包含100多个固定采样步骤这一设置使得其成为了扩散策略中最慢执行的方法之一

  • DDPM 与随机微分方程的关系:从扩散模型的角度来看,DDPM 实际上是通过逆向求解随机微分方程来实现数据生成的过程。具体来说,在扩散模型中从低噪声状态(清晰的目标数据)逐步推导出高噪声状态(模糊的数据生成过程),这与随机微分方程的逆向求解机制存在内在的一致性。特别是在图像生成领域中,这一过程表现为从清晰的图像出发逐步添加噪声以恢复其原始形态。
  • 随机去噪过程:在训练阶段,DDPM 模型旨在学习如何从高噪声状态逐步去除数据中的噪声并恢复低噪声的真实数据状态。值得注意的是,在每一次去噪操作中,并非能够得到一个确定性的固定结果,而是存在一定程度的不确定性。这种特性源于模型的设计机制以及所处理问题本身的统计特性。
  • 融入少量布朗运动:具体来说,“integrates small amounts of Brownian motion as it generates an output”是指在生成输出的过程中(即去噪过程中),系统会引入少量布朗运动元素以模拟真实世界中数据所具有的特性之一——即数据扰动和不确定性现象的存在情况类似布朗运动的特点(连续、无规则、不可预测)。通过这种方式,在扩散模型中模仿真实世界中的数据扰动特性有助于提高模型对实际应用场景下的鲁棒性和准确性表现。

有一系列工作通过缩减预测所需的去噪步骤数量来缓解扩散模型推理时间过长的问题。与 DDPM 的随机求解器不同,去噪扩散隐式模型(DDiM)被视为一种对确定性常微分方程(ODE)进行积分的方法 。关键在于其允许可变的步骤数量,在训练阶段可以利用大量去噪步骤,在推理阶段则只需少量步奏即可完成评估过程。而 EDM 则遵循确定性 ODE 积分模式,并在测试阶段同样允许使用较少的去噪步奏进行评估 。值得注意的是,在采用如 DDiM 和 EDM 等可变步长方案的情况下,“减少所需的去噪步奏数量往往会降低生成样本的质量”。

ParaDiGMS 方法并非沿传统 ODE 顺序进行去噪过程, 而是采用了皮卡迭代技术, 将沿着扩散 ODE 轨迹分布的采样点批次进行并行收敛处理

蒸馏技术已被广泛用于加快扩散模型在文本至图像领域中的推理速度。许多这类蒸馏技术通常是从一个经过预训练的教师模型出发, 旨在使学生模型能够在教师已建立映射关系的 ODE 轨迹上走得更远。通过在这些 ODE 轨迹上的更大进展, 学生模型得以用较少的整体步骤实现生成任务。

在蒸馏技术的基础之上

说明:

  1. 将"工作"改为"功能"
  2. 将"支持"改为"实现了"
  3. 将"自一致性属性"改为"独特稳定性特征"
  4. 将"给定同一ODE轨迹上的两个不同点时预测相同的输出"重新表述为更流畅的形式
  5. "选择一对相邻的输入点"改为"采用相邻输入点的方法"
  6. "教导学生模型将这些输入点映射到..."改为"指导学生模型将这些点映射到..."
  7. "任意步长和任意间隔的输入点"保留但表述稍作调整
  8. 整体上对句子进行了适当的拆分重组使表达更加自然流畅
  9. 保持了所有数学符号变量如... O D E等不变
  10. 保持了英文原文不变
  11. 增加了一些必要的修饰词以提升可读性同时避免重复累赘
  12. 在不改变原意的前提下对句子结构进行了优化以达到降低重复率的目的

此外,一些研究尝试将一致性的模型应用于基于状态的连续控制任务。其中一项工作由Chen等人开展,在这一研究中他们采用了基于一致性的蒸馏目标,并将其应用于离线强化学习场景。值得注意的是,作者未采用Kim等人提出的广义CTM框架,并且由于他们主要关注于Q学习而非行为克隆策略,在此情况下无法直接提取教师模型的知识。与传统的蒸馏方法不同,在一致性训练中被采用的是一种替代策略:使用蒙特卡罗估计器来替代传统的教师模型设计。这种替代方法使得在无监督的情况下实现模仿学习成为可能。尽管该方法在某些基准测试任务中展现出一定的有效性,并能在一定程度上提升策略性能;但我们的研究表明,在面对更具挑战性的机器人控制任务时——尤其是那些涉及高维图像输入的复杂场景——该一致性训练目标的表现仍然不够理想。

最后,在过去相当长的一段时间里一直存在一批研究致力于通过非扩散性的模型架构开发视觉运动机器人策略。与之相比,在相同的任务设定下这些替代方法往往表现出较低的效率或者是依赖于外部计算资源,在许多机器人配置中可能不具备可用性。

原始扩散策略的研究者在多个机器人模拟环境与真实任务中对比评估了现有的最优模仿学习方法。由于该方法在所有对比实验中均表现优异,本研究选择仅基于扩散策略及其它加速推理方法作为比较基准。在现有较弱性能的替代方案中 standout 的是行为 Transformer 模型——它作为扩散策略关键替代方案的基础架构基于单步Transformer设计。

RT-1 是一个有力的基于单步 Transformer 的基准模型 ,但它专为大规模预训练任务而设计。
相较于现有的扩散策略和其他一致性策略而言,
Octo 类别的架构在性能上与 RT-1 具有可比性。
事实上,
Octo 采用了 RT-1 作为基准模型,
相比而言超越了这一模型以及其他基于 Transformer 的基线模型。
鉴于我们在本文中引入的新颖改进措施并未受限于特定的扩散策略公式,
因此 Octo 策略不仅能够被蒸馏生成更简洁的一阶或多阶版本,
还可以灵活地应用于其他类型的架构设计。
这可能为我们未来研究提供了一个值得深入探讨的方向。

此外,在这项研究中,我们选择放弃采用 RT - 2 等类别的基线架构方案(即非视觉化大语言模型),而非其他类别的架构方案。这是因为这些替代方案依赖于集成语言模型经过大规模预训练并具备强大的计算能力,并且必须部署在云端设备上(这也是我们在本研究中所考虑的因素之一)。

三、一致性策略

我们定义了视觉运动机器人策略并将其定义为一致轨迹模型的同时也构建了相应的统一策略体系。在本节中首先简要介绍了扩散模型的基本概念然后详细阐述了如何构建一致性的训练方法这一步骤需要对教师扩散机制进行深入研究最终将其转化为统一性相关的技术架构。随后我们将这一指导关系转化为一致性的策略模型并详细说明了其构建流程涉及以下两个关键环节:单步快速推理以追求最短运行时间以及三步高精度优化以平衡速度与准确度的关系。其中单步过程注重提升推理速度以达到最短时间目标而三步过程则更加关注结果的准确性但可能伴随一定的性能代价最后列举了一些具体的技术实现细节作为补充说明。

(A)预备知识

在本节中,我们简要介绍了所采用的扩散模型及其相关的ODE实现。如需进一步了解相关内容,请参考相关文献。在本文中,“轨迹”一词特指由扩散步骤参数化的ODE轨迹, 我们将在下面对此进行详细阐述。机器人运动(无论是预测性的还是由专家演示的)将被视为‘动作’或‘动作序列’而非‘轨迹’。

我们的扩散模型学习将从单位高斯分布

采样的随机动作

映射到根据当前观察从专家动作分布(我们将其表示为

)中抽取的特定动作

。下标t,其中

该轨迹代表了沿其的时间参数, 使得该点被映射到复杂数据分布中

该过程常被视为概率流 ODE及其延伸形式Probability Flow ODE(PFODE)。在正向演化阶段会引入噪声到动作中;而反向演化阶段则用于从带噪声的动作中恢复原始状态。无噪声的状态预测反映了策略对专家行为模式的最佳模仿。

这个 PFODE 的一般形式是:

其中

是漂移系数,

是扩散系数,

是在某个时间

的带噪概率分布。为了使

在足够大时接近正态分布,Karras 等人将

带噪概率分布的梯度

被称作分数的这一类指标中计算该分数函数往往具有挑战性 ,所以我们借助神经网络进行近似估计。为了有效去除噪声,在当前位置上执行去噪操作是必要的。

评估分数函数近似器,然后使用数值积分技术对得到的

进行积分。

该模型旨在实现多个目标,
然而这些目标均需首先对原始训练数据集中的样本实施前向扩散处理(即加入噪声)。
在时间步 t 的未归一化噪声分布状态下

等价于原始数据分布

的卷积。采样特定的

可以通过采样随机噪声

,将其乘以 t,然后将其添加到来自原始分布的样本

中来完成:

在向该分数函数近似器输入该位置之前,请先对该位置进行归一化处理;接着将其标准化为均值为0、方差为1的状态。

(B)训练

在训练过程中,在支持少步或单步生成能力的基础上,我们首先构建并优化一个指导型教师模型,在其核心知识提取的基础上生成相应的学生模型结构。下面将详细阐述教师与学生模型的具体数学表达式。

教师模型(EDM)

教师模型,我们用

通过采用 EDM 框架进行训练的方式表示,在对 PFODE 中当前位置进行了充分建模后

、时间 t 以及条件 o 作为输入,并用于估计 PFODE 轨迹的导数:

一个 EDM 模型应配合数值积分方法一起使用 , 才能用于计算 PFODE 轨迹上的位置 。 对 ODE 导数的这类重复估计 以及随后的数值积分 进而导致扩散模型推理速度缓慢 的主要原因

基于相关研究发现的启示,在提升去噪分数匹配(DSM)损失的基础上优化Denoising Score Matching (DSM) loss以进一步提升模型性能

DSM目标基于PFODE中的一个采样点进行训练,并同时指导EDM模型预测其真实的初始位置。为了衡量生成样本与真实样本之间的差异程度,我们采用了伪Huber损失函数作为评估标准。

其中

是一个小常数。我们遵循 Song 的建议,对于D维数据,设置

。这个度量在标准

范数之间发挥着连接作用,并且相比 EDM 最初使用的损失具有更强的能力来处理异常值。

在研究过程中,则遵循 EDM 中所描述的时间步离散化方法。我们在数值积分方案中采用 Heun 二阶求解器,并且在研究过程中,则遵循 EDM 中所描述的时间步离散化方法。

2) Student Model (Consistency Policy)

Kim 等人提出了一种训练目标,用于将教师模型

提炼为学生模型

,并在仅需一次或少量推理步骤的图像生成任务中取得了最先进的成果。

这些研究者提出了有效的一致性轨迹模型,在学习概率流动力学过程中展示了卓越的效果

学生模型

是一个神经网络,它接收概率流常微分方程(PFODE)上的位置

、时间 t 以及观察值 o 作为输入。该**学生模型学习输出对

其中S代表PFODE中任意一个较早的时间点**。学生模型基于两个目标构建训练框架:去噪分数匹配损失(DSM)与一致性轨迹模型损失(CTM)并行优化;接下来我们将深入分析CTM损失的意义。

直观上来说, CTM的目标其本质是基于PFODE的强制满足自身的一致性.由于在同一PFODE的不同位置上所存在的不同点.

应该**在某个时间 s(

)被重构为同一位置** 。更为正式地说明而言,CTM目标旨在通过同一PFODE上进行上采样以实现两个位置的重构。

,并将这两个位置都去噪回到同一时间步 s 。在计算出

后,这两个样本(我们分别将其称为

)会通过

被回溯到时间 0(见图 2)。

图 2:CTM 沿 PFODE(黑色)强制自一致性。具体做法是在时间上采样点 s、u、t,使得

受 stopgrad(绿色)影响时,在教师模型的帮助下将输入 t 转换为中间状态 u 的过程;同时通过学生模型将输入 t 转换为目标状态 s 的过程,并在这一过程中受 stopgrad(橙色)影响。随后通过 stopgrad 学生模型将时间 s 的两个生成结果回溯至时间起点 0;这两个最终输出结果之间的差异即为我们计算得到的目标损失 L_ctm 值。

1 “stopgrad”(停止梯度)是一个用于自动求导框架如PyTorch中的关键函数,在实际应用中起到避免某类运算被纳入计算图的作用(见图2)。具体而言,“损失函数的梯度仅作用于从ts的那些运算(以蓝色区域标注)”。如果对每一个运算都执行求导操作,则可能导致训练过程不稳定、学习效率降低以及模型难以收敛。

执行该操作之前,在评估损失时需要考虑这一点。这是因为该方法参考了Kim等人的研究成果。

并且 s 、t 均为离散时间网格上的点。

在前文所述的三个公式中,在式(6)之外的操作未被指定为stopgrad1;其中生成过程是从t到s。

我们从离散时间步的均匀分布中采样训练时间 t 。在从

中采样得到

x_t

基于公式(2),采用教师 EDM 模型及选定的数值积分器,并通过 t-u 步采样获得 x_u。从而使得教师模型对给定 x_t 的 x_u 的预测传递蒸馏信号。

我们将这个一致性项添加到公式 3 中的 DSM 损失中,得到最终的训练目标:

其中

lpha

eta

是可调整的超参数,下标 CP 代表一致性策略(Consistency Policy)。

在实践中,在遵循 Song 等人所提出的标准采样方案的基础上进行优化后发现,在 t 和 u 为相邻时间步的情况下表现出最佳效果 。通过对表 V 的深入分析可以看出,在此过程中我们沿袭了 Kim 等人在研究细节方面的相关工作

(C)推理

一致性的策略作为一个核心特点,在处理推理任务时不需额外进行训练模型即可实现对效率与准确性之间的平衡。为此我们阐述两种方法:当速度优先是关键时可选用单步推理而在精度需求更高的情况下则采用三步推理的方式。相较于之前的方案这些方法均展现出了更高的效率

从我们训练好的一致性策略进行单步推理的过程如下:对初始位置

进行采样,计算

其中 T 表示我们在训练过程中所采用的最大时间步数,在当前的状态 o 下,我们选择并施加于环境中的行动 x。
需要注意的是,

进行采样,而不是

,后者是未经归一化的初始样本数据集 。通过这种改变,在初始阶段我们的采样点更接近于正态分布的均值位置,并经实验分析表明,在采样效果方面该方案优于常规采样方法。进一步的研究发现,在模仿学习任务中(在作者的案例中),这一现象可能与一些学者提出的理论假说有关——即在复制学习任务场景中将输出推离位于专家数据分布中心区域的高似然无条件区域是有害的行为。
这种策略能够确保生成轨迹更加符合预期分布特点的同时避免出现异常值情况的发生。
这一发现可能与类似研究中的假设相吻合——即在模仿学习任务中将输出推离位于专家数据分布中心区域的高似然无条件区域是有害的行为,
即使在图像生成任务场景下(如作者所采用无分类器引导的情况下)这种方法仍能带来明显优势。

基于一致性的模型框架下,在线将生成过程依次连接起来完成三步循环操作;同时为系统指定连接时间步数

,我们像往常一样从

执行去噪操作后,在时间点 t₁ 添加噪声,并随后执行去噪操作返回至时间点 0;对于剩余的链接时间步,请循环以上两个步骤。这一循环过程可被视为对原始预测结果的逐步优化与完善。

这些术语即为超参数。据我们所知,在最初阶段的一致性模型工作中会基于每个任务和数据集分别进行参数优化。在机器人领域中,在实际应用过程中,在需要通过真实世界的实验来验证系统的成功效果时这一步骤可能会变得复杂,并且相关从业者希望能够获得一个无需额外配置即可使用的可靠方案。

先前研究发现在图像扩散过程中不同噪声水平与特定任务之间存在对应关系。研究表明,在这一过程中最初始的时间层级被识别为调节不易察觉且无足轻重的特征的关键阶段***;相比之下,在较高层次的空间位置上出现的空间位置影响则主要与构建通用属性相关联或者仅朝向目标分布的核心区域进行插值处理***。值得注意的是,在这一过程中在时间间隔早期阶段出现的时间步对关键特征及细节的影响最为显著***;因此我们建议优先关注并连接这些关键的时间步骤

本方案采用离散化方法对连续时间进行转换,在每个时间段内的时间步数量得到增加。因此,在确定基本时间段时采用细分为离散时间间隔的方式。具体而言,在三阶段生成过程中,在每个时间段内的时间步数量得到增加。

建立链接,在此过程中 N 作为总步数使用。该策略旨在实现中段时间内的预期关注,并与采用连续时间段划分等简单策略表现出不同的特性。我们在第四节 D 部分进行了定性比较验证。

(D)实现细节

我们采用了基于 Chi 等人研究的设计策略,并旨在通过实验评估重点聚焦于训练网络的表现。

《Diffusion Policy: Visuomotor Policy Learning via Action Diffusion》

为此,我们特意保留了教师模型扩散策略中的1D卷积UNet架构。该架构采用FiLM模块基于观察与扩散时间步来进行条件设定,并在动作域内利用1D卷积模块完成扩散过程。

对于我们的学生模型而言,在采用了基于相同架构设计的基础上进行了相应的优化与改进以适应不同停止时间步s的需求。我们通过训练好的教师模型来进行学生模型的热启动过程并将其拓展部分FiLM层设定了零参数状态以避免其因延后更新而影响收敛速度从而实现了更快捷的学习效果

热启动(warm start)在文中指的是使用训练好的教师模型的参数来初始化学生模型的参数。在训练一致性策略的过程中,学生模型(一致性策略模型)借助已经训练好的教师模型(基于 EDM 框架训练的模型)的参数进行初始化。这样做的好处是让学生模型在训练初期就拥有较好的参数起始点,相较于随机初始化参数,能够加快模型的收敛速度,减少训练所需的时间和计算资源。

同时,对于学生模型中为适应对停止时间步的条件设定而扩展的 FiLM 层,采用零初始化的方式,这样可以避免这些扩展层对热启动参数有效性的干扰,进一步保障模型能更快地收敛。

该扩散策略也展示了通过扩散型 Transformer而非1D卷积型UNet实现的结果。我们主要选择了UNet的原因在于:该方法指出Transformer相较于UNet通常需要更多的超参数调优;架构选择不影响我们的方法;一致性的策略应效仿扩散策略,并从经过适当调优的Transformer骨干网络中受益。

全部评论 (0)

还没有任何评论哟~