Diffusion Policies
扩散策略通过对以观测为条件的随机噪声进行去噪来生成机器人动作,从而保留回归平均化会抹去的多模态行为。
扩散策略是一种机器人控制策略,其生成动作的方式与扩散模型生成图像类似:它从随机噪声开始,根据机器人当前的观察结果逐步去噪,直到剩下一系列电机指令。它根据人类演示进行训练,属于行为克隆的一种,可以表示完成任务的多种不同方式,而不是将它们平均成一种方式。Columbia University、Toyota Research Institute 和 MIT 的 Cheng Chi、Shuran Song 及其同事在 2023 年 RSS 上发表了 Diffusion Policy: Visuomotor Policy Learning via Action Diffusion,在 4 个机器人技术操作基准测试的 12 项任务中,平均超过早期模仿学习方法 46.9%。
回归为何无法处理多模态演示#
同一观察结果通常对应多种有效的人类演示动作。一个人推方块时,可能在一段演示中从左侧靠近,在下一段演示中从右侧靠近;驾驶员遇到障碍物时,也可能从任一侧通过。使用均方误差损失函数进行回归训练的策略会预测这些选择的平均值,而这个平均值可能并非任何演示者采取过的动作,例如直接开向障碍物。早期的改进方法会将动作离散化到不同区间、拟合高斯混合模型,或使用隐式行为克隆等基于能量的模型;后者需要在训练期间使用负样本,而且通常不稳定。扩散模型可以表示任意动作分布,而其噪声预测损失与普通回归一样稳定。
这个 PyTorch示例使用演示数据训练一个小型扩散策略;这些演示中,绕障碍物向左或向右转向的次数相同:
import torch
from torch import nn
torch.manual_seed(0)
# Demonstrations at one observation: facing an obstacle, experts steer left (-1) or right (+1) equally often
actions = torch.cat([torch.randn(500, 1) * 0.05 - 1, torch.randn(500, 1) * 0.05 + 1])
print(f"Regression policy (mean squared error) outputs: {actions.mean():.2f}")
# Diffusion policy: train a network to predict the noise mixed into an action at denoising step k
K, betas = 50, torch.linspace(1e-4, 0.2, 50)
alpha_bar = torch.cumprod(1 - betas, 0)
net = nn.Sequential(nn.Linear(2, 64), nn.SiLU(), nn.Linear(64, 64), nn.SiLU(), nn.Linear(64, 1))
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
for _ in range(3000):
k, noise = torch.randint(0, K, (1000, 1)), torch.randn(1000, 1)
noisy = alpha_bar[k].sqrt() * actions + (1 - alpha_bar[k]).sqrt() * noise
loss = (net(torch.cat([noisy, k / K], 1)) - noise).pow(2).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Sampling: start from pure noise and remove the predicted noise one step at a time
a = torch.randn(1000, 1)
with torch.no_grad():
for k in reversed(range(K)):
eps = net(torch.cat([a, torch.full_like(a, k / K)], 1))
a = (a - betas[k] / (1 - alpha_bar[k]).sqrt() * eps) / (1 - betas[k]).sqrt()
a += betas[k].sqrt() * torch.randn_like(a) if k > 0 else 0
print(
f"Diffusion policy samples: {(a < -0.5).float().mean():.0%} steer left, {(a > 0.5).float().mean():.0%} steer right"
)回归策略会输出 0.00,径直撞向障碍物。扩散策略生成的样本中,44% 向左转向,55% 向右转向,只有约 1% 落在两种模式之间,因此每个样本都会选择一种有效的机动方式。
扩散策略的工作原理#
训练时,模型从演示中取出一个动作序列,在随机选定的 100 个步骤之一加入高斯噪声,并训练神经网络,使其根据加噪后的动作和观察结果预测噪声。运行时,机器人先采样一个纯噪声序列,再根据最新观察结果对其去噪。原始论文在实体机器人上使用了 10 个 DDIM 去噪步骤,在 NVIDIA RTX 3080 GPU 上的推理延迟为 0.1 秒。
该策略会分块规划,并在滚动时域循环中重新规划。在论文的大多数任务中,策略观察最近 2 个时间步,预测接下来的 16 个动作,执行前 8 个动作,然后根据最新观察结果对新序列去噪,这样既能保持动作平滑,也能对变化作出反应:
作者构建了两种噪声预测网络:一种是 1D 时序卷积神经网络,通过 FiLM 条件注入观察特征;作者建议从这种网络开始尝试。另一种是用于动作变化迅速任务的 Transformer。摄像头图像会输入从头训练的 ResNet-18,使用空间 softmax 池化代替全局平均池化,并使用组归一化代替批归一化。与预测速度相比,预测目标位置的效果始终更好,也让策略对延迟更加稳健。
策略也可以观察低维状态,而非像素。在 Push-T 基准测试中,机器人要将一个 T 形方块推到目标位置;基于状态的变体会观察方块的 9 个二维关键点以及末端执行器的位置。使用自定义关键点布局训练的 Ultralytics YOLO26 姿态估计模型,会根据数据集 YAML 中的 kpt_shape 设置,从摄像头画面中返回这类关键点,并通过 result.keypoints.xy 输出。
速度更快、规模更大的后续模型#
去噪需要对每个动作块进行多次网络推理,因此许多后续研究都聚焦于实时推理。Aaditya Prasad 及其同事于 2024 年提出的 Consistency Policy,将扩散策略蒸馏为一致性模型,只需一步或几步即可完成推理,将速度提高了一个数量级。Yanjie Ze 及其同事于 2024 年提出的 3D Diffusion Policy(DP3)使用紧凑点云而非图像作为条件,在大多数 72 项模拟任务中只需 10 次演示,并在 4 项真实机器人任务中以每项 40 次演示达到 85% 的成功率。
同样的方法如今也用于机器人基础模型。Physical Intelligence 的 π0(2024)将一个通过流匹配训练的动作专家(它与扩散方法关系密切)接入一个拥有 30 亿参数的视觉语言模型,并生成包含 50 个动作的动作块,以最高 50 Hz 的频率控制机器人。Hugging Face 的开源 LeRobot库实现了扩散策略及其他模仿学习策略,用于具身 AI研究。










