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

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

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

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

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

論文創(chuàng)新點：

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

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

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

技術(shù)路線

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

A. 世界模型學(xué)習(xí)：

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

(b) 交互擴(kuò)散：同時對兩種隱空間表示實施擴(kuò)散過程，并在擴(kuò)散過程中讓二者充分交互，學(xué)習(xí)幾何與語義表征之間的依賴關(guān)系，從而促進(jìn)兩種表示的準(zhǔn)確動態(tài)預(yù)測。

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

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

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

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

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

實驗結(jié)果

虛擬實驗：

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

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

表 1: LIBERO-LONG 性能對比

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

表 2: CALVIN D-D 性能對比

團(tuán)隊進(jìn)一步驗證了所提出框架的可擴(kuò)展性，如圖 4 所示。

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

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

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

圖 4 : 可擴(kuò)展性實驗

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

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

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

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

真機(jī)實驗：

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

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

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

表 4: 真實場景性能對比

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

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

總結(jié)

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