交叉扩散：通过自监督学习改进基于扩散的视觉运动策略

24年1月来自纽约州 Stony Brook U 的论文“Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning”。

序列建模方法在机器人模仿学习中已显示出良好的效果。最近，扩散模型已被用于以序列建模方式进行行为克隆，这得益于它们在建模复杂数据分布方面的卓越能力。基于扩散的标准策略以输入状态为条件，从随机噪声中迭代生成动作序列。尽管如此，扩散策略模型在视觉表示方面还可以进一步改进。这项工作提出 交叉扩散（Crossway Diffusion） ，一种简单有效的方法，通过设计的状态解码器和辅助的自监督学习 (SSL) 目标来增强基于扩散的视觉运动策略学习。状态解码器从反向扩散过程的中间表示中重建原始图像像素和其他状态信息。整个模型由 SSL 目标和原始扩散损失共同优化。实验证明交叉扩散在各种模拟和现实世界机器人任务中的有效性，证实它相对于基于标准扩散的策略具有优势，并且比基线有显着改进。

行为克隆 (BC) [1] 是一种用于机器人动作策略学习的监督学习公式。给定由一系列状态-动作对组成的专家演示数据，训练模型预测给定输入状态（例如图像）的正确动作向量。事实证明，该框架非常有效，尤其是在提供足够数量的训练数据后 [2]。

最近，序列建模方法 [3]、[4]、[5] 经常用于行为克隆，因为它们能够对多步信息进行建模。在这种公式中，目标是对多步状态-动作轨迹的概率分布进行建模。这使得 BC 能够考虑单步回归以外，更好地利用历史。

鉴于在建模人类语言 [6] 和图像 [7] 方面的成功，Transformers [8] 已被广泛用于基于序列建模的策略 [3]、[4]、[5]、[2]、[9]。

扩散模型 [10]、[11]、[12]、[13] 在建模多模态数据分布和从分布中生成新样本方面具有出色的能力，这使得它们适合生成轨迹来模仿行为。最近的研究 [14]、[15]、[16] 已成功将扩散模型应用于使用低维状态的顺序建模。对于视觉运动控制任务，[17] 使用包括视觉观察在内的多模态状态作为扩散模型的条件，表现出良好的性能。

在 机器人策略学习 中，扩散模型已用于数据增强 [23] 和样本合成 [24]、[25]、[26]。最近的研究 [3]、[4]、[27] 将机器人策略学习表述为动作序列生成问题，扩散模型也已成功改编为序列生成模型 [16]、[23]、[15]、[14]。Diffuser [14] 将低维状态和动作连接起来，并将多个状态-动作对建模为矩阵（图像）生成问题。然而，由于视觉观察的高维性，这种设置对于视觉运动策略学习是不可行的。扩散策略 [17] 仅生成动作序列来应对挑战，同时以视觉和其他状态为条件。具体来说，给定一个 Ts 状态序列 St = {s/t−Ts +1 , s/t−Ts +2 , ..., s/t }，其中每个 s 包含视觉状态和低维状态，扩散模型根据状态序列生成一个 Ta 动作序列 At = {a/t, a/t+1, ..., a/t+Ta−1}。当智体完成执行 At 时，扩散策略在给定 S/t+Ta 的情况下生成后续动作序列 A/t+Ta，从而形成闭环控制，如图所示：Ts = 2，Ta = 4。

给定一个同时具有视觉和低维状态的状态序列 St = {S/t,img , S/t,low-dim }，状态编码器从图像 h/t,img = E/S (S/t,img ) 中提取视觉嵌入。然后将视觉嵌入 h/t,img 与其他低维状态 S/t,low-dim 连接起来，形成观察条件 h/t = h/t,img ⊕ S/t,low-dim。

动作编码器 E/A 采用扩散步骤 k 处的噪声动作序列 Ak/t 和观察条件 h/t，并从最深层生成表示 Xk/t，从较浅层生成用于跳连接 Xk/t,skip 的其他张量：Xk/t, Xk/t,skip = E (Ak/t, h/t , k)，其中 t 是状态或动作轨迹的时间戳。

动作解码器采用 Xk/t，Xk/t,skip 和条件 ht 来估计前向扩散过程中应用于 Ak/t 的噪声 ε。条件 ht 应用于残差块中的两个卷积层之间，使用逐特征线性调制 (FiLM) [28]。然后从估计的噪声 εθ 和 Ak/t 得出（略微）去噪的动作序列 Ak−1/t。

在推理过程中，上述去噪过程迭代重复 K 次，最终生成无噪声的动作序列 A0t。整个模型使用与 DDPM [11] 相同的均方误差 (MSE) 进行优化，以预测应用于 A0/t 去构建 Ak/t 的噪声 ε。

交叉扩散扩展了现有的扩散策略 [17]，引入了 1) 状态解码器和 2) 辅助目标，两者均用于重建输入状态。交叉扩散的整体架构如图所示。具体来说，状态解码器采用扩散过程 Xt 的中间表示来重建输入状态。重建目标在训练期间与扩散损失 L/DDPM 联合优化。

这个表示 Xk/t 被称为 相交 ，因为动作去噪流和状态重建流都经过这个张量，然后前往它们相应的目的地。出于同样的原因，方法命名为 交叉扩散（Crossway Diffusion） ，其主要特征是在训练阶段在扩散模型中让两个信息流相交。通过这样做，明确规范 Xk/t 以捕获来自两个流的信息，这有利于表示学习。设计包括两种不同类型的状态解码器：一种用于视觉状态，另一种用于低维状态。对于每个状态源（多个摄像头、关节角度等），分配一个专用解码器以获得最佳重建结果。为了重建视觉状态，视觉状态解码器由一系列 2D 残差卷积块、用于上采样的转置卷积层 (ConvTranspose) 和原始卷积层组成，如图所示。

其他低维状态由三层 MLP 回归。隐藏层宽度与低维状态的宽度之比为 4:2:1。请注意，使用状态解码器 DS 的重构仅在训练期间使用，充当一个“解释器”生成额外的监督信号，训练更好的中间表示。状态解码器将在训练后被丢弃。

相交变换 g(·) 调整中间表示 Xk/t，以便状态解码器执行重建任务。假设 Xk/t 沿时间轴拆分为向量列表 Xk/t := Xk/t,i |i = 1, 2, ..., T 。默认情况下，只有列表中的第一个向量 Xk/t,1 被发送到状态解码器进行状态重建。

低维状态解码器 g/low-dim(·) 的变换是一个恒等（identity）函数，这意味着状态解码器直接将 Xk/t,1 作为输入。