【论文阅读】ChainedDiffuser: Unifying Trajectory Diffusion and Keypose Prediction for Robotic Manipulation

Abstract

我们提出了chaineddiffuser，这是一种policy架构，它结合了动作键预测和轨迹扩散生成，用于从演示中学习机器人操作。我们的主要创新是使用全局基于转换器的动作预测器来预测关键帧的动作，这是一项需要多模态语义场景理解的任务，并使用局部轨迹扩散器来预测连接预测宏动作的轨迹段。chaineddiffuser 在已建立的操作基准上设置了新的记录，并且优于使用运动规划器进行轨迹预测的 state-of-the-art 关键姿势（宏观动作）预测模型，以及不预测关键帧宏动作的轨迹扩散 policies。我们在模拟和现实世界环境中进行了实验，并展示了 chaineddiffuser 解决涉及不同对象交互的广泛操作任务的能力。

Keywords: Manipulation, Imitation Learning, Transformers, Diffusion Models

1 Introduction

虽然从演示中学习操作policies是一个监督学习问题，但动作轨迹的多模态和多样性给机器学习方法带来了巨大的挑战。一些任务，例如将杯子放在橱柜中，可以由policy处理，policy只为杯子提供所需的目标姿态[1,2,3]，而其他任务，如地板上擦掉污垢，需要policy为抓取的mop生成连续的动作轨迹[4,5]。

一行操作学习方法从演示中建模动作轨迹。这些方法要么反应性地将视觉和语言映射到密集的时间动作 temporal actions[6，7，8，5，9]，要么使用基于能量的模型对输入动作兼容性进行建模[10，11，12，13]。尽管最近取得了进展，但这些方法可能会遇到多模态动作轨迹分布或体验训练稳定性 experience training stabilities[14,13,15]。基于扩散模型[16，17，18]的成功，最近的一项工作建议训练基于扩散的policies[14，4，19]来生成动作轨迹。这些方法已经证明了稳定的训练行为和令人印象深刻的捕捉多模态动作轨迹分布的能力。然而，它们还没有在长期操作任务上进行测试。 

另一项工作将机器人操作问题转化为预测关键帧上离散末端执行器动作序列 [1, 20, 12, 21]。该范式从连续演示中提取关键帧并预测这些关键帧中的末端执行器动作[2,22,3,21]。随后，低级路径规划器连接预测的关键姿势（宏观动作），并返回符合环境和任务约束的完整轨迹。利用基于注意力的架构[23]的最新进展，许多方法将关键帧动作预测扩展到6自由度的语言指令操作任务[2,3,22,24,21]。

关键帧预测背后的假设阻碍了它对超出拾取和放置类型的操作任务的适用性。许多任务，例如擦拭桌子，打开门，同时尊重运动学约束等，只能通过与环境的连续交互来解决。此外，对低级路径规划的依赖进一步限制了这些方法的能力:虽然一系列任务需要无碰撞的轨迹，但其他任务，如物体推[24,3,25]，需要运动规划者忽略避碰。尽管在模拟数据集[26]中很容易获得这种额外推理的监督，但现实世界的人类演示通常缺乏这样的数据，更不用说现实世界中的无碰撞运动规划需要精确的状态估计，这带来了其自身的挑战。

鉴于上述情况，我们提出了 chaineddiffuser，这是一种统一上述两种范式的神经架构。chaineddiffuser 是一个 policy 架构，它以输入视觉信号，可选地，语言指令并输出时间密集的末端执行器动作。在粗略级别，它使用全局基于转换器的动作预测器预测宏观步骤末端执行器动作（我们称之为宏观动作 macro-actions），这是一项需要全局理解视觉环境和完成任务的高级任务。然后，低级轨迹扩散器生成局部轨迹段来连接预测的宏动作。与基于transformer的宏观步骤预测方法[2,3,22,24]相比，我们的模型预测平滑轨迹以适应需要连续交互和无碰撞动作的任务。与仅扩散轨迹生成方法 [14, 15, 4, 19] 相比，我们的分层方法以更结构化的方式处理长视距任务，并允许不同的模块专注于它们表现出色的任务。

我们在RLBench[26]上测试ChainedDiffuser，这是一个已建立的从演示中学习操作基准。我们在先前文献[22,24]研究的各种任务和场景中评估我们的模型。ChainedDiffuser 设置了新的技术水平，并且优于没有预测宏动作或使用回归或运动规划器进行关键帧到关键帧轨迹预测的消融版本。此外，我们使用少量人类演示进行训练，在具有许多长期操作任务的现实场景中验证了我们的模型。

2 Related Work

Learning from Demonstrations

[27,28]一直是机器人的常见范式，但需要在现实世界中进行演示数据收集[6,29,30]或模拟[26,31,32]。为了提高数据效率，有几种方法在预先训练的视觉表示之上学习policy，这些视觉表示利用了大型的纯视觉数据集[33,34,35,36,37,38]。与此正交，其他方法将每个任务抽象为一系列子目标，表示为拾取和放置图元[1,2]或关键帧[39,40]。在这种情况下，采用手工设计的低级控制器来规划末端执行器在中间子目标之间的运动。虽然数据高效，但这种抽象并不能充分推广到只有少数尊重所有物理约束的特定轨迹有效的场景[41]，例如可变形[42,43]或铰接[44]物体的操纵、闭环机器人系统的运动[45,46]，或在混乱环境中穿过障碍物的轨迹[47]。因此，最近的研究对每个额外的约束（例如，碰撞避免、轨迹平滑度[4]）都采用了半手动成本规范。更接近我们的方法，james和abbeel[41]学会了对手工设计或基于学习的planners提出的轨迹进行评分。相反，我们训练场景条件扩散模型来生成连接预测keyposes的轨迹。

Transformers for Robotics

继它们在自然语言处理[23,48,49]和计算机视觉[50,51]方面的成功之后，最近的许多工作使用基于transformer的架构进行机器人和控制[52,53,6,54,55,21]。一个主要的动机是在结合来自多个感官流的信息时，注意力对长期预测的灵活性，例如视觉观察和语言指令[56,22]。与我们最相关的是基于多任务Transformer的模型流，这些模型在不同的数据集上进行训练，以实现更高的分布内[3,24]或分布外泛化[6,57,58,59]。我们的模型由两个基于注意力的模块组成，一个用于宏观步骤动作预测，一个用于局部轨迹优化，它可以利用不同的输入模式并在不同的抽象上运行。

Diffusion Models

[60，18，16，17]学习通过迭代去噪过程来逼近数据分布，并在无条件和条件图像生成方面显示出令人印象深刻的结果[61，62，63，64]。在机器人领域，扩散模型在规划[15,65,66]、场景重排[67,68]、可控运动优化[69,70]、视频生成[71]和模仿学习[14,19]等方面找到了应用。它们的主要优点是与以前的生成模型相比，它们可以更好地捕获动作轨迹分布。最近的工作使用扩散模型来预测完整的轨迹，通常是自回归的[14,72]。相反，我们使用扩散模型来生成与宏动作链接的局部轨迹。

3 ChainedDiffuser

3.1 Overview

chaineddiffuser的体系结构如图1所示。chaineddiffuser将宏动作预测与条件轨迹扩散相结合。它的输入包括环境的视觉观察和任务的自然语言描述 l。在每一步，chaineddiffuser 使用全局 policy πglobal 预测宏动作 ^at，然后将 ^at 及其当前末端执行器状态 qt 馈送到低级局部轨迹生成器 πlocal(qt, ^at) 以生成连接 qt 和 ^at 的密集微动作，如图 1 所示。 at 和 qt 共享相同的空间 a = {apos, arot, agrip}，由末端执行器的 3d 位置 apos、旋转表示为 4d 四元数的旋转 arrot 和指示夹具是否打开的二进制标志 aggrip 组成。对于每个任务，我们假设可以访问数据集 d = {ζ1, ζ2,。.., ζm} 的 m 个专家演示，其中 ζi 包含演示中所有时间步的语言指令 l、视觉观察 o 和末端执行器状态 qt。

图1:chaineddiffuser是一种机器人操纵policy架构，它预测一组机器人按键，并使用预测的轨迹段将它们链接起来。它使用预训练的2d图像主干对输入的多视图图像进行特征化处理，并使用感测的深度将生成的2d特征图提升到3d。在（b）中，我们使用pca可视化3d特征云，并保留3个主成分，将其映射到rgb。然后，该模型使用粗到细的注意力操作来预测末端执行器的keypose，以估计末端执行器3d位置的3d动作图，并回归机器人的3d方向，类似于[21]（d）。然后，它将当前的末端执行器姿态与预测的姿态联系起来，并使用基于3d场景特征云和预测关键点的扩散模型预测轨迹（e）。

Input Encoding

chaineddiffuser在3d空间中运行，以在不断变化的相机视点上实现鲁棒性，这是与假设固定相机视点的现有2d方法相比的一个重要优势[22,24,14]。与依赖于基于体素的3d表示的现有机器人架构（例如[3,40]）相比，chaineddiffuser采用了基于点的表示，这有助于稀疏计算，并避免了体素化过程中的精度损失。chaineddiffuser使用冻结的clip[73]将语言指令l和rgb图像ot分别编码为一组语言和视觉特征tokens。然后，它使用深度通道信息将2d图像特征tokens投影到3d特征云中（图1（b）），其中每个视觉token都有2d外观信息和3d位置信息。我们还使用简单的mlp对本体感觉信息qt进行编码。

3.2 Macro-Action Predictor

我们的宏动作预测器基于 act3d [21]，这是一种 state-of-the-art 宏动作预测方法，它使用基于点的 transformer，将末端执行器动作预测转换为 3d 动作图预测。为了完整起见，我们在这里包括它的主要 pipeline。act3d迭代地采样3d点候选，并使用场景3d特征云的相对位置注意力对其进行特征化。然后，使用可训练查询token 来对场景中的n个点候选的池进行评分，并选择下一个宏动作的位置。首先在机器人的空工作空间内对候选点进行均匀采样，并且仅包含3d位置信息和可训练特征嵌入。查询token和候选点分别处理语言tokens-、视觉特征标记和本体感觉token-的连接（跨序列维度）（图1（d））：

其中 attn(x, y) 是一个注意力操作 [23, 74]，其中查询来自 x、键和值来自 y 组成。在这个上下文化contextualization 步骤之后，查询 token 和点候选已经捕获了任务和场景信息。我们将上下文化查询嵌入与所有候选点进行点积，并为预测宏动作的位置选择最匹配的候选点:        

一旦我们获得了最佳点候选，我们在查询之上使用简单的 MLP 预测旋转和抓取打开标志：

3.3 Local Trajectory Diffuser

一旦我们获得了当前步骤t的宏动作，我们调用基于扩散的局部轨迹生成器来用微动作填补中间的空白。我们将这种轨迹生成建模为去噪过程[18,14,4]:首先在归一化SE(3)空间中绘制S个随机高斯样本序列，然后进行K次去噪迭代，将有噪声的轨迹转换为无噪声的路径点序列。每次去噪迭代描述为:

其中 εθ 是噪声预测网络，k 是去噪步骤，是每次迭代添加的高斯噪声，λk, γk, σk 是依赖于 k 的标量噪声调度函数（附录 7.1）。

噪声预测网络(图1 (e))也是一个基于注意力的模型，它吸收与宏动作选择器相似的输入，即语言指令l、rgb-d观察值ot和当前末端执行器状态qt，但还附加了目标宏动作和去噪时间步k的条件。语言tokens zins、视觉tokens zvis和当前末端执行器状态zrobotas的特征与宏动作选择器相似。我们使用mlp将目标宏动作编码为。我们使用正弦位置嵌入将去噪时间步编码为[23]，并使用mlp将采样噪声编码为tokens 序列。我们让这个序列首先迭代地交叉关注所有编码输入:

然后自注意力获得最终确定的（请注意，为了表示清晰，我们重用相同的符号）：

同样，我们使用相对位置嵌入来编码所有 tokens 的空间位置。对于轨迹噪声tokens，我们还使用正弦位置嵌入对每个样本的时间位置s进行编码。这些被添加到各自的噪声 tokens zk s 中。然后将上下文化的噪声样本馈送到另一个 mlp 进行噪声回归：

在将等式 6 代入 5 的 K 个去噪步骤之后，我们通过不归一化 unnormalizing 将去噪样本转换回实际的micro-actions：。有关更多实现和培训细节，请参阅附录 3.4。

Noise schedulers

我们将局部轨迹优化建模为离散时间扩散过程，我们使用DDPM采样器[18]实现。DDPM 使用非参数时变噪声方差调度器 βk，它定义了在每个时间步添加多少噪声。我们采用缩放的线性时间表来确定位置，采用平方余弦时间表来确定每个轨迹步的旋转。

3.4 Implementation and Training Details

chaineddiffuser 将场景的 m 个多视图 rgb-d 图像作为输入。对于模拟的实验，我们使用 m = 3（左、右、手腕）或 m = 4（带有额外的前视图），具体取决于我们比较的基线设置。对于真实世界的实验，我们使用 k = 1，使用单个前视图相机。每个rgb-d图像为256 × 256，用clip的resnet50视觉编码器编码为64 × 64的视觉tokens[73]。演示数据包含所有时间步的末端执行器状态。为了提取宏观动作来监督动作选择transformer，我们在以前的文献[3,22,24]之后使用简单的启发式:如果夹持器打开或关闭，则认为时间步是一个包含宏观动作的关键帧，或者如果机械臂不移动(当所有关节速度接近零时)。演示中存在的所有密集动作都用于监督局部轨迹扩散器。我们将提取的宏动作之间的密集轨迹重新采样到固定长度S = 50的轨迹 。我们在实践中发现，对固定数量的micro-actions进行去噪会导致更稳定的训练，并且比学习具有预测轨迹长度的变长轨迹扩散效果更好。我们联合训练动作检测器和轨迹扩散器，使用交叉熵（CE）损失通过预测池中所有候选点的概率分布q来监督候选点的选择，并使用均方误差（MSE）损失来监督四元数、夹具开口和轨迹噪声回归：

其中 * 表示预测值，, k 是随机采样的去噪步骤，εk 是采样的地面实况噪声。为了在实践中加快训练速度，我们训练前 2 个项直到收敛，然后添加第 3 个项进行联合优化，而不是使用地面实况宏动作来训练轨迹扩散器。这允许轨迹扩散器结合一定的误差恢复能力来处理不准确的宏观动作预测。此外，在测试时，我们对预测的四元数进行归一化，以确保它在将其馈送到机器人之前尊重归一化约束。我们在所有实验中使用 B = 24 和 AdamW [75] 优化器的批量大小，学习率为 1e-4。我们的单任务模型在一个 A100 GPU 上训练 1 天，多任务模型在 4 个 A100 GPU 上训练 5 天。

Control

我们的控制算法在宏观动作 macro-action 层面是闭环的，这意味着宏观动作预测器对周围环境进行推理并预测处理环境变化的动作。在低层，我们的控制器是开环的，遵循笛卡尔空间末端执行器轨迹，由使用位置控制的预测微动作 micro-actions 组成，控制频率为10Hz。对于我们的真实机器人设置，我们使用带有 ROS 的开源 frankapy 包 [76]，它使用 1kHz 的低级 PID 控制器。

4 Experiments

我们在模拟和现实世界环境中的各种操作任务中测试 ChainedDiffuser。我们的实验旨在回答以下问题：

• 与以前的 SOTA 2D 和 3D 操作方法相比，ChainedDiffuser 如何。

• 宏观动作预测是否有助于指导轨迹生成。

• ChainedDiffuser 是否适用于只有单个摄像头和有限数量的演示可用的现实世界？

4.1 Simulation Experiments

我们使用 RLBench [26] 在模拟中进行实验，这是一个广泛采用的操作基准，具有不同的任务，涉及与广泛对象的交互，如图 2 所示。我们遵循先前工作 [22, 24] 中使用的相同设置，其中每个任务有多个变化，包含 100 个演示。我们报告了每个任务的成功率，平均超过 100 个看不见的测试集。对于基线，在可能的情况下，我们使用他们论文中报告的官方数字。

Baselines

我们将 ChainedDiffuser 与以下基线进行比较：

1. auto-λ[77]和hiveformer[24]，policy学习者在多视图2.5d图像上操作，并通过偏移输入图像中检测到的点来预测动作。

2. instructrl [22]，一种 policy，它使用预训练的视觉和语言编码器对多视图 2d 图像进行操作，并直接预测 6-dof 末端执行器动作。

3.act3d [21]， policy，它使用 3d 动作检测 transformer 预测关键帧末端执行器宏动作，并依靠低级运动规划器连接宏动作。

4. 开环轨迹扩散，即没有宏动作检测器的chaineddiffuser，使其成为轨迹扩散模型。

5. act3d+轨迹回归，用确定性轨迹回归代替了chaineddiffuser中的局部轨迹扩散器

Dataset

我们考虑以下单任务实验设置：

• Auto-λ[77]实验设置中考虑的10个任务。这些任务由许多先前的工作考虑，这使我们能够将我们的性能与它们进行比较。
• 我们在RLbench中确定了10项需要与环境持续交互的任务，例如擦拭桌子，需要擦拭轨迹来清除桌子上的污垢，以及打开冰箱，当机器人抓住门把手时，局部轨迹需要遵守运动学约束。RLBench中的大多数任务仅需宏动作预测和运动规划器即可合理解决。我们考虑的这组任务突显了这些方法的局限性。

Results

我们训练单任务 chaineddiffuser 和基线。对于 auto-λ、hiveformer 和 instructrl，我们使用相应论文中报告的数字。我们在表 1 和表 2 中展示了定量结果。 chaineddiffuser 在所有任务类别上始终比以前的方法获得更好的性能。在运动规划者的挑战性任务中，chaineddiffuser 平均显着提高了 60%。chaineddiffuser改进了开环轨迹扩散模型，这表明将全局宏观动作预测委托给高级policy来指导局部轨迹扩散有帮助。act3d+trajectory regression在演示中存在多模态轨迹的地方挣扎，例如橱柜中的杯子，其中存在多模态轨迹来抓取杯子并输入训练集中的橱柜。这表明将轨迹生成建模为多步去噪过程优于基于回归的模型，这与先前文献[14]的结论一致。

4.2 Real-world Experiments

我们使用带有平行爪抓手的 Franka Emika Panda 机器人与真实世界的设置进行实验。我们使用单个Azure Kinect摄像头来收集前视图RGB-D图像输入。有关我们的硬件和数据收集设置的更多详细信息，请参见附录 7.4。我们设计了7个任务，涉及多步动作和与场景的连续交互(如图2所示)，为每个任务收集了10 - 20个演示，并训练多任务ChainedDiffuser进行实际部署。我们将读者推荐给我们的补充视频，用于机器人的定性执行。我们在每个任务的 10 集上评估它，并在表 3 中报告了成功率。 ChainedDiffuser 能够在大多数任务上表现得相当好，即使对于具有多种动作模式和技能的任务。最常见的失败案例是由噪声深度图像引起的：我们利用点选择进行宏观动作预测，这将在现实世界中遭受不正确的深度估计。这可以通过使用多视图相机设置和学习从噪声输入中恢复更准确的相机校准来解决，我们将其留作我们未来的工作。

4.3 Limitations

我们的方法目前有以下局限性:

1) 我们的轨迹扩散器取决于SE（3）空间中的末端执行器姿态。理想的情况是将其扩展到全关节配置空间，以实现更灵活的轨迹预测。

2）我们的模型在宏观行动层面进行闭环控制，这限制了它在高度动态环境中的灵活性。也就是说，我们的框架可以很容易地通过微观行动层面的闭环重新规划进行扩展，使policy对环境动态更具鲁棒性，我们将其作为未来的工作。

3) 根据RLBench中的标准设置，我们的方法假设可以访问校准的相机。我们认为这一假设是有效的，因为未来为人类执行家务的移动机器人应该在机器人上安装摄像头，这些摄像头可以在出厂时进行校准。

5 Conclusion

我们提出了chaineddiffuser，这是一种神经policy架构，用于从演示中学习6自由度机器人操作。我们的模型在各种任务设置（在模拟和现实世界）上都实现了具有竞争力的性能。我们的实验表明，通过统一基于变压器的宏观动作检测和基于扩散的轨迹生成，chaineddiffuser 在两个家族中都取得了最好的结果，并解决了它们各自的局限性。chaineddiffuser 优于单独的关键帧预测方法和轨迹扩散，这证明了它们在我们的框架中的统一。它在 rlbench 中设置了一个新的 state-of-the-art，并特别提高了涉及铰接对象的接触丰富任务和任务的性能，其中依赖手工设计的规划器的方法通常会遇到困难。