Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets

作者: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · RSS 2025

授權條款

開放取用
論文: CC BY 4.0

以下內容皆依據論文全文，依其 CC BY 4.0 授權條款使用；圖表已標示作者。

一句話摘要

單一擴散 transformer 為動作與未來影像分別設置獨立的雜訊時間步，可擔任策略、動力學模型或影片預測器，並能以沒有動作的影片訓練。

為何重要

一個模型既能以附有動作標註的機器人資料訓練，也能以取得容易得多的無動作影片訓練。在採用 DROID 式 Franka 設備的五項真實機器人任務測試中，預訓練的 UWM 在每項任務上的成功率都最高，而以額外影片共同訓練讓成功率進一步提升，在視覺干擾下也是如此。

問題

大型機器人模型的模仿學習仰賴高品質的專家示範，而透過遠端操作蒐集這些示範既昂貴又緩慢。呈現真實世界動態的影片非常豐富，但目前大多數模仿學習方法無法使用，因為這些影片沒有動作標籤；而且行為複製並不會明確捕捉軌跡中存在的時間動態。

提出的方法

UWM 在一個以目前觀測為條件的 transformer 中，耦合了動作擴散過程與未來影像擴散過程，每種模態各有獨立的擴散時間步。把某個模態的時間步設為 T 即將其邊際化，設為 0 即以其為條件，因此同一個模型可以取樣策略 p(a|o)、正向動力學 p(o′|o,a)、逆向動力學 p(a|o,o′) 或影片預測 p(o′|o)。訓練時使用無動作影片的方法，是把動作時間步固定為 T，並以高斯雜訊填補缺少的動作。

關鍵見解

每種模態獨立的擴散時間步具有部分遮罩的作用，因此單一網路就能涵蓋動作與未來畫格上的邊際分布與條件分布。
帶來提升的是預測未來畫格，而不只是重建目前畫格：在 Stack-Bowls/Block-Cabinet 上成功率為 0.86/0.76，重建目前畫格為 0.70/0.66，Diffusion Policy 為 0.48/0.60。
暫存器 token（預設 8 個）有助於在動作 token 與影像 token 之間交換資訊，而在 LIBERO 消融實驗中，AdaLN 條件調制優於交叉注意力。
從頭訓練時，UWM 與 Diffusion Policy 表現大致相同；UWM 從預訓練中獲益更多。
以域內機器人影片共同訓練，比以網際網路人類活動影片（Kinetics-400 與 Something-Something v2）共同訓練更有幫助，不過兩者都勝過只用機器人資料。

主要結果

真實機器人，以 2,000 條 DROID 軌跡預訓練，分布內成功率：UWM 在 Stack-Bowls、Block-Cabinet、Paper-Towel、Hang-Towel 與 Rice-Cooker 上分別為 0.86、0.76、0.78、0.82 與 0.60，Diffusion Policy 為 0.48、0.60、0.52、0.64、0.35，GR1 為 0.66、0.66、0.60、0.66、0.40。
以另外 2,000 條當作無動作影片使用的 DROID 軌跡共同訓練後，UWM 的分布內成功率提高到 0.92、0.84、0.86、0.86 與 0.65，分布外成功率（前四項任務）從 0.76/0.60/0.78/0.64 提高到 0.84/0.72/0.84/0.76。
分類後的 OOD 場景（光線、背景、雜物）：共同訓練的 UWM 在 Stack-Bowls 與 Block-Cabinet 上為 21/30 與 15/30，預訓練的 UWM 為 15/30 與 8/30，Diffusion Policy 為 12/30 與 6/30。
LIBERO（以 LIBERO-90 預訓練，在五項具分布偏移的 LIBERO-10 任務上微調）：UWM 平均成功率為 0.79 ± 0.11，DP 為 0.71 ± 0.12，PAD 為 0.57 ± 0.19，GR1 為 0.58 ± 0.14。
逆向動力學模式在軌跡長度內追蹤 50 條 LIBERO 軌跡：在 Book-Caddy 與 Soup-Cheese 上成功率為 0.73 與 0.45，策略模式為 0.37 與 0.33。
在 DROID 上訓練 UWM 100K 步，需要 4 張 NVIDIA A100 GPU 耗時 24 小時。

限制

此模型尚未從大規模人類影片中學習，因此並未處理人與機器人之間的形態差距。
正向動力學的預測可能含有瑕疵，可能降低其在規劃上的用處。
作者預期更密集的影片預測會帶來進一步提升，而此版本並未使用。
在 LIBERO 分布外設定中，相對於基準方法的提升遠小於真實機器人上的提升；作者認為模擬的動態較為簡單。
Rice-Cooker 只在 20 個接近分布內的初始化上評估，沒有在 OOD 設定下評估，而所有真實機器人測試都使用同一套 Franka 設備。

對領域的影響

這篇論文被引用 244 次。論文表示影片與程式碼可在其專案頁面取得，該頁面連結了 GitHub 上的程式碼儲存庫 WEIRDLabUW/unified-world-model。

摘要

模仿學習已成為打造通用機器人的一種有前景的途徑。然而，由於仰賴高品質的專家示範，為大型機器人基礎模型擴展模仿學習仍充滿挑戰。同時，描繪各式環境與多樣行為的大量影片資料隨手可得。這些資料提供了關於真實世界動態與代理-環境互動的豐富資訊來源。然而，由於缺少當今多數方法所需的動作標註，直接將這些資料用於模仿學習已證實相當困難。在本研究中，我們提出統一世界模型（Unified World Models，UWM），一個能同時利用影片與動作資料進行策略學習的框架。具體而言，UWM 在統一的 transformer 架構中整合了動作擴散過程與影片擴散過程，每種模態由各自獨立的擴散時間步控制。我們展示，只要控制各個擴散時間步，UWM 就能彈性地表示策略、正向動力學、逆向動力學與影片生成器。透過模擬與真實世界實驗，我們展示：(1) UWM 能在大規模多任務機器人資料集上，同時以動力學預測與動作預測進行有效的預訓練，得到比模仿學習更具泛化性與穩健性的策略；(2) 透過獨立控制各模態的擴散時間步，UWM 自然地支援從沒有動作的影片資料中學習，進一步提升微調後策略的效能。我們的結果顯示，UWM 朝著運用大型異質資料集實現可擴展機器人學習邁出了有前景的一步，並在常彼此分立的模仿學習與世界建模兩種範式之間，提供了一種簡單的統一。影片與程式碼見 https://weirdlabuw.github.io/uwm/

連結

arXiv 2504.02792: https://arxiv.org/abs/2504.02792
DOI https://doi.org/10.15607/rss.2025.xxi.015: https://doi.org/https://doi.org/10.15607/rss.2025.xxi.015
PDF: https://arxiv.org/pdf/2504.02792v1
程式碼: https://github.com/WEIRDLabUW/unified-world-model
專案頁面: https://weirdlabuw.github.io/uwm/