世界模型（World Model）作為近年來(lái)機(jī)器學(xué)習(xí)和強(qiáng)化學(xué)習(xí)的研究熱點(diǎn)，通過(guò)建立智能體對(duì)其所處環(huán)境的一種內(nèi)部表征和模擬，能夠加強(qiáng)智能體對(duì)于世界的理解，進(jìn)而更好地進(jìn)行規(guī)劃和決策。在強(qiáng)化學(xué)習(xí)領(lǐng)域中，世界模型通常被建模為一個(gè)神經(jīng)網(wǎng)絡(luò)，通過(guò)歷史狀態(tài)和動(dòng)作，預(yù)測(cè)未來(lái)可能出現(xiàn)的狀態(tài)。其中，Dreamer 算法在多種模擬環(huán)境的成功表現(xiàn)讓我們看到了世界模型優(yōu)秀的表征和泛化能力。如果將世界模型應(yīng)用于復(fù)雜真實(shí)場(chǎng)景，是否能夠?qū)崿F(xiàn)更好的控制決策呢？

對(duì)此，ByteDance Research 研究團(tuán)隊(duì)成功將世界模型應(yīng)用于四足機(jī)器人視覺(jué)控制領(lǐng)域，提出了基于世界模型的感知算法 WMP（World Model-based Perception），WMP 通過(guò)在模擬器中學(xué)習(xí)世界模型和策略，其中世界模型通過(guò)歷史感知信息（包括視覺(jué)感知和本體感知）預(yù)測(cè)未來(lái)的感知，策略以世界模型提取的特征作為輸入，輸出具體控制動(dòng)作。

WMP 將模擬器中訓(xùn)練的世界模型和策略 Zero-Shot 遷移到宇樹(shù) A1 機(jī)器人進(jìn)行驗(yàn)證，在多種環(huán)境下取得了出色的成績(jī)，達(dá)到了目前為止 A1 機(jī)器人 SOTA 的通過(guò)性能。同時(shí)，使用模擬數(shù)據(jù)訓(xùn)練的世界模型可以準(zhǔn)確預(yù)測(cè)真實(shí)軌跡，展示出卓越的泛化性能，有望成為一種機(jī)器人控制的新范式。

• 項(xiàng)目主頁(yè)：https://wmp-loco.github.io/
• 論文地址：https://arxiv.org/abs/2409.16784

研究背景

近年來(lái)，強(qiáng)化學(xué)習(xí) (RL) 通過(guò)在物理模擬器中訓(xùn)練策略，然后將其轉(zhuǎn)移到現(xiàn)實(shí)世界（Sim-to-Real transfer），在足式機(jī)器人的運(yùn)動(dòng)控制領(lǐng)域被廣泛應(yīng)用。在足式機(jī)器人的運(yùn)動(dòng)控制中，視覺(jué)圖像信息對(duì)于諸如越障等復(fù)雜環(huán)境是不可或缺的。

憑借強(qiáng)化學(xué)習(xí)的獎(jiǎng)勵(lì)信號(hào)從長(zhǎng)序列、高維的視覺(jué)信息中學(xué)習(xí)策略難度極大。為了更好地引入視覺(jué)信息，傳統(tǒng)特權(quán)學(xué)習(xí)（privileged learning）框架首先學(xué)習(xí)一個(gè)教師策略，教師策略的輸入包含只能在模擬器中得到的特權(quán)信息，如高度圖掃標(biāo)點(diǎn)（heightmap scandots），各類(lèi)障礙物的幾何參數(shù)等。之后再學(xué)習(xí)一個(gè)以深度圖序列為輸入的學(xué)生策略以模仿教師策略的動(dòng)作。

由于特權(quán)學(xué)習(xí)兩階段的訓(xùn)練模式，學(xué)生策略的性能往往落后于教師策略；并且特權(quán)信息的選擇需要人工設(shè)計(jì)，較為繁瑣，同時(shí)易受到具體環(huán)境的限制。如下圖所示，使用 scandots 作為特權(quán)信息無(wú)法處理需要精確距離的環(huán)境以及存在空中障礙物的環(huán)境。

相比之下，動(dòng)物能夠在沒(méi)有特權(quán)信息的情況下僅憑借視覺(jué)感官信息通過(guò)各種非結(jié)構(gòu)化的地形，并且在不熟悉的環(huán)境中依然能夠做出合理的決策。認(rèn)知科學(xué)的一種解釋是動(dòng)物會(huì)利用腦中建立的心智模型（mental model）對(duì)外部環(huán)境進(jìn)行理解和預(yù)測(cè)以便更好地理解和應(yīng)對(duì)周?chē)沫h(huán)境，從而做出合理的動(dòng)作和決策。

ByteDance Research 的機(jī)器人研究團(tuán)隊(duì)在此啟發(fā)下設(shè)計(jì)了一種更加通用的機(jī)器人控制框架 WMP（World Model-based Perception）。WMP 通過(guò)構(gòu)建世界模型來(lái)處理復(fù)雜的感知信息，并將世界模型提取的環(huán)境信息輸入給策略，解決了特權(quán)學(xué)習(xí)中特權(quán)信息難以設(shè)計(jì)的局限性。通過(guò)訓(xùn)練的世界模型和策略可以直接遷移到真實(shí)環(huán)境的 A1 機(jī)器人上，在多種復(fù)雜任務(wù)中達(dá)到了目前該領(lǐng)域的領(lǐng)先水平效果。例如，在世界模型的幫助下，A1 機(jī)器人可以跳過(guò) 85cm 的間隙，跳上 55cm 的高臺(tái)，穿過(guò) 22cm 高的橋洞。這些結(jié)果證明了世界模型對(duì)于決策的正向作用，為之后世界模型在機(jī)器人等領(lǐng)域的研究提供了重要參考。

方法

WMP 采用經(jīng)典的 RSSM 框架作為世界模型的結(jié)構(gòu)，RSSM 包括編碼模塊 encoder，解碼模塊 decoder，以及循環(huán)模塊 recurrent model。encoder 將感知信息以及循環(huán)狀態(tài)編碼為一個(gè)隨機(jī)變量，decoder 通過(guò)循環(huán)狀態(tài)和隨機(jī)變量恢復(fù)出原始的感知信息，而循環(huán)模塊則通過(guò)循環(huán)狀態(tài)、隨機(jī)變量以及動(dòng)作序列預(yù)測(cè)下一個(gè)循環(huán)狀態(tài)。為了滿(mǎn)足真機(jī)運(yùn)行的算力要求，WMP 將世界模型的運(yùn)行頻率設(shè)定為策略運(yùn)行頻率的 k 分之一。由于世界模型主要處理更高層級(jí)的信息，較低的控制頻率同樣能滿(mǎn)足底層控制的需求，這與人體大腦和小腦展現(xiàn)出的不同的控制頻率情況有一定的相似之處。

一個(gè)訓(xùn)練有素的世界模型的循環(huán)狀態(tài)包含足夠多的信息從而預(yù)測(cè)未來(lái)的時(shí)間步，也有助于策略執(zhí)行動(dòng)作。因此在 WMP 框架中，策略會(huì)接受來(lái)自世界模型的循環(huán)狀態(tài)作為輸入。并使用強(qiáng)化學(xué)習(xí)算法 PPO 進(jìn)行訓(xùn)練。此外，策略的訓(xùn)練和世界模型的訓(xùn)練使用模擬數(shù)據(jù)同步進(jìn)行，簡(jiǎn)化了特權(quán)學(xué)習(xí)中的兩階段訓(xùn)練。訓(xùn)練后的策略和世界模型可以無(wú)需微調(diào)直接遷移到真實(shí)機(jī)器人設(shè)備。

實(shí)驗(yàn)結(jié)果

模擬實(shí)驗(yàn)：

WMP 算法使用 Issacgym 模擬器構(gòu)建的 6 種地形上進(jìn)行訓(xùn)練：Slope、Stair、Gap、Climb、Crawl、Tilt。由于 scandot 特權(quán)信息的局限性，使用特權(quán)學(xué)習(xí)訓(xùn)練的 Student baseline 只使用前四個(gè)地形進(jìn)行訓(xùn)練。在模擬器的定量對(duì)比實(shí)驗(yàn)中，WMP 在絕大多數(shù)任務(wù)中獲得了比 Baseline 更高的回報(bào)獎(jiǎng)勵(lì)以及更小的速度追蹤誤差。

真機(jī)實(shí)驗(yàn)：

在真機(jī)實(shí)驗(yàn)中，WMP 繼承了模擬器中的優(yōu)秀表現(xiàn)，相比 baseline 能以更高的成功率通過(guò)更難的地形，并且在室內(nèi)和室外環(huán)境中表現(xiàn)保持一致，進(jìn)一步展現(xiàn)出世界模型優(yōu)秀的泛化能力。

一鏡到底視頻：

驗(yàn)證實(shí)驗(yàn)：

使用模擬數(shù)據(jù)訓(xùn)練需要考慮的一個(gè)問(wèn)題是世界模型對(duì)真實(shí)軌跡預(yù)測(cè)的準(zhǔn)確性如何。驗(yàn)證實(shí)驗(yàn)表明，世界模型對(duì)于真實(shí)軌跡的圖像能給出準(zhǔn)確的預(yù)測(cè)，尤其是對(duì)于關(guān)鍵的部分。例如，世界模型對(duì)橋洞障礙物整體形狀的預(yù)測(cè)存在偏差，但對(duì)機(jī)器人需通過(guò)的窄縫的位置角度的預(yù)測(cè)十分準(zhǔn)確。這驗(yàn)證了世界模型有利于解構(gòu)和提取不同域中的關(guān)鍵要素，從而有助于模擬到真實(shí)的泛化。

總結(jié)

本研究提出了一種新的運(yùn)動(dòng)控制框架，通過(guò)構(gòu)筑的世界模型來(lái)處理視覺(jué)信息和輔助決策，在四足機(jī)器人運(yùn)動(dòng)控制領(lǐng)域取得了不錯(cuò)的效果。WMP 揭示了世界模型在 Sim2Real 以及機(jī)器人控制領(lǐng)域的巨大潛力，為之后世界模型在現(xiàn)實(shí)世界更廣泛的應(yīng)用提供了樣例和寶貴的經(jīng)驗(yàn)。