25年10月来自香港大学、西湖大学、纽卡斯尔大学和优必选科技公司的论文“On-Device Diffusion Transformer Policy for Efficient Robot Manipulation”。
扩散策略通过模仿学习显著提升机器人操作任务的性能，但由于计算效率低下和内存占用过大，其在资源受限的移动平台上的应用仍然面临挑战。本文提出一种名为 LightDP 的框架，旨在加速扩散策略在移动设备上的实时部署。LightDP 通过两项核心策略解决计算瓶颈问题：一是压缩去噪模块的网络，二是减少所需的采样步数。首先对现有的扩散策略架构进行广泛的计算分析，发现去噪网络是造成延迟的主要原因。为了克服传统剪枝方法通常导致的性能下降，引入一个统一的剪枝和重训练流程，显式地优化模型剪枝后的可恢复性。此外，将剪枝技术与一致性蒸馏相结合，在有效减少采样步数的同时，保持动作预测的准确性。在标准数据集（即 PushT、Robomimic、CALVIN 和 LIBERO）上的实验评估表明，LightDP 在移动设备上实现具有竞争力的实时动作预测性能，这标志着基于扩散的策略在资源受限环境中的实际部署迈出重要一步。
一致性模型是一类扩散生成模型，旨在通过一步将噪声输入直接映射到对应的干净输入，从而高效地生成数据。它们强制执行自洽性，确保模型的输出在不同的噪声水平下保持不变，即 f_θ(x_t,t) = f_θ(x_t′,t′)，其中 x_t 和 x_t′ 是沿 ODE 轨迹在不同时间步 t 和 t′ 采集的样本。在 EDM 框架中，一致性模型采用边界条件 c_skip(0) = 1 和 c_out(0) = 0。一种训练这些模型的方法称为一致性蒸馏，它通过最小化一致性损失L_CD来改进预训练的扩散模型。
问题描述
近年来，模仿学习技术的进步使得机器人能够从人类专家收集的演示中学习复杂的操作任务。给定演示 T，轨迹 τ ∈ T 是观测值 o 和机器人动作 a 的序列，记为 τ = {(o_1, a_1), ..., (o_N_τ , a_N_τ )}。扩散策略 π_φ(a|o, g) 的训练目标是通过最大化给定观测值 o 和目标 g 时动作 a 的对数似然，来模仿专家的行为。在多模态设置下，目标 g 是一条高级指令，用于指定任务的预期结果，可以是语言指令或目标观测值。通常，由 φ 参数化的扩散策略由观测编码器 E、扩散transformer D 和目标编码器 G 组成。观测编码器 E 从观测值 o 中提取特征，而扩散transformer D 生成基于观测值 o 和目标 g 的动作 a。扩散策略通过得分匹配（score matching）估计时间步 t 的得分函数 ∇_a log p(a|o, g)，其中 π_φ = a + σ^2^_t∇_a log p(a|o, g) 是神经网络，a_t 是时间步 t 的带噪声动作，α(σ_t) 是损失权重。扩散模型通过最小化得分匹配损失 LDM 进行训练，这鼓励模型生成与专家演示一致的动作。在这项工作中，本文专注于通过剪枝和蒸馏算法加速预训练策略模型，然后将这些模型部署到移动设备上，以实现实时机器人操作。
扩散策略的延迟分析
由于扩散策略是为实时机器人操作而设计的，因此评估策略模型的设备端延迟至关重要。鉴于这些模型之间的结构相似性，以 MDT-V 【36】模型为例。如图所示，MDT-V 模型支持多种输入模态，包括用于提取图像特征的观察编码器（即 MDT-V 模型中的 Voltron 网络 [21]）、用于处理高级指令的目标编码器（即 CLIP 文本编码器）以及用于生成机器人动作的扩散transformer。
如表所示，在 iPhone 13 上评估 Policy Transformer (DP-T)【8】和 MDT-V 【36】模型的延迟。对于 DP-T，该网络由两个主要组件构成：图像编码器采用 ResNet18 模型将输入图像转换为嵌入，作为扩散Transformer的条件，这部分延迟仅占总延迟的一小部分（1.28 毫秒）。扩散Transformer是一个8层Transformer，它是模型的主要瓶颈（耗时90.6毫秒），需要100次迭代去噪步骤才能得到最终的动作预测。在MDT-V模型中也观察到了类似的现象，其中Voltron网络的耗时（7.56毫秒）远低于扩散Transformer（22.25毫秒），这减缓设备端的生成过程。分析策略模型的架构，得到模型的瓶颈在于扩散Transformer，在两个模型中都是如此。扩散Transformer的架构可以表示为N个Transformer模块的堆叠，每个模块包含一个多头注意层（MHA）和一个前馈网络（FFN）层，表示为φ_i = FFN(MHA(·))。由于扩散Transformer需要多次去噪步骤才能生成动作预测，这导致模型的延迟较高。为了解决这个问题，建议通过剪枝和蒸馏来加速模型。
通过学习进行模型剪枝
为了获得更小的模型，采用层剪枝技术来移除扩散transformer中的冗余层。给定一个 N 层扩散transformer，目标是找到一个二值掩码 M(N) = {m_1, m_2, ..., m_N }，用于标识需要剪枝的层，其中 m_i ∈ {0, 1} 表示该层是保留还是剪枝。通常，剪枝过程被表述为一个优化问题，旨在最小化剪枝后的损失 L，该问题可以表述为min_M,πφˆ E_x [L(x,π ,M)]，其中 π_φ =Π(φ_i) 是原始模型，π_φˆ 是剪枝后的模型。
然而，由于掩码 M 和权重 φˆ 需要联合优化，因此该剪枝问题是 NP-难 [3, 14]。为了解决这个问题，一种常见的方法是两阶段剪枝过程：首先确定掩码 M（通过最小化给定准则下的损失 L），然后微调剪枝后的模型以恢复性能。然而，这种两步法可能并非最优，因为模型在剪枝后可能无法完全恢复性能。为了解决这个问题，提出使用单阶段剪枝方法 [10]，其中掩码 M 和权重 φˆ 联合优化以最小化剪枝后的损失 L。
具体来说，M 被建模为概率分布 M_i ∼ Bernoulli((p_i))，其中 p_i 是在训练过程中优化的门得分。用奇异值分解 (SVD) 来估计层重要性，因为 SVD 是模型压缩中常用的技术 [17,25]。与规范多项式分解或克罗内克积分解等其他方法相比，奇异值分解 (SVD) 能够提供奇异值，从而捕捉权重矩阵中最重要的分量。用 SVD 分解来初始化门控得分。具体来说，SVD 分解应用于每个 Transformer 模块的权重矩阵，包括注意层和 FFN 模块中 MLP 层的Q权重矩阵、K权重矩阵和V权重矩阵。然后，通过 p_i = I(φ_i)/sum(I(φ_i)) 初始化门控得分，其中 φ_i 是扩散 Transformer 第 i 个模块中的权重矩阵。
如图所示，该模型采用基于 Gumbel-Softmax 技巧 [11, 20] 的可学习门选择机制进行训练，该机制可用于选择待剪枝的块。如果在训练过程中丢弃第 i 个块，则将其输出与其输入相同（恒等映射）。门得分在训练过程中更新，可用于选择待剪枝的块。训练结束后，为了获得 N 层扩散transformer，选择门得分最高的 N 层。为了进一步恢复剪枝后的性能，继续对模型进行微调，但不采用掩码选择过程。
单步蒸馏
使用剪枝后的模型，单步推理速度可以显著提高。然而，该模型仍然需要多次去噪步骤才能获得高质量的动作预测，这会带来不可忽略的计算成本。为了解决这个问题，采用一致性蒸馏来训练模型，使其成为一个一致性模型，这样可以在更少的去噪步骤下达到与原始模型相当的性能。
一致性蒸馏旨在训练模型π_φ，使其在不同噪声水平下满足一致性属性，记为π_φ(a_t, o, g, σ_t) = π_φ(a_t′, o, g, σ_t′)。蒸馏后的模型被重新参数化为EDM。
如上图所示，学生模型f_φ使用教师模型f_ψ进行初始化，然后通过可学习剪枝技术进行剪枝。给定采样演示数据 (o, a, g)，首先添加噪声以获得时间步 t + k 的噪声动作 a_t+k，然后使用教师模型 f_ψ 预测时间步 t 的噪声动作 at。接着，将两个噪声动作 a_t+k 和 at 输入到学生模型 f_φ 和目标模型 f_φ⋆ 中，以计算一致性损失 L_CD。目标模型 f_φ⋆ 使用参数 f_φ 的指数移动平均值 (EMA) 进行更新，定义为 f_φ⋆ ← sg(μf_φ⋆ +(1−μ)f_φ)，其中 sg(·) 表示stopgrad运算，μ 满足 0 ≤ μ < 1。学生模型和目标模型均使用教师模型进行初始化。
实现细节
基础模型。本文提到 Diffusion Policy Transformer 和 MDT-V，因为它们在模仿学习中应用广泛，尤其是在目标操作任务中。目标是压缩模型，使其在移动设备上更高效、更快速地运行。因此，选择这两个模型作为基础模型。Diffusion Policy Transformer 是一个基于 Transformer 的策略网络，仅支持图像输入。该模型由一个扩散 Transformer 和一个视觉编码器组成。
MDT 是一个多模态策略网络，它集成预训练的多模态特征提取器 Voltron。还实现 MoDE，这是一个基于 MoE 的策略网络，在 CALVIN 和 LIBERO 基准测试中取得最先进的性能。在实验中，考虑压缩广泛使用的扩散策略，包括 Diffusion-Policy-T [8] 和 MDT [36]。 Diffusion-Policy-T [8] 是一个基于 Transformer 的策略网络，用于模仿学习，仅支持图像输入。MDT [36] 通过集成预训练的多模态特征提取器 Voltron [21]，在 CALVIN 数据集上取得了良好的结果。
实现细节。实现基于 PyTorch。在 NVIDIA RTX 3090 和 H800 GPU 上进行训练。然后，将 GPU 上训练的模型转换为 Core ML 模型格式（mlpackage，基于 Apple 的 ml-stable-diffusion），并在 iPhone 13（A15 Bionic，iOS 18.3.1）上使用 Xcode Instruments 测量延迟。对于网络剪枝，采用 TinyFusion [10] 中的局部块剪枝方案，构建一个方案为 N:M 的局部块。在这个 N:M 方案中，每组连续的 M 层（一个“块”）被剪枝到 N 层。例如，当从一个总共有 M = 4 层的局部块中保留 N = 3 层时，我们有 4 种选择，对应于 M = [[1,1,1,0],[1,1,0,1],[1,0,1,1],[0,1,1,1]]。按照惯例，将一致性蒸馏应用于模型的 x_0 预测（预测去噪后的动作），并将 EMA 衰减率初始设置为 0.95，并在训练过程中逐渐增加到 0.999，以稳定目标模型的更新。用 DDIM 求解器 [40] 进行蒸馏，跳跃间隔为 10 步（即每 10 个扩散步进行一次蒸馏）。尽可能保持超参数与基础模型的原始实现一致。对于 DP-T 模型，输入是 RGB 图像和低维状态的混合，图像尺寸为 84 × 84，观测序列长度设置为 2，DiT的 Transformer 模块隐藏层大小为 256，头数为 4，DP-T 模型的层数为 8。对于 MDT 模型，输入是多模态的，包括两幅不同视角的 RGB 图像作为观测值，以及一条语言指令作为目标。采用 AdamW 作为优化器，学习率为 1e-4，批大小设置为 64。在 CALVIN 数据集上训练模型 30 个 epoch，在最后几个 epoch 中，根据第 20 个 epoch 的 gate 分数对 Student Model f_φ 进行剪枝。
基于 DP-T 和 MoDE 两种模型，将 LightDP 部署在两个机械臂上（DP-T 使用 Inovo 机器人，MoDE 使用 Lebai 机器人），每个任务执行 20 次。如图和如表所示，精简后的模型在这些实际任务中取得相当的成功率。考虑到大多数家庭用户通常没有必要购买高级设备，选择最易获取且便携的设备（即 iPhone）作为机器人开发平台。此外，还基于 Jetson Orin NX（16 GB，Jetpack 5.1.1）评估方法，基于 DP-T 的延迟为 244.68 毫秒（基于 DP-T-D6/6-8 的延迟为 37.69 毫秒）。