galilai-group/stable-worldmodel

解決的問題

提供一個統一且可重現的世界模型研究平台，簡化工作流程的三個主要階段：資料收集、模型訓練，以及使用模型預測控制（MPC）進行評估。它消除了研究人員撰寫資料處理與規劃求解器樣板程式碼的需求，讓他們能專注於模型架構與目標函數。

如何運作

該平台將多個關鍵元件整合至單一介面中：

• 資料管理：格式註冊表支援多種後端（LanceDB、HDF5、資料夾、影片與LeRobot），用於記錄與載入資料集，並提供內建工具實現格式間的一次性遷移。
• 環境套件：提供來自 DeepMind Control Suite、Gymnasium、OGBench 與 Atari 等來源的大量標準化環境，其中許多環境包含「變異因素」（factors of variation），用於測試模型在不同視覺或物理設定下的泛化能力。
• 規劃求解器：內建採樣型求解器（如 CEM、MPPI）與基於梯度的求解器（如 SGD、Adam）的參考實作，用於將世界模型預測轉換為動作。
• 統一 API：提供簡單的 Python API，讓使用者能使用專家策略收集資料、訓練模型，並搭配 WorldModelPolicy 與求解器進行評估。

適用對象

需要在多樣化環境中進行結果基準測試與重現的世界模型、具身智能與模型預測控制領域的 AI 研究人員與工程師。

主要亮點

• 豐富的環境資料庫：涵蓋從簡單控制（CartPole）到複雜機器人（Fetch）與 Atari 遊戲的廣泛任務。
• 高效率資料 I/O：利用 LanceDB 實現快速索引讀取與可追加儲存，優化效能。
• 內建基線：包含基於 JEPA 的模型（LeWM、DINO-WM）以及其他行為克隆或強化學習基線的參考實作。
• 泛化能力測試：整合的變異因素使評估模型在分佈偏移下的零樣本泛化能力變得輕鬆。