在機器人操作任務中，預測性策略近年來在具身人工智能領域引起了廣泛關注，因為它能夠利用預測狀態(tài)來提升機器人的操作性能。然而，讓世界模型預測機器人與物體交互的精確未來狀態(tài)仍然是一個公認的挑戰(zhàn)，尤其是生成高質(zhì)量的像素級表示。

為解決上述問題，國防科大、北京大學、深圳大學團隊提出LaDi-WM（Latent Diffusion-based WorldModels），一種基于隱空間擴散的世界模型，用于預測隱空間的未來狀態(tài)。

具體而言，LaDi-WM 利用預訓練的視覺基礎模型 (Vision Fundation Models) 來構(gòu)建隱空間表示，該表示同時包含幾何特征（基于 DINOv2 構(gòu)造）和語義特征（基于 Siglip 構(gòu)造），并具有廣泛的通用性，有利于機器人操作的策略學習以及跨任務的泛化能力。

基于 LaDi-WM，團隊設計了一種擴散策略，該策略通過整合世界模型生成的預測狀態(tài)來迭代地優(yōu)化輸出動作，從而生成更一致、更準確的動作結(jié)果。通過在虛擬和真實數(shù)據(jù)集上的大量實驗，LaDi-WM 能夠顯著提高機器人操作任務的成功率，尤其是在 LIBERO-LONG 數(shù)據(jù)集上提升27.9%，超過之前的所有方法。

• 論文地址：https://arxiv.org/abs/2505.11528
• 項目主頁：https://guhuangai.github.io/LaDiWM.github.io/

論文創(chuàng)新點：

1.一種基于隱空間擴散的世界模型：使用視覺基礎模型構(gòu)建隱空間的通用表示，并在隱空間學習可泛化的動態(tài)建模能力。

2.一種基于世界模型預測迭代優(yōu)化的擴散策略：利用世界模型生成未來預測的狀態(tài)，將預測的狀態(tài)反饋給策略模型，迭代式地優(yōu)化策略輸出。

圖 1 :（左）通過任務無關的片段學習隱擴散世界模型；（右）通過世界模型的未來狀態(tài)預測來優(yōu)化策略模型

技術路線

該團隊提出一種利用世界模型優(yōu)化策略學習的框架，以學習機器人抓取操作相關的技能策略。該框架可分為兩大階段：世界模型學習和策略學習。

A. 世界模型學習：

(a)隱空間表示：通過預訓練的視覺基礎模型對觀測圖像提取幾何表征與語義表征，其中幾何表征利用 DINOv2 提取，而語義表征則使用 Siglip 提取。

(b)交互擴散：同時對兩種隱空間表示實施擴散過程，并在擴散過程中讓二者充分交互，學習幾何與語義表征之間的依賴關系，從而促進兩種表示的準確動態(tài)預測。

圖 2 : 基于交互擴散的世界模型架構(gòu)

B. 策略模型訓練與迭代優(yōu)化推理

（a）結(jié)合世界模型的未來預測引導策略學習：將世界模型給出的未來預測作為額外的輸入，引導策略模型的準確動作預測；模型架構(gòu)基于擴散策略模型，有利于學習多模態(tài)動作分布。

（b）迭代優(yōu)化策略輸出：策略模型可以在一個時間步多次利用世界模型的未來預測作為引導，從而不斷優(yōu)化自身的動作輸出。實驗顯示，該方案可以逐漸降低策略模型的輸出分布熵，達到更準確的動作預測。

圖 3 : 基于未來預測引導的策略模型架構(gòu)

實驗結(jié)果

虛擬實驗：

在公開的虛擬數(shù)據(jù)集（LIBERO-LONG，CALVIN D-D）中，團隊驗證了所提出框架在機器人抓取相關的操作任務上的性能。在實驗中，世界模型的訓練數(shù)據(jù)會與策略模型的訓練數(shù)據(jù)區(qū)分開，從而驗證世界模型的泛化能力。對于 LIBERO-LONG，給定語言指令，多次執(zhí)行并統(tǒng)計機器人完成各項任務的成功率。對于 CALVIN D-D，連續(xù)給定五個語言指令，多次執(zhí)行并統(tǒng)計平均完成任務的數(shù)量。

在 LIBERO-LONG 數(shù)據(jù)集，為了驗證世界模型對策略模型的引導作用，團隊僅使用 10 條軌跡去訓練各任務，對比結(jié)果如表 1 所示。相比于其他方法，LaDi-WM 能夠提供精確的未來預測，并將預測反饋給策略模型，不斷優(yōu)化動作輸出，僅需少量訓練數(shù)據(jù)即可達到 68.7% 的成功率，顯著優(yōu)于其他方法。

表 1: LIBERO-LONG 性能對比

在 CALVIN D-D 數(shù)據(jù)集上，LaDi-WM 同樣展示了在長時任務中的強大性能（表 2）。

表 2: CALVIN D-D 性能對比

團隊進一步驗證了所提出框架的可擴展性，如圖 4 所示。

（a）逐漸增大世界模型的訓練數(shù)據(jù)，模型的預測誤差逐漸降低且策略性能逐漸提升；

（b）逐漸增大策略模型的訓練數(shù)據(jù)，抓取操作的成功率逐漸提升；

（c）逐漸增大策略模型的參數(shù)量，抓取操作的成功率逐漸提升。

圖 4 : 可擴展性實驗

為了驗證 LaDi-WM 的跨場景泛化能力，團隊在 LIBERO-LONG 上訓練世界模型，并直接應用于 CALVIN D-D 的策略學習中，實驗結(jié)果如表 3 所示。若是使用在 LIBERO-LONG 訓練的原始策略模型，直接應用到 CALVIN D-D 是不工作的（表第一行）；而使用在 LIBERO-LONG 訓練的世界模型來引導 CALVIN 環(huán)境下的策略學習，則可以比在 CALVIN 環(huán)境訓練的原始策略的性能高 0.61（表第三行）。這表明，世界模型的泛化能力要優(yōu)于策略模型的泛化能力。

表 3: 跨場景實驗結(jié)果。L 代表 LIBERO-LONG，C 代表 CALVIN D-D

團隊進一步探索了利用世界模型迭代優(yōu)化的工作原理。團隊收集不同迭代輪次下策略模型的輸出動作并繪制其分布，如圖 5 所示。迭代優(yōu)化的過程中，輸出動作分布的熵在逐漸降低，這表明策略模型每一步的輸出動作更加穩(wěn)定，從而提升整體的抓取成功率。

圖 5 : 迭代優(yōu)化的動作分布對比

真機實驗：

團隊也在真實場景中驗證了所提出框架的性能，具體操作任務包括「疊碗」、「開抽屜」、「關抽屜」以及「抓取物體放入籃子」等，如圖 6 所示。

圖 6 : （左）真實場景環(huán)境；（右）機器人實際操作樣例

在真實場景中，LaDi-WM 將原始模仿學習策略的成功率顯著提升 20%（表 4）。

表 4: 真實場景性能對比

圖 7 展示了最終所得策略模型在不同任務上的執(zhí)行軌跡，從圖中可以發(fā)現(xiàn)，提出的策略能夠在不同光照條件以及不同初始位置的情況下有魯棒的泛化性。

圖 7 : 真實場景機器人執(zhí)行軌跡

總結(jié)

國防科大、北京大學、深圳大學團隊提出了一種隱空間擴散的世界模型 LaDi-WM（Latent Diffusion-based World Models），利用視覺基礎模型提取通用的隱空間表示，并在隱空間學習可泛化的動態(tài)建模。同時，團隊提出基于世界模型的未來預測來引導策略學習，在推理階段通過迭代式地優(yōu)化策略輸出，從而進一步提高策略輸出動作的準確度。團隊通過虛擬與真機上廣泛的實驗證明了 LaDi-WM 的有效性，所提出的方法顯著提升了機器人抓取操作技能的性能。