港科大&哈佛最新!TesserAct:4D具身世界模型全新问世~

点击下方卡片,关注“具身智能之心”公众号

作者丨具身智能之心

本文只做学术分享,如有侵权,联系删文


>>点击进入→具身智能之心技术交流群

更多干货,欢迎加入国内首个具身智能全栈学习社区具身智能之心知识星球(戳我)这里包含所有你想要的。

写在前面&出发点

这里提出了一种学习新型4D具身世界模型的有效方法,该模型能根据具身智能体的动作预测3D场景随时间的动态演变,保证空间和时间上的一致性。通过对RGB-DN(RGB、深度和法线)视频进行训练来学习4D世界模型。这不仅通过在预测中纳入详细的形状、结构和时间变化,超越了传统的2D模型,还能够有效地为具身智能体学习精确的逆动力学模型。具体来说,首先利用现成的模型,用深度和法线信息扩展现有的机器人操作视频数据集。接着,在这个带注释的数据集上微调视频生成模型,该模型联合预测每一帧的RGB-DN。然后提出一种算法,将生成的RGB、深度和法线视频直接转换为高质量的4D世界场景。方法确保了从具身场景中预测的4D场景在时间和空间上的连贯性,实现了具身环境的新视角合成,并促进了策略学习,其性能显著优于之前基于视频的世界模型。

项目链接:https://TesserActWorld.github.io

背景介绍

学习世界模型在实现具身智能体中起着关键作用,它可以模拟环境动态,支持灵活的策略合成、数据模拟与生成以及长期规划。然而,现有世界模型大多在2D像素空间中运行,而真实物理世界是三维的,这种差异导致空间关系表示不完整,影响了需要精确深度和姿态信息的任务,如机器人难以准确确定物体位置和方向,且2D模型还会产生不真实的结果,限制了其在数据驱动模拟和稳健策略学习中的应用。

这里旨在探索学习4D具身世界模型TesserAct,它能模拟3D世界的动态,生成传统2D模型无法实现的逼真3D交互,为机器人操作提供关键的深度和姿态信息。但学习4D具身世界模型面临挑战,世界动态的训练和学习计算成本高昂,需要模型在三维空间和时间中生成输出。为有效表示和预测世界动态,我们采用了一种更轻量级的4D世界表示方法,即预测场景的RGB、深度和法线图序列,这种表示能准确捕捉场景的外观、几何和表面信息,且与现有视频模型相似,便于利用其生成能力构建4D世界模型。

此外,还提出了一种从生成的RGB-DN视频重建精确4D场景的高效算法,通过结合深度和法线预测来整合场景的平滑3D表面,并利用光流区分背景和动态区域,引入两个新的损失函数来保证场景随时间的一致性。训练TesserAct的一个关键挑战是缺乏高质量4D注释的大规模数据集,为此我们收集了包含合成数据和真实世界数据的4D具身视频数据集。

主要贡献包括:

1)收集了带有紧凑高质量RGB、深度和法线注释的4D具身视频数据集,并学习了4D具身世界模型;

2)提出了将生成的RGB-DN视频转换为时空连贯的高质量4D场景的算法;

3)大量实验表明,与传统基于视频的世界模型相比,我们的4D具身世界模型能预测高保真的4D场景,并在下游具身任务中表现更优。

相关工作一览

1)具身基础模型

近期许多工作聚焦于构建通用智能体的基础模型,一类是构建多模态语言模型,处理图像或3D输入并输出描述智能体动作的文本;另一类是构建视觉 - 语言 - 动作(VLA)模型,直接输出动作tokens。而我们旨在构建具身智能体的基础4D世界模型,用于下游的规划或策略合成等应用。

2)学习世界模型

学习给定控制输入下的世界动力学模型一直是系统识别、基于模型的强化学习和最优控制领域的挑战。早期工作多在低维状态空间学习世界模型,泛化性较差;后来有研究在像素空间图像上学习,但多基于简单游戏环境。随着生成模型的发展,大量研究将视频模型作为基础世界模型,但这些模型在2D像素空间运行,无法完全模拟3D世界。与我们最相似的工作中,有的仅预测机器人任务的3D目标状态,有的是纯基于合成数据训练的几何感知世界模型,缺乏语言基础和下游机器人任务支持,而我们的方法从RGB-DN视频建模4D场景,并支持语言条件控制。

3)4D视频生成

近年来,4D视频生成受到越来越多关注,但现有方法常因混合框架导致优化速度慢,以及SDS损失的收敛问题。我们使用RGB-DN视频表示4D场景,训练效率更高,能提供对具身任务至关重要的高精度3D信息,且首次从当前帧和文本描述的具身智能体动作直接预测4D场景。

预备知识基础

1)潜在视频扩散模型

扩散模型可以通过正向过程逐渐向数据中添加噪声,直至其近似高斯分布,从而学习数据分布 。在推理阶段,训练一个去噪器 从噪声状态恢复数据。潜在视频扩散模型利用变分自编码器(VAE)在潜在空间中编码数据,在保持高质量输出的同时更有效地建模数据分布。

我们将RGB、深度和法线视频生成任务定义为条件去噪生成任务,即建模分布 ,其中 分别表示预测的未来RGB、深度和法线图的潜在序列, 是RGB图像、深度和法线图的潜在表示以及具身智能体的文本动作。

正向扩散过程在 个时间步内向潜在变量 添加高斯噪声,定义为:

其中 表示扩散步骤, 是控制每个步骤噪声影响的参数, 是单位矩阵。在反向过程中,模型旨在从噪声中恢复原始潜在变量。令 表示 的连接,训练一个带有学习参数 的去噪网络 来预测每个时间步添加的噪声。反向过程定义为:

得到去噪后的潜在变量 后,模型将其映射回像素空间,得到最终的RGB-DN视频。

2)通过法线积分优化深度

法线图提供了关于表面方向的重要信息,在深度优化过程中,对施加几何约束、保证表面平滑和连续性至关重要,有助于获得更准确可靠的深度估计,捕捉精细的表面细节。

我们使用透视相机模型对深度和表面法线设置约束。在第 帧的2D图像坐标系中,像素位置为 ,其对应的深度标量、法向量为 。在焦距为 、主点为 的透视相机假设下,对数深度 应满足以下方程: ,其中

此外,假设所有位置都是平滑表面,可将上述约束转换为二次损失函数,以找到优化的深度图:

可以将上述目标转换为迭代优化的损失目标。在迭代步骤 ,计算矩阵 ,并迭代优化得到更精确的深度预测

其中 由法线和相机内参定义。

学习4D具身世界模型

学习3D场景如何基于当前观察和动作随时间变化,对具身智能体至关重要。我们提出通过对RGB-DN视频进行训练并从中重建4D场景,来学习4D具身世界模型TesserAct。

1)4D具身视频数据集

学习4D具身世界模型需要大规模4D数据集,而在现实世界中收集成本高昂。我们提出一种数据收集和注释流程,可从现有视频数据集自动构建4D数据集。

我们从RLBench中选择20个难度较高的任务,为每个任务从4个不同视角生成1000个实例,共生成80k个合成4D具身视频。虽然模拟器提供了度量深度信息,但缺乏表面法线数据,我们使用DSINE中的depth2normal函数估计法线,并采用场景随机化技术增强泛化性。

合成数据的多样性有限,与现实场景存在差距。因此,我们纳入了现实世界视频数据集,利用RollingDepth为其标注仿射不变深度,使用Temporal-Consistent Marigold-LCM-normal1标注法线图。我们还选择了OpenX中的Fractal data和Bridge数据集,并纳入了人类 - 物体交互数据集Something Something V2,以增加指令的多样性。

2)模型架构和训练策略

训练扩散模型生成时间相关的RGB-DN数据具有挑战性。由于RGB-DN数据集规模相对较小,无法从头训练世界模型。因此修改并微调CogVideoX作为我们的RGBDN预测模型,利用其预训练知识启动4D模型。

模型架构如图2所示,首先使用CogVideoX中的3D VAE分别对RGB、深度和法线视频进行编码,不额外微调VAE。将这些潜在表示添加噪声得到 ,与相应的图像潜在头 一起输入模型。输入设计上,为每个模态引入三个单独的投影仪提取嵌入: ,其中 。DiT以这些嵌入的和为输入,结合文本输入 和去噪步骤 ,得到隐藏状态: 。输出方面,保留原始RGB输出方法,同时为深度和法线预测引入额外模块,通过Conv3D层编码输入潜在和预测的RGB去噪输出的连接,与DiT骨干网络产生的隐藏状态结合,经过输出投影仪得到深度和法线的去噪预测。为保留预训练知识,模型用CogVideoX权重初始化,其他模块用零初始化,训练时从4D具身数据集中随机选择样本,添加噪声并最小化目标函数:

3)4D场景重建

获得RGB-DN视频后,我们进一步优化深度并重建4D场景。以往工作中,深度表示常为相对图,无法直接重建整个场景,且基于深度的重建往往粗糙,会导致平面或墙壁倾斜。

利用法线图 ,通过法线积分优化深度图 ,得到更精确的深度图 ,并引入空间一致性损失项 。但这种方法逐帧优化深度,缺乏动态场景的时间一致性。为此,我们计算帧间光流 ,并强制深度在帧间保持一致。定义每帧的静态区域为光流幅度小于阈值 的像素,其掩码为 ,动态部分为 ,背景为跨帧固定的静态区域,即

根据光流,可以从先前帧检索当前帧对应位置的深度,以施加一致性约束。对于帧 中的每个像素 ,光流提供位移 ,可找到帧 中对应位置 的像素。基于此映射,定义 ,并引入图像动态和背景区域的一致性损失

此外,还纳入正则化损失 ,确保优化后的深度与生成的深度图 相似:

整体优化的损失目标为:

从第一帧开始,通过优化上述损失迭代细化深度图,用生成的深度图 初始化第 帧的深度图。利用细化的深度图 和RGB图像 ,可以重建在空间和时间上一致的代表世界的4D点云。

4)基于4D场景的具身动作规划

生成4D场景后,提取其中的几何细节,这可以显著增强机器人下游任务的执行能力,如在机器人抓取任务中,这些场景捕捉的详细几何信息至关重要。

我们使用基于4D点云构建的逆动力学模型,根据当前状态 、预测的未来状态 和指令 预测合适的机器人动作 ,数学关系为 。使用PointNet编码点云并提取3D特征,将这些特征与指令文本嵌入相结合,再通过MLP生成最终的7自由度动作。

实验分析

1)4D场景预测
  • 数据集:在真实领域使用RT1 Fractal和Bridge数据集中400个未见样本进行实验;在合成领域使用RLBench中的200个未见样本,该数据集中可直接获取真实的深度和法线图。

  • 评估指标:使用FVD、SSIM和PSNR评估视频质量;使用AbsRel、 评估深度质量;使用Mean、Median和11.25°评估法线图质量;使用  Chamfer距离评估重建点云的质量。对每个样本生成10次,并报告平均值。

  • 基线对比:将我们的方法与两个视频扩散模型和一个4D点云扩散模型进行比较。

    • OpenSora是一个视频扩散模型,在相同数据集上使用LoRA进行微调,但该数据集没有深度和法线注释。为构建完整的4D场景,在给定预测视频的情况下,使用Rolling Depth和Marigold额外估计深度和法线。

    • CogVideoX同样是一个视频扩散模型,在相同数据集上使用LoRA微调,且无深度和法线注释。其4D场景的获取方式与OpenSora类似。

    • 4D Point-E,由于此前没有工作能直接从第一帧和文本输入生成动态场景,我们将4D点云扩散模型作为基线。具体做法是修改Point-E模型,使其基于从文本和图像输入中提取的CLIP特征的均值进行条件设定,输出T个大小为n的点云。由于计算限制,T设置为4,n设置为8192。

  • 实施细节:使用多分辨率训练方法,在收集的4D具身视频数据集上训练模型,每次预测49帧。

TesserAct能预测高质量的4D场景:如表2所示,与进行后估计的视频扩散模型相比,模型能更准确地预测深度和法线图,验证了所学习模型的有效性。凭借更好的深度和法线图,方法重建的4D点云在真实和合成数据集上均实现了最低的Chamfer距离。4D Point-E方法在RLBench上的表现优于视频扩散模型,但仍落后于我们的方法。此外,直接使用点云进行训练计算成本高昂,限制了所用的帧数。相比之下,我们的模型利用RGB-DN视频的高效表示生成更精确的4D场景,尤其在捕捉动态场景的细粒度细节方面表现出色。我们在图4(a)中展示了定性结果。

TesserAct能高效合成新视角:我们的方法还可以通过预测深度和法线序列并生成点云,执行单目视频到4D的任务。我们在RLBench上对新视角合成任务进行实验,并与Shape of Motion进行比较,后者是一种利用高斯溅射的前沿视频重建方法。输入为单目前置摄像头视频,我们比较了来自头顶和左肩摄像头的结果。在表4中,报告了PSNR(重建精度)、SSIM(结构相似性)、LPIPS(感知差异)、CLIP Score(语义匹配)、CLIP aesthetic(视觉质量)和时间成本。结果表明,我们的方法可以在显著更短的时间内合成视觉质量更高、对齐更好的新视角。在Bridge数据集上的定性结果如图4(c)所示。

一致性和正则化损失有效:对新设计的损失项进行了消融实验。结果如图3所示,前两行展示了一致性损失的效果,我们从不同时间步的同一相机视角渲染帧。结果显示,应用一致性损失后,机器人手臂的运动更加连贯。最后一行突出了正则化损失的作用,展示了同一帧从三个不同视角的图像,表明该项有助于提高重建的几何精度。

TesserAct在场景和实体间具有泛化性:模型展现出强大的泛化能力,受益于CogVideoX的知识,该模型在未见场景和未见物体上实现了良好的生成效果。此外,它在跨实体场景中也表现出色,比如在RT-1环境中使用Bridge数据集的机械臂。图4(b)展示了在未见场景和物体上的泛化结果。

2)具身动作规划

数据集:从RLBench中挑选了9个具有挑战性的任务,其中包含需要高精度抓取的任务。

指标:以100个情节的平均成功率作为衡量指标。

基线:将我们的方法与一个行为克隆智能体以及一个基于视频的世界模型进行对比。

  • Image-BC:一种行为克隆智能体,它接收图像和任务指令,并输出7自由度动作。

  • UniPi∗:该方法接收任务指令和当前图像,预测未来视频,并使用基于2D的逆动态策略来预测动作。为保证公平比较,我们用微调后的CogVideoX替换其骨干网络,重新实现了这个基线模型。

实现细节:为每个任务收集了500个样本,用于训练逆动力学模型。在推理时,给定初始状态,我们首先预测并记录所有未来关键帧。在后续动作中,仅根据当前状态和预测的未来状态,通过查询逆动力学模型来获得相应动作。最大帧长度为13,固定分辨率为512×512。

结果与分析

结果如表3所示,在大多数任务中,我们的方法优于视频扩散模型和图像行为克隆智能体。这是因为在大多数任务中,4D点云能够揭示物体的几何形状,为机器人规划提供更好的空间引导,比如在关闭盒子和打开罐子等任务中就有所体现。同时,3D信息有助于工具的使用,像在扫地入簸箕和给植物浇水这类任务中发挥作用。然而,在打开微波炉和称重等任务中,我们的方法性能不如基线,这可能是因为这些任务在2D前视图中就已经包含了足够的信息。总体而言,这些结果凸显了将4D场景预测与逆动力学模型相结合,对提升机器人任务执行能力的潜力。

参考

[1] TesserAct: Learning 4D Embodied World Models

论文辅导计划

具身智能干货社区

具身智能之心知识星球是国内首个具身智能开发者社区,也是最专业最大的交流平台,近1000人。主要关注具身智能相关的数据集、开源项目、具身仿真平台、大模型、视觉语言模型、强化学习、具身智能感知定位、机器臂抓取、姿态估计、策略学习、轮式+机械臂、双足机器人、四足机器人、大模型部署、端到端、规划控制等方向。星球内部为大家汇总了近40+开源项目、近60+具身智能相关数据集、行业主流具身仿真平台、强化学习全栈学习路线、具身智能感知学习路线、具身智能交互学习路线、视觉语言导航学习路线、触觉感知学习路线、多模态大模型学理解学习路线、多模态大模型学生成学习路线、大模型与机器人应用、机械臂抓取位姿估计学习路线、机械臂的策略学习路线、双足与四足机器人开源方案、具身智能与大模型部署等方向,涉及当前具身所有主流方向。

全栈技术交流群

具身智能之心是国内首个面向具身智能领域的开发者社区,聚焦大模型、视觉语言导航、VLA、机械臂抓取、双足机器人、四足机器人、感知融合、强化学习、模仿学习、规控与端到端、机器人仿真、产品开发、自动标注等多个方向,目前近60+技术交流群,欢迎加入!扫码添加小助理微信邀请入群,备注:学校/公司+方向+昵称(快速入群方式)。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值