论文标题:Lift3D Foundation Policy: Lifting 2D Large-Scale Pretrained Models for Robust 3D Robotic Manipulation
项目主页: https://lift3d-web.github.io/
代码链接: https://github.com/PKU-HMI-Lab/LIFT3D
1. 摘要
3D 几何信息对于机器人操纵任务至关重要,因为机器人需要感知三维环境,推理空间关系,并与复杂的几何物体进行交互。近年来,研究逐渐聚焦于显式提取 3D 特征,但仍面临诸如:缺乏大规模机器人 3D 数据和潜在的空间几何信息丢失等挑战。为了解决这些问题,我们提出了 Lift3D 框架,通过隐式和显式的 3D 机器人表示逐步增强 2D 大规模预训练模型的 3D 空间感知能力,从而构建一个鲁棒的 3D 操纵策略。具体来说,我们首先设计了一个任务感知的掩码自编码器,该自编码器通过 Mask 与任务相关的 Affordance token 并重建深度几何信息,增强了 2D 基础模型的隐式 3D 机器人表达能力。在自监督微调之后,我们引入了一种 2D 基础模型 Lifting 策略,该策略在输入 3D 点和 2D 模型的位置编码之间建立了位置映射关系。基于该映射,Lift3D 利用 2D 基础模型直接显式的编码点云数据,借助大规模预训练知识提升 3D 模仿学习效率,同时最小化空间信息丢失。
图 1: 为了赋予大规模 2D 预训练模型 3D 操控能力, Lift3D 首先通过 mask 任务相关的 affordance 区域和深度重建,以提升隐式 3D 机器人表示;随后通过利用预训练的 2D 位置嵌入来直接编码 3D 点云数据,进行显式 3D 操纵模仿学习。Lift3D 在各种仿真(i.e., 机械臂,灵巧手)和现实世界任务中展现出强大的鲁棒性和令人惊讶的效果
基于视觉的操纵策略的核心目标之一是理解场景并预测相应的 3D 位姿。一些现有的方法利用 2D 图像作为输入,通过强化学习或模仿学习直接预测 3D 末端执行器的位姿。尽管这些方法能够有效地处理一系列操纵任务,但它们未能完全理解物理世界中的空间关系和 3D 结构。在机器人操纵中,3D 几何信息对于应对复杂任务至关重要,因为机器人必须感知 3D 环境、推理几何关系并与复杂的空间物体进行交互。
近年来,研究越来越多地集中在机器人操纵任务中显式提取 3D 特征表示,这些方法可分为两类:1)一些方法直接编码点云数据,通过从头开始训练 3D 策略模型或微调预训练的点云编码器(例如 PointNet++ 和 PointNext)。然而,由于缺乏大规模机器人 3D 数据和基础模型,限制了它们的泛化能力。此外,处理 3D 或体素特征需要大量计算资源,这限制了其可扩展性和在现实世界应用中的可行性。2)另一方面,一些方法涉及模态转换,例如将预训练的 2D 特征提升到 3D 空间,或将 3D 点云投影到多视角图像中,以供 2D 预训练模型输入。尽管在一些后续的操纵任务中表现出色,但这些模态转换不可避免地导致空间信息的丢失,并不能完全激发模型对 3D 空间关系的理解。
基于上述 3D 策略中的挑战,我们提出了一个问题:“我们能否开发一个即有大规模预训练知识,又融合完整 3D 空间数据输入的 3D 策略基础模型?”
3. Lift3D 方法
3.1 贡献概述
我们提出了 Lift3D,通过系统地改进隐式和显式的 3D 机器人表示,将 2D 大规模预训练模型提升为鲁棒的 3D 操纵策略模型。
对于隐式 3D 机器人表示,我们设计了一种任务感知 MAE(Mask Autoencoder)自监督方法,它对任务相关的可操作区域进行掩蔽,并重建深度几何信息,从而增强 2D 基础模型的 3D 空间认知。
对于显式 3D 机器人表示,我们提出了一种 2D 基础模型 lifting 策略,利用 2D 基础模型的预训练位置编码(PE)来编码 3D 点云数据,用于 3D 操纵模仿学习。
3.2 Task-aware Masked Autoencoder (隐式 3D 机器人表示)
图 2: Lift3D 的整体流程。a) 对于隐式 3D 机器人表示,我们利用 CLIP 离线提取基于任务描述的图像注意力图,并将其反向投影到 2D 输入图像上,以指导 MAE 的掩码。然后,我们将可见的 tokens 输入 2D 基础模型进行特征提取。掩码 tokens 和编码后的可见 tokens 经过 MAE 解码器进行深度重建,从而增强 3D 空间感知。同时,编码后的可见 tokens 还会通过与原模型的特征进行蒸馏,以减轻灾难性遗忘问题。b) 对于显式 3D 机器人表示,我们首先将点云数据投影到多个虚拟平面上,建立 3D 输入点与每个虚拟平面上的 2D 位置嵌入(PEs)之间的位置映射。映射完成后,我们将每个 3D 补丁对应的 2D PEs 进行平均,形成统一的 3D 位置嵌入(3D PEs),然后将其与 3D tokens 结合。这些 3D tokens 是通过将点云输入轻量级 3D 分词器生成的。最后,来自 2D 基础模型的输出特征通过策略头进行处理,以预测模仿学习的位姿。
一些研究表明,2D 基础模型在各种下游机器人任务中展示了强大的操纵和泛化能力。在此基础上,Lift3D 首先增强了 2D 基础模型中的隐式 3D 机器人表示。现有的机器人 MAE 重建方法大多采用随机掩码。然而掩码部分可能主要包含无关的背景信息,这会阻碍对于前景物体表示的有效学习。不同于之前的方法,Lift3D 的目标是掩码与任务相关的 Affordance token,并重建深度几何信息,从而增强 2D 基础模型的 3D 空间感知能力。
具体而言,利用来自机器人操作的大规模数据集来构建 MAE 训练数据集,该数据集包括从视频中随机抽取的 100 万个训练样本,包含配对的图像和深度数据。如下图 2 a) 所示,具体的操作步骤如下:
在赋予 2D 基础模型隐式的 3D 机器人感知能力后,我们引入了一种 Lifting 策略,使 2D 模型能够显式地理解点云数据。近期的工作,无论是将 3D 点云投影到多视图图像中,还是将 2D 特征提升到 3D 空间中,都面临由于模态转换而丢失空间信息的挑战。因此,高效地编码 3D 数据一直是 3D 机器人领域的一个关键研究方向。
对于基于 Transformer 的 2D 模型,位置嵌入(PEs)起着重要作用,因为它们为注意力机制中的输入 token 提供了位置信息。然而,直接创建新的 3D 位置嵌入来编码 3D token 可能会引入语义差异,这些差异源于预训练的 2D 基础模型与新添加的 3D 位置嵌入之间的不同,从而可能导致大规模预训练知识的丧失。因此,受 [23, 24, 71] 启发,将 3D token 投影到多个虚拟平面上。与之前的工作不同,我们的投影过程并不是为了构建模型的输入,而是为了在输入的 3D 点和每个虚拟平面的预训练 2D 位置嵌入之间建立位置对应关系。然后,这些 2D 位置嵌入被用于直接编码 3D 标记。
如图 2 b) 所示,我们将原始的点云数据转化为高维空间 (B x 128 x 768),通过轻量级的 3D 标记器获得 k 个(k = 128)3D token。3D tokenizer 由最远点采样、k - 最近邻聚合算法和可学习的线性层构成,用于特征编码。随后,每个 3D 坐标投影到 n 个虚拟平面上,得到相应的 3D 到 2D 的坐标对应。投影机制无可学习参数且高效,我们采用带有 6 个面的立方体投影方法,有效捕获空间信息。这 n 个虚拟平面对应于 n 个原始的 2D 位置嵌入。
使用 3D 到 2D 坐标对应,将每个 3D token 分配到 n 个原始 2D 位置嵌入。在将每个 3D 标记与 n 个 2D 位置嵌入对齐后,我们通过平均值的方式创建一个统一的 3D 位置信号,表示为:
将 3D 位置嵌入与 3D tokens 结合,输入到 2D 基础模型中。通过这种方式,我们利用 n 个结合的原始 2D PEs 来编码 3D tokens,从而有效地提供了 2D 空间内的多样位置关系,并减轻了 3D 集合信息的丢失。来自 2D 基础模型的输出特征将通过一个简单的 policy head 进行处理,以预测用于模仿学习的位姿。我们使用三层多层感知机(MLP)来构建策略头。需要注意的是,我们的 Lift3D 编码器可以轻松适配不同的解码器或策略头,例如 Diffusion Transformer, Large Language Model;这里使用 MLP 头部是为了简单验证。最后,模仿学习的损失函数如下所示,其中 T、R 和 G 分别代表末端执行器位姿中的平移、旋转和夹持器状态。