机器人扩散Transformer的成分

24年10月来自CMU和UC Berkeley的论文“The Ingredients for Robotic Diffusion Transformers”。

近年来,机器人专家利用高容量 Transformer 网络架构和生成扩散模型,在解决灵巧机器人硬件上日益普遍的任务方面取得了显著进展。不幸的是,将这两项正交改进结合起来已被证明出奇地困难,因为没有明确且易于理解的流程来做出重要的设计选择。本文确定、研究和改进高容量扩散Transformer策略的关键架构设计决策。由此产生的模型可以有效地解决多个机器人实施例上的不同任务,而无需每次设置超参调整的痛苦。将研究结果与改进的模型组件相结合,提出一种名为 DiT-Block Policy 的新型架构,在解决双手 ALOHA 机器人上的长范围(1500+ 时间步)灵巧任务方面的表现,明显优于最先进的技术。此外,在 10 小时的高度多模态、语言注释的 ALOHA 演示数据上进行训练,该策略表现出更好的扩展性能。

现代机器学习利用高度表达的深度神经网络从大量离线模仿数据集 [1]、[2]、[3] 中生成和建模样本,取得了显著的成功。受这些进步的启发,机器人领域正在采用类似的技术来开发用于操控 [4]、[5] 和运动任务 [6]、[7] 的通用策略和控制器。然而,机器人任务存在多重挑战,阻碍了这些方法的直接应用。

首先,策略必须学会处理来自多个摄像头的高维观察流,而不会过拟合数据中的虚假相关性。例如,策略可以学会直接从本体感受信号或特定的摄像头视图回归动作。因此,在测试期间,它会完全忽略来自其他模态(例如腕式摄像头)的信号,而这些信号对于解决潜遮挡的高灵巧任务至关重要。这通常会导致部署期间发生灾难性故障。

其次,由于目标操控的容错率低,机器人必须做出极其精确的动作预测。这在解决长期任务时尤其重要,因为机器人可能需要在轨迹结束之前按顺序实现多个子目标。例如,负责准备寿司的机器人需要达到多个“分解”子目标,每个子目标都有毫米级的误差阈值,如图所示。

请添加图片描述

最后,策略学习需要应对多模态动作分布,即解决同一任务的不同方法。简单地从这个分布中学习平均动作通常会导致犹豫不决且容易出错的行为。随着数据集大小的增加,处理动作多模态变得尤为重要,因为不同的专家自然会表现出不同的行为。如果不能解决这些挑战,将导致部署期间的策略不可靠且不安全。

最近的进展已经开始解决这些问题,通过为灵巧任务开发更高容量的网络架构 [8],并利用改进的生成建模框架(如扩散 [9])进行有效的多模态动作学习 [10]。结合这两种正交改进可以产生非常有效的策略,但迄今为止已被证明极具挑战性。例如,原始扩散策略论文 [10] 提出了一种简单的交叉注意 Transformer [11] 实现策略网络,根据他们自己的分析,该网络极难训练。因此,大多数后续工作 [12] 都以他们的 U-Net 架构 [13] 为基础,该架构更易于调整,但对任务设置有严格的要求(例如,动作信号必须足够平滑 [10])。因此,高容量扩散建模对于广泛的机器人应用来说仍然难以实现。

编码高维观测值。为了感知周围环境,机器人通常利用多种感官观测值。因此,如何最好地组合来自多个传感器的信息是机器人技术和计算机视觉领域的一个古老问题 [18]、[19]、[20]、[21]、[22]。例如,像 ALOHA [8] 这样的双手机器人必须结合来自全局摄像头(可查看整个场景)和腕部摄像头(可近距离查看操作本身)的信息。处理此问题最直接的方法是学习一个可同时在所有输入模态上运行的单个共享网络/编码器 [4]、[23]。然而,这些系统通常会学习易变的特征,这些特征会过拟合特定输入(例如本体感受数据和全局摄像头),而完全忽略其他输入。先前研究提出的可能解决方案,包括为每个视觉流使用单独的高容量图像编码器 [8]、[10]、将 3D 感知空间偏差注入表示网络 [24]、[25],以及使用丢弃观测适当地正则化特征 [26]、[27]。

预测多模态动作分布。对多模态动作分布进行建模(即机器人可能根据同一观察/目标采取多个完全不同动作的场景)是 行为克隆(BC)方法的一项众所周知的挑战 [28]。随着专家数据量的增加,这一挑战往往会加剧,因为不同的演示可能会展示同一任务的不同解决方案。潜在的解决方案包括动作空间离散化 [29]、[30]、[31]、[32]、[33]、修改 BC 的目标策略 π 以预测更高容量的动作分布 [23]、[34]、[35]、隐式建模动作分布 [36]、[37]、[38]、[39]、[40],以及使用像扩散这样的生成建模目标 [41]、[42]、[10]、[43]、[44]、[45]。扩散尤其在机器人技术中展现出了最先进的成果 [10]:它可以学习复杂的 3D 感知策略 [46]、[47],同时开展的工作甚至在双手机械臂上展现出了最先进的操作结果 [48]。然而,模型架构/超参非常敏感且难以调整 [10]、[12]。这是扩展的主要障碍,因为更高容量的网络架构(如 Transformers [11])对于拟合大型和更多样化的数据集至关重要。

考虑通过模仿学习 [51]、[52]、[53]、[54]、[55]、[56] 获取机器人控制器的问题,该控制器可在语言指令提示时执行具有挑战性的灵巧操作行为。具体而言,机器人必须学习一个目标条件策略 πθ (at | ot, g),该策略在给定新输入观测 (ot) 和期望语言目标 (g) 的情况下,预测在环境动态 T : S × A → S 下,其中 ot ∈ S 和 at ∈ A 时,在 ∼ π(·|ot,g) 处的动作分布。通过行为克隆 [57]、[28]、[58] (BC) 优化策略 π,以匹配给定演示数据集 D = {τ1, …, τn} 的最佳动作分布,其中每个轨迹 τi = {g, o0, a0, o1, … } 是从专家智体收集的(例如,人类远程操作数据)。在测试期间,从 π 中采样动作并在机器人上执行。ot 是来自机器人的一组图像观察;at 是 H 关节/笛卡尔状态动作的块,g 是任务的文本描述。这种设置能够最大限度地提高灵活性和通用性,应对广泛的机器人任务,其中精确的状态难以推断,目标是自由形式的自然语言指令。
策略 π 被表述为条件去噪扩散概率模型 [41] (DDPM),这是一种生成模型,其中输出使用去噪过程进行采样 [59]。给定初始高斯噪声 x/k 和噪声预测网络 εθ(x/k, k, ot, g),DDPM 过程产生 x/k−1 = α(x/k − γεθ(x/k, k, ot, g) + N(0, σ2I)),其中 k 是扩散时间指数,α、γ、σ 是与扩散噪声计划相关的参数 [41]。当 εθ 经过适当训练后,该过程将产生一个以最佳动作终止的序列:x/k, x/k−1, . . . , x/0 ≃ at。因此,目标是通过梯度下降[60],[61]学习εθ,遵循MSE目标:L = ||ε/k−εθ(at+ε/k,k,ot,g)||^2。注:在训练期间使用k=100个扩散步,一个余弦噪声调度[62]和标准确定性采样过程,减少测试期间所需的样本数量(至k=10)[63]。

如图所示策略架构:DiT-Block 策略架构可为各种机器人任务提供可扩展、目标条件的策略学习;使用单独的 ResNet-26 [16] 编码器对图像观测进行token化。使用预训练的 Distill BeRT 模型 [49] 对文本目标进行token化并将其编码为嵌入向量。使用 FiLM 层 [50] 将此向量合并到观测tokens中。观测tokens被传递到编码器-解码器Transformer网络(中间),该网络负责预测用于扩散的噪声 epsilon (ε)。为了进行稳定的训练,解码器块利用自定义的 adaLN-Zero 架构(右),使Transformer能够可扩展地优化扩散目标。

请添加图片描述

受到先前数据扩展研究 [68]、[69]、[17]、[5]、[29]、[70] 的启发,试图了解 DiT-Block 策略在越来越多样化的演示数据上进行训练时的表现。然而,少数(开源)双手机器人数据集 [8]、[71] 仅包含使用相同受控场景/目标收集的少数任务。因此,它们对于测试双手设置中的泛化能力没有用处。为了解决这一缺点,收集并注释 BiPlay,这是一个更加多样化的双手操作数据集,具有随机目标和背景设置,如图所示。将 BiPlay 收集为一系列 3.5 分钟长的episodes。对于每一episode,构建一个包含各种目标的随机场景,并解决了该场景中的一系列任务。收集后,episodes被分解成片段,然后用适当的语言任务描述进行注释。最终的数据集包含跨越 10 小时机器人数据收集的 7023 个剪辑。

请添加图片描述

为了训练模型,为每个评估任务收集一组固定的演示(100 多个演示)。此外,还汇编前期工作的开源数据(ALOHA [8] 数据集和 YaY [71] 的最佳策略展开),并将其添加到训练组合中进行正则化。完整的数据组合如表所示。

请添加图片描述

所有 DiT-Block 策略都在此数据组合上训练了 250K 次迭代,使用 AdamW [61] 优化器和余弦学习调度 [72]。最后,不是在每个步骤预测单个动作,而是训练 DiT-Block 策略模型,预测一组 H = 100 个动作。这在训练期间起到了正则化的作用,并允许用时间集成(temporal emsembling) [8] 来提高运行时的稳定性。

第一个任务集考虑的是双手动、低成本的 ALOHA 机器人 [8],这样能够研究具有高度灵巧、精确行为的挑战性场景。描述这些任务及其成功标准如下:
(1)拾取放置:给定一个文本指令(如 g =“拿起玉米并将其放入碗中”),机器人必须找到目标物体,抓住它,然后将其放入目标盘子/碗中。场景中总是有两个物体和两个可能的目标,因此机器人必须根据文本指令正确地确定其行为。如果物体最终落入正确的容器中,则试验标记为成功。
(2)拔笔帽:此任务评估机器人的精确抓取能力及其同时控制双臂的能力。机器人必须用一只手臂拿起记号笔,将其带到另一只手臂,然后从笔上取下笔帽。如果试验结束时笔帽已拔出,则试验标记为成功。
(3) 切寿司:这项任务评估机器人在较长时间内连续执行精确、灵巧的操作任务能力。这是最具挑战性的任务,因为即使是 2 分钟内的小错误也可能破坏策略。机器人必须将寿司放在砧板上,从杯子里拿起刀,用另一只手重新抓住它,然后将寿司切成四块。如果任务以寿司分成四块结束,则该任务标记为成功,但部分得分将根据成功切割的比例给出(例如,一次切割获得 1/3)。

下一个任务集使用单臂 Franka FR3 机器人。虽然灵活性更有限,但 Franka 允许测试对全新形态和控制空间(笛卡尔速度)的泛化。考虑以下任务:
(1) 烘烤:在这个长期任务中,机器人必须拿起目标物体,将其放入烤面包机中,然后关闭烤面包机。如果烤面包机完成了整个过程,则试验成功;如果物体仅放在烤面包机中,则试验成功一半。
(2)擦拭:机器人必须找到海绵,抓住它,然后将碎屑垃圾推入簸箕。如果在运行结束时擦掉所有碎屑垃圾,则试验成功。

猜你喜欢

转载自blog.csdn.net/yorkhunter/article/details/143205653