鉸鏈物體的通用世界模型,超越擴(kuò)散方法,入選CVPR 2025
基于當(dāng)前觀察,預(yù)測鉸鏈物體的的運(yùn)動(dòng),尤其是 part-level 級別的運(yùn)動(dòng),是實(shí)現(xiàn)世界模型的關(guān)鍵一步。盡管現(xiàn)在基于 diffusion 的方法取得了很多進(jìn)展,但是這些方法存在處理效率低,同時(shí)缺乏三維感知等問題,難以投入真實(shí)環(huán)境中使用。
清華大學(xué)聯(lián)合北京大學(xué)提出了第一個(gè)基于重建模型的 part-level 運(yùn)動(dòng)的建?!狿artRM。用戶給定單張輸入圖像和對應(yīng)的 drag ,PartRM 能生成觀測物體未來狀態(tài)的三維表征 ,使得生成數(shù)據(jù)能夠真正服務(wù)于機(jī)器人操縱等任務(wù)。實(shí)驗(yàn)證明 PartRM 在生成結(jié)果上都取得了顯著的提升。該研究已入選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
研究動(dòng)機(jī)
世界模型是一種基于當(dāng)前觀察和動(dòng)作來預(yù)測未來狀態(tài)的函數(shù)。該模型的研發(fā)使得計(jì)算機(jī)能夠理解物理世界中的復(fù)雜規(guī)律,在機(jī)器人等領(lǐng)域得到了廣泛應(yīng)用。近期,對 part-level 的動(dòng)態(tài)建模的興趣日益增長,給定當(dāng)前時(shí)刻的觀察并給與用戶給定的拖拽,預(yù)測下一時(shí)刻的鉸鏈物體各個(gè)部件的運(yùn)動(dòng)受到越來越多的關(guān)注,這種類型的世界模型對于需要高精度的任務(wù),例如機(jī)器人的操縱任務(wù)等,具有重要的意義。
然而,我們對這個(gè)充滿前景的領(lǐng)域的調(diào)研表明,目前的前沿研究(如 Puppet-Master)通過對預(yù)訓(xùn)練的 大規(guī)模視頻擴(kuò)散模型進(jìn)行微調(diào),以實(shí)現(xiàn)增加拖拽控制的功能。盡管這種方法有效地利用了預(yù)訓(xùn)練過程中 學(xué)習(xí)到的豐富運(yùn)動(dòng)模式,但在實(shí)際應(yīng)用中仍顯不足。其中一個(gè)主要局限是它僅輸出單視角視頻作為表示,而模擬器需要三維表示來從多個(gè)視角渲染場景。此外,擴(kuò)散去噪過程可能需要幾分鐘來模擬單個(gè)拖 拽交互,這與為操作策略(Manipulation Policies)提供快速試錯(cuò)反饋的目標(biāo)相悖。
因此,我們需要采用三維表征,為了實(shí)現(xiàn)從輸入單視角圖像的快速三維重建,我們利用基于三維高斯?jié)姙R(3DGS)的大規(guī)模重建模型,這些模型能以前饋方式從輸入圖像預(yù)測三維高斯?jié)姙R,使重建時(shí)間從傳 統(tǒng)優(yōu)化方法所需的幾分鐘減少到僅需幾秒鐘。同時(shí),通過將用戶指定的拖拽信息加入到大規(guī)模三維重建 網(wǎng)絡(luò)中,我們實(shí)現(xiàn)了部件級別的動(dòng)態(tài)建模。在這個(gè)問題中,我們認(rèn)為聯(lián)合建模運(yùn)動(dòng)和幾何是至關(guān)重要的,因?yàn)椴考夁\(yùn)動(dòng)本質(zhì)上與每個(gè)部件的幾何特性相關(guān)聯(lián)(例如,抽屜在打開時(shí)通常沿其法線方向滑動(dòng))。這種集成使我們能夠?qū)崿F(xiàn)更真實(shí)和可解釋的部件級動(dòng)態(tài)表示。
同時(shí),由于我們是第一個(gè)做這個(gè)任務(wù)的,在這個(gè)任務(wù)上缺少相關(guān)的數(shù)據(jù)集,因此我們基于 PartNet- Mobility 構(gòu)建了 PartDrag-4D 數(shù)據(jù)集,并在這個(gè)數(shù)據(jù)集上建立了衡量對部件級別動(dòng)態(tài)建模的基準(zhǔn)(Benchmark),實(shí)驗(yàn)結(jié)果表明,我們的方法在定量和定性上都取得了最好的效果。
PartDrag-4D 數(shù)據(jù)集的構(gòu)建
我們首先定義 PartRM 需要完成的任務(wù),給定單張鉸鏈物體的圖像 ot 和用戶指定的拖拽 at ,我們需要設(shè)計(jì) 一個(gè)模型,完成
其中 是 Ot 在 at 作用下的三維表征。
現(xiàn)有的數(shù)據(jù)集分為兩種, 一種是只含有數(shù)據(jù)對,但是缺乏對應(yīng)的三維表征(比如 DragAPart 中提出的 Drag-a-Move 數(shù)據(jù)集)。還有一種是通用數(shù)據(jù)集,比如 Objaverse 中的動(dòng)態(tài)數(shù)據(jù),這種數(shù)據(jù)不止還有部件級別的運(yùn)動(dòng),還會(huì)含有物體整體的變形等運(yùn)動(dòng),不適合我們訓(xùn)練。
因此,我們基于提供鉸鏈物體部件級別標(biāo)注的 PartNet-Mobility 數(shù)據(jù)集構(gòu)建了 PartDrag-4D 數(shù)據(jù)集。我們選取了 8 種鉸鏈物體(其中 7 種用于訓(xùn)練, 1 種用于測試),共 738 個(gè) mesh。對于每個(gè) mesh,如圖所示,我們使其中某個(gè)部件在兩種極限狀態(tài)(如完全閉合到完全開啟)間運(yùn)動(dòng)至 6 個(gè)狀態(tài),同時(shí)將其他部分狀態(tài) 設(shè)置為隨機(jī),從而產(chǎn)生共 20548 個(gè)狀態(tài),其中 20057 個(gè)用于訓(xùn)練,491 個(gè)用于測試。為渲染多視角圖像,我們利用 Blender 為每個(gè) mesh 渲染了 12 個(gè)視角的圖像。對于兩個(gè)狀態(tài)之間拖拽數(shù)據(jù)的采樣,我們在鉸鏈物體運(yùn)動(dòng)部件的 Mesh 表面選取采樣點(diǎn),并將兩個(gè)狀態(tài)中對應(yīng)的采樣點(diǎn)投影至 2D 圖像空間,即可獲得對應(yīng)的拖拽數(shù)據(jù)。
PartRM 方法
方法概覽
上圖提供了 PartRM 方法的概述,給定一個(gè)單視角的鉸鏈物體的圖像 ot 和對應(yīng)的拖拽 at,我們的目標(biāo)是生 成對應(yīng)的 3D 高斯?jié)姙R。我們首先會(huì)利用多視角生成模型 Zero123++ 生成輸入的多視角圖像,然后對輸入的拖拽在用戶希望移動(dòng)的 Part 上進(jìn)行傳播。這些多視角的圖像和傳播后的拖拽會(huì)輸入進(jìn)我們設(shè)計(jì)的網(wǎng)絡(luò)中,這個(gè)網(wǎng)絡(luò)會(huì)對輸入的拖拽進(jìn)行多尺度的嵌入,然后將得到的嵌入拼接到重建網(wǎng)絡(luò)的下采樣層中。在訓(xùn)練過程中,我們采用兩階段訓(xùn)練方法,第一階段學(xué)習(xí) Part 的運(yùn)動(dòng),利用高斯庫里的 3D 高斯進(jìn)行 監(jiān)督,第二階段學(xué)習(xí)外觀,利用數(shù)據(jù)集里的多視角圖像進(jìn)行監(jiān)督。
圖像和拖拽的預(yù)處理
圖像預(yù)處理:由于我們的主網(wǎng)絡(luò)是基于 LGM 設(shè)計(jì)的, LGM 需要多視角的圖像作為輸入,所以我們需要將 輸入的單視角圖像變成多視角,我們利用多視角圖像生成網(wǎng)絡(luò) Zero123++,為了使得 Zero123++ 生成的 圖像質(zhì)量更高,我們會(huì)在訓(xùn)練集上對其進(jìn)行微調(diào)。
拖拽傳播:如果用戶只輸入一個(gè)拖拽,后續(xù)網(wǎng)絡(luò)可能會(huì)對拖拽的區(qū)域產(chǎn)生幻覺從而出錯(cuò),因此我們需要 對拖拽進(jìn)行傳播到需要被拖拽部分的各個(gè)區(qū)域,使得后續(xù)網(wǎng)絡(luò)感知到需要被拖拽的區(qū)域,為此我們設(shè)計(jì)了一個(gè)拖拽傳播策略。如圖所示,我們首先拿用戶給定的拖拽的起始點(diǎn)輸入進(jìn) Segment Anything 模型中得到對應(yīng)的被拖拽區(qū)域的掩碼,然后在這個(gè)掩碼區(qū)域內(nèi)采樣一些點(diǎn)作為被傳播拖拽的起始點(diǎn),這些被傳播的拖拽的強(qiáng)度和用戶給定的拖拽的強(qiáng)度一樣。盡管在拖動(dòng)強(qiáng)度大小的估計(jì)上可能存在不準(zhǔn)確性,我們后續(xù)的模型仍然足夠穩(wěn)健,能夠以數(shù)據(jù)驅(qū)動(dòng)的方式學(xué)習(xí)生成預(yù)期的輸出。
拖拽嵌入
PartRM 重建網(wǎng)絡(luò)的 UNet 部分沿用了 LGM 的網(wǎng)絡(luò)架構(gòu),為了將上一步處理好的拖拽注入到重建網(wǎng)絡(luò)中, 我們設(shè)計(jì)了一個(gè)多尺度的拖拽嵌入。具體地,對于每一個(gè)拖拽,我們會(huì)將它的起始點(diǎn)和終止點(diǎn)先過一個(gè) Fourier 嵌入,然后過一個(gè)三層的 MLP:
其中代表第 i 個(gè)拖拽在第 l 層的嵌入,其余部分設(shè)為 0。F 代表 Fourier 嵌入和 MLP ,
代表在 channel 維度上連接。得到第 l 層的嵌入后,我們將
和網(wǎng)絡(luò)第 l 層的輸出 Ol 在 channel 維度上連接,并過一個(gè)卷積層,作為 Ol 的殘差加到 Ol 上作為下一層的輸入,具體地:
其中卷積層的參數(shù)全零初始化,為第 l + 1 層的輸入。
兩階段訓(xùn)練流程
為了保證對靜態(tài) 3D 物體外觀和幾何的建模能力,我們在預(yù)訓(xùn)練的 LGM 基礎(chǔ)上構(gòu)建了 PartRM。但直接在新數(shù)據(jù)集上微調(diào)會(huì)導(dǎo)致已有知識(shí)災(zāi)難性遺忘,從而降低對真實(shí)場景數(shù)據(jù)的泛化能力。為此,我們提出了 兩階段學(xué)習(xí)方法:先專注于學(xué)習(xí)之前未涉及的運(yùn)動(dòng)信息,再訓(xùn)練外觀、幾何和運(yùn)動(dòng)信息,以確保更好的性能。
運(yùn)動(dòng)學(xué)習(xí)階段:在運(yùn)動(dòng)學(xué)習(xí)階段,我們期望模型能夠?qū)W到由輸入的拖拽引起的運(yùn)動(dòng)。我們首先利用在我 們的數(shù)據(jù)集上微調(diào)好的 LGM 去推理每個(gè)狀態(tài) Mesh 對應(yīng)的 3D 高斯?jié)姙R表征,拿這些作為監(jiān)督數(shù)據(jù)我們第 一階段的訓(xùn)練。對于兩個(gè) 3D 高斯之間的對應(yīng),我們利用 LGM 輸出的是一個(gè) splatter image 這一優(yōu)勢,即 LGM 會(huì)對 2D 圖像的每一個(gè)像素點(diǎn)學(xué)一個(gè)高斯?jié)姙R,我們可以直接對監(jiān)督數(shù)據(jù)和 PartRM 網(wǎng)絡(luò)預(yù)測的輸出 做像素級別的 L2 損失,即:
其中 i 代表在 splatter image 里的坐標(biāo), GSi 和 GSj 均為每個(gè)像素點(diǎn)對應(yīng)的 14 維高斯球參數(shù)。
外觀學(xué)習(xí)階段: 在運(yùn)動(dòng)學(xué)習(xí)階段之后,我們引入了一個(gè)額外的階段來聯(lián)合優(yōu)化輸出的外觀,幾何以及部 件級別的運(yùn)動(dòng)。這個(gè)階段我們會(huì)渲染我們輸出的 3D 高斯,利用數(shù)據(jù)集中提供的多視角圖像計(jì)算一個(gè)損失,具體地:
實(shí)驗(yàn)結(jié)果
實(shí)驗(yàn)設(shè)置
我們在兩個(gè)數(shù)據(jù)集上來衡量我們提出的 PartRM 方法,這兩個(gè)數(shù)據(jù)集包括我們提出的 PartDrag-4D 數(shù)據(jù)集 以及通用數(shù)據(jù)集 Objaverse-Animation-HQ。因?yàn)?Objaverse-Animation-HQ 數(shù)據(jù)量比較大,我們只從其中采樣 15000 條數(shù)據(jù),然后手動(dòng)拆分訓(xùn)練集和測試集。驗(yàn)證時(shí),我們對輸出的 3D 高斯渲染 8 個(gè)不同的視角,在這 8 個(gè)視角上算 PSNR ,SSIM 和 LPIPS 指標(biāo)。
我們選用 DragAPart , DiffEditor 和 Puppet-Master 作為我們的 baseline。對于不需要訓(xùn)練的 DiffEditor 方法,我們直接拿它官方的訓(xùn)練權(quán)重進(jìn)行推理。對于需要訓(xùn)練的 DragAPart 和 Puppet-Master,我們在訓(xùn)練 集上對他們進(jìn)行微調(diào)。
由于現(xiàn)有的方法只能輸出 2D 圖像,不能輸出 3D 表征,為了和我們的任務(wù)對齊,我們設(shè)計(jì)了兩種方法。第一種稱為 NVS-First,即我們首先對輸入的單視角圖像利用 Zero123++ 生成多視角圖像,再分別對每個(gè)視角結(jié)合每個(gè)視角對應(yīng)的拖拽進(jìn)行推理,生成對應(yīng)的圖像后再進(jìn)行 3D 高斯重建;第二種稱為 Drag-First,
即我們首先先對輸入視角進(jìn)行拖拽,然后對生成的結(jié)果利用 Zero123++ 進(jìn)行多視角生成,最后進(jìn)行 3D 高斯重建。我們采用了兩種 3D 高斯重建方法,第一種為直接用 LGM (下圖中兩個(gè)時(shí)間的第一個(gè))進(jìn)行重建,第二種利用基于優(yōu)化的 3D 高斯?jié)姙R進(jìn)行重建(下圖中兩個(gè)時(shí)間的第二個(gè))。
定性比較
在視覺效果方面, PartRM 通過對外觀,幾何和運(yùn)動(dòng)的聯(lián)合建模,能夠在抽屜開合等場景中生成物理合理的三維表征。相比之下, DiffEditor 由于缺乏三維感知,導(dǎo)致部件形變錯(cuò)位; DragAPart 雖然能夠處理簡 單的關(guān)節(jié)運(yùn)動(dòng),但在生成微波門板時(shí)出現(xiàn)了明顯的偽影等問題,同時(shí)在通用數(shù)據(jù)集上表現(xiàn)不佳;Puppet- Master 在外觀的時(shí)間連續(xù)性和運(yùn)動(dòng)部分的建模方面表現(xiàn)不佳。
在 in the wild 質(zhì)量方面,我們從互聯(lián)網(wǎng)上采了一些數(shù)據(jù),手動(dòng)設(shè)置拖拽,利用我們在 PartDrag-4D 上訓(xùn)練 好的 PartRM 進(jìn)行推理。圖中可以看到,我們的方法在一些和訓(xùn)練數(shù)據(jù)分布差別不大的數(shù)據(jù)上可以取得較 好的效果;但是在一些分布差別較大的數(shù)據(jù)上效果欠佳。
定量比較
定量評估中, PartRM 在 PSNR、SSIM、 LPIPS 指標(biāo)上較基線模型均有提升;同時(shí)大幅提升了生成效率, PartRM 僅需 4 秒即可完成單次生成,而傳統(tǒng)方案需分步執(zhí)行 2D 形變與三維重建。
總結(jié)
本文介紹了 PartRM ,一種同時(shí)建模外觀、幾何和部件級運(yùn)動(dòng)的新方法。為了解決 4D 部件級運(yùn)動(dòng)學(xué)習(xí)中的數(shù)據(jù)稀缺問題,我們提出了 PartDrag-4D 數(shù)據(jù)集,提供了部件級動(dòng)態(tài)的多視角圖像。實(shí)驗(yàn)結(jié)果表明,我們的方法在部件運(yùn)動(dòng)學(xué)習(xí)上優(yōu)于以往的方法,并且可應(yīng)用于具身 AI 任務(wù)。然而,對于與訓(xùn)練分布差異較大的關(guān)節(jié)數(shù)據(jù),可能會(huì)遇到挑戰(zhàn)。