扩散策略

概述

扩散策略（Diffusion Policy, Chi et al., 2023）将去噪扩散概率模型（DDPM）引入机器人行为克隆，解决了传统 BC 难以处理的多模态动作分布问题。在接触丰富的操作任务中，扩散策略展示了显著优于传统方法的性能。

---

背景：去噪扩散概率模型

DDPM 回顾

DDPM（Ho et al., 2020）定义了两个过程：

前向过程（加噪）：逐步向数据 \(x_0\) 添加高斯噪声：

\[ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I) \]

其中 \(\beta_t\) 是噪声调度。利用累积参数 \(\bar{\alpha}_t = \prod_{s=1}^t (1 - \beta_s)\)，可直接从 \(x_0\) 采样任意时刻 \(t\) 的噪声数据：

\[ q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1 - \bar{\alpha}_t) I) \]

即：

\[ x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I) \]

反向过程（去噪）：学习从噪声 \(x_T \sim \mathcal{N}(0, I)\) 逐步去噪恢复数据：

\[ p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \sigma_t^2 I) \]

DDPM 训练目标

通过简化的变分下界，DDPM 的训练目标等价于噪声预测：

\[ \mathcal{L}_{\text{DDPM}} = \mathbb{E}_{t \sim U[1,T], \, x_0, \, \epsilon \sim \mathcal{N}(0,I)} \left[ \| \epsilon - \epsilon_\theta(x_t, t) \|^2 \right] \]

其中 \(x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon\)，\(\epsilon_\theta\) 是噪声预测网络。

推导：

DDPM 的反向过程后验为 \(q(x_{t-1}|x_t, x_0)\)，其均值为：

\[ \tilde{\mu}_t = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} x_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t \]

将 \(x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \sqrt{1-\bar{\alpha}_t}\epsilon)\) 代入，得到用 \(\epsilon_\theta\) 参数化的均值：

\[ \mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right) \]

最小化 \(\|\tilde{\mu}_t - \mu_\theta\|^2\) 等价于最小化 \(\|\epsilon - \epsilon_\theta\|^2\)（忽略系数）。

DDIM 加速采样

DDPM 需要 \(T\) 步（通常 1000 步）去噪，太慢。DDIM（Song et al., 2021）提供非马尔可夫采样过程，可用 \(S \ll T\) 步完成：

\[ x_{t-1} = \sqrt{\bar{\alpha}_{t-1}} \underbrace{\left( \frac{x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon_\theta(x_t, t)}{\sqrt{\bar{\alpha}_t}} \right)}_{\text{predicted } x_0} + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \cdot \epsilon_\theta(x_t, t) + \sigma_t \epsilon \]

当 \(\sigma_t = 0\) 时为确定性采样（DDIM），\(\sigma_t = \sqrt{\beta_t}\) 时退化为 DDPM。

---

Diffusion Policy 公式

问题设定

给定观测 \(\mathbf{o}_t\)（图像 + 机器人状态），预测未来 \(T_a\) 步动作序列：

\[ \mathbf{A}_t = (a_t, a_{t+1}, \ldots, a_{t+T_a-1}) \sim p_\theta(\mathbf{A}_t | \mathbf{O}_t) \]

其中 \(\mathbf{O}_t = (o_{t-T_o+1}, \ldots, o_t)\) 是观测历史（\(T_o\) 步）。

条件去噪

将动作序列 \(\mathbf{A}_t\) 视为扩散模型要生成的"数据"。条件扩散模型以观测 \(\mathbf{O}_t\) 为条件：

训练：

\[ \mathcal{L} = \mathbb{E}_{k, \mathbf{A}_t^0, \epsilon} \left[ \left\| \epsilon - \epsilon_\theta\left(\sqrt{\bar{\alpha}_k} \mathbf{A}_t^0 + \sqrt{1-\bar{\alpha}_k} \epsilon, \, k, \, \mathbf{O}_t \right) \right\|^2 \right] \]

其中 \(k\) 是扩散时间步（区别于环境时间步 \(t\)），\(\mathbf{A}_t^0\) 是专家动作序列。

推理：

1. 采样 \(\mathbf{A}_t^K \sim \mathcal{N}(0, I)\)
2. For \(k = K, K-1, \ldots, 1\):

\[ \mathbf{A}_t^{k-1} = \frac{1}{\sqrt{\alpha_k}} \left( \mathbf{A}_t^k - \frac{\beta_k}{\sqrt{1-\bar{\alpha}_k}} \epsilon_\theta(\mathbf{A}_t^k, k, \mathbf{O}_t) \right) + \sigma_k z \]

1. 执行 \(\mathbf{A}_t^0\) 的前 \(T_e\) 步（\(T_e \leq T_a\)）

动作空间参数

参数	典型值	说明
观测历史 \(T_o\)	2	当前帧 + 1 帧历史
预测长度 \(T_a\)	16	预测未来 16 步动作
执行长度 \(T_e\)	8	实际执行 8 步后重新预测
去噪步数 \(K\)	100（训练）/ 10-16（推理用 DDIM）	推理加速

---

网络架构

观测编码器

视觉编码器：使用预训练的 ResNet-18 提取图像特征：

\[ h_{\text{img}} = \text{ResNet18}(I_t) \in \mathbb{R}^{512} \]

状态编码器：对机器人本体状态做线性投影：

\[ h_{\text{state}} = W_s \cdot [q_t, \dot{q}_t, p_t^{\text{ee}}] + b_s \]

融合：将视觉和状态特征拼接后投影：

\[ h_t = \text{MLP}([h_{\text{img}}, h_{\text{state}}]) \in \mathbb{R}^{d} \]

噪声预测网络

Diffusion Policy 提供两种架构选择：

1. CNN-based（1D 时序卷积）：

将动作序列视为 1D 信号，使用 ResNet 风格的 1D 卷积网络。观测通过 FiLM 条件化注入：

\[ \text{FiLM}(x, h) = \gamma(h) \odot x + \beta(h) \]

2. Transformer-based：

使用 Transformer Decoder，观测编码作为 cross-attention 的 key/value：

\[ \hat{\epsilon} = \text{TransformerDecoder}(\mathbf{A}_t^k, k, \mathbf{O}_t) \]

其中扩散时间步 \(k\) 通过正弦位置编码注入。

两种架构对比

维度	CNN-based	Transformer-based
推理速度	更快	较慢
参数量	~25M	~40M
长序列建模	受感受野限制	全局注意力
调参难度	较易	较难
典型任务	通用	长时序任务

---

为什么扩散策略优于传统 BC

多模态分布建模

传统 BC（MSE 损失 + 确定性网络）只能输出单一动作。当数据中存在多种合理动作时，MSE 回归会预测均值——这通常不是任何一种合理动作。

例子：桌上有一个杯子，可以从左边或右边绕过去。

• 传统 BC：预测中间路径（可能撞到杯子）
• 高斯混合模型：需要预设模态数量
• 扩散策略：自然建模任意复杂的多模态分布

接触丰富任务

在涉及接触的任务（如插入、折叠、倒水）中，微小的位置差异导致截然不同的力学行为。扩散策略的迭代去噪过程能精确生成满足接触约束的动作序列。

实验对比（Chi et al., 2023）

任务	BC (MSE)	BC (GMM)	IBC	Diffusion Policy
Push-T	56.5%	65.1%	62.3%	86.2%
Can	78.5%	82.1%	80.4%	94.6%
Square	45.2%	51.8%	48.9%	84.7%
Transport	35.1%	42.3%	38.7%	73.8%

---

3D Diffusion Policy（DP3）

动机

标准 Diffusion Policy 以 RGB 图像为输入，但图像缺乏 3D 几何信息。DP3（Ze et al., 2024）使用点云作为输入，提供显式的 3D 空间理解。

架构

点云编码：使用 PointNet++ 或类似网络将点云 \(P \in \mathbb{R}^{N \times 3}\) 编码为紧凑表征：

\[ h_{\text{3D}} = \text{PointNet++}(P) \in \mathbb{R}^{d} \]

条件扩散：与标准 Diffusion Policy 相同，但以 \(h_{\text{3D}}\) 替代 \(h_{\text{img}}\) 作为条件。

优势

• 视角不变性：3D 表示天然对相机视角变化鲁棒
• 空间推理：显式编码物体间的空间关系
• 深度信息：精确的距离信息对操作任务至关重要

实验结果

DP3 在多个基准上超越 RGB-based Diffusion Policy，尤其在需要精确空间推理的任务（如装配、对齐）上优势明显。

---

Consistency Policy：快速推理

问题

标准 Diffusion Policy 推理需要 10-100 步去噪迭代，每步需要一次神经网络前向传播。在 50Hz 控制频率下：

• 16 步 DDIM：每次推理 ~50ms（仅勉强满足 20Hz）
• 100 步 DDPM：~300ms（远不能满足实时要求）

Consistency Model 回顾

Consistency Model（Song et al., 2023）学习一个函数 \(f_\theta\)，将扩散轨迹上任意点直接映射到起点 \(x_0\)：

\[ f_\theta(x_t, t) = x_0 \quad \forall t \in [0, T] \]

自一致性条件：对同一轨迹上的任意两点 \(x_t\) 和 \(x_{t'}\)，有 \(f_\theta(x_t, t) = f_\theta(x_{t'}, t')\)。

Consistency Policy

将 Consistency Model 应用于机器人策略（Prasad et al., 2024）：

训练目标（蒸馏方式）：

\[ \mathcal{L}_{\text{CD}} = \mathbb{E}_{k, \mathbf{A}^0, \epsilon} \left[ d\left( f_\theta(\mathbf{A}^{t_{k+1}}, t_{k+1}, \mathbf{O}) , \, f_{\theta^-}(\hat{\mathbf{A}}^{t_k}, t_k, \mathbf{O}) \right) \right] \]

其中 \(\hat{\mathbf{A}}^{t_k}\) 是用预训练扩散模型从 \(\mathbf{A}^{t_{k+1}}\) 一步去噪的结果，\(\theta^-\) 是 EMA 目标网络。

推理：仅需 1-2 步即可生成高质量动作序列：

\[ \mathbf{A}_t^0 = f_\theta(\mathbf{A}_t^T, T, \mathbf{O}_t), \quad \mathbf{A}_t^T \sim \mathcal{N}(0, I) \]

速度对比

方法	去噪步数	推理时间	控制频率
DDPM	100	~300 ms	~3 Hz
DDIM	16	~50 ms	~20 Hz
Consistency Policy	1-2	~5 ms	~200 Hz

Consistency Policy 实现了 10-60x 的推理加速，使扩散策略能够应用于高频控制场景。

---

实现要点

动作归一化

扩散模型假设数据在 \([-1, 1]\) 范围内。需要对动作进行归一化：

\[ a_{\text{norm}} = \frac{2(a - a_{\min})}{a_{\max} - a_{\min}} - 1 \]

观测窗口与动作窗口

实践中的重要超参数：

• 观测窗口：使用最近 \(T_o = 2\) 步观测
• 预测窗口：预测 \(T_a = 16\) 步未来动作
• 执行窗口：执行 \(T_e = 8\) 步后重新规划（receding horizon）

重叠执行提供了平滑的动作过渡。

训练技巧

1. EMA 模型：使用指数滑动平均模型用于推理
2. 学习率调度：余弦退火
3. 数据增强：随机裁剪、颜色抖动
4. 梯度裁剪：防止训练不稳定

---

与其他章节的联系

• 模仿学习：模仿学习 中的 BC 是扩散策略的基础对照方法
• VLA 模型：VLA模型 中的部分模型也采用扩散解码器
• 遥操作：遥操作与数据收集 提供扩散策略所需的高质量演示数据

---

参考文献

1. Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. NeurIPS.
2. Song, J., Meng, C., & Ermon, S. (2021). Denoising Diffusion Implicit Models. ICLR.
3. Chi, C., et al. (2023). Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. RSS.
4. Ze, Y., et al. (2024). 3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations. RSS.
5. Song, Y., et al. (2023). Consistency Models. ICML.
6. Prasad, V., et al. (2024). Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation. RSS.

---

评论