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）作为损失函数，来衡量网络的预测值