Diffusion Policy:基于扩散模型的机器人动作生成策略

项目地址：

Diffusion Policy (columbia.edu)

一、摘要

本文介绍了 "扩散策略"，这是一种生成机器人行为的新方法，它将机器人的视觉运动策略（visuomotor policy）表示为条件去噪扩散过程（conditional denoising diffusion process）。在 4 个不同的机器人操纵基准中的 15 个不同任务中对 Diffusion Policy 进行了基准测试，发现它始终优于现有的最先进的机器人学习方法，平均提高了 46.9%。Diffusion Policy 可以学习动作分布得分函数的梯度（gradient of the action-distribution score），并在推理过程中通过一系列随机朗之文动力学（stochastic Langevin dynamics）步骤对该梯度场进行迭代优化。我们发现，扩散公式在用于机器人策略时具有强大的优势，包括可以优雅地处理多模态动作分布、适用于高维动作空间，以及表现出令人印象深刻的训练稳定性。为了充分释放扩散模型在物理机器人视觉运动策略学习方面的潜力，本文提出了一系列关键技术贡献，包括递减视界控制receding horizon control（MPC控制）、视觉调节（visual conditioning）和时间序列扩散（time-series diffusion）transformer。我们希望这项工作将有助于推动新一代策略学习技术的发展，从而充分利用扩散模型强大的生成建模能力。

图 1. a) 具有不同行动表示类型的显式策略。 b) 隐式策略学习以行动和观测为条件的能量函数，并优化使能量最小化的行动。 c) 扩散策略通过学习的梯度场将噪声细化为行动。这种方法能提供稳定的训练，使学习到的策略能准确模拟多模态动作分布，并适应高维动作序列。

二、介绍

从演示中学习策略，其最简单的形式可以表述为学习将观察结果映射到行动的监督回归任务。然而，在实践中，预测机器人动作的独特性--如存在多模态分布、连续相关性和高精度要求--使得这项任务与其他监督学习问题相比具有独特性和挑战性。之前的研究试图通过探索不同的动作表示法（图 1 a）来应对这一挑战--使用高斯混合、量化动作的分类表示法，或者将策略表示法（图 1 b）从显式转换为隐式，以更好地捕捉多模态分布。在这项工作中，我们试图通过引入一种新形式的机器人视觉运动策略来应对这一挑战，即通过 "机器人动作空间上的条件去噪扩散过程"（Diffusion Policy）来生成行为。在这种方法中，策略不是直接输出动作，而是根据视觉观察结果，在 K 次去噪迭代中推导出动作分数梯度（图 1 c）。这种表述方式让机器人策略继承了扩散模型的几个关键特性，大大提高了性能。

- 多模态动作分布。通过学习动作得分函数的梯度，并在此梯度场上执行随机朗文动力学采样，扩散策略可以表达任意可归一化分布，其中包括多模态动作分布，这是策略学习面临的一个著名挑战。

- 高维输出空间。正如其令人印象深刻的图像生成结果所证明的那样，扩散模型在高维输出空间中表现出卓越的可扩展性。这一特性允许策略联合推断一系列未来行动，而不是单步行动，这对于鼓励时间行动一致性和避免近视规划至关重要。

- 稳定的训练。训练基于能量的策略通常需要负采样来估计一个难以处理的归一化常数，而这已知会导致训练的不稳定性。扩散策略通过学习能量函数的梯度绕过了这一要求，从而在保持分布表达性的同时实现了稳定的训练。

我们的主要贡献是将上述优势引入机器人领域，并在复杂的真实世界机器人操纵任务中展示其有效性。为了成功地将扩散模型用于视觉运动策略学习，我们提出了以下技术贡献，以提高扩散策略的性能，释放其在物理机器人上的全部潜力：

- 闭环动作序列（Closed-loop action sequences）。我们将该策略预测高维动作序列的能力与MPC控制相结合，以实现稳健的执行。这种设计允许策略以闭环方式不断重新规划其行动，同时保持时间行动的一致性，在长视距规划和响应速度之间取得平衡。

- 视觉调节（Visual conditioning）。我们引入了一种视觉条件扩散策略，在这种策略中，视觉观察结果被视为条件，而不是联合数据分布的一部分。在这种方法中，无论去噪迭代多少次，该策略都能提取一次视觉表征，从而大大减少了计算量，实现了实时行动推理。

- 时间序列扩散变换器（Time-series diffusion transformer）。我们提出了一种新的基于transformer的扩散网络，它能最大限度地减少典型的基于 CNN 模型的过度平滑效应，并在需要高频率动作变化和速度控制的任务中实现最先进的性能。

三、扩散策略

我们将视觉运动机器人策略制定为去噪扩散概率模型（DDPMs）。最重要的是，扩散策略能够表达复杂的多模态动作分布，并具有稳定的训练行为--几乎不需要针对特定任务进行超参数调整。下文将更详细地介绍 DDPMs，并解释如何将其用于表示视觉运动策略。

3.1 去噪扩散概率模型

DDPM 是一类生成模型，其输出生成被模拟为一个去噪过程，通常称为随机朗文动力学（Stochastic Langevin Dynamics）。

从高斯噪声采样的 xK 开始，DDPM 执行 K 次迭代去噪，产生一系列噪声水平递减的中间动作 xk、xk-1...x0，直到形成所需的无噪声输出 x0。该过程遵循公式

其中，εθ 是噪声预测网络，将通过学习进行优化，N(0, σ2)是每次迭代时添加的高斯噪声。上述公式 1 也可以解释为一个单一的噪声梯度下降步骤：

其中噪声预测网络 εθ (x, k) 有效地预测了梯度场 ∇E(x)，γ 是学习率。作为迭代步数 k 的函数，α、γ、σ 的选择也称为噪声调度(noise schedule)，可以解释为梯度赋值过程中的学习率调度。事实证明，α 略小于 1 可以提高稳定性。

3.2 DDPM训练

训练过程的第一步是从数据集中随机抽取未经修改的样本 x0。对于每个样本，我们随机选择一个去噪迭代 k，然后为迭代 k 采样一个具有适当方差的随机噪声 εk。噪声预测网络需要从添加了噪声的数据样本中预测噪声。

使用均方误差（MSE, Mean Squared Error）作为损失函数，来衡量网络的预测值 𝜖𝜃(𝑥0+𝜖𝑘,𝑘) 和实际噪声 𝜖𝑘 之间的差异。

图 2：扩散策略概述 a) 一般表述。在时间步长 t，该策略将最新 To 步的观测数据 Ot 作为输入，并输出 Ta 步的行动 At。 b) 在基于 CNN 的扩散策略中，Perez 等人（2018 年）将观测特征 Ot 的 FiLM（Feature-wise Linear Modulation）调理应用于每个卷积层，按通道进行调理。c) 在基于transformer的 Vaswani 等人（2017）扩散策略中，观测特征 Ot 的嵌入被传递到每个变换器解码块的多头交叉注意层。利用图示的注意力掩码，每个动作嵌入只关注自身和之前的动作嵌入（因果注意力）。

3.3 视觉运动策略学习中的扩散

DDPM 通常用于生成图像（x 是一幅图像），而我们则使用 DDPM 来学习机器人视觉运动策略。这需要对公式进行两大修改： 1. 将输出 x 改为代表机器人动作。2. 使去噪过程以输入观测 Ot 为条件。以下段落将逐一讨论这些修改，图 2 显示了相关概述。

闭环行动序列预测：

有效的行动表述应鼓励长视距规划的时间一致性和平稳性，同时允许对意外观察做出迅速反应。为了实现这一目标，我们在重新规划之前，会在固定时间内使用扩散模型生成的行动序列预测。具体来说，在时间步长 t 处，策略将最新的 To 步观测数据 Ot 作为输入，并预测出 Tp 步行动，其中 Ta 步行动将在机器人身上执行，而无需重新规划。在这里，我们将 To 定义为观察期，Tp 定义为行动预测期，Ta 定义为行动执行期。这样既能保证行动的时间一致性，又能保持反应灵敏。我们还使用receding horizon control通过使用之前的行动序列预测来warm-starting下一个推理设置，从而进一步提高行动的平滑性。

视觉观测条件：

我们使用 DDPM 来近似条件分布 p(At |Ot ) 而不是 联合分布 p(At , Ot ) 。这种表述方式使模型能够根据观察结果预测行动，而无需付出推断未来状态的代价，从而加快了扩散过程，提高了生成行动的准确性。为了捕捉条件分布 p(At |Ot )，我们将公式 1 修改为

将公式 3 中的训练损失修改为：

从去噪过程的输出中排除观测特征 Ot 能显著提高推理速度，更好地适应实时控制。这也有助于使视觉编码器的端到端训练变得可行。

四、关键设计

在本节中，我们将介绍 "扩散策略 "的关键设计决策以及神经网络架构对 εθ 的具体实现。

4.1 网络架构选择

第一个设计决策是为εθ 选择神经网络架构。在这项工作中，我们研究了两种常见的网络架构类型：CNNs和Transformers，并比较了它们的性能和训练特性。

在实践中，我们发现基于 CNN 的骨干网络在大多数任务中都能很好地运行，无需进行过多的超参数调整。然而，当所需的动作序列随时间发生快速而剧烈的变化时（如速度指令动作空间），它的表现就会很差，这可能是由于时间卷积偏好低频信号的归纳偏差 。

在我们基于状态的实验中，大多数表现最佳的策略都是通过transformer主干实现的，尤其是在任务复杂度和行动变化率较高的情况下。然而，我们发现transformer对超参数更为敏感。未来有可能通过改进transformer训练技术或扩大数据规模来解决。

一般来说，我们建议从基于 CNN 的扩散策略实施开始，作为新任务的首次尝试。如果由于任务复杂或动作变化率高而导致性能低下，则可以使用时间序列扩散变换器（Time-series diffusion transformer）来提高性能，但需要付出额外的调整代价。

4.2 视觉编码器

视觉编码器将原始图像序列映射为潜在嵌入 Ot，并使用扩散策略进行端到端训练。不同的摄像机视角使用不同的编码器，每个时间步中的图像独立编码，然后串联形成 Ot。我们使用标准 ResNet-18（无预训练）作为编码器，并做了以下修改： 1) 用spatial softmax pooling取代全局平均池化（global average pooling），以保持空间信息。2) 用 GroupNorm 取代 BatchNorm ，以获得稳定的训练。当归一化层与指数移动平均（Exponential Moving Average）结合使用时，这一点非常重要 （常用于 DDPMs）。

4.3 Noise Schedule

噪声表由 σ、α、γ 和加性高斯噪声 εk 定义，是 k 的函数，Ho 等人（2020 年）、Nichol 和 Dhariwal（2021 年）对噪声表进行了积极研究。基本噪声表控制着扩散策略捕捉动作信号高频和低频特征的程度。在我们的控制任务中，我们根据经验发现，iDDPM中提出的平方余弦时间表（Square Cosine Schedule）最适合我们的任务。

4.4 加速推理

我们使用扩散过程作为机器人的策略；因此，对于闭环实时控制来说，拥有快速的推理速度至关重要。Song 等人（2021 年）提出的去噪扩散隐含模型（DDIM）方法将训练和推理中的去噪迭代次数分离开来，从而允许算法使用较少的推理迭代次数来加快推理过程。在我们的实际实验中，使用 100 次训练迭代和 10 次推理迭代的 DDIM，可以在 Nvidia 3080 GPU 上实现 0.1 秒的推理延迟。

五、局限性与未来工作

尽管我们在模拟和真实世界系统中都证明了扩散策略的有效性，但仍存在一些局限性，未来的工作可以加以改进。首先，我们的实现继承了行为克隆的局限性，例如在演示数据不足的情况下，我们的实现效果并不理想。扩散策略可应用于其他范式，如强化学习 Wang 等人（2023 年）；Hansen-Estruch 等人（2023 年），以利用次优数据和负面数据。其次，与 LSTM-GMM 等简单方法相比，扩散策略具有更高的计算成本和推理延迟。我们的动作序列预测方法部分缓解了这一问题，但对于需要高速率控制的任务来说可能还不够。未来的工作可以利用扩散模型加速方法的最新进展来减少所需的推理步骤，如新的噪声调度、推理求解器和一致性模型。

六、结论

在这项工作中，我们评估了基于扩散的机器人行为策略的可行性。通过对仿真和真实世界中的 15 项任务进行综合评估，我们证明了基于扩散的视觉运动策略能够持续、明确地优于现有方法，而且稳定、易于训练。我们的研究结果还强调了关键的设计因素，包括receding-horizon action prediction、末端执行器位置控制和高效视觉调节，这些因素对于释放基于扩散的策略的全部潜力至关重要。影响行为克隆策略最终质量的因素有很多，包括演示的质量和数量、机器人的物理能力、策略结构以及所使用的预训练机制，但我们的实验结果强烈表明，策略结构是行为克隆过程中的一个重要性能瓶颈。我们希望这项工作能推动该领域进一步探索基于扩散的策略，并强调在行为克隆过程中，除了考虑用于策略训练的数据外，还要考虑所有方面的因素。