PartRM:清华北大提出铰链物体Part-Level运动建模新方法,入选CVPR 2025

清华北大提出PartRM,一种铰链物体Part-Level运动建模新方法,入选CVPR 2025。通过单张图像和拖拽信息预测物体运动,为机器人操纵提供新思路。

原文标题:铰链物体的通用世界模型,超越扩散方法,入选CVPR 2025

原文作者:机器之心

冷月清谈:

清华大学和北京大学联合提出了PartRM,一种用于铰链物体part-level运动建模的新方法。该方法旨在解决现有基于diffusion的方法在处理效率和三维感知方面的不足,以更好地服务于机器人操纵等任务。PartRM基于重建模型,通过单张输入图像和拖拽信息,生成观测物体未来状态的三维表征。研究团队还构建了PartDrag-4D数据集,用于评估部件级别动态建模的基准。实验结果表明,PartRM在生成结果上取得了显著提升,且效率更高,入选CVPR 2025。

怜星夜思:

1、PartRM方法中提到的“拖拽传播策略”具体是如何实现的?在实际应用中,如果用户给定的拖拽信息不准确或者存在偏差,会对最终的生成结果产生什么影响?
2、PartRM的两阶段训练方法,先学习Part的运动,再学习外观,这种训练方式的优势是什么?为什么不采用端到端的训练方式?
3、PartRM在PartDrag-4D数据集上表现良好,但在与训练数据差异较大的数据上效果欠佳,如何提高模型在真实世界场景中的泛化能力?

原文内容


基于当前观察,预测铰链物体的的运动,尤其是 part-level 级别的运动,是实现世界模型的关键一步。尽管现在基于 diffusion 的方法取得了很多进展,但是这些方法存在处理效率低,同时缺乏三维感知等问题,难以投入真实环境中使用。


清华大学联合北京大学提出了第一个基于重建模型的 part-level 运动的建模——PartRM。用户给定单张输入图像和对应的 drag ,PartRM 能生成观测物体未来状态的三维表征 ,使得生成数据能够真正服务于机器人操纵等任务。实验证明 PartRM 在生成结果上都取得了显著的提升。该研究已入选CVPR 2025。



  • 论文题目:PartRM: Modeling Part-Level Dynamics with Large Cross-State Reconstruction Model

  • 论文主页:https://partrm.c7w.tech/

  • 论文链接:https://arxiv.org/abs/2503.19913

  • 代码链接:https://github.com/GasaiYU/PartRM


研究动机

                                                                         

世界模型是一种基于当前观察和动作来预测未来状态的函数。该模型的研发使得计算机能够理解物理世界中的复杂规律,在机器人等领域得到了广泛应用。近期,对 part-level 的动态建模的兴趣日益增长,给定当前时刻的观察并给与用户给定的拖拽,预测下一时刻的铰链物体各个部件的运动受到越来越多的关注,这种类型的世界模型对于需要高精度的任务,例如机器人的操纵任务等,具有重要的意义。


然而,我们对这个充满前景的领域的调研表明,目前的前沿研究(如 Puppet-Master)通过对预训练的 大规模视频扩散模型进行微调,以实现增加拖拽控制的功能。尽管这种方法有效地利用了预训练过程中 学习到的丰富运动模式,但在实际应用中仍显不足。其中一个主要局限是它仅输出单视角视频作为表示,而模拟器需要三维表示来从多个视角渲染场景。此外,扩散去噪过程可能需要几分钟来模拟单个拖 拽交互,这与为操作策略(Manipulation Policies)提供快速试错反馈的目标相悖。


因此,我们需要采用三维表征,为了实现从输入单视角图像的快速三维重建,我们利用基于三维高斯泼溅(3DGS)的大规模重建模型,这些模型能以前馈方式从输入图像预测三维高斯泼溅,使重建时间从传 统优化方法所需的几分钟减少到仅需几秒钟。同时,通过将用户指定的拖拽信息加入到大规模三维重建   网络中,我们实现了部件级别的动态建模。在这个问题中,我们认为联合建模运动和几何是至关重要的,因为部件级运动本质上与每个部件的几何特性相关联(例如,抽屉在打开时通常沿其法线方向滑动)。这种集成使我们能够实现更真实和可解释的部件级动态表示。


同时,由于我们是第一个做这个任务的,在这个任务上缺少相关的数据集,因此我们基于 PartNet- Mobility 构建了 PartDrag-4D 数据集,并在这个数据集上建立了衡量对部件级别动态建模的基准(Benchmark),实验结果表明,我们的方法在定量和定性上都取得了最好的效果。



PartDrag-4D 数据集的构建


我们首先定义 PartRM 需要完成的任务,给定单张铰链物体的图像 ot 和用户指定的拖拽 at ,我们需要设计 一个模型,完成


图片


其中 图片 是 Ot 在 at 作用下的三维表征。


现有的数据集分为两种, 一种是只含有图片数据对,但是缺乏对应的三维表征(比如 DragAPart 中提出的 Drag-a-Move 数据集)。还有一种是通用数据集,比如 Objaverse 中的动态数据,这种数据不止还有部件级别的运动,还会含有物体整体的变形等运动,不适合我们训练。


因此,我们基于提供铰链物体部件级别标注的 PartNet-Mobility 数据集构建了 PartDrag-4D 数据集。我们选取了 8 种铰链物体(其中 7 种用于训练, 1 种用于测试),共 738 个 mesh。对于每个 mesh,如图所示,我们使其中某个部件在两种极限状态(如完全闭合到完全开启)间运动至 6 个状态,同时将其他部分状态 设置为随机,从而产生共 20548 个状态,其中 20057 个用于训练,491 个用于测试。为渲染多视角图像,我们利用 Blender 为每个 mesh 渲染了 12 个视角的图像。对于两个状态之间拖拽数据的采样,我们在铰链物体运动部件的 Mesh 表面选取采样点,并将两个状态中对应的采样点投影至 2D 图像空间,即可获得对应的拖拽数据。


PartRM 方法


方法概览



上图提供了 PartRM 方法的概述,给定一个单视角的铰链物体的图像 ot 和对应的拖拽 at,我们的目标是生 成对应的 3D 高斯泼溅图片。我们首先会利用多视角生成模型 Zero123++ 生成输入的多视角图像,然后对输入的拖拽在用户希望移动的 Part 上进行传播。这些多视角的图像和传播后的拖拽会输入进我们设计的网络中,这个网络会对输入的拖拽进行多尺度的嵌入,然后将得到的嵌入拼接到重建网络的下采样层中。在训练过程中,我们采用两阶段训练方法,第一阶段学习 Part 的运动,利用高斯库里的 3D 高斯进行 监督,第二阶段学习外观,利用数据集里的多视角图像进行监督。


图像和拖拽的预处理


图像预处理:由于我们的主网络是基于 LGM 设计的, LGM 需要多视角的图像作为输入,所以我们需要将 输入的单视角图像变成多视角,我们利用多视角图像生成网络 Zero123++,为了使得 Zero123++ 生成的  图像质量更高,我们会在训练集上对其进行微调。



拖拽传播:如果用户只输入一个拖拽,后续网络可能会对拖拽的区域产生幻觉从而出错,因此我们需要   对拖拽进行传播到需要被拖拽部分的各个区域,使得后续网络感知到需要被拖拽的区域,为此我们设计了一个拖拽传播策略。如图所示,我们首先拿用户给定的拖拽的起始点输入进 Segment Anything 模型中得到对应的被拖拽区域的掩码,然后在这个掩码区域内采样一些点作为被传播拖拽的起始点,这些被传播的拖拽的强度和用户给定的拖拽的强度一样。尽管在拖动强度大小的估计上可能存在不准确性,我们后续的模型仍然足够稳健,能够以数据驱动的方式学习生成预期的输出。


拖拽嵌入



PartRM 重建网络的 UNet 部分沿用了 LGM 的网络架构,为了将上一步处理好的拖拽注入到重建网络中,  我们设计了一个多尺度的拖拽嵌入。具体地,对于每一个拖拽,我们会将它的起始点和终止点先过一个 Fourier 嵌入,然后过一个三层的 MLP:


图片


其中图片代表第 i 个拖拽在第 l 层的嵌入,其余部分设为 0。F 代表 Fourier 嵌入和 MLP ,图片代表在 channel 维度上连接。得到第 l 层的嵌入后,我们将 图片 和网络第 l 层的输出 Ol 在 channel 维度上连接,并过一个卷积层,作为 Ol 的残差加到 Ol 上作为下一层的输入,具体地:


图片


其中卷积层的参数全零初始化,图片为第 l + 1 层的输入。


两阶段训练流程


为了保证对静态 3D 物体外观和几何的建模能力,我们在预训练的 LGM 基础上构建了 PartRM。但直接在新数据集上微调会导致已有知识灾难性遗忘,从而降低对真实场景数据的泛化能力。为此,我们提出了 两阶段学习方法:先专注于学习之前未涉及的运动信息,再训练外观、几何和运动信息,以确保更好的性能。


运动学习阶段:在运动学习阶段,我们期望模型能够学到由输入的拖拽引起的运动。我们首先利用在我 们的数据集上微调好的 LGM 去推理每个状态 Mesh 对应的 3D 高斯泼溅表征,拿这些作为监督数据我们第 一阶段的训练。对于两个 3D 高斯之间的对应,我们利用 LGM 输出的是一个 splatter image 这一优势,即 LGM 会对 2D 图像的每一个像素点学一个高斯泼溅,我们可以直接对监督数据和 PartRM 网络预测的输出  做像素级别的 L2 损失,即:



其中 i 代表在 splatter image 里的坐标, GSi 和 GSj 均为每个像素点对应的 14 维高斯球参数。


外观学习阶段: 在运动学习阶段之后,我们引入了一个额外的阶段来联合优化输出的外观,几何以及部 件级别的运动。这个阶段我们会渲染我们输出的 3D 高斯,利用数据集中提供的多视角图像计算一个损失,具体地:


图片


实验结果


实验设置


我们在两个数据集上来衡量我们提出的 PartRM 方法,这两个数据集包括我们提出的 PartDrag-4D 数据集 以及通用数据集 Objaverse-Animation-HQ。因为 Objaverse-Animation-HQ 数据量比较大,我们只从其中采样 15000 条数据,然后手动拆分训练集和测试集。验证时,我们对输出的 3D 高斯渲染 8 个不同的视角,在这 8 个视角上算 PSNR ,SSIM 和 LPIPS 指标。


我们选用 DragAPart , DiffEditor 和 Puppet-Master 作为我们的 baseline。对于不需要训练的 DiffEditor 方法,我们直接拿它官方的训练权重进行推理。对于需要训练的 DragAPart 和 Puppet-Master,我们在训练 集上对他们进行微调。


由于现有的方法只能输出 2D 图像,不能输出 3D 表征,为了和我们的任务对齐,我们设计了两种方法。第一种称为 NVS-First,即我们首先对输入的单视角图像利用 Zero123++ 生成多视角图像,再分别对每个视角结合每个视角对应的拖拽进行推理,生成对应的图像后再进行 3D 高斯重建;第二种称为 Drag-First,

即我们首先先对输入视角进行拖拽,然后对生成的结果利用 Zero123++ 进行多视角生成,最后进行 3D 高斯重建。我们采用了两种 3D 高斯重建方法,第一种为直接用 LGM (下图中两个时间的第一个)进行重建,第二种利用基于优化的 3D 高斯泼溅进行重建(下图中两个时间的第二个)。


定性比较




在视觉效果方面, PartRM 通过对外观,几何和运动的联合建模,能够在抽屉开合等场景中生成物理合理的三维表征。相比之下, DiffEditor 由于缺乏三维感知,导致部件形变错位; DragAPart 虽然能够处理简  单的关节运动,但在生成微波门板时出现了明显的伪影等问题,同时在通用数据集上表现不佳;Puppet- Master 在外观的时间连续性和运动部分的建模方面表现不佳。



在 in the wild 质量方面,我们从互联网上采了一些数据,手动设置拖拽,利用我们在 PartDrag-4D 上训练 好的 PartRM 进行推理。图中可以看到,我们的方法在一些和训练数据分布差别不大的数据上可以取得较 好的效果;但是在一些分布差别较大的数据上效果欠佳。



定量比较



定量评估中, PartRM 在 PSNR、SSIM、 LPIPS 指标上较基线模型均有提升;同时大幅提升了生成效率, PartRM 仅需 4 秒即可完成单次生成,而传统方案需分步执行 2D 形变与三维重建。


总结


本文介绍了 PartRM ,一种同时建模外观、几何和部件级运动的新方法。为了解决 4D 部件级运动学习中的数据稀缺问题,我们提出了 PartDrag-4D 数据集,提供了部件级动态的多视角图像。实验结果表明,我们的方法在部件运动学习上优于以往的方法,并且可应用于具身 AI 任务。然而,对于与训练分布差异较大的关节数据,可能会遇到挑战。



© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

PartRM中的拖拽传播策略是利用Segment Anything模型,根据用户给定的拖拽起始点,得到对应的被拖拽区域的掩码,然后在该区域内采样多个点作为被传播拖拽的起始点,这些点的拖拽强度与用户给定的相同。这样做的目的是为了让网络更好地感知需要被拖拽的整个区域,避免只关注用户指定的单一拖拽点。

如果用户给定的拖拽信息不准确或存在偏差,肯定会对结果产生影响。但文中提到,即使拖拽强度大小估计不准确,模型也足够稳健,可以通过数据驱动的方式学习生成预期的输出。不过,如果偏差过大,例如起始点或方向完全错误,模型可能就无法正确预测运动了。所以,拖拽信息的质量对最终结果还是至关重要的。

拖拽传播感觉有点像图像处理里的“区域生长”算法,从一个点出发,把相似的区域都算作是被影响的部分。但是如果用户给的拖拽方向反了,或者力度不对,那结果肯定会歪掉吧?感觉这个“稳健性”应该是有个限度的。

楼上说的有道理,感觉这个拖拽传播策略有点像是在“脑补”用户意图。如果用户给的信息靠谱,那模型就能“脑补”得很好;反之,如果用户给的信息太离谱,那模型可能就“脑补”出一个四不像了。所以,使用的时候还是要尽量给模型提供准确的引导。不过这个策略确实比只用一个拖拽点要好,至少考虑到了物体是一个整体。

我觉得可以考虑加入一些对抗训练,让模型能够更好地识别和处理真实世界中的噪声和干扰。另外,可以尝试用一些无监督或自监督的学习方法,让模型能够从大量无标签的数据中学习到更通用的特征表示。

提高泛化能力是个老生常谈的问题了,针对PartRM,可以考虑以下几个方向:
1. 数据增强: 引入更多样化的数据,例如不同光照、不同角度、不同材质的铰链物体图像。可以采用数据增强技术,如旋转、缩放、裁剪、颜色抖动等,来增加数据的多样性。
2. 领域自适应: 考虑使用领域自适应技术,例如对抗训练,来减小训练数据和真实世界数据之间的差异。
3. 元学习: 采用元学习方法,让模型学习如何快速适应新的铰链物体类型。
4. 模型结构改进: 考虑引入注意力机制或Transformer结构,让模型更好地关注重要的特征,从而提高泛化能力。
5. 增加数据集规模: 简单粗暴但是有效

此外,还可以考虑结合先验知识,例如物理约束,来提高模型的鲁棒性。

两阶段训练的优势主要在于避免灾难性遗忘。PartRM是在预训练的LGM基础上构建的,LGM已经具备了对静态3D物体外观和几何的建模能力。如果直接在新数据集上微调,可能会导致模型忘记之前学习到的知识,从而降低对真实场景数据的泛化能力。

先学习运动信息,可以使模型专注于学习新的动态特性,然后再联合训练外观、几何和运动信息,可以更好地平衡新知识和旧知识,从而获得更好的性能。

至于为什么不采用端到端训练,可能是因为端到端训练需要大量的数据,而且训练难度较大。两阶段训练可以分解问题,降低训练难度,并且可以更好地利用预训练模型的知识。

个人感觉,两阶段训练有点像是在“分步走”。第一步先把骨架搭起来(运动),第二步再往骨架上贴皮(外观)。如果一步到位,可能骨架就搭歪了,皮也贴不好。而且,LGM本身就是一个很强的模型,直接用它的能力,比从头开始训练要高效得多。

我觉得除了避免灾难性遗忘,可能还有一个原因是运动和外观的loss尺度不一样。运动可能用高斯距离,外观用图像loss,直接一起训容易出现梯度冲突,导致训练不稳定。分开训可能更容易调参。

除了数据和模型,我觉得还可以从任务本身入手。PartRM现在是直接预测3D高斯泼溅,这个任务可能太难了,对模型的泛化能力要求很高。可以考虑把任务分解成几个更简单的子任务,例如先预测物体的姿态,再预测部件的运动,最后再生成3D模型。这样可以降低每个子任务的难度,从而提高整体的泛化能力。