世界模型（World Model）作為近年來機器學習和強化學習的研究熱點，透過建立智慧體對其所處環境的一種內部表徵和模擬，能夠加強智慧體對於世界的理解，進而更好地進行規劃和決策。在強化學習領域中，世界模型通常被建模為一個神經網路，透過歷史狀態和動作，預測未來可能出現的狀態。其中，Dreamer 演算法在多種模擬環境的成功表現讓我們看到了世界模型優秀的表徵和泛化能力。如果將世界模型應用於複雜真實場景，是否能夠實現更好的控制決策呢？
對此，ByteDance Research 研究團隊與上海交通大學合作，成功將世界模型應用於四足機器人視覺控制領域，提出了基於世界模型的感知演算法 WMP（World Model-based Perception），WMP 透過在模擬器中學習世界模型和策略，其中世界模型透過歷史感知資訊（包括視覺感知和本體感知）預測未來的感知，策略以世界模型提取的特徵作為輸入，輸出具體控制動作。
WMP 將模擬器中訓練的世界模型和策略 Zero-Shot 遷移到宇樹 A1 機器人進行驗證，在多種環境下取得了出色的成績，達到了目前為止 A1 機器人 SOTA 的透過效能。同時，使用模擬資料訓練的世界模型可以準確預測真實軌跡，展示出卓越的泛化效能，有望成為一種機器人控制的新正規化。
專案主頁：https://wmp-loco.github.io/
論文地址：https://arxiv.org/abs/2409.16784
研究背景
近年來，強化學習 (RL) 透過在物理模擬器中訓練策略，然後將其轉移到現實世界（Sim-to-Real transfer），在足式機器人的運動控制領域被廣泛應用。在足式機器人的運動控制中，視覺影像資訊對於諸如越障等複雜環境是不可或缺的。
憑藉強化學習的獎勵訊號從長序列、高維的視覺資訊中學習策略難度極大。為了更好地引入視覺資訊，傳統特權學習（privileged learning）框架首先學習一個教師策略，教師策略的輸入包含只能在模擬器中得到的特權資訊，如高度圖掃標點（heightmap scandots），各類障礙物的幾何引數等。之後再學習一個以深度圖序列為輸入的學生策略以模仿教師策略的動作。
由於特權學習兩階段的訓練模式，學生策略的效能往往落後於教師策略；並且特權資訊的選擇需要人工設計，較為繁瑣，同時易受到具體環境的限制。如下圖所示，使用 scandots 作為特權資訊無法處理需要精確距離的環境以及存在空中障礙物的環境。
相比之下，動物能夠在沒有特權資訊的情況下僅憑藉視覺感官資訊透過各種非結構化的地形，並且在不熟悉的環境中依然能夠做出合理的決策。認知科學的一種解釋是動物會利用腦中建立的心智模型（mental model）對外部環境進行理解和預測以便更好地理解和應對周圍的環境，從而做出合理的動作和決策。
ByteDance Research 的機器人研究團隊在此啟發下設計了一種更加通用的機器人控制框架 WMP（World Model-based Perception）。WMP 透過構建世界模型來處理複雜的感知資訊，並將世界模型提取的環境資訊輸入給策略，解決了特權學習中特權資訊難以設計的侷限性。透過訓練的世界模型和策略可以直接遷移到真實環境的 A1 機器人上，在多種複雜任務中達到了目前該領域的領先水平效果。例如，在世界模型的幫助下，A1 機器人可以跳過 85cm 的間隙，跳上 55cm 的高臺，穿過 22cm 高的橋洞。這些結果證明了世界模型對於決策的正向作用，為之後世界模型在機器人等領域的研究提供了重要參考。
方法
WMP 採用經典的 RSSM 框架作為世界模型的結構，RSSM 包括編碼模組 encoder，解碼模組 decoder，以及迴圈模組 recurrent model。encoder 將感知資訊以及迴圈狀態編碼為一個隨機變數，decoder 透過迴圈狀態和隨機變數恢復出原始的感知資訊，而迴圈模組則透過迴圈狀態、隨機變數以及動作序列預測下一個迴圈狀態。為了滿足真機執行的算力要求，WMP 將世界模型的執行頻率設定為策略執行頻率的 k 分之一。由於世界模型主要處理更高層級的資訊，較低的控制頻率同樣能滿足底層控制的需求，這與人體大腦和小腦展現出的不同的控制頻率情況有一定的相似之處。
一個訓練有素的世界模型的迴圈狀態包含足夠多的資訊從而預測未來的時間步，也有助於策略執行動作。因此在 WMP 框架中，策略會接受來自世界模型的迴圈狀態作為輸入。並使用強化學習演算法 PPO 進行訓練。此外，策略的訓練和世界模型的訓練使用模擬資料同步進行，簡化了特權學習中的兩階段訓練。訓練後的策略和世界模型可以無需微調直接遷移到真實機器人裝置。
實驗結果
模擬實驗：
WMP 演算法使用 Issacgym 模擬器構建的 6 種地形上進行訓練：Slope、Stair、Gap、Climb、Crawl、Tilt。由於 scandot 特權資訊的侷限性，使用特權學習訓練的 Student baseline 只使用前四個地形進行訓練。在模擬器的定量對比實驗中，WMP 在絕大多數任務中獲得了比 Baseline 更高的回報獎勵以及更小的速度追蹤誤差。
真機實驗：
在真機實驗中，WMP 繼承了模擬器中的優秀表現，相比 baseline 能以更高的成功率透過更難的地形，並且在室內和室外環境中表現保持一致，進一步展現出世界模型優秀的泛化能力。
一鏡到底影片：
驗證實驗：
使用模擬資料訓練需要考慮的一個問題是世界模型對真實軌跡預測的準確性如何。驗證實驗表明，世界模型對於真實軌跡的影像能給出準確的預測，尤其是對於關鍵的部分。例如，世界模型對橋洞障礙物整體形狀的預測存在偏差，但對機器人需透過的窄縫的位置角度的預測十分準確。這驗證了世界模型有利於解構和提取不同域中的關鍵要素，從而有助於模擬到真實的泛化。
總結
本研究提出了一種新的運動控制框架，透過構築的世界模型來處理視覺資訊和輔助決策，在四足機器人運動控制領域取得了不錯的效果。WMP 揭示了世界模型在 Sim2Real 以及機器人控制領域的巨大潛力，為之後世界模型在現實世界更廣泛的應用提供了樣例和寶貴的經驗。
參考文獻：
1. Ha, David, and Jürgen Schmidhuber. "World models." arXiv preprint arXiv:1803.10122 (2018).
2. Hafner, Danijar, et al. "Learning latent dynamics for planning from pixels." International conference on machine learning. PMLR, 2019.
3. Hafner, Danijar, et al. "Learning latent dynamics for planning from pixels." International conference on machine learning. PMLR, 2019.
4. Hafner, Danijar, et al. "Mastering atari with discrete world models." arXiv preprint arXiv:2010.02193 (2020).
5. Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).
6. Lee, Joonho, et al. "Learning quadrupedal locomotion over challenging terrain." Science robotics 5.47 (2020): eabc5986.
7. Miki, Takahiro, et al. "Learning robust perceptive locomotion for quadrupedal robots in the wild." Science robotics 7.62 (2022): eabk2822.
8. Agarwal, Ananye, et al. "Legged locomotion in challenging terrains using egocentric vision." Conference on robot learning. PMLR, 2023.
9. Wu, Philipp, et al. "Daydreamer: World models for physical robot learning." Conference on robot learning. PMLR, 2023.
10. Zhuang, Ziwen, et al. "Robot parkour learning." arXiv preprint arXiv:2309.05665 (2023).
11. Cheng, Xuxin, et al. "Extreme parkour with legged robots." 2024 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2024.