ManiCM

基于扩散模型的 Policy 不可避免地面临运行效率低下的问题，因为在推理阶段需要大量采样步骤以生成高质量的动作，即使采用了一些采样加速方法 [36]。

一致性模型（Consistency Models, CMs）[37, 22] 被引入图像生成领域，能够通过极少的推理步骤生成高保真图像。在本文中，提出了 ManiCM 方法，将一致性模型的概念应用于实时机器人操作。

具体而言，首先在点云条件下构建了一个扩散过程，以去噪机器人动作，并确保一致性属性，即原始动作可以从 ODE 轨迹上的任意点直接去噪。然后，设计了一种一致性蒸馏技术直接预测动作样本，而不是像常见的那样预测噪声，因为预测去噪动作收敛更快。

在 Adroit 和 Metaworld 的 31 个机器人操作任务上的大量实验表明，该方法在推理速度上显著加速了 (10x) 当前最先进的方法，同时保持了具有竞争力的成功率。

---

Problem Formulation

Consistency Model, CM

CM 是一种高效的生成模型，能够在单步或少步推理中生成高质量的结果，同时保持与扩散模型相当的性能。

给定一个概率流常微分方程（PF-ODE），该方程可以将数据转换为噪声，一致性模型的目标是学习该 PF-ODE 的解函数 $f(\cdot, \cdot)$，使得 ODE 轨迹上的任意点都可以通过 $f(\cdot, \cdot)$ 映射回原始分布。

给定 PF-ODE 的解轨迹 $\{x_t\}_{t \in [\epsilon, T]}$，一致性函数 $f$ 定义为 $f : (x_t, t) \to x_\epsilon$，其中 $t \in [0, T]$，$T > 0$ 是一个固定常数。为了避免数值不稳定性，$\epsilon$ 是一个非常小的正数，$\hat{x}_\epsilon$ 可以被视为数据分布的近似值（$\hat{x}_\epsilon \sim p_{\text{data}}(x)$）。

根据 [37]，一致性函数需要满足以下自一致性属性（self-consistency property）：

$$f(x_t, t) = f(x_{t'}, t'), \quad \forall t, t' \in [\epsilon, T].$$

如上式所示，自一致性属性的含义是，对于同一 PF-ODE 轨迹上的任意输入对 $(x_t, t)$，它们的输出 $f(x_t, t)$ 应保持一致。所有一致性模型还需满足边界条件 $f(x_\epsilon, \epsilon) = x_\epsilon$，这一条件在一致性模型训练的成功中起着重要作用。

为了将自一致性属性引入深度神经网络，将函数 $f(\cdot, \cdot)$ 实现为：

$$f_\theta(x, t) = c_{\text{skip}}(t)x + c_{\text{out}}(t)F_\theta(x, t),$$

其中 $c_{\text{skip}}(t)$ 和 $c_{\text{out}}(t)$ 是可微函数，满足 $c_{\text{skip}}(\epsilon) = 1$ 和 $c_{\text{out}}(\epsilon) = 0$，$F_\theta(\cdot, \cdot)$ 是一个参数化的深度神经网络，用于学习自一致性。

Consistency Distillation

一致性蒸馏是一种通过从预训练的扩散模型（即“教师网络”）中提取知识来训练一致性模型的方法。一致性损失定义为：

$$L(\theta, \theta^-; \Phi) = \mathbb{E}\left[d\left(f_\theta(x_{t_{n+1}}, t_{n+1}), f_{\theta^-}(\hat{x}_\Phi^{t_n}, t_n)\right)\right],$$

其中 $d(\cdot, \cdot)$ 是用于衡量两个样本之间距离的度量函数，$f_\theta(\cdot, \cdot)$ 和 $f_{\theta^-}(\cdot, \cdot)$ 分别被称为“在线网络”（online network）和“目标网络”（target network）[37]。当在线网络的优化收敛时，目标网络最终会匹配在线网络，因为 $\theta^-$ 是 $\theta$ 的指数移动平均（EMA）。只要 ODE 求解器的步长足够小且一致性蒸馏损失趋于零，估计的一致性模型就可以变得任意精确。重要的是，边界条件 $f_\theta(x, \epsilon) \equiv x$ 确保模型在训练中不会收敛到无意义的解 $f_\theta(x, t) \equiv 0$。

ManiCM: Manipulation Consistency Model

选择 DP3 [48] 作为底层扩散模型，将其知识蒸馏到单步采样器中，从而实现实时推理，同时保持操作质量。

3D Conditional Generation

将点云表示注入扩散策略作为条件。具体来说，给定从单视角 RGB-D 摄像头接收到的深度图像，我们首先通过摄像头的内参和外参获取点云。然后，通过 FPS 对点云进行下采样，以聚焦于感兴趣的区域（如机器人手臂和目标物体）。接着，稀疏点云通过 MLP 点云编码器进行编码，生成紧凑的 3D 表示，对原始扩散模型进行条件化：

$$a_{t-1} = \alpha_t \left(a_t - \gamma_t \epsilon_\theta(a_t, t, p, q)\right) + \sigma_t N(0, I), \quad t > 0,$$

其中 $\epsilon_\theta$ 是预测的噪声，$p$ 和 $q$ 分别表示点云条件和机器人位姿。上述公式也被称为反向过程 [14]。

Manipulation Self-Consistency Function

$$a_0 = c_{\text{skip}}(t)a_t + c_{\text{out}}(t)F_\theta(a_t, t, p, q), \quad t \geq 0,$$

其中 $F_\theta$ 是一个操作一致性函数，包含了沿轨迹累积的 $\epsilon_\theta$。在实践中，$c_{\text{skip}}(t) = \sigma_d^2 / (\beta_t^2 + \sigma_d^2)$，$c_{\text{out}}(t) = \beta_t / \sqrt{\beta_t^2 + \sigma_d^2}$，其中 $\beta_t$ 表示时间步缩放值，$\sigma_d$ 是平衡值。$F_\theta(\cdot, \cdot)$ 可以被参数化为预测动作样本或噪声 $\epsilon$。在一般情况下，这两种预测方法都是可行的，因为它们本质上优化的是相同的扩散过程。

操作自一致性函数 $F_\theta(\cdot, \cdot)$ 参数化为预测动作样本而非噪声。

Manipulation Consistency Distillation

为了训练 ManiCM，采用一致性蒸馏方法。这里，$a_0$ 表示机器人手臂动作序列，包含多个动作，其维度由具体任务决定。$p_0$ 表示从图像中获取并下采样的点云，$q_0$ 表示从离线数据集中获取的机器人手臂当前位姿。这些输入通过编码器处理，生成紧凑的表示，随后作为条件输入到教师网络、在线网络和目标网络中。

给定一个真实的专家动作 $a_0$，首先通过 $n + k$ 步的前向扩散操作对 $a_0$ 添加噪声，得到噪声动作 $a_{n+k}$，其中 $k$ 是跳步间隔。随后，噪声动作 $a_{n+k}$ 被输入到冻结的教师网络和在线网络中，分别解码出干净的 $\hat{a}_0^*$ 和 $\hat{a}_0$。为了确保自一致性属性，设计了一个从在线网络克隆的目标网络，并期望在线网络和目标网络的去噪输出保持一致。

具体来说，目标网络利用通过 $k$-步 ODE 求解器 $\Phi$（例如 DDIM [36]）处理 $\hat{a}_0^*$ 得到的 $a_n$ 来生成 $\hat{a}_0^-$。因此，在线网络和目标网络的输出 $\hat{a}_0$ 和 $\hat{a}_0^-$ 由一致性函数 $f(\cdot, \cdot)$ 生成。一致性蒸馏损失表达为：

$$L_{\text{MCD}}(\theta, \theta^-) = \mathbb{E}\left[\|f_\theta(a_{n+k}, t_{n+k}) - f_{\theta^-}(\hat{a}_\Phi^n, t_n)\|_2^2\right],$$

其中 $\|\cdot\|_2$ 是 $\ell_2$ 范数。此外，$\theta^-$ 使用在线网络参数 $\theta$ 的指数移动平均（EMA）更新，定义为：

$$\theta^- \leftarrow \text{sg}(\mu \theta^- + (1 - \mu)\theta),$$

其中 $\text{sg}(\cdot)$ 表示停止梯度操作，$\mu$ 满足 $0 \leq \mu < 1$。在公式 (6) 中，$\hat{a}_\Phi^n$ 是从 $a_{n+k}$ 通过 $k$-步 ODE 求解器得到的估计值。因此，$\hat{a}_\Phi^n$ 可表示为：

$$\hat{a}_\Phi^n \leftarrow a_{n+k} + (t_n - t_{n+k})\Phi(a_{n+k}, t_{n+k}),$$

其中 $\Phi(\cdot, \cdot)$ 是应用于 PF-ODE 的 $k$-步 ODE 求解器的更新函数。

利用 EMA 从在线网络的参数更新目标网络的参数。在实验中，将预训练的 3D 操作扩散模型（如 DP3 [48]）作为“教师网络”。实际上，在线网络和目标网络均以教师网络的参数初始化。