第12章 行为克隆:从扩散策略到可扩展的视觉运动控制
原课程: Robots That Learn, UC Berkeley CS294-291, Spring 2026 主讲: Jitendra Malik 本章基于 Lecture 8A: Behavioral Cloning 精读论文: Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, Shuran Song. “Diffusion Policy: Visuomotor Policy Learning via Action Diffusion,” RSS 2023. 代码仓库: https://diffusion-policy.cs.columbia.edu
一、本章概要
核心问题: 假设你有一个机器人臂,你通过遥操作演示了50次”拿起杯子倒水”这个任务。给定这些演示数据,如何让机器人自主学会这个策略,并且在面对未曾见过的初始条件时(杯子位置偏左5cm、桌面光照变了)依然成功执行?
这就是行为克隆(Behavioral Cloning, BC)试图回答的问题——它是最直接的模仿学习范式:将专家演示数据视为有监督学习样本 \(\{(\mathbf{o}_i, \mathbf{a}_i)\}_{i=1}^N\),直接学习从观测到动作的映射 \(\pi: \mathcal{O} \to \mathcal{A}\)。
然而,将BC应用于机器人操作面临三个独特挑战,使其与传统的图像分类或语言建模截然不同:
- 动作分布的多模态性(Multimodality)。对于同一个观测状态,人类专家可能选择从左侧绕行或右侧绕行——这两种策略都是”正确”的,但它们的平均值(如果使用MSE损失)可能是完全错误的动作。
- 序列相关性(Sequential Correlation)。机器人动作不是独立的——\(t\) 时刻的动作直接影响 \(t+1\) 时刻的观测。一个微小误差可能被逐步放大,导致分布偏移(distribution shift)的灾难性累积。
- 高精度要求(High Precision)。许多操作任务(如插销、拧螺丝)需要亚厘米级定位精度,这对策略的预测精度提出了极高要求。
本章以 Chi et al. (2023) 的 Diffusion Policy 为核心,深入剖析如何利用去噪扩散概率模型(DDPM)作为机器人视觉运动策略的表示形式。Diffusion Policy 不直接输出动作,而是学习动作评分函数的梯度场(the gradient of the action-score function),在推理时通过一系列随机 Langevin 动力学步骤迭代优化。这一范式带来了三项关键优势:(1)天然支持多模态动作分布的表达;(2)适用于高维动作序列的输出空间;(3)训练过程异常稳定。
与前后章节的逻辑关系: 第11章(强化学习)介绍了通过试错学习策略的方法论,本章引入其互补范式——从专家演示中直接学习策略。第13章(视觉模仿学习)将 BC 扩展到更复杂的视觉输入场景(如 UMI 的3D表示和跨具身迁移)。本章建立的扩散策略框架(去噪过程、receding-horizon 控制、视觉条件化)是理解后续章节中各类模仿学习方法的数学基础。
二、核心概念与数学基础
2.1 行为克隆的基本形式化
定义 1(行为克隆,Behavioral Cloning): 给定一个专家演示数据集 \(\mathcal{D} = \{(\mathbf{o}_i, \mathbf{a}_i)\}_{i=1}^N\),其中 \(\mathbf{o}_i \in \mathcal{O}\) 是观测(通常是图像序列),\(\mathbf{a}_i \in \mathcal{A} \subseteq \mathbb{R}^{d_a}\) 是对应的动作,行为克隆的目标是找到一个参数化策略 \(\pi_{\boldsymbol{\theta}}: \mathcal{O} \to \mathcal{A}\),最小化经验风险:
\[\pi^* = \arg\min_{\boldsymbol{\theta}} \frac{1}{N} \sum_{i=1}^N \ell(\pi_{\boldsymbol{\theta}}(\mathbf{o}_i), \mathbf{a}_i)\]对于连续动作空间,最常用的损失是均方误差(MSE):
\[\mathcal{L}_{\text{BC}}(\boldsymbol{\theta}) = \frac{1}{N} \sum_{i=1}^N \|\pi_{\boldsymbol{\theta}}(\mathbf{o}_i) - \mathbf{a}_i\|_2^2\]但这个看似简单的公式隐藏着一个巨大的陷阱:训练分布 \(\neq\) 部署分布。在训练时,\(\mathbf{o}_i\) 来自专家的轨迹,策略只见过专家访问过的状态。在部署时,策略自己的动作会偏离专家轨迹,将智能体带入全新的、训练中从未出现过的状态——这就是分布偏移(Distribution Shift)问题的根源。
2.2 BC的复合误差分析
本节给出 BC 中分布偏移导致复合误差的严格分析框架(Ross et al., 2011; Syed & Schapire, 2010)。
设定:考虑一个确定性环境,状态转移为 \(s_{t+1} = T(s_t, a_t)\),专家策略为 \(\pi^E\)(假设为确定性),BC 学到的策略为 \(\hat{\pi}\)。定义单步动作误差:
\[\epsilon = \mathbb{E}_{s \sim d_{\pi^E}}[\|\hat{\pi}(s) - \pi^E(s)\|]\]其中 \(d_{\pi^E}\) 是专家策略下的状态访问分布。
关键假设:环境的动力学是Lipschitz 连续的。即存在常数 \(L_T > 0\) 使得 \(\forall s, s', a\):
\[\|T(s, a) - T(s', a)\| \leq L_T \|s - s'\|\]这意味着相邻状态在相同动作下的演化不会发散太快。机器人动力学通常满足此假设——因为物理定律对状态变化施加了平滑性约束(惯性、阻尼)。
定理 1(BC 复合误差界): 在上述设定下,\(\hat{\pi}\) 在第 \(t\) 步的状态偏差满足:
\[\boxed{\|s_t^{\hat{\pi}} - s_t^{\pi^E}\| \leq \epsilon \cdot \frac{L_T^t - 1}{L_T - 1} \quad (\text{当 } L_T \neq 1)}\]其中 \(s_t^{\hat{\pi}}\) 是 \(\hat{\pi}\) 执行 \(t\) 步后的状态,\(s_t^{\pi^E}\) 是专家策略执行 \(t\) 步后的状态。当 \(L_T > 1\) 时,上界随 \(t\) 呈指数增长。
证明:使用归纳法。
基础情形 (\(t=1\)): 从相同初始状态 \(s_0\) 出发: \(\begin{aligned} \|s_1^{\hat{\pi}} - s_1^{\pi^E}\| &= \|T(s_0, \hat{\pi}(s_0)) - T(s_0, \pi^E(s_0))\| \\ &\text{(动作空间上的 Lipschitz 连续性)} \\ &\leq L_a \|\hat{\pi}(s_0) - \pi^E(s_0)\| \leq L_a \epsilon \end{aligned}\) 其中 \(L_a\) 是转移函数对动作输入的 Lipschitz 常数。
归纳步 (\(t \to t+1\)): 假设 \(\|s_t^{\hat{\pi}} - s_t^{\pi^E}\| \leq \delta_t\),则:
\[\begin{aligned} \|s_{t+1}^{\hat{\pi}} - s_{t+1}^{\pi^E}\| &= \|T(s_t^{\hat{\pi}}, \hat{\pi}(s_t^{\hat{\pi}})) - T(s_t^{\pi^E}, \pi^E(s_t^{\pi^E}))\| \\ &\leq \|T(s_t^{\hat{\pi}}, \hat{\pi}(s_t^{\hat{\pi}})) - T(s_t^{\pi^E}, \hat{\pi}(s_t^{\hat{\pi}}))\| \\ &\quad + \|T(s_t^{\pi^E}, \hat{\pi}(s_t^{\hat{\pi}})) - T(s_t^{\pi^E}, \pi^E(s_t^{\pi^E}))\| \\ &\leq L_T \|s_t^{\hat{\pi}} - s_t^{\pi^E}\| + L_a \|\hat{\pi}(s_t^{\hat{\pi}}) - \pi^E(s_t^{\pi^E})\| \end{aligned}\]其中第二项可以通过在 \(s_t^{\pi^E}\) 处插入一项来界住:
\[\begin{aligned} \|\hat{\pi}(s_t^{\hat{\pi}}) - \pi^E(s_t^{\pi^E})\| &\leq \|\hat{\pi}(s_t^{\hat{\pi}}) - \pi^E(s_t^{\hat{\pi}})\| + \|\pi^E(s_t^{\hat{\pi}}) - \pi^E(s_t^{\pi^E})\| \\ &\leq \epsilon + L_\pi \|s_t^{\hat{\pi}} - s_t^{\pi^E}\| \end{aligned}\]合并得到递推关系:
\[\delta_{t+1} \leq (L_T + L_a L_\pi) \delta_t + L_a \epsilon\]定义有效膨胀因子 \(L_{\text{eff}} = L_T + L_a L_\pi\)。当 \(L_{\text{eff}} \leq 1\) 时,误差保持有界;当 \(L_{\text{eff}} > 1\) 时,求解递推关系得(取 \(L_T\) 为膨胀因子以简化记号):
\[\delta_t \leq \epsilon \cdot \frac{L_T^t - 1}{L_T - 1}\]特别地,当 \(L_T\) 严格大于 1 时,\(\delta_t = \Theta(\epsilon \cdot L_T^t)\)——误差以指数速度放大。当 \(L_T = 1\) 时,\(\delta_t \leq t \cdot \epsilon\)——线性累积。\(\blacksquare\)
直觉:复合误差的本质是 BC 策略在”它自己的世界”里运行,而这个”自己的世界”与训练分布渐行渐远。当 \(L_T\) 较大时(例如,操作任务中的接触动力学——物体被推动后状态会快速改变),即使是很小的单步误差也会在几十步内导致灾难性失败。这就是为什么 BC 策略在初始阶段通常表现良好(训练分布中见过的状态),但随后性能急剧下降。
2.3 DAgger:通过在线纠正抑制分布偏移
Ross et al. (2011) 提出的 DAgger(Dataset Aggregation)算法是解决分布偏移的经典策略。
DAgger 的核心洞察: 如果 BC 训练数据只覆盖专家的状态分布 \(d_{\pi^E}\),则需要通过”查询专家在学到的策略访问的状态上的动作”来扩充训练数据。
DAgger 算法:
- 初始化数据集 \(\mathcal{D}_0 \leftarrow\) 专家演示数据
- 训练策略 \(\hat{\pi}_0\) 在 \(\mathcal{D}_0\) 上
- For \(i = 1, 2, \ldots, M\):
- 执行 \(\hat{\pi}_{i-1}\),收集其访问的状态 \(\{s_1, s_2, \ldots, s_T\}\)
- 对每个 \(s_t\),查询专家动作 \(\mathbf{a}_t^E = \pi^E(s_t)\)
- \[\mathcal{D}_i \leftarrow \mathcal{D}_{i-1} \cup \{(s_t, \mathbf{a}_t^E)\}_{t=1}^T\]
- 在 \(\mathcal{D}_i\) 上重新训练 \(\hat{\pi}_i\)
DAgger 的误差界: Ross et al. 证明,在 DAgger 下,策略的期望误差为:
\[\mathbb{E}\left[\sum_{t=1}^T \ell(\hat{\pi}(s_t), \pi^E(s_t))\right] \leq T \epsilon_N + O\left(\frac{1}{N}\right)\]其中 \(\epsilon_N\) 是 BC 在 \(N\) 个训练样本上的分类/回归误差。关键在于:误差不再以 \(T^2\) 增长,而是以 \(T\) 线性增长——DAgger 成功阻止了指数级复合。
为什么 DAgger 在机器人操作中受限: DAgger 假设可以实时查询专家策略 \(\pi^E(s_t)\)。在机器人操作中,这意味着人类专家需要持续监督并随时介入——这在纯遥操作实验设置中是不现实的。这也是为什么后续工作(如 Diffusion Policy)选择通过动作序列预测+ receding-horizon 执行来在没有在线专家的情况下抑制分布偏移。
2.4 去噪扩散概率模型(DDPM)核心数学
在深入 Diffusion Policy 之前,必须先理解 DDPM 的数学机制。DDPM 是一种生成模型,通过模拟从噪声到数据的逐步去噪过程来生成样本。
2.4.1 前向扩散过程
给定来自真实数据分布的样本 \(\mathbf{x}_0 \sim q(\mathbf{x}_0)\),前向过程在 \(K\) 步内逐步向数据添加高斯噪声,产生一系列噪声版本 \(\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_K\):
\[q(\mathbf{x}_k \mid \mathbf{x}_{k-1}) = \mathcal{N}\left(\mathbf{x}_k; \sqrt{1 - \beta_k} \mathbf{x}_{k-1}, \beta_k \mathbf{I}\right)\]其中 \(\beta_k \in (0, 1)\) 是噪声调度(noise schedule),控制每一步的噪声增量。
关键性质——封闭形式的任一步采样: 利用高斯分布的卷积性质,可以直接从 \(\mathbf{x}_0\) 采样 \(\mathbf{x}_k\) 而不需要迭代:
\[\boxed{q(\mathbf{x}_k \mid \mathbf{x}_0) = \mathcal{N}\left(\mathbf{x}_k; \sqrt{\bar{\alpha}_k} \mathbf{x}_0, (1 - \bar{\alpha}_k) \mathbf{I}\right)}\]其中 \(\alpha_k = 1 - \beta_k\),\(\bar{\alpha}_k = \prod_{i=1}^k \alpha_i\)。这个性质对训练效率至关重要——我们不需要从 \(k=1\) 迭代到 \(k\),只需一步采样。
推导: 记 \(\tilde{\mathbf{x}} \sim \mathcal{N}(\sqrt{\alpha} \mathbf{x}, (1-\alpha)\mathbf{I})\) 为”一步衰减+加噪”,反复应用:
\[\begin{aligned} \mathbf{x}_1 &= \sqrt{\alpha_1} \mathbf{x}_0 + \sqrt{1 - \alpha_1} \boldsymbol{\epsilon}_1 \\ \mathbf{x}_2 &= \sqrt{\alpha_2} \mathbf{x}_1 + \sqrt{1 - \alpha_2} \boldsymbol{\epsilon}_2 \\ &= \sqrt{\alpha_2 \alpha_1} \mathbf{x}_0 + \sqrt{\alpha_2 (1 - \alpha_1) + (1 - \alpha_2)} \bar{\boldsymbol{\epsilon}} \\ &= \sqrt{\bar{\alpha}_2} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_2} \bar{\boldsymbol{\epsilon}} \end{aligned}\]其中 \(\bar{\boldsymbol{\epsilon}}\) 是独立高斯噪声的加权和(仍然是高斯)。推广即得上述封闭形式。
2.4.2 逆向去噪过程
生成过程逆转前向扩散:从纯噪声 \(\mathbf{x}_K \sim \mathcal{N}(0, \mathbf{I})\) 开始,逐步去噪以得到 \(\mathbf{x}_0\):
\[p_{\boldsymbol{\theta}}(\mathbf{x}_{k-1} \mid \mathbf{x}_k) = \mathcal{N}\left(\mathbf{x}_{k-1}; \boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_k, k), \sigma_k^2 \mathbf{I}\right)\]其中 \(\boldsymbol{\mu}_{\boldsymbol{\theta}}\) 是由神经网络预测的均值,\(\sigma_k\) 通常设为 \(\sigma_k = \sqrt{\frac{1 - \bar{\alpha}_{k-1}}{1 - \bar{\alpha}_k} \beta_k}\) 或简化为 \(\sigma_k = \sqrt{\beta_k}\)。
Ho et al. (2020) 的关键简化:与其直接预测 \(\boldsymbol{\mu}_{\boldsymbol{\theta}}\),不如预测添加的噪声分量 \(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_k, k)\)。通过将前向过程的后验 \(q(\mathbf{x}_{k-1} \mid \mathbf{x}_k, \mathbf{x}_0)\) 写为高斯形式:
\[q(\mathbf{x}_{k-1} \mid \mathbf{x}_k, \mathbf{x}_0) = \mathcal{N}\left(\mathbf{x}_{k-1}; \frac{\sqrt{\bar{\alpha}_{k-1}} \beta_k}{1 - \bar{\alpha}_k} \mathbf{x}_0 + \frac{\sqrt{\alpha_k} (1 - \bar{\alpha}_{k-1})}{1 - \bar{\alpha}_k} \mathbf{x}_k, \frac{1 - \bar{\alpha}_{k-1}}{1 - \bar{\alpha}_k} \beta_k \mathbf{I}\right)\]再利用 \(\mathbf{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_k}}(\mathbf{x}_k - \sqrt{1 - \bar{\alpha}_k} \boldsymbol{\epsilon})\) 替换,得到反向采样公式:
\[\boxed{\mathbf{x}_{k-1} = \frac{1}{\sqrt{\alpha_k}}\left(\mathbf{x}_k - \frac{1 - \alpha_k}{\sqrt{1 - \bar{\alpha}_k}} \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_k, k)\right) + \sigma_k \mathbf{z}}\]其中 \(\mathbf{z} \sim \mathcal{N}(0, \mathbf{I})\)(当 \(k > 1\) 时,\(k=1\) 时 \(\mathbf{z}=0\))。
2.4.3 训练目标
DDPM 的训练目标简化为预测噪声与真实噪声之间的 MSE:
\[\boxed{\mathcal{L}_{\text{DDPM}}(\boldsymbol{\theta}) = \mathbb{E}_{k \sim \mathcal{U}(1,K), \mathbf{x}_0 \sim q(\mathbf{x}_0), \boldsymbol{\epsilon} \sim \mathcal{N}(0,\mathbf{I})}\left[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\sqrt{\bar{\alpha}_k} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_k} \boldsymbol{\epsilon}, k)\|_2^2\right]}\]Ho et al. 证明,这个简化目标等价于最小化数据分布 \(q(\mathbf{x}_0)\) 与模型分布 \(p_{\boldsymbol{\theta}}(\mathbf{x}_0)\) 之间 KL 散度的变分下界。
Langevin 动力学视角: 反向过程可以被解释为随机 Langevin 动力学在学到的评分函数(score function)梯度场上运行:
\[\mathbf{x}_{k-1} = \mathbf{x}_k - \gamma \nabla_{\mathbf{x}} E_{\boldsymbol{\theta}}(\mathbf{x}_k) + \mathcal{N}(0, \sigma^2 \mathbf{I})\]其中 \(\nabla_{\mathbf{x}} E_{\boldsymbol{\theta}}(\mathbf{x}) \approx -\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}, k)\),即噪声预测网络近似负评分函数 \(\nabla_{\mathbf{x}} \log p(\mathbf{x})\)。这一洞察是 Diffusion Policy 训练稳定性的理论根源(详见第 2.6 节)。
2.5 Diffusion Policy:条件去噪过程
2.5.1 从图像生成到策略生成
Diffusion Policy 将 DDPM 从图像生成迁移到视觉运动策略学习,需要两项关键修改:
-
输出 \(\mathbf{x}\) 变为动作序列(而非图像像素)。具体来说,\(\mathbf{x}_0\) 是 \(T_p\) 步未来动作的拼接向量 \(\mathbf{A}_t = [\mathbf{a}_t, \mathbf{a}_{t+1}, \ldots, \mathbf{a}_{t+T_p-1}] \in \mathbb{R}^{T_p \times d_a}\)。
-
去噪过程以观测 \(\mathbf{O}_t\) 为条件。不再建模无条件分布 \(p(\mathbf{A}_t)\),而是建模条件分布 \(p(\mathbf{A}_t \mid \mathbf{O}_t)\),其中 \(\mathbf{O}_t\) 是最近 \(T_o\) 步观测数据的序列。
条件去噪: 反向步骤(公式 1)的条件化版本为(Chi et al., Eq. 4):
\[\boxed{\mathbf{A}_t^{k-1} = \frac{1}{\sqrt{\alpha_k}}\left(\mathbf{A}_t^k - \frac{1 - \alpha_k}{\sqrt{1 - \bar{\alpha}_k}} \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{O}_t, \mathbf{A}_t^k, k)\right) + \sigma_k \mathbf{z}}\]条件训练损失(Chi et al., Eq. 5):
\[\boxed{\mathcal{L}_{\text{cond}}(\boldsymbol{\theta}) = \mathbb{E}_{k, \mathbf{A}_t^0, \boldsymbol{\epsilon}}\left[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{O}_t, \mathbf{A}_t^0 + \boldsymbol{\epsilon}^k, k)\|_2^2\right]}\]其中 \(\boldsymbol{\epsilon}^k\) 是方差为 \((1 - \bar{\alpha}_k)\) 的噪声(即直接使用封闭形式的加噪公式)。
关键设计选择——只对动作空间建模: Chi et al. 选择建模 \(p(\mathbf{A}_t \mid \mathbf{O}_t)\) 而非联合分布 \(p(\mathbf{A}_t, \mathbf{O}_t)\)(与 Janner et al. 的 Diffuser 对比)。这一选择带来了两个关键优势:
- 推理速度: 视觉特征在去噪循环外提取一次即可,无需在每个去噪步骤中重新编码图像。
- 实时性: 这使得 0.1 秒级别的推理延迟成为可能(使用 DDIM 加速),满足闭环控制需求。
2.5.2 Receding-Horizon 控制
给定 DDPM 生成的 \(T_p\) 步动作序列,机器人并非执行全部 \(T_p\) 步,而是只执行前 \(T_a\) 步(\(T_a \leq T_p\)),然后重新规划。这就是 Receding-Horizon Control(RHC)的核心思想。
形式化: 在时间步 \(t\):
- 观测 \(\mathbf{O}_t\)(最近 \(T_o\) 帧)
- 生成动作序列 \(\hat{\mathbf{A}}_t = [\hat{\mathbf{a}}_t, \ldots, \hat{\mathbf{a}}_{t+T_p-1}]\)(通过 DDPM 去噪)
- 执行 \(\hat{\mathbf{a}}_t, \ldots, \hat{\mathbf{a}}_{t+T_a-1}\)(\(T_a\) 步)
- 推进到 \(t + T_a\),重复
设计超参数: \(T_o\)(观测窗口)、\(T_p\)(预测窗口)、\(T_a\)(执行窗口)。
\(T_a\) 的权衡(如图 6 实验所验证):
- \(T_a = 1\)(每步重规划): 反应性最强,但计算开销最大,且丧失了时间一致性(每步可能产生不一致的动作)。
- \(T_a = T_p\)(完全开环): 时间一致性最强,但对延迟和意外观测缺乏鲁棒性。
- 实验结论:\(\mathbf{T_a = 8}\) 在大多数任务上是最优平衡点——它提供了足够的时间一致性来避免抖动,同时保持足够的反应性来应对扰动。
延迟鲁棒性: RHC 的另一个关键优势是对系统延迟的鲁棒性。因为策略预测未来动作序列,在观测延迟(如图像处理、网络传输)存在时,策略可以”预判”未来的需求。实验表明,Diffusion Policy 可以在最多 4 步延迟的情况下维持峰值性能。
2.5.3 视觉条件化机制
Diffusion Policy 将视觉观测 \(\mathbf{O}_t\) 作为去噪过程的条件信息。两种不同的网络架构使用了不同的条件化机制:
CNN-based: FiLM(Feature-wise Linear Modulation)
FiLM 对卷积层的输出特征 \(\mathbf{f}\) 进行通道级别的仿射变换(Perez et al., 2018):
\[\text{FiLM}(\mathbf{f} \mid \mathbf{o}, k) = \boldsymbol{\gamma}(\mathbf{o}, k) \odot \mathbf{f} + \boldsymbol{\beta}(\mathbf{o}, k)\]其中 \(\boldsymbol{\gamma}, \boldsymbol{\beta}\) 是由观测嵌入和扩散步 \(k\) 的条件网络 \(\psi\) 生成的缩放和偏移向量:
\[[\boldsymbol{\gamma}, \boldsymbol{\beta}] = \psi(\text{Enc}_{\text{obs}}(\mathbf{O}_t), \text{Emb}(k))\]FiLM 被应用到 1D 时序 CNN 的每一个卷积层,使得条件信息能够在网络的所有深度层次上影响动作生成。
Transformer-based: Cross-Attention
Transformer 架构使用多头交叉注意力(multi-head cross-attention)将观测信息注入动作标记:
\[\text{CrossAttn}(\mathbf{Q}_{\text{action}}, \mathbf{K}_{\text{obs}}, \mathbf{V}_{\text{obs}}) = \text{softmax}\left(\frac{\mathbf{Q}_{\text{action}} \mathbf{K}_{\text{obs}}^T}{\sqrt{d_k}}\right) \mathbf{V}_{\text{obs}}\]其中 \(\mathbf{Q}_{\text{action}}\) 来自动作标记(含噪声的 \(\mathbf{A}_t^k\) 的每个时间步的动作向量 + 扩散步 \(k\) 的正弦嵌入),\(\mathbf{K}_{\text{obs}}, \mathbf{V}_{\text{obs}}\) 来自观测嵌入(观测帧通过共享 MLP 提取的特征序列)。动作标记还应用了因果注意力掩码(causal attention mask),确保每个动作只关注自身和之前的动作。
两种架构的比较:
| 特性 | CNN-based (1D U-Net) | Transformer-based |
|---|---|---|
| 基础模型 | 1D 时间卷积 + FiLM | minGPT 解码器 + Cross-Attn |
| 归纳偏置 | 平滑性(频域低通) | 无(注意力完全数据驱动) |
| 高频动作 | 差(temporal convolution 倾向平滑) | 优(注意力可捕获快速变化) |
| 超参数敏感性 | 低 | 高(dropout、weight decay 需精调) |
| 推荐使用场景 | 大多数任务的首选 baseline | 高频动作变化 / 高速控制任务 |
| 训练稳定性 | 高 | 中等 |
2.6 隐式策略的不稳定性与 Diffusion Policy 的解
这一节是理解 Diffusion Policy 训练稳定性优势的数学关键。
2.6.1 隐式策略的训练不稳定性来源
隐式策略(Implicit Policy, e.g., IBC, Florence et al. 2021)将动作分布表示为基于能量的模型(Energy-Based Model, EBM):
\[\boxed{p_{\boldsymbol{\theta}}(\mathbf{a} \mid \mathbf{o}) = \frac{e^{-E_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a})}}{Z(\mathbf{o}, \boldsymbol{\theta})}}\]其中:
\[Z(\mathbf{o}, \boldsymbol{\theta}) = \int_{\mathcal{A}} e^{-E_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a})} d\mathbf{a}\]是不可解的归一化常数(intractable normalization constant)。
隐式策略的训练使用 InfoNCE 风格的损失(等价于 Chi et al. Eq. 6 的负对数似然):
\[\mathcal{L}_{\text{InfoNCE}} = -\log\left(\frac{e^{-E_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a})}}{e^{-E_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a})} + \sum_{j=1}^{N_{\text{neg}}} e^{-E_{\boldsymbol{\theta}}(\mathbf{o}, \tilde{\mathbf{a}}^j)}}\right)\]其中 \(\{\tilde{\mathbf{a}}^j\}_{j=1}^{N_{\text{neg}}}\) 是负样本集合,用于估计 \(Z(\mathbf{o}, \boldsymbol{\theta})\)。
不稳定性根源: \(Z(\mathbf{o}, \boldsymbol{\theta})\) 的估计需要从高维动作空间 \(\mathcal{A} \subseteq \mathbb{R}^{d_a}\) 中采样负样本。当 \(\mathcal{A}\) 的维度很高时(例如 \(T_p \times d_a = 16 \times 7 = 112\) 维),负样本几乎不可能覆盖能量低的区域——这导致 \(Z\) 的估计存在系统性偏差,在训练中表现为梯度波动和收敛失败。
2.6.2 Diffusion Policy 的评分函数视角
DDPM 的噪声预测网络 \(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a}^k, k)\) 实际上在逼近动作评分函数(action score function)(Chi et al. Eq. 8):
\[\nabla_{\mathbf{a}} \log p(\mathbf{a} \mid \mathbf{o}) = -\nabla_{\mathbf{a}} E_{\boldsymbol{\theta}}(\mathbf{a}, \mathbf{o}) - \underbrace{\nabla_{\mathbf{a}} \log Z(\mathbf{o}, \boldsymbol{\theta})}_{= 0} \approx -\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{a}, \mathbf{o})\]关键洞察: \(Z(\mathbf{o}, \boldsymbol{\theta})\) 对 \(\mathbf{a}\) 的导数恒为零!因此评分函数独立于归一化常数。Diffusion Policy 通过让网络学习评分函数(即噪声)而非能量函数,完全规避了对 \(Z\) 的估计问题。这意味着:
- 训练时:不需要负样本采样,不需要估计 \(Z\)。只需从均匀分布采样 \(k\) 和从标准高斯采样 \(\boldsymbol{\epsilon}\)。
- 推理时:不需要求解优化问题 \(\arg\min_{\mathbf{a}} E_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a})\),而是运行 \(K\) 步评分函数引导的 Langevin 动力学。
稳定性实证: 论文图 7(训练稳定性对比)显示,IBC 在训练过程中虽然训练 Loss 平滑下降(EBM 的能量函数在改善),但评估成功率剧烈振荡。相比之下,Diffusion Policy 的训练 Loss 和评估成功率都呈现稳定收敛趋势。
2.6.3 动作分布多模态性的来源
Diffusion Policy 表达多模态性的能力来自两个随机源:
- 随机初始化: 每个推理过程中,\(\mathbf{A}_t^K \sim \mathcal{N}(0, \mathbf{I})\) 是不同的。不同的初始样本自然地收敛到评分函数梯度场的不同吸引盆(basins of attraction)。
- 随机优化轨迹: 在 \(K\) 步 Langevin 动力学过程中,每步添加的 \(\mathcal{N}(0, \sigma_k^2 \mathbf{I})\) 噪声使样本可以在不同模式之间”跳跃”。
论文图 4(Push-T 任务的多模态行为)是这一能力的经典例证:对于从下方推 T 形块的同一状态,Diffusion Policy 以约 50% 的概率选择从左绕行、50% 的概率选择从右绕行,并且在每次 rollout 中始终如一地遵循选中的模式——不像 BC-RNN 那样在两个模式间摇摆,也不像 IBC 那样偏向某一个模式。
2.7 DDIM 加速推理:方程与直觉
在 Diffusion Policy 的部署中,推理速度是硬性约束——机器人需要以 10-125 Hz 的频率生成控制指令。标准 DDPM 在训练中使用 \(K=100\) 步去噪,如果推理也使用 100 步,则单次推理可能需要数秒——这对于实时控制完全不可接受。
DDIM(Denoising Diffusion Implicit Models)(Song et al., 2021)通过将去噪过程从马尔可夫链推广为非马尔可夫过程,解耦了训练和推理的去噪步数。
2.7.1 DDIM 的采样公式
DDIM 的逆向采样不使用噪声 \(\mathbf{z}\),而是通过隐式求解器直接从 \(\mathbf{x}_k\) 跳到 \(\mathbf{x}_{k-1}\):
\[\boxed{\mathbf{x}_{k-1} = \sqrt{\bar{\alpha}_{k-1}} \cdot \underbrace{\left(\frac{\mathbf{x}_k - \sqrt{1 - \bar{\alpha}_k} \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_k, k)}{\sqrt{\bar{\alpha}_k}}\right)}_{\text{"预测的 } \mathbf{x}_0\text{"}} + \sqrt{1 - \bar{\alpha}_{k-1}} \cdot \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_k, k)}\]这个公式的直觉解读:
- 第一项:用当前噪声样本 \(\mathbf{x}_k\) 和噪声预测 \(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\) 推断”如果这是真实数据,\(\mathbf{x}_0\) 在哪里”
- 第二项:朝着”预测的 \(\mathbf{x}_0\)“方向,但以 \(\bar{\alpha}_{k-1}\) 确定的比例添加噪声(保持分布的一致性)
关键性质——确定性生成: 当 \(\sigma_k = 0\) 时,DDIM 是完全确定性的(给定 \(\mathbf{x}_K\),\(\mathbf{x}_0\) 唯一确定)。这意味着给定相同的初始噪声,DDIM 始终产生相同的动作序列——这是策略部署中可重复性需求的理想性质。
2.7.2 子序列加速
给定训练时使用 \(K\) 步,推理时可以选择 \(\mathbf{\{1, 2, \ldots, K\}}\) 的任意子序列 \(\{\tau_1, \tau_2, \ldots, \tau_S\}\)(\(S < K\))。DDIM 采样在这些子序列点上运行:
\[\mathbf{x}_{\tau_{i-1}} = \sqrt{\bar{\alpha}_{\tau_{i-1}}} \cdot \frac{\mathbf{x}_{\tau_i} - \sqrt{1 - \bar{\alpha}_{\tau_i}} \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_{\tau_i}, \tau_i)}{\sqrt{\bar{\alpha}_{\tau_i}}} + \sqrt{1 - \bar{\alpha}_{\tau_{i-1}}} \cdot \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_{\tau_i}, \tau_i)\]在 Diffusion Policy 的实践中:
- 训练使用 \(K = 100\) 步去噪
- 推理使用 \(S = 10\text{-}16\) 步(DDIM 子序列加速)
- 在 Nvidia RTX 3080 上实现约 0.1 秒的推理延迟
- 输出动作序列在 10 Hz 控制频率下以 125 Hz 线性插值
加速比:\(\frac{K}{S} \approx 6\text{-}10\times\) 的推理加速,且性能退化微乎其微。
2.7.3 噪声调度的选择
噪声调度 \(\{\beta_k\}_{k=1}^K\) 控制去噪过程中先后去除的噪声量。Chi et al. 实证发现 iDDPM 的平方余弦调度(Square Cosine Schedule, Nichol & Dhariwal, 2021)在机器人控制任务中表现最佳:
\[\bar{\alpha}_k = \frac{f(k)}{f(0)}, \quad f(k) = \cos\left(\frac{k/K + s}{1 + s} \cdot \frac{\pi}{2}\right)^2\]其中 \(s = 0.008\) 是一个小的偏移量(防止 \(\beta_1\) 过小)。
平方余弦调度的影响:相比线性调度,平方余弦调度在早期步骤(\(k\) 接近 \(K\) 时)使用更小的 \(\beta_k\)(即更慢地添加噪声),在后期步骤使用更大的 \(\beta_k\)。这使得模型在学习过程中更多地关注高噪声(即困难)的去噪任务,提升了整体质量。
2.8 位置控制 vs. 速度控制的数学分析
Diffusion Policy 的一个反直觉发现是:位置控制(预测末端执行器的绝对位姿)显著优于速度控制(预测位姿增量)。这一发现与当时的主流实践(多数 BC 工作使用速度控制)形成鲜明对比。
定义:
- 位置控制: \(\mathbf{a}_t^{\text{pos}} = \mathbf{p}_t^{\text{des}}\)(绝对目标末端执行器位姿)
- 速度控制: \(\mathbf{a}_t^{\text{vel}} = \mathbf{p}_t^{\text{des}} - \mathbf{p}_{t-1}^{\text{des}}\)(相对于上一步目标位姿的增量)
速度控制的复合误差问题: 速度控制面临更严重的分布偏移累积。设真实位姿 \(\mathbf{p}_t = \mathbf{p}_{t-1} + \hat{\mathbf{a}}_t^{\text{vel}} + \boldsymbol{\xi}_t\),其中 \(\boldsymbol{\xi}_t\) 是策略的速度预测误差。经过 \(t\) 步,累积位姿误差为:
\[\Delta \mathbf{p}_t = \sum_{i=1}^t \boldsymbol{\xi}_i\]位置的方差以 \(\mathcal{O}(t)\) 线性增长(假设 \(\boldsymbol{\xi}_i\) 独立同分布)。在操作任务中,这对应着末端执行器逐渐偏离目标位置,直至完全失败。
位置控制的自纠正性质: 位置控制不存在这个累积问题——每一步的动作是绝对位姿 \(\mathbf{p}_t^{\text{des}}\),上一步的误差不会传播到下一步。即使某一步的位姿未完全达到目标,下一步的策略会基于当前真实状态重新规划绝对目标位姿。
位置控制与动作序列预测的协同: 在动作序列预测中(\(T_p > 1\)),位置控制自然编码了轨迹的几何形状(一系列绝对目标点),而速度控制编码的是轨迹的差分(一系列增量)。前者对噪声更鲁棒,因为绝对位姿的内部一致性强于速度增量的一致性。
实验(论文图 5)证实:在 Push-T 真实世界任务中,位置控制的 Diffusion Policy 成功率为 95%,速度控制仅 10%(CNN-based)。
三、论文精读
3.1 策略表示的分类学
Chi et al. 将机器人策略表示分为三大类(论文图 1),这是理解 Diffusion Policy 所处技术上下文的关键框架:
显式策略(Explicit Policy) — \(\pi: \mathcal{O} \to \mathcal{A}\) 直接映射:
- 标量回归(Scalar Regression): 最简单但无法表达多模态性。典型工作如 ALVINN (Pomerleau, 1988),使用神经网络直接从图像预测方向盘转角。
- 高斯混合(Mixture of Gaussians, MDN): 用多个高斯分量表达多模态性。典型工作如 BC-RNN / LSTM-GMM (Mandlekar et al., 2021) 使用 LSTM + GMM 输出动作概率分布。局限性在于需要预先指定模式数量 \(M\)(超参数),且容易出现模式崩塌(mode collapse)。
- 离散分类(Categorical / Discretized): 将连续动作空间离散化为 \(B\) 个格子,转化为分类问题。典型工作如 BET (Shafiullah et al., 2022) 使用 k-means 聚类动作 + Transformer 预测聚类中心偏移。问题:格子数随动作维度指数增长(\(B^{d_a}\)),高维动作空间不可行。
隐式策略(Implicit Policy) — 通过能量函数 \(E_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a})\) 定义分布:
- 典型工作如 IBC (Florence et al., 2021)。推理时需要在能量函数上进行梯度优化 \(\mathbf{a}^* = \arg\min_{\mathbf{a}} E_{\boldsymbol{\theta}}(\mathbf{o}, \mathbf{a})\)。理论上可表达任意多模态分布,但实践中训练不稳定(详见 2.6 节)。
扩散策略(Diffusion Policy) — 通过去噪过程隐式定义策略:
- 不直接输出动作,而是学习评分函数 \(\nabla_{\mathbf{a}} \log p(\mathbf{a} \mid \mathbf{o})\)。
- 推理时从噪声通过 Langevin 动力学渐进收敛到动作样本。
- 继承了扩散模型的全部优势(多模态表达、高维输出、稳定训练)同时规避了隐式策略的不稳定性。
3.2 实验设计与关键结果解读
3.2.1 仿真基准测试(Robomimic + Push-T + BlockPush + Kitchen)
Diffusion Policy 在 12 个任务上与三种 baseline 对比:LSTM-GMM (BC-RNN), IBC, BET。
核心结果(论文表 I / 表 II):
- 平均提升 46.9%: 在所有任务和变体上,Diffusion Policy 一致且显著地超越了所有 baseline。
- 复杂任务优势更为明显: 在 Transport(2 臂 3 物 14DoF)和 ToolHang(高精度挂工具)等最困难的任务上提升幅度最大(有时数倍于 baseline)。
- 多阶段长程任务: 在 BlockPush(将两个块推入两个目标框,顺序任意)和 Kitchen(操作 7 个物体,完成任意 4 个子任务)中,Diffusion Policy 的优势尤为突出——BlockPush 的 p2 指标提升 32%,Kitchen 的 p4 指标提升 213%。这体现了 Diffusion Policy 处理长程多模态性(long-horizon multimodality)的能力。
检查点选择方法: 论文使用”最后 10 个检查点的平均值”而非”最佳检查点”,这是因为前者更公平地反映了方法的训练稳定性和实用性(真实部署无法通过后验选择最佳检查点)。Diffusion Policy 在两个指标上均显著领先。
3.2.2 真实世界任务
四项真实世界实验验证了 Diffusion Policy 的实用性:
- Push-T(UR5): 95% 成功率,接近人类水平(0.8 vs 0.84 IoU)。baseline 最佳仅 20%。
- Mug Flipping(Franka): 6DoF 下翻转随机放置的杯子,90% 成功率。策略自发学会了推把手对齐、杯子掉落重抓等从未被演示过的行为。
- Sauce Pouring(Franka): 倾倒番茄酱到披萨面团,0.74 vs 0.79 IoU(人类)。
- Sauce Spreading(Franka): 周期性螺旋涂抹,0.77 vs 0.79 覆盖率(人类)。周期性动作是 BC 中的经典难题。
关键发现——End-to-End 训练优于预训练编码器: 在 Push-T 任务上,使用预训练编码器(R3M, ImageNet)的 Diffusion Policy 性能显著低于端到端训练的版本。特别是 ImageNet 预训练的版本”预测抖动动作”,表明 ImageNet 特征可能过度关注静态语义信息而忽略了与控制相关的动态视觉线索。
3.3 消融实验中的关键洞察
3.3.1 动作空间:位置 vs. 速度
这是整篇论文最反直觉的实验结果(论文图 5)。当大多数 baseline 切换到速度控制时表现下降,而 Diffusion Policy 切换到位置控制时表现反而提升。原因分析(详见 2.8 节):
- Diffusion Policy 更好地表达了位置控制动作空间中的多模态性
- 位置控制不受速度控制复合误差的影响
- 位置控制与动作序列预测天然契合
3.3.2 动作序列预测的益处
动作序列预测(\(T_p > 1\))带来三项好处(论文第 IV-C 节讨论):
- 时间一致性: 在序列中,连续动作被迫保持一致性,避免了单步策略在不同模式间来回跳跃导致的抖动动作。
- 对空闲动作的鲁棒性: 遥操作演示中常有暂停(产生连续的相同位置/零速度动作)。单步策略容易过拟合到这种暂停行为。动作序列预测通过建模时间上下文来区分”应当暂停”和”偶然空闲”。
- 延迟补偿: 预测未来动作使得策略可以”预判”系统延迟的影响。
3.3.3 数据效率
论文图 13 显示 Diffusion Policy 在所有训练数据集规模下都优于 LSTM-GMM。即使只有极少量演示数据(如全数据集的 25%),Diffusion Policy 的优势依然保持。这表明扩散策略的表示优势独立于数据规模。
3.4 局限性分析(论文第四节中的诚实讨论)
论文在 VIII 节坦率地指出了 Diffusion Policy 的局限性:
- 行为克隆的继承局限: 如果演示数据质量不足,BC 学到的策略天然是次优的。Diffusion Policy 可以被扩展到其他范式(如 RL,后续工作如 Wang et al. 2023 已实现)。
- 计算开销: 相比 LSTM-GMM 的前向推理,DDPM 的多步去噪推理需要更多计算。DDIM 加速部分缓解了此问题,但对于需要 100+ Hz 高速控制的任务可能仍不足。
- Transformer 超参数敏感性: Transformer-based Diffusion Policy 的性能高度依赖 dropout rate 和 weight decay 的调节——不同任务的最优值差异很大,这增加了实际应用中的调参负担。
- CNN 的高频动作限制: CNN-based Diffusion Policy 由于时间卷积的归纳偏置(偏好低频信号),在需要快速、尖锐动作变化的任务上表现不佳。
四、算法与代码实现
本节提供 Diffusion Policy 的完整 PyTorch 实现。我们将实现两个核心组件:
- DDPM/DDIM 调度器 — 管理前向加噪和逆向去噪过程
- CNN-based 和 Transformer-based 扩散策略
环境要求: PyTorch >= 2.0, torchvision >= 0.15。所有代码均以教学为目的进行了简化,完整实现请参考官方仓库 https://github.com/columbia-ai-robotics/diffusion_policy。
4.1 噪声调度与 DDPM/DDIM 采样器
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Tuple, Optional
class DiffusionScheduler:
"""DDPM/DDIM 去噪扩散调度器
管理噪声调度 {β_k}、系数 {α_k, ᾱ_k}、以及前向加噪和逆向去噪操作。
论文引用: Ho et al. (2020) DDPM §2 + Nichol & Dhariwal (2021) iDDPM §3.1
"""
def __init__(self, num_train_steps: int = 100, num_infer_steps: int = 16,
beta_start: float = 0.0001, beta_end: float = 0.02,
schedule_type: str = 'squaredcosine'):
"""
Args:
num_train_steps: K_train, 训练时的去噪步骤数 (论文使用 100)
num_infer_steps: K_infer, 推理时的 DDIM 加速步数 (论文使用 10-16)
beta_start, beta_end: 线性调度的起止值
schedule_type: 'linear' | 'squaredcosine' — 论文推荐后者
论文引用: Noise Schedule §III-C, 推荐 iDDPM 的 Squared Cosine Schedule
"""
self.K_train = num_train_steps
self.K_infer = num_infer_steps
# 计算噪声调度 {β_k}_{k=1}^K
if schedule_type == 'squaredcosine':
# 公式: ᾱ_k = cos²((k/K + s)/(1 + s) * π/2) / cos²(s/(1 + s) * π/2)
# 其中 s = 0.008 (防止 β_1 过小)
s = 0.008
steps = torch.arange(num_train_steps + 1, dtype=torch.float32)
f = torch.cos(((steps / num_train_steps + s) / (1 + s)) * (np.pi / 2)) ** 2
alphas_cumprod = f / f[0] # ᾱ_k
betas = 1 - alphas_cumprod[1:] / alphas_cumprod[:-1]
betas = torch.clamp(betas, max=0.999) # 数值稳定性限制
else:
betas = torch.linspace(beta_start, beta_end, num_train_steps)
# 预计算 DDPM 所需的所有系数
alphas = 1.0 - betas # α_k = 1 - β_k
alphas_cumprod = torch.cumprod(alphas, dim=0) # ᾱ_k = ∏_{i=1}^k α_i
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) # ᾱ_{k-1}
# 注册为 buffer (不参与梯度但随模型保存)
self.register('betas', betas)
self.register('alphas', alphas)
self.register('alphas_cumprod', alphas_cumprod)
self.register('alphas_cumprod_prev', alphas_cumprod_prev)
# 从 ᾱ_k 推导的中间系数
self.register('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
self.register('sqrt_one_minus_alphas_cumprod', torch.sqrt(1.0 - alphas_cumprod))
self.register('sqrt_recip_alphas', torch.sqrt(1.0 / alphas))
# DDPM 逆向采样公式中的系数: 1/√α_k * (1-α_k)/√(1-ᾱ_k)
self.register('posterior_variance',
betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod))
# DDIM 推理子序列: 从 K_train 中均匀采样 K_infer 个时间步
if num_infer_steps < num_train_steps:
step_ratio = num_train_steps // num_infer_steps
self.ddim_steps = torch.arange(0, num_train_steps, step_ratio).long()
# 确保包含最后一步
if self.ddim_steps[-1] != num_train_steps - 1:
self.ddim_steps = torch.cat([
self.ddim_steps, torch.tensor([num_train_steps - 1])
])
else:
self.ddim_steps = torch.arange(num_train_steps)
self.register('ddim_steps', self.ddim_steps)
def register(self, name: str, tensor: torch.Tensor):
"""模拟 nn.Module.register_buffer(用于非 nn.Module 类)"""
setattr(self, name, tensor)
def add_noise(self, x0: torch.Tensor, k: torch.Tensor,
eps: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
"""前向加噪: x_k = √ᾱ_k * x_0 + √(1 - ᾱ_k) * ε
论文引用: DDPM 训练 §II-B, Eq. (3)/(5) 的前向步骤
"""
if eps is None:
eps = torch.randn_like(x0)
sqrt_ac = self.sqrt_alphas_cumprod[k].view(-1, 1, 1) # [B, 1, 1]
sqrt_omac = self.sqrt_one_minus_alphas_cumprod[k].view(-1, 1, 1)
xk = sqrt_ac * x0 + sqrt_omac * eps
return xk, eps
def predict_x0_from_noise(self, xk: torch.Tensor, k: torch.Tensor,
noise_pred: torch.Tensor) -> torch.Tensor:
"""从噪声预测反推干净样本: x̂_0 = (x_k - √(1-ᾱ_k) * ε_θ) / √ᾱ_k
这是 DDIM 采样中的"预测的 x_0"步骤
论文引用: DDIM §3.1, Eq. (12) 的 x0 预测项
"""
sqrt_ac = self.sqrt_alphas_cumprod[k].view(-1, 1, 1)
sqrt_omac = self.sqrt_one_minus_alphas_cumprod[k].view(-1, 1, 1)
return (xk - sqrt_omac * noise_pred) / sqrt_ac
@torch.no_grad()
def ddim_sample(self, noise_pred_net: nn.Module, cond: torch.Tensor,
shape: Tuple[int, ...], device: torch.device) -> torch.Tensor:
"""DDIM 加速采样 — 从纯噪声渐变到干净动作序列
论文引用: DDIM §3.1 (Song et al., 2021), Diffusion Policy §III-D
Args:
noise_pred_net: 噪声预测网络 ε_θ(O, A^k, k)
cond: 观测条件 O_t, shape [B, T_o, d_obs_feat]
shape: 动作张量形状 [B, T_p, d_a]
device: 计算设备
Returns:
去噪后的动作序列 A^0_t, shape [B, T_p, d_a]
"""
b = shape[0]
x = torch.randn(shape, device=device) # A^K_t ~ N(0, I)
# 沿 DDIM 子序列逆向迭代 (从 K_train-1 到 0)
ddim_steps = self.ddim_steps.to(device)
for i in reversed(range(len(ddim_steps))):
k = ddim_steps[i]
k_tensor = torch.full((b,), k, device=device, dtype=torch.long)
# 预测噪声: ε_θ(O, A^k, k)
noise_pred = noise_pred_net(x, cond, k_tensor)
# 推断 x̂_0
x0_pred = self.predict_x0_from_noise(x, k_tensor, noise_pred)
if i > 0:
k_prev = ddim_steps[i - 1]
# DDIM 逆向步骤:
# x_{k-1} = √ᾱ_{k-1} * x̂_0 + √(1-ᾱ_{k-1}) * ε_θ
sqrt_ac_prev = self.sqrt_alphas_cumprod[k_prev].view(-1, 1, 1)
sqrt_omac_prev = self.sqrt_one_minus_alphas_cumprod[k_prev].view(-1, 1, 1)
x = sqrt_ac_prev * x0_pred + sqrt_omac_prev * noise_pred
else:
x = x0_pred # 最后一步直接输出 x̂_0
return x
@torch.no_grad()
def ddpm_sample(self, noise_pred_net: nn.Module, cond: torch.Tensor,
shape: Tuple[int, ...], device: torch.device) -> torch.Tensor:
"""DDPM 随机采样 (训练/评估模式)
论文引用: DDPM §II-A, Eq. (1); Diffusion Policy §II-C, Eq. (4)
与 DDIM 不同: DDPM 的每步添加 σ_k * z 随机噪声
"""
b = shape[0]
x = torch.randn(shape, device=device) # A^K_t ~ N(0, I)
K = self.K_train
for k in reversed(range(K)):
k_tensor = torch.full((b,), k, device=device, dtype=torch.long)
# 预测噪声
noise_pred = noise_pred_net(x, cond, k_tensor)
# DDPM 逆向步骤 (§II-A, Eq. (1)):
# x_{k-1} = 1/√α_k * (x_k - (1-α_k)/√(1-ᾱ_k) * ε_θ) + σ_k * z
sqrt_recip_alpha = self.sqrt_recip_alphas[k].view(-1, 1, 1)
one_minus_alpha = 1.0 - self.alphas[k]
sqrt_one_minus_ac = self.sqrt_one_minus_alphas_cumprod[k].view(-1, 1, 1)
x = sqrt_recip_alpha * (x - one_minus_alpha / sqrt_one_minus_ac * noise_pred)
if k > 0:
sigma = torch.sqrt(self.posterior_variance[k]).view(-1, 1, 1)
x = x + sigma * torch.randn_like(x)
return x
4.2 视觉编码器
import torchvision.models as models
class SpatialSoftmax(nn.Module):
"""空间 Softmax 池化 — 保留空间位置信息
替代 ResNet 的全局平均池化 (GAP),提取特征点的加权空间位置。
对于操作任务,GAP 会丢弃"物体在图像哪个位置"的关键信息。
论文引用: Visual Encoder §III-B, 取自 Mandlekar et al. (2021)
"""
def __init__(self, temperature: float = 1.0):
super().__init__()
self.temperature = temperature
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""x: [B, C, H, W] → 输出: [B, 2C] (每个通道的 (x,y) 软位置)"""
B, C, H, W = x.shape
x_flat = x.view(B, C, -1) # [B, C, H*W]
# 对每个通道计算 softmax 权重
weights = F.softmax(x_flat / self.temperature, dim=-1) # [B, C, H*W]
# 创建坐标网格
y_grid, x_grid = torch.meshgrid(
torch.linspace(-1, 1, H, device=x.device),
torch.linspace(-1, 1, W, device=x.device),
indexing='ij'
)
x_grid, y_grid = x_grid.reshape(-1), y_grid.reshape(-1) # [H*W]
# 加权坐标期望
expected_x = (weights * x_grid).sum(dim=-1) # [B, C]
expected_y = (weights * y_grid).sum(dim=-1) # [B, C]
return torch.cat([expected_x, expected_y], dim=-1) # [B, 2C]
class VisualEncoder(nn.Module):
"""基于 ResNet-18 的视觉编码器
论文引用: Visual Encoder §III-B
- ResNet-18 基础架构
- 全局平均池化 → 空间 Softmax 池化 (保留空间信息)
- BatchNorm → GroupNorm (与 DDPM 的 EMA 兼容,稳定训练)
- 无 ImageNet 预训练 (端到端训练表现最佳)
"""
def __init__(self, in_channels: int = 3, output_dim: int = 64,
spatial_softmax_temp: float = 1.0):
super().__init__()
# ResNet-18 backbone (无原始分类头)
resnet = models.resnet18(weights=None) # 无预训练
self.backbone = nn.Sequential(
resnet.conv1, # 7x7 conv, stride 2
resnet.bn1,
resnet.relu,
resnet.maxpool, # 3x3 maxpool, stride 2
resnet.layer1, # 2x BasicBlock [64, H/4, W/4]
resnet.layer2, # 2x BasicBlock [128, H/8, W/8]
resnet.layer3, # 2x BasicBlock [256, H/16, W/16]
resnet.layer4, # 2x BasicBlock [512, H/32, W/32]
)
self.spatial_softmax = SpatialSoftmax(temperature=spatial_softmax_temp)
self.projection = nn.Sequential(
nn.Linear(1024, output_dim), # 512*2 (spatial softmax)
nn.ReLU(),
nn.Linear(output_dim, output_dim)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""x: images [B, T_o, C, H, W]
→ 独立编码每帧 → 按时间拼接 → 线性投影
→ output: [B, T_o, output_dim]
论文引用: §II-C — "different camera views use separate encoders,
and images in each timestep are encoded independently and then concatenated"
"""
B, T_o, C, H, W = x.shape
x = x.view(B * T_o, C, H, W) # 将时间维合并到 batch 维
features = self.backbone(x) # [B*T_o, 512, h, w]
features = self.spatial_softmax(features) # [B*T_o, 1024]
features = self.projection(features) # [B*T_o, output_dim]
features = features.view(B, T_o, -1) # [B, T_o, output_dim]
return features
4.3 CNN-based Diffusion Policy
class FiLMBlock(nn.Module):
"""Feature-wise Linear Modulation 块
对卷积特征进行通道级别的条件仿射变换:
FiLM(f | c) = γ(c) ⊙ f + β(c)
论文引用: FiLM conditioning §III-A, 基于 Perez et al. (2018)
应用到 1D 时间 CNN 的每个卷积层
"""
def __init__(self, in_channels: int, cond_dim: int):
super().__init__()
# 条件网络: 观测嵌入 + 扩散步 → 仿射参数 (γ, β)
self.condition_mlp = nn.Sequential(
nn.Linear(cond_dim, cond_dim * 2),
nn.ReLU(),
nn.Linear(cond_dim * 2, in_channels * 2)
)
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
"""x: [B, C, L], cond: [B, cond_dim]
→ output: [B, C, L] (条件调制后的特征)
"""
gamma_beta = self.condition_mlp(cond) # [B, 2C]
gamma, beta = gamma_beta.chunk(2, dim=-1) # [B, C], [B, C]
gamma = gamma.unsqueeze(-1) # [B, C, 1]
beta = beta.unsqueeze(-1) # [B, C, 1]
return gamma * x + beta
class ResidualConvBlock(nn.Module):
"""1D 残差卷积块: Conv1d → GroupNorm → ReLU → Conv1d → GroupNorm"""
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
self.conv1 = nn.Conv1d(in_channels, out_channels, 3, padding=1)
self.norm1 = nn.GroupNorm(8, out_channels) # GroupNorm for stable EMA training
self.conv2 = nn.Conv1d(out_channels, out_channels, 3, padding=1)
self.norm2 = nn.GroupNorm(8, out_channels)
self.shortcut = nn.Conv1d(in_channels, out_channels, 1) \
if in_channels != out_channels else nn.Identity()
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = self.shortcut(x)
out = F.relu(self.norm1(self.conv1(x)))
out = self.norm2(self.conv2(out))
return F.relu(out + residual)
class SinusoidalPositionEmbedding(nn.Module):
"""扩散步 k 的正弦位置嵌入 (Transformer 风格)
论文引用: DDPM §3.3 (Positional Embeddings for diffusion step)
"""
def __init__(self, dim: int):
super().__init__()
self.dim = dim
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
device = timesteps.device
half_dim = self.dim // 2
emb = torch.exp(torch.arange(half_dim, device=device, dtype=torch.float32) *
(-np.log(10000) / (half_dim - 1)))
emb = timesteps.float().unsqueeze(-1) * emb.unsqueeze(0)
emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
return emb
class CNNDiffusionPolicy(nn.Module):
"""CNN-based 扩散策略 (1D 时间 U-Net + FiLM 条件化)
架构: 1D 时间卷积 encoder-decoder,每个卷积层应用 FiLM 条件化。
输入观测 O_t 和扩散步 k,在每个 FiLM 块中调制卷积特征。
论文引用: §III-A, Fig 3(b)
- 基于 Janner et al. (2021) 的 1D 时序 CNN
- 仅建模 p(A_t | O_t),不建模联合分布 p(A_t, O_t)
- FiLM 条件化观测 O_t 和扩散步 k
- 预测动作轨迹 (非拼接的观测-动作轨迹)
"""
def __init__(self, obs_dim: int, act_dim: int, pred_horizon: int = 16,
hidden_dim: int = 256, num_layers: int = 6):
"""
Args:
obs_dim: 观测嵌入维度 (视觉编码器输出)
act_dim: 动作维度 (e.g., 7 for 6DoF + gripper)
pred_horizon: T_p — 动作预测窗口
hidden_dim: FiLM 条件嵌入维度
num_layers: 卷积层数
"""
super().__init__()
self.pred_horizon = pred_horizon
self.act_dim = act_dim
# 扩散步 k 的正弦位置嵌入
self.k_embedding = SinusoidalPositionEmbedding(hidden_dim)
# 观测嵌入 (将不同时间步的观测特征聚合为条件向量)
self.obs_encoder = nn.Sequential(
nn.Linear(obs_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim)
)
# FiLM 条件融合 MLP: 观测嵌入 + 扩散步嵌入 → FiLM 条件
self.cond_fusion = nn.Sequential(
nn.Linear(hidden_dim * 2, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim)
)
# 动作嵌入: 将含噪动作 A^k_t 提升到隐空间
self.action_embedding = nn.Conv1d(act_dim, hidden_dim, kernel_size=1)
# Encoder: 下采样 + FiLM 条件化的残差卷积块
self.encoder_blocks = nn.ModuleList()
self.downsample_layers = nn.ModuleList()
current_dim = hidden_dim
for i in range(num_layers):
self.encoder_blocks.append(
nn.ModuleList([
ResidualConvBlock(current_dim, current_dim),
FiLMBlock(current_dim, hidden_dim)
])
)
if i < num_layers - 1:
self.downsample_layers.append(
nn.Conv1d(current_dim, current_dim * 2, 3, stride=2, padding=1)
)
current_dim *= 2
# Decoder: 上采样 + 残差卷积块
self.decoder_blocks = nn.ModuleList()
self.upsample_layers = nn.ModuleList()
for i in range(num_layers):
self.decoder_blocks.append(ResidualConvBlock(current_dim, current_dim))
if i < num_layers - 1:
self.upsample_layers.append(
nn.ConvTranspose1d(current_dim, current_dim // 2, 3,
stride=2, padding=1, output_padding=1)
)
current_dim //= 2
# 输出投影: 隐空间 → 动作空间 (预测噪声 ε_θ)
self.output_proj = nn.Conv1d(current_dim, act_dim, kernel_size=1)
def forward(self, noisy_actions: torch.Tensor, obs_emb: torch.Tensor,
k: torch.Tensor) -> torch.Tensor:
"""前向传播 — 预测噪声
Args:
noisy_actions: A^k_t, shape [B, T_p, d_a] — 含噪动作序列
obs_emb: O_t 嵌入, shape [B, T_o, d_o] — 编码后的观测序列
k: [B] — 当前扩散步编号
Returns:
noise_pred: [B, T_p, d_a] — 预测的噪声 ε_θ(O, A^k, k)
论文引用: Diffusion Policy §II-C, Eq. (4)/(5)
"""
B = noisy_actions.shape[0]
# 1. 条件编码
k_emb = self.k_embedding(k) # [B, hidden_dim]
obs_pooled = obs_emb.mean(dim=1) # [B, d_o] — 时间池化
obs_feat = self.obs_encoder(obs_pooled) # [B, hidden_dim]
cond = self.cond_fusion(torch.cat([obs_feat, k_emb], dim=-1)) # [B, hidden_dim]
# 2. 动作嵌入 (1D Conv: 时间维度不变, 通道维度提升)
x = noisy_actions.transpose(1, 2) # [B, d_a, T_p]
x = self.action_embedding(x) # [B, hidden_dim, T_p]
# 3. Encoder (每个残差块 → FiLM 调制 → 下采样)
skip_connections = []
for i, (res_block, film_block) in enumerate(self.encoder_blocks):
x = res_block(x)
x = film_block(x, cond)
skip_connections.append(x)
if i < len(self.downsample_layers):
x = self.downsample_layers[i](x)
# 4. Decoder (上采样 → 残差块)
for i, res_block in enumerate(self.decoder_blocks):
if i < len(self.upsample_layers):
x = self.upsample_layers[i](x)
# Skip connection (从 encoder 同层的输出)
skip = skip_connections[-(i + 2)] # 从后往前取
if x.shape[-1] != skip.shape[-1]:
x = F.interpolate(x, size=skip.shape[-1])
x = x + skip
x = res_block(x)
# 5. 输出投影
noise_pred = self.output_proj(x) # [B, d_a, T_p]
noise_pred = noise_pred.transpose(1, 2) # [B, T_p, d_a]
return noise_pred
4.4 Transformer-based Diffusion Policy
class TransformerDiffusionPolicy(nn.Module):
"""Transformer-based 扩散策略 (minGPT 解码器 + Cross-Attention)
架构: 动作标记 (含噪 A^k_t 的每个时间步) 作为 transformer decoder 的输入 tokens,
扩散步 k 的正弦嵌入作为第一个 token 前置,观测嵌入 O_t 通过共享 MLP 转换为
encoder 特征序列,注入每个 decoder 层的交叉注意力。
动作标记之间使用因果注意力掩码 (只能 attend 到自身和之前的动作)。
论文引用: §III-A, Fig 3(c) — Time-series Diffusion Transformer
"""
def __init__(self, obs_dim: int, act_dim: int, pred_horizon: int = 16,
n_layers: int = 8, n_heads: int = 8, emb_dim: int = 256,
attn_dropout: float = 0.1):
"""
Args:
obs_dim: 观测嵌入维度
act_dim: 动作维度
pred_horizon: T_p — 动作预测窗口
n_layers: Transformer decoder 层数
n_heads: 多头注意力头数
emb_dim: 标记嵌入维度
attn_dropout: 注意力 dropout (不同任务差异大, 论文表 VII)
"""
super().__init__()
self.pred_horizon = pred_horizon
self.act_dim = act_dim
# 扩散步 k 的正弦位置嵌入 → 第一个 token
self.k_embedding = SinusoidalPositionEmbedding(emb_dim)
self.k_proj = nn.Linear(emb_dim, emb_dim)
# 动作嵌入: 含噪动作 A^k_t 的每个时间步 → token
self.action_tokenizer = nn.Linear(act_dim, emb_dim)
# 动作时间步位置编码 (Transformer 需要的序列位置信息)
self.action_pos_embedding = nn.Parameter(
torch.randn(1, pred_horizon, emb_dim) * 0.02
)
# 观测编码器: O_t 的每个时间步 → token
self.obs_tokenizer = nn.Linear(obs_dim, emb_dim)
self.obs_pos_embedding = nn.Parameter(
torch.randn(1, 1, emb_dim) * 0.02
)
# Transformer Decoder 层 (minGPT 风格 + cross-attention)
decoder_layer = nn.TransformerDecoderLayer(
d_model=emb_dim,
nhead=n_heads,
dim_feedforward=emb_dim * 4,
dropout=attn_dropout,
activation='gelu',
batch_first=True
)
self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_layers)
# 输出投影: token → 预测噪声
self.output_proj = nn.Linear(emb_dim, act_dim)
# 因果注意力掩码 (动作标记只能看到自身和之前的动作)
self.register_buffer(
'causal_mask',
torch.triu(torch.ones(pred_horizon, pred_horizon) * float('-inf'), diagonal=1)
)
def forward(self, noisy_actions: torch.Tensor, obs_emb: torch.Tensor,
k: torch.Tensor) -> torch.Tensor:
"""前向传播 — 预测噪声
Args:
noisy_actions: [B, T_p, d_a]
obs_emb: [B, T_o, d_o]
k: [B]
Returns:
noise_pred: [B, T_p, d_a]
论文引用: §III-A, Fig 3(c) — "Actions with noise A^k_t are passed in as
input tokens for the transformer decoder blocks, with the sinusoidal
embedding for diffusion iteration k prepended as the first token"
"""
B, T_p, _ = noisy_actions.shape
# 1. 观测 token 序列 (作为 cross-attention 的 memory)
obs_tokens = self.obs_tokenizer(obs_emb) # [B, T_o, emb_dim]
# 2. 动作 token 序列
act_tokens = self.action_tokenizer(noisy_actions) # [B, T_p, emb_dim]
act_tokens = act_tokens + self.action_pos_embedding[:, :T_p, :] # 位置编码
# 扩散步前置 token
k_emb = self.k_proj(self.k_embedding(k)) # [B, emb_dim]
k_token = k_emb.unsqueeze(1) # [B, 1, emb_dim]
# 动作序列 = [k_token] + [action tokens]
decoder_input = torch.cat([k_token, act_tokens], dim=1) # [B, 1+T_p, emb_dim]
# 3. Transformer Decoder (Self-Attn: causal mask, Cross-Attn: obs_tokens)
# causal mask 扩展到 1+T_p (额外的一个 k_token)
causal_mask = torch.zeros(1 + T_p, 1 + T_p, device=noisy_actions.device)
causal_mask[1:, 1:] = self.causal_mask[:T_p, :T_p] # 动作部分使用因果掩码
decoded = self.decoder(
tgt=decoder_input,
memory=obs_tokens,
tgt_mask=causal_mask
) # [B, 1+T_p, emb_dim]
# 4. 输出投影 — 去掉第一个 token (k_token),只保留 T_p 个动作 token
action_outputs = decoded[:, 1:, :] # [B, T_p, emb_dim]
noise_pred = self.output_proj(action_outputs) # [B, T_p, d_a]
return noise_pred
4.5 完整训练与推理流程
class DiffusionPolicyAgent:
"""Diffusion Policy 完整 Agent — 整合所有组件
论文引用: 完整 Diffusion Policy 流程 §II-C, §III
"""
def __init__(self, config: dict):
"""
config keys:
obs_horizon (T_o), pred_horizon (T_p), action_horizon (T_a)
obs_dim, act_dim
arch: 'cnn' | 'transformer'
num_train_steps, num_infer_steps
lr, weight_decay
"""
self.T_o = config['obs_horizon'] # 观测窗口
self.T_p = config['pred_horizon'] # 预测窗口
self.T_a = config['action_horizon'] # 执行窗口
# 扩散调度器
self.scheduler = DiffusionScheduler(
num_train_steps=config.get('num_train_steps', 100),
num_infer_steps=config.get('num_infer_steps', 16),
schedule_type='squaredcosine'
)
# 视觉编码器
self.visual_encoder = VisualEncoder(
in_channels=3,
output_dim=config['obs_dim']
)
# 噪声预测网络
if config['arch'] == 'cnn':
self.noise_pred_net = CNNDiffusionPolicy(
obs_dim=config['obs_dim'],
act_dim=config['act_dim'],
pred_horizon=config['pred_horizon']
)
elif config['arch'] == 'transformer':
self.noise_pred_net = TransformerDiffusionPolicy(
obs_dim=config['obs_dim'],
act_dim=config['act_dim'],
pred_horizon=config['pred_horizon']
)
else:
raise ValueError(f"Unknown architecture: {config['arch']}")
# 优化器 (论文对 CNN 和 Transformer 使用不同权重衰减)
if config['arch'] == 'cnn':
self.optimizer = torch.optim.AdamW(
list(self.noise_pred_net.parameters()) +
list(self.visual_encoder.parameters()),
lr=config.get('lr', 1e-4),
weight_decay=config.get('weight_decay', 1e-6)
)
else:
self.optimizer = torch.optim.AdamW(
list(self.noise_pred_net.parameters()) +
list(self.visual_encoder.parameters()),
lr=config.get('lr', 1e-4),
weight_decay=config.get('weight_decay', 1e-3) # Transformer 更高的 wd
)
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.to(self.device)
def to(self, device):
self.visual_encoder.to(device)
self.noise_pred_net.to(device)
self.device = device
def compute_loss(self, batch: dict) -> torch.Tensor:
"""训练损失 — Diffusion Policy 条件 MSE
论文引用: §II-C, Eq. (5)
L = MSE(ε, ε_θ(O_t, A^0_t + ε^k, k))
batch keys:
'images': [B, T_o, C, H, W] — 观测图像序列
'actions': [B, T_p, d_a] — 目标动作序列
"""
images = batch['images'].to(self.device)
actions = batch['actions'].to(self.device) # A^0_t (干净的动作序列)
B = actions.shape[0]
# 1. 编码视觉观测
obs_emb = self.visual_encoder(images) # [B, T_o, d_o]
# 2. 随机采样扩散步 k ~ Uniform(0, K-1)
k = torch.randint(0, self.scheduler.K_train, (B,), device=self.device)
# 3. 前向加噪: A^k_t = √ᾱ_k * A^0_t + √(1-ᾱ_k) * ε
noise = torch.randn_like(actions) # ε ~ N(0, I)
noisy_actions, _ = self.scheduler.add_noise(actions, k, eps=noise)
# 4. 前向传播: 预测噪声 ε_θ(O_t, A^k_t, k)
noise_pred = self.noise_pred_net(noisy_actions, obs_emb, k)
# 5. MSE 损失
loss = F.mse_loss(noise_pred, noise)
return loss
def training_step(self, batch: dict) -> float:
"""单步训练"""
self.visual_encoder.train()
self.noise_pred_net.train()
loss = self.compute_loss(batch)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
@torch.no_grad()
def infer_action(self, images: torch.Tensor) -> np.ndarray:
"""推理 — 生成动作序列 (DDIM 加速)
将最近的 T_o 帧图像输入,输出 T_a 步动作。
论文引用: Receding-horizon control §II-C + DDIM acceleration §III-D
Args:
images: [1, T_o, C, H, W] — 单条观测图像序列
Returns:
action: [T_a, d_a] — 待执行的前 T_a 步动作
"""
self.visual_encoder.eval()
self.noise_pred_net.eval()
images = images.to(self.device)
B = images.shape[0]
# 1. 编码观测 (在去噪循环外,一次完成)
obs_emb = self.visual_encoder(images) # [1, T_o, d_o]
# 2. DDIM 采样
actions = self.scheduler.ddim_sample(
self.noise_pred_net,
cond=obs_emb,
shape=(B, self.T_p, self.noise_pred_net.act_dim),
device=self.device
) # [1, T_p, d_a]
# 3. Receding-horizon: 只取前 T_a 步
actions = actions[0, :self.T_a, :].cpu().numpy() # [T_a, d_a]
return actions
4.6 数据归一化
class ActionNormalizer:
"""动作数据归一化器
论文引用: Appendix A — Normalization
- 每个动作维度独立缩放到 [-1, 1] (通过 min-max normalization)
- 零均值归一化会造成 DDPM 预测范围外区域不可达 (因为 DDPM 每步 clip 到 [-1,1])
- 旋转表示 (e.g., 6D rotation) 不做归一化
- 方差极小的维度 (近常量) 做零均值居中但不缩放
"""
def __init__(self):
self.mins = None
self.maxs = None
self.has_scale = None # 哪些维度需要缩放
def fit(self, action_data: np.ndarray):
"""action_data: [N, T, d_a] — 所有演示动作数据"""
actions_flat = action_data.reshape(-1, action_data.shape[-1])
self.mins = actions_flat.min(axis=0)
self.maxs = actions_flat.max(axis=0)
# 如果 min ≈ max (近常量),不做缩放
self.has_scale = (self.maxs - self.mins) > 1e-5
def normalize(self, actions: np.ndarray) -> np.ndarray:
"""归一化到 [-1, 1]"""
actions = actions.copy()
for d in range(actions.shape[-1]):
if self.has_scale[d]:
actions[..., d] = 2.0 * (actions[..., d] - self.mins[d]) / \
(self.maxs[d] - self.mins[d]) - 1.0
else:
actions[..., d] = actions[..., d] - self.mins[d] # 仅居中
return actions
def denormalize(self, actions: np.ndarray) -> np.ndarray:
"""反归一化: [-1, 1] → 原始范围"""
actions = actions.copy()
for d in range(actions.shape[-1]):
if self.has_scale[d]:
actions[..., d] = (actions[..., d] + 1.0) / 2.0 * \
(self.maxs[d] - self.mins[d]) + self.mins[d]
else:
actions[..., d] = actions[..., d] + self.mins[d]
return actions
五、核心 Takeaway 与延伸阅读
核心洞察
-
评分函数是连接扩散模型与策略学习的数学桥梁。Diffusion Policy 通过学习动作评分函数 \(\nabla_{\mathbf{a}} \log p(\mathbf{a} \mid \mathbf{o})\) 而非直接输出动作,从根本上规避了隐式策略中归一化常数的估计问题。这个设计使得训练异常稳定——因为评分函数独立于 \(Z(\mathbf{o}, \boldsymbol{\theta})\),不需要负样本采样,不需要复杂的训练技巧。
-
多模态性不是 BC 的 bug,而是 feature。人类演示数据天然是多模态的——同一个任务有多条正确的执行路径。Diffusion Policy 通过随机初始化 + Langevin 动力学随机优化,能够自然地表达和采样多模态动作分布。在每个 rollout 中,它承诺于一个模式(而非在模式间摇摆),实现了时间一致性和多模态表达的统一。
-
位置控制 + 动作序列预测 + Receding Horizon = 鲁棒的闭环控制配方。这三个设计决策的协同效应是 Diffusion Policy 实用性的基础:位置控制避免了速度控制的复合误差;动作序列预测提供了时间一致性;Receding-horizon control 在一致性和反应性之间找到了平衡点。\(T_a = 8\) 是大多数任务的经验最优值。
-
端到端训练优于预训练视觉编码器。对于操作任务,任务特定的视觉特征(如接触点、物体边缘)比 ImageNet 的语义特征(如”这是一个杯子”)更重要。端到端训练的 ResNet-18 + 空间 Softmax 池化 + GroupNorm 是 Diffusion Policy 中最有效的视觉编码组合。
-
CNN 先尝试,Transformer 后优化。CNN-based Diffusion Policy 超参数鲁棒、训练稳定,是大多数新任务的推荐起点。当任务需要高频动作变化(如速度控制空间)或表现瓶颈明确来自 CNN 的低频偏置时,再转向 Transformer-based 架构——尽管它的调参成本更高。
开放问题与研究方向
- 超越行为克隆: Diffusion Policy 天然适合扩展为 RL 策略(后续工作 Wang et al. 2023, Hansen-Estruch et al. 2023 已探索),使其能从次优甚至负面数据中学习。
- 一致性模型加速: Song et al. (2023) 的 Consistency Models 可以实现单步生成,将推理延迟从 100ms 降至 1ms 量级,对高速控制任务具有革命性意义。
- 多模态条件化: 如何将语言指令、目标图像、3D 场景理解同时作为 Diffusion Policy 的条件输入?
- 跨具身泛化: 在一种机器人上训练的 Diffusion Policy 能否迁移到不同形态的机器人?这需要解耦策略表示中的具身特定组件。
推荐延伸阅读
| 文献 | 关键贡献 | 与本章的关联 |
|---|---|---|
| Ho et al. (2020), “DDPM” | 去噪扩散概率模型的原始公式 | DDPM 数学基础,噪声预测简化 |
| Nichol & Dhariwal (2021), “iDDPM” | 改进的噪声调度 + 学习方差 | 平方余弦调度 (本章使用的噪声调度) |
| Song et al. (2021), “DDIM” | 非马尔可夫加速采样 | 推理加速 (100步→16步),确定性生成 |
| Ross et al. (2011), “DAgger” | 数据集聚合解决分布偏移 | BC 复合误差的理论分析框架 |
| Florence et al. (2021), “IBC” | 隐式行为克隆 | 隐式策略的对比方法论 |
| Mandlekar et al. (2021), “Robomimic” | 大规模 BC 基准 | 本章实验评估的主要平台 |
| Shafiullah et al. (2022), “BET” | 行为 Transformer | BC 中多模态建模的替代方案 |
| Janner et al. (2022), “Diffuser” | 扩散规划 | 联合建模 p(A, O) 的对比方案 |
| Song et al. (2023), “Consistency Models” | 单步扩散生成 | Diffusion Policy 的未来加速方向 |
| Wang et al. (2023), “Diffusion-QL” | 扩散策略 + RL | Diffusion Policy 从 BC 到 RL 的扩展 |
六、练习与思考
基础练习:
-
DDPM 前向过程验证: 假设 \(\beta_k = 0.001\) 对所有 \(k\)(恒定噪声调度),\(K = 1000\)。计算 \(\bar{\alpha}_{500}\) 和 \(\bar{\alpha}_{1000}\)。当 \(k=1000\) 时,信号与噪声的比例(SNR = \(\frac{\sqrt{\bar{\alpha}_k}}{\sqrt{1-\bar{\alpha}_k}}\))是多少?解释为什么 \(k=K\) 时 \(\mathbf{x}_K\) 近似为纯噪声。
-
BC 复合误差模拟: 在一维环境中模拟 BC 的复合误差:状态转移为 \(s_{t+1} = 1.05 s_t + a_t\)(\(L_T = 1.05\)),专家策略为 \(\pi^E(s) = -0.8s\),BC 学到的策略为 \(\hat{\pi}(s) = \pi^E(s) + \epsilon\)(\(\epsilon = 0.1\))。从 \(s_0 = 10\) 开始,运行 50 步,绘制 \(\vert s_t^{\hat{\pi}} - s_t^{\pi^E}\vert\) 随 \(t\) 的变化曲线,验证指数增长趋势。
-
DDIM 加速实现: 修改
DiffusionScheduler.ddim_sample方法,支持不同的推理步数 \(S \in \{5, 10, 16, 25, 50, 100\}\)。在一个简单的 2D 数据分布(如 Swiss Roll)上训练 DDPM,比较不同 \(S\) 下的生成质量和推理时间。
进阶练习:
-
位置 vs. 速度控制实验: 在一个简单的 2D 到达任务(reach a target point)中,实现位置控制和速度控制两种动作空间。对每种控制模式,分别训练一个简单的 BC 策略(MLP),添加单步噪声扰动 \(\sim \mathcal{N}(0, \sigma^2)\) 并测量 100 步后的终点位置误差。改变 \(\sigma\),绘制位置控制和速度控制的误差曲线,验证速度控制的累积误差效应。
- 多模态建模能力对比: 在二维 double-reach 任务(target 在 \((1,1)\) 和 \((-1,-1)\) 两处随机出现,每个初始状态有两种有效路径)上,比较以下方法的”模式切换频率”(连续两次 rollout 中策略选择不同模式的概率):
- MSE 损失的 MLP(应表现为模式平均)
- 混合密度网络(MDN, \(M=3\))
- 简化的 Diffusion Policy(1D CNN backbone) 分析每种方法成功/失败的原因。
- FiLM vs. Cross-Attention 消融: 在同一任务(如 Push-T 仿真版)上训练 CNN-based (FiLM) 和 Transformer-based (Cross-Attention) Diffusion Policy。对比:
- 收敛速度(达到相同成功率的训练轮数)
- 超参数敏感性(改变隐藏维度、学习率对性能的影响)
- 高频动作跟踪能力(测量动作序列的逐帧变化量的方差) 结论是否与论文第 III-A 节的分析一致?
开放研究问题:
-
扩散策略 vs. 流匹配: 流匹配(Flow Matching, Lipman et al. 2023)是一种替代扩散的生成建模范式,使用 ODE 而非 SDE 来连接噪声和数据。比较 DDPM 和流匹配在策略学习中的理论优势和劣势:训练效率、推理速度、分布覆盖质量。如果训练一个基于流匹配的视觉运动策略,receding-horizon 控制的设计是否可以直接沿用?
-
RL + Diffusion Policy: Wang et al. (2023) 将 Diffusion Policy 扩展到离线 RL (Diffusion-QL)。阅读该论文,分析 DDPM 的训练目标如何从 BC 的 MSE 损失修改为包含 Q 值引导的损失。扩散策略的随机性如何自然地实现了 RL 中的探索?在在线 RL 设置中,扩散策略的推理延迟是否会影响样本效率?