鉸鏈物體的通用世界模型,超越擴散方法,入選CVPR 2025
基于當前觀察,預測鉸鏈物體的的運動,尤其是 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 任務。然而,對于與訓練分布差異較大的關節數據,可能會遇到挑戰。