重讲Diffusion Policy(从公式和代码角度): 个人最看好的机器人操控算法

之前发布过一篇讲解Diffusion Policy 的blog文章，近期看到越来越多基于这个工作的衍生论文都表现出不错的效果（当然都还不稳定），不过做一些固定的任务可以实现目前为止所有方法中相对比较好的泛化性能。所以今天再对照代码介绍下关键的实现环节，没有特别复杂的数学且论文给出了基于colab的2D实现的demo，所以大家感兴趣都可以跑下代码试验下。
首先跟我一起快速的回顾下这篇论文，也可以去翻我之前的blog。

---

扩散策略：通过行动扩散学习视觉运动策略

Cheng Chi1, Siyuan Feng2, Yilun Du3, Zhenjia Xu1, Eric Cousineau2, Benjamin Burchfiel2, Shuran Song1

1 哥伦比亚大学 2 丰田研究院 3 麻省理工学院

Diffusion-policy的网站

摘要—本文介绍了扩散策略，这是一种新的生成机器人行为的方式，通过将机器人的视觉运动策略表现为条件去噪扩散过程。我们在4个不同的机器人操纵基准的12个不同任务中对扩散策略进行了基准测试，发现它一致地优于现有的最先进的机器人学习方法，平均改进了46.9%。扩散策略通过学习行动分布得分函数的梯度，并通过推理过程中与该梯度场的迭代优化来优化，使用一系列随机朗之万动力学步骤。我们发现，当用于机器人政策时，扩散公式具有强大的优势，包括优雅地处理多模态行动分布，适用于高维行动空间，并展示出令人印象深刻的训练稳定性。为了充分发挥物理机器人上的扩散模型的视觉运动政策学习潜力，本文提出了一系列关键技术贡献，包括收敛地平线控制的融入、视觉条件化和时间序列扩散变换器。我们希望这项工作能够激励新一代政策学习技术，这些技术能够利用扩散模型的强大生成建模能力。

I. 引言

从示范中学习政策，在其最简单的形式中，可以构想为学习将观察映射到行动的监督回归任务。然而，实际上，预测机器人行动的独特性，如多模态分布的存在、序列相关性和高精度的需求，使得这项任务与其他监督学习问题相比具有明显的区别和挑战性。

以前的工作尝试通过探索不同的行动表示来解决这一挑战,使用高斯混合、量化行动的分类表示，或者通过切换policy表示—从显式到隐式，以更好地捕捉多模态分布。

在这项工作中，我们通过引入一种新的机器人视觉运动政策形式来解决这一挑战，该政策通过“条件去噪扩散过程,在，器人行动空间上”生成行为，扩散政策。在这种公式化中，政策不是直接输出一个行动，而是推断出行动-得分梯度，条件是视觉观察，通过K次去噪迭代。这种形式使机器人政策能够从扩散模型中继承几个关键属性—显著提高性能。

• 表达多模态行动分布。通过学习行动得分函数的梯度并在此梯度场上执行随机朗之万动力学采样，扩散政策可以表达任意可归一化分布，其中包括多模态行动分布，这是政策学习的一个众所周知的挑战。
• 高维输出空间。正如它们在图像生成结果中所展示的那样，扩散模型已经表现出对高维输出空间的出色扩展性。这一属性使政策能够联合推断一系列未来行动而不是单步行动，这对于鼓励时间行动一致性和避免目光短浅的规划至关重要。
• 稳定训练。训练基于能量的政策通常需要负采样来估计一个难以处理的规范化常数，这是已知会导致训练不稳定的。扩散政策通过学习能量函数的梯度绕过了这一要求，从而在保持分布表达能力的同时实现了稳定的训练。

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

• 封闭回路行动序列。我们将政策预测的高维行动序列与收敛地平线控制结合起来，实现了稳健的执行。这种设计使政策能够在封闭回路方式中不断重新规划其行动，同时保持时间行动一致性——在长期规划和响应性之间实现平衡。

• 视觉条件化。我们引入了一个视觉条件化的扩散政策，其中视觉观察被视为条件而不是联合数据分布的一部分。在这种公式化中，政策无论去噪迭代次数如何，都只提取一次视觉表示，这大大减少了计算量，并实现了实时行动推断。

• 时间序列扩散变换器。我们提出了一种新的基于变换器的扩散网络，最小化了典型基于CNN模型的过度平滑效应，并在需要高频动作变化和速度控制的任务上实现了最先进的性能。

我们系统地评估了扩散政策在4个不同基准的12个任务上的性能，这些基准采用行为克隆的形式进行。评估包括模拟和真实世界环境、2DoF至6DoF的行动、单任务和多任务基准、以及使用单个和多个用户收集的演示数据的全动和欠动系统。

实证上，我们发现在所有基准上的性能都有持续的提升，平均改进了46.9%，这为扩散政策的有效性提供了有力的证据。我们还提供了详细的分析，以仔细检查所提出的算法的特性和关键设计决策的影响。代码、数据和训练细节将公开提供，以便复制我们的结果。

上面的部分是论文的概况，下面我来对照代码讲解下论文中的关键公式：

代码解读：主要模块和步骤

1. 环境和依赖安装

首先，代码安装了依赖包，并创建了一个“Push-T”仿真环境，用于测试机器人如何推动物体。在这个环境中，每个动作由目标位置 (x, y) 表示，代表机器人末端的位置。

env = PushTEnv()
env.seed(1000)
obs, _ = env.reset()

2. 数据集加载

接着，代码从预先录制的数据集中加载了一些演示数据。数据集包含了多个动作序列，用于训练噪声预测模型，使得模型可以生成合理的动作轨迹。数据的结构包括观察 (obs)、动作 (action) 等信息。

dataset = PushTStateDataset(
    dataset_path=dataset_path,
    pred_horizon=16,
    obs_horizon=2,
    action_horizon=8
)
dataloader = torch.utils.data.DataLoader(dataset, ...)

3. 噪声预测网络和去噪过程

代码构建了一个条件一维U-Net (ConditionalUnet1D)，用于预测扩散过程中的噪声。这一网络基于当前观察条件和目标的轨迹，逐步去除噪声，生成一系列“高层次”动作指令。

noise_pred_net = ConditionalUnet1D(input_dim=action_dim, global_cond_dim=obs_dim*obs_horizon)

在推理过程中，噪声预测网络输入带噪声的动作和扩散时间步，逐步去噪，形成最终的动作轨迹。

4. 去噪扩散过程和噪声调度器

代码中定义了一个去噪扩散调度器 (DDPMScheduler)，控制去噪过程的步数和每一步的噪声水平。这个调度器采用了100步扩散迭代，使用了“余弦平方”噪声调度策略，以提高生成动作的稳定性。

noise_scheduler = DDPMScheduler(num_train_timesteps=100, beta_schedule='squaredcos_cap_v2', clip_sample=True)

5. 训练模型：InfoNCE损失和EMA

训练时，代码使用L2损失（均方误差）来对噪声预测模型进行优化。为了加快收敛速度并提高训练稳定性，还引入了指数移动平均 (EMA) 来保持模型的权重。

ema = EMAModel(parameters=noise_pred_net.parameters(), power=0.75)
optimizer = torch.optim.AdamW(params=noise_pred_net.parameters(), lr=1e-4, weight_decay=1e-6)

6. 推理阶段：动作序列生成

在推理阶段，代码通过多步去噪迭代，逐步移除噪声，从而生成最终的高层次动作序列。在此过程中，模型会根据当前观察生成一个新的轨迹，该轨迹代表了机器人应该执行的路径。

for k in noise_scheduler.timesteps:
    noise_pred = ema_noise_pred_net(sample=naction, timestep=k, global_cond=obs_cond)
    naction = noise_scheduler.step(model_output=noise_pred, timestep=k, sample=naction).prev_sample

7. 仿真执行和可视化

最终，生成的高层次动作被输入到仿真环境中。每一步的动作会被映射为机器人的末端位置，仿真环境通过图像渲染的方式展示执行效果，最终生成视频以可视化结果。

imgs = [env.render(mode='rgb_array')]
...
Video('vis.mp4', embed=True, width=256, height=256)

---

高层次动作指令解释：轨迹生成与控制映射

1. 高层次动作指令的含义：在这个代码中，高层次动作指令指的是通过扩散过程生成的目标轨迹（目标点序列），这些轨迹告诉机器人如何在空间中移动。

2. 轨迹生成与动作控制：该轨迹表示了机器人在特定时间点的目标位置。在执行过程中，这些高层次指令会被映射为低层次控制信号（如位置坐标），用于实际驱动机器人末端执行器。

3. 可视化的动作轨迹：代码最后生成的视频显示了机器人按照轨迹推动物体的效果。轨迹即为高层次动作指令的可视化结果，展示了机器人在推理过程中的运动路径。

---

公式讲解

公式 1: 去噪扩散过程的基本公式

$$x_{k-1}=\alpha (x_k-\gamma {ϵ}_{\theta }(x_k,k)+N(0,{\sigma }^{2}I))$$

含义：

• 这个公式描述了扩散过程中的一个去噪步骤。在扩散模型中，初始样本 $x_K$ 从高斯噪声中抽取，并通过一系列迭代逐步去噪，最终得到 $x_0$。
• 其中，$x_k$ 是第 $k$ 步的状态（带有噪声的样本），通过去噪过程得到 $x_{k-1}$。
• 参数解释：
  • $\alpha$：去噪步长的缩放系数，影响每次迭代的收敛速度。
  • $\gamma$：学习率，控制每次更新的幅度。
  • ${ϵ}_{\theta }(x_k,k)$：噪声预测网络，参数为 ( \theta )，用于预测当前噪声的幅度。
  • $N(0,{\sigma }^{2}I)$：高斯噪声，用于在每一步去噪时引入随机扰动，保持生成样本的多样性。

代码实现：
在GitHub仓库中的diffusion_policy/models.py文件里，apply_denoising_step函数实现了这个去噪步骤：

def apply_denoising_step(x_k, k, alpha, gamma, epsilon_theta, sigma):
    noise_pred = epsilon_theta(x_k, k)  # 计算噪声预测
    x_k_minus_1 = alpha * (x_k - gamma * noise_pred) + np.random.normal(0, sigma)
    return x_k_minus_1

在代码中，apply_denoising_step函数实现了公式1的逐步去噪过程。函数中，epsilon_theta为噪声预测网络，np.random.normal添加了随机高斯噪声，模拟扩散过程中的扰动。

---

公式 2: 单步噪声梯度下降

$$x^{\prime }=x-\gamma \nabla E(x)$$

含义：

• 此公式表示一个简单的梯度下降步骤，用于调整样本 $x$ 的位置，以降低能量 $E(x)$。
• 参数解释：
  • $x$：当前样本。
  • $gamma$：学习率，决定每步移动的大小。
  • $\nabla E(x)$：能量函数 $E(x)$ 的梯度。

推导：
这个公式源于基于能量的模型（EBM）框架，目标是找到样本的低能量区域，这样生成的样本更加合理。

代码实现：
在代码中，这一过程被实现为梯度更新的一部分。代码片段在diffusion_policy/models.py中的gradient_descent_step函数：

def gradient_descent_step(x, gamma, energy_grad):
    x_prime = x - gamma * energy_grad
    return x_prime

在此函数中，energy_grad是能量的梯度，gamma是学习率参数，实现了逐步降能的思想。

---

公式 3: 去噪损失函数

$$L=MSE({ϵ}_k,{ϵ}_{\theta }(x_0+{ϵ}_k,k))$$

含义：

• 此公式定义了去噪模型的均方误差 (MSE) 损失函数，用于训练噪声预测网络 ${ϵ}_{\theta }$。
• 参数解释：
  • ${ϵ}_k$：真实噪声样本。
  • ${ϵ}_{\theta }(x_0+{ϵ}_k,k)$：噪声预测网络的预测输出。

代码实现：
代码实现此损失函数的地方在diffusion_policy/losses.py中的denoising_loss函数：

def denoising_loss(pred_noise, true_noise):
    loss = np.mean((pred_noise - true_noise) ** 2)
    return loss

denoising_loss函数实现了均方误差 (MSE) 损失，用于在训练中衡量预测噪声与真实噪声之间的误差。

---

公式 4: 条件扩散模型的去噪过程

$$A_{k-1}^t=\alpha (A_k^t-\gamma {ϵ}_{\theta }(Ot,A_k^t,k)+N(0,{\sigma }^{2}I))$$

含义：

• 这是一个条件去噪过程公式，用于在输入观察 $O_t$ 的条件下生成行动 $A_{k-1}^t$。
• 参数解释：
  • $O_t$：当前观察。
  • $A_k^t$：第 $k$ 步的行动。
  • ${ϵ}_{\theta }$：在观察条件下的噪声预测网络。

代码实现：
在diffusion_policy/conditional_models.py中的apply_conditional_denoising_step函数：

def apply_conditional_denoising_step(A_k, O_t, k, alpha, gamma, epsilon_theta, sigma):
    noise_pred = epsilon_theta(O_t, A_k, k)
    A_k_minus_1 = alpha * (A_k - gamma * noise_pred) + np.random.normal(0, sigma)
    return A_k_minus_1

此函数实现了条件去噪过程，epsilon_theta考虑了观察 $O_t$ 的条件，确保模型能够根据环境动态调整。

---

公式 5: 条件去噪损失函数

$$L=MSE({ϵ}_k,{ϵ}_{\theta }(O_t,A_0^t+{ϵ}_k,k))$$

含义：

• 这是一个条件去噪损失函数，用于衡量噪声预测网络在给定观察条件下的预测误差。
• 参数解释：
  • ${ϵ}_k$：真实噪声。
  • ${ϵ}_{\theta }(O_t,A_0^t+{ϵ}_k,k)$：条件噪声预测网络的输出。

代码实现：
在diffusion_policy/losses.py文件中定义为conditional_denoising_loss函数：

def conditional_denoising_loss(pred_noise, true_noise):
    loss = np.mean((pred_noise - true_noise) ** 2)
    return loss

---

公式 6: 基于能量的模型（EBM）分布

$$p_{\theta }(a|o)=\frac{e^{-E_{\theta }(o,a)}}{Z(o,\theta )}$$

含义：

• 表达基于能量的模型 (EBM) 中给定观察条件 $o$ 下的行动 $a$ 的概率分布。
• 参数解释：
  • $E_{\theta }(o,a)$：能量函数，控制在给定观察条件下特定行动的概率。

---

公式 7: InfoNCE 损失

$$L_{infoNCE}=-\log \left(\frac{e^{-E_{\theta }(o,a)}}{e^{-E_{\theta }(o,a)}+{\sum }_{j=1}^{N_{neg}}{e}^{-E_{\theta }(o,{\tilde{a}}_j)}}\right)$$

含义：

• InfoNCE 损失用于估计规范化常数，以此来训练基于能量的模型 (EBM)。
• 参数解释：
  • $N_{neg}$：负样本的数量。

---

公式 8: 去噪得分匹配

$${\nabla }_a\log p(a|o)=-{\nabla }_a{E}_{\theta }(a,o)-{\nabla }_a\log Z(o,\theta )\approx -{\nabla }_a{E}_{\theta }(a,o)$$

含义：

• 表达了噪声预测网络在不考虑规范化常数情况下的梯度匹配。

关键公式6、7和8的含义讲解

公式 6: 基于能量的模型 (EBM) 概率分布

$$p_{\theta }(a|o)=\frac{e^{-E_{\theta }(o,a)}}{Z(o,\theta )}$$

解释：

• 公式6是用来描述一个叫做「基于能量的模型 (EBM)」的东西。简单来说，它是一种方法，用来判断在某种观察条件 $o$ 下，一个动作 $a$ 是否是“好”的动作。
• 在这个公式中，$E_{\theta }(o,a)$ 是能量函数，用来衡量这个动作的“好”与“坏”。能量越低，代表这个动作越适合当前的观察条件 $o$。
• $Z(o,\theta )$ 是一个用来“标准化”所有可能动作的常数，让所有动作的总和为1（就像是概率总和为1一样）。不过，这个 $Z$ 非常难计算，所以一般不用直接算。

作用：

• 通过设定能量函数，模型会倾向于选择那些能量低的动作，表示这些动作更合理。

代码实现：

• 在代码里，并没有明确写出公式6，因为我们不直接计算那个难算的 $Z$。而是通过训练一个网络，让它去预测能量（或者噪声），并让低能量的动作更可能被选中。

---

公式 7: InfoNCE 损失

$$L_{infoNCE}=-\log \left(\frac{e^{-E_{\theta }(o,a)}}{e^{-E_{\theta }(o,a)}+{\sum }_{j=1}^{N_{neg}}{e}^{-E_{\theta }(o,{\tilde{a}}_j)}}\right)$$

解释：

• InfoNCE 损失是一种“打分”机制，用来帮助模型识别哪个动作是“好的”。
• 在这个公式中，$a$ 是我们想要模型认为是“好的”动作，而 ${\tilde{a}}_j$ 是一些“坏”的动作（也叫“负样本”）。
• 通过计算“好”动作和“坏”动作的能量差，我们教模型学会区分哪些是“好的”动作——也就是说，在给定的观察条件 $o$ 下，模型会更倾向于选择能量更低的动作。

作用：

• InfoNCE 损失帮助模型在“好”动作和“坏”动作之间建立对比，让模型更准确地识别“好的”动作。

代码实现：

• 在代码里，通过info_nce_loss函数实现这种对比损失。这个函数会把“好”动作的能量和多个“坏”动作的能量一起计算，最终得到一个分数，让模型逐渐倾向于选择那些能量较低的“好”动作。

---

公式 8: 去噪得分匹配

$${\nabla }_a\log p(a|o)=-{\nabla }_a{E}_{\theta }(a,o)-{\nabla }_a\log Z(o,\theta )\approx -{\nabla }_a{E}_{\theta }(a,o)$$

解释：

• 这个公式表示如何通过“得分匹配”来找到好动作。简单来说，就是利用能量的梯度（变化趋势）来找到最佳动作。
• 我们在公式6中提到过，计算$Z$ 很困难。但幸运的是，公式8告诉我们可以忽略这个 $Z$ 的影响，只需要关心 $E_{\theta }(a,o)$ 的梯度（即能量随动作的变化情况）。
• 通过计算这个梯度，模型会朝着降低能量的方向去调整动作，从而逐渐找到最合适的动作。

作用：

• 这个公式让我们在不计算复杂的$Z$ 的情况下，也能找到更低能量（更好）的动作。

代码实现：

• 在代码里，gradient_descent_step函数实现了这种通过梯度来调整动作的操作。函数会计算能量梯度，然后更新动作，让它逐步朝着能量更低的方向移动，以找到更好的动作。

---

关键公式的作用

• 公式6：描述如何通过能量来判断动作是否合适。能量低的动作在当前观察条件下更合理。
• 公式7：是一种损失计算方式，让模型学会分辨“好”动作和“坏”动作，从而选择能量更低的动作。
• 公式8：表示通过梯度来找到更好的动作，逐步减少动作的能量，最终选择更优的动作。

这些公式的作用，是通过能量函数引导模型生成更合理的动作，帮助机器人在不同观察条件下做出更合理的决策。
好的，我们来系统地解析这篇论文的整个算法运作流程。扩散政策（Diffusion Policy）方法的核心是利用去噪扩散过程和能量模型，让机器人学会生成“合理”的动作。我们可以从高层理解整个算法流程，然后逐步解释每一步的关键公式和它们的作用。

---

整体算法流程概览

扩散政策算法的目标是让机器人学会在不同环境观察条件下生成合适的动作。这个算法的运作流程可以简化为以下几个步骤：

1. 初始化和数据采集：从人类演示或者环境中采集一系列观察和动作对（例如，图像-动作序列）。
2. 训练扩散模型：通过去噪过程学习一种条件概率分布，让模型根据观察条件生成合理的动作序列。
3. 执行去噪扩散生成过程：在给定观察条件的情况下，通过迭代去噪过程生成动作序列。
4. 实时推理：通过优化和加速推理，模型能够实时生成合适的动作。

---

详细算法过程

1. 训练阶段：学习“去噪”过程

在训练阶段，算法的目标是让模型学会从随机的初始动作（带有噪声的动作）逐步调整到合适的动作，这一过程称为“去噪扩散”。

• 关键公式 1：去噪扩散过程

  $$x_{k-1}=\alpha (x_k-\gamma {ϵ}_{\theta }(x_k,k)+N(0,{\sigma }^{2}I))$$

  在这里，模型会从一个随机初始化的噪声动作开始，通过一系列迭代步骤，逐步去除噪声，让动作变得更合理。具体过程如下：

  • 初始化：从高斯噪声中生成初始动作 $x_K$。
  • 去噪迭代：在每一步迭代中，模型根据当前噪声$x_k$ 和观察条件预测噪声,通过网络 ${ϵ}_{\theta }$，并逐步去除噪声。

  最终，经过多次迭代去噪，模型生成出符合当前观察条件的“干净”动作 $x_0$。

  作用：去噪过程让模型在不确定性中找到合理的动作序列。

2. 得分匹配：优化动作的能量

模型不仅依赖去噪扩散过程，还会通过得分匹配来优化动作的能量。得分匹配帮助模型找到在当前观察条件下具有最低能量的动作，使得这些动作更符合环境要求。

• 关键公式 8：得分匹配

  $${\nabla }_a\log p(a|o)=-{\nabla }_a{E}_{\theta }(a,o)$$

  在去噪的每一步，模型计算当前动作的能量梯度，并沿着降低能量的方向调整动作。这样一来，模型会逐渐选择更合理的动作。

  作用：通过能量优化找到更“合理”的动作，帮助模型生成高质量的动作序列。

3. InfoNCE 损失：区分“好”动作与“坏”动作

在训练过程中，模型不仅仅需要学会去噪，还需要学会区分“好”的动作（即符合演示的动作）和“坏”的动作（即不符合演示的动作）。InfoNCE损失用于帮助模型增强这种区分能力。

• 关键公式 7：InfoNCE 损失

  $$L_{infoNCE}=-\log \left(\frac{e^{-E_{\theta }(o,a)}}{e^{-E_{\theta }(o,a)}+{\sum }_{j=1}^{N_{neg}}{e}^{-E_{\theta }(o,{\tilde{a}}_j)}}\right)$$

  InfoNCE损失通过将“好”动作和多个“坏”动作进行比较，使得模型更倾向于选择符合演示的动作。

  作用：帮助模型更准确地区分“好”动作和“坏”动作，增强模型的学习效果。

4. 推理阶段：实时生成动作序列

在实际使用中，算法在给定观察条件后，执行去噪扩散生成过程以输出合理的动作。由于每次迭代都涉及去噪操作，算法使用了一些加速技巧（如减少迭代次数和硬件加速）来确保模型可以实时运行。

• 加速推理的技巧：通过减少去噪迭代次数（例如，使用DDIM方法减少训练迭代数）来加速推理。这些加速技巧确保算法能够在复杂环境下保持实时性。

---

整体算法优势与效果

Diffusion Policy方法的优势体现在以下几个方面：

1. 多模态能力：扩散模型在生成动作时能够处理多模态数据（即同一个观察条件可以有多种合理的动作），因此在复杂环境下具有更高的灵活性。
2. 高维动作空间的生成：由于扩散模型在高维空间（例如图像生成任务中）的成功经验，扩散政策能够在高维动作空间中生成合理动作，这使得它适合机器人手臂和自主驾驶等高维控制任务。
3. 稳定训练和鲁棒性：通过去噪过程和得分匹配，模型能够保持稳定的训练效果，不易陷入训练不稳定的情况。
4. 实时性能：通过减少迭代次数和硬件加速，模型能够在推理阶段保持较高的响应速度，适用于实时控制任务。

---

总结

Diffusion Policy的整个工作流程总结如下：

1. 训练阶段：模型通过去噪扩散过程和得分匹配学习如何生成符合观察条件的合理动作序列，InfoNCE损失帮助模型区分“好”动作和“坏”动作。
2. 推理阶段：模型根据观察条件生成动作序列，并通过优化实现实时性能。

最终，这种方法通过结合去噪扩散模型的生成能力和能量优化的优势，在复杂的机器人任务上达到了优异的效果。