next-state/open-dreamer

Open Dreamer – 基於 JAX/Flax 的即時 Dreamer 4 世界模型

是什麼

• 使用 JAX/Flax（NNX 庫）撰寫的 Dreamer 4 世界模型管道之開源實作。
• 提供完整的 訓練堆疊，包含動畫分詞器、動作條件的潛在動力學模型，以及評估工具（滾動生成、Fréchet Video Distance）。
• 模型在 Minecraft/VPT 風格的遊戲影片 上訓練，可透過雲端主機演示實時運行。

核心元件

元件	作用
dreamer/models.py	因果動畫分詞器與潛在動力學模型的定義
dreamer/training.py	訓練迴圈、損失函數與檢查點處理（透過 Orbax）
dreamer/generation.py	去噪排程、產生滾動、渲染影片的工具
dreamer/fvd/	特徵提取與 Fréchet Video Distance 計算，用於品質評估
scripts/	可直接執行的入口點：train_tokenizer.py、tokenize_minecraft_dataset.py、train_dynamics.py、eval_fvd.py
configs/	描述資料集、模型超參數與評估設定的 Hydra/OmegaConf YAML 檔案
site/	托管互動式演示與文件的 Next.js 網站

典型工作流程

1. 準備資料 – 將原始 Minecraft MP4 录影轉換為 ArrayRecord 分片（包含影片位元組、動作等的 pickle 字典）。
2. 訓練分詞器 – 學習影片畫格的潛在表示（scripts/train_tokenizer.py）。
3. 分詞資料集 – 將每個原始劇集編碼為潛在標記（scripts/tokenize_minecraft_dataset.py）。
4. 訓練動力學模型 – 在分詞資料上擬合動作條件的潛在動力學模型（scripts/train_dynamics.py）。
5. 生成與評估 – 滾動生成模型以合成新影片畫格並計算 FVD（scripts/eval_fvd.py）。

如何試用

• 即時演示 – 無需安裝；網站托管瀏覽器內演示，可串流 Minecraft 世界，並在真實遊戲與模型預測之間切換。
• 本地推理 – 附屬倉儲 reactor-team/open-dreamer 包含輕量級推理腳本，可在您自己的影片/動作序列上執行訓練好的檢查點。

安裝（訓練端）

# 需要 Python 3.11 和相容 CUDA-12 的 JAX 建構
pip install uv               # UV 是此處使用的快速 Python 套件管理器
uv sync                       # 將固定相依性安裝至虛擬環境
source .venv/bin/activate     # 激活環境

若需不同 JAX 輪子（例如用於不同 GPU 或 CPU），可在 uv sync 步驟後安裝。

執行訓練腳本（分詞器範例）

# 首先編輯 configs/tokenizer.yaml 和 configs/dataset/minecraft_vpt.yaml
uv run scripts/train_tokenizer.py

類似指令適用於分詞、動力學訓練與評估，均由 configs/ 下的 YAML 設定檔驅動。

路線圖

• 目前僅支援世界模型訓練管道。
• 未來工作目標是加入完整的 Dreamer 4 行為克隆 / 強化學習迴圈。

授權與引用

• 倉儲目前帶有「保留所有權利」的占位符通知；正式開源授權計畫於後續版本發布。
• 若在研究中使用程式碼，請引用 README 中提供的 Zenodo 項目與原始 Dreamer 4 論文。

誰可能用得上？

• 探索可擴展世界模型學習的研究人員，特別是在 Minecraft 等視訊豐富環境中。
• 需要快速 JAX 實作的即時生成代理的工程師。
• 尋找完整視訊分詞器 + 潛在動力學管道具體範例的學生。

---

以上所有資訊均直接取自倉儲的 README；未推斷任何額外功能。