認識 Open Dreamer：基於 JAX/Flax 復現 Dreamer 4 世界模型管道，完整訓練配方已釋出

一群 AI 研究人員（Reactor）釋出了 Open Dreamer，這是 Dreamer 4 世界模型管道的開源實現，使用 JAX 和 Flax NNX 編寫。

實際釋出的內容

釋出了兩個倉庫。next-state/open-dreamer 包含訓練管道：因果影片分詞器、動作條件的潛動態模型、 rollout 生成和 FVD 評分。reactor-team/open-dreamer 包含一個最小的本地 rollout 工具，可從 MP4 和匹配的動作檔案生成幀。

第三個成果是託管在 Reactor 執行時上的瀏覽器演示。它即時流式傳輸生成的 Minecraft 世界，並公開一個 Game ⟷ Dream 切換開關，可將流從真實遊戲逐幀切換到世界模型。

明確的目標是復現 Dreamer 4 研究。研究團隊特意避免使用該研究論文之外的方法，以保持搜尋空間狹窄。他們從 CoinRun 開始，這是一個可以在單個 GPU 上訓練的程式化生成的 2D 平臺遊戲，然後將工作管道擴充套件到 Minecraft/VPT 風格的遊戲影片。

架構：一個骨幹，兩個模型

分詞器和動態模型都使用相同的塊因果變換器骨幹。該骨幹交替使用兩種注意力型別。空間層在同一幀的元素之間傳播資訊。因果時間層在幀之間傳播資訊。

分詞器是基於變換器的掩碼自編碼器，而不是 VAE。團隊報告大約 100 倍的壓縮率，並指出該設計不需要 KL 或對抗性損失。他們認為掩碼使潛在空間更易於擴散。

動態模型執行下一幀預測，並使用擴散強制、流匹配和 shortcut 模型進行訓練。它還預測下一個動作。rollout 被摺疊到每個時間步的塊（前一個動作、狀態、策略）中，而不是在單獨的轉換模組和策略之間交替。空間注意力在每個塊內部執行；因果時間注意力連線跨時間的塊。

關鍵地，世界模型 token 不能讀取智慧體 token。因此，任務和策略資訊只能透過下一個動作影響未來狀態。

訓練配方配置

釋出的 Minecraft 配置使配方具體化。

動態模型為 1.6B 引數：30 個塊因果層，d_model 1920，30 個注意力頭，3 個 KV 頭用於分組查詢注意力。每四層中有一層是時間注意力層。每個時間步攜帶 32 個學習到的暫存器 token，packing_factor: 2 將相鄰的分詞器潛在變數打包到每個動態空間 token 中。時間注意力使用 192 步滑動視窗。

訓練執行 200,000 步，使用 Muon、WSD 排程和峰值學習率 3e-4。Shortcut/bootstrap 樣本在第 100,000 步以 0.25 的批次分數開啟。EMA 衰減為 0.999。

分詞器配置每幀輸出 512 個潛在 token，瓶頸寬度為 16。原始 360×640 幀被填充到 368×640，以便兩個維度都能被 16×16 塊整除。編碼器深度為 12，d_model 1536；解碼器深度為 8，d_model 1024。MAE 掩碼機率最高為 0.9，LPIPS 以權重 0.2 應用在半個時間步上。

VPT 動作被解析為 27 個二進位制動作通道加上 121 個分類滑鼠類，沒有連續通道。

計算最大化與記憶體牆

研究團隊報告模型 FLOPs 利用率為 57-58%，而健康變換器訓練的基準為 60%。原因是 roofline 論點。在 B200 上，頻寬受限和計算受限之間的交叉點為 292 FLOP/位元組。每 GPU 提供 256 幀將工作負載推過該峰值點。

分片方式與預期相反。1.6B 引數的模型狀態——引數、梯度、最佳化器狀態和 EMA——大約佔用 24 GiB，適合 B200。啟用才是真正的成本。研究團隊嘗試了資料並行、FSDP、張量並行和序列並行，最終選擇了普通資料並行加啟用檢查點。

資料載入透過將整個資料集預分詞為 .arrayrecord 檔案，然後使用 Grain 和 GPU 端預取緩衝區解決。ffmpeg 解碼速度不足以保持 GPU 滿載。

穩定性部分才是真正的重點

研究團隊明確指出穩定性消耗了他們大部分時間。他們的關鍵觀察：大多數穩定性問題發生在損失下降的情況下。MSE 平滑改善而生成質量下降。

記錄了六種修復方法。Muon 取代了 LaProp，後者隨機且頻繁地尖峰，在兩次約 400 B200 小時的執行中。EMA 權重被視為擴散推斷的必要條件。混合精度對邊界敏感：引數保持 float32，BF16 覆蓋大多數 matmul 啟用和注意力輸入，而 float32 保留用於歸一化和動態流輸出頭。

在損失加權方面，他們使用 x-prediction 和 v-space 損失，這減少到與 Dreamer 4 類似的加權項，但分母是平方的。他們報告了微小但顯著的改進。噪聲和潛在序列之間的 Minibatch barycentric optimal transport 使 rollout 生成更加穩定。μ-引數化經過測試並被判斷為不必要，部分原因是 Muon 在不同模型大小上保持超引數更穩定。

CoinRun 階段的一個進一步結果：iso-FLOPs 掃描將計算最優縮放置於大約 N∝C^0.56 和 D∝C^0.44。

不包含的內容

倉庫不包括行為克隆或 RL 訓練迴圈；完整的 Dreamer 4 BC/RL 智慧體迴圈被列為開放路線圖專案。帖子中描述的 CoinRun 策略工作未用於 Minecraft 且未釋出。

該帖子也未釋出 FVD 分數，儘管 scripts/eval_fvd.py 附帶了一個基於 I3D 的工具，配置了 4 個上下文幀和 240 幀的視野。

關鍵要點

Open Dreamer 在 JAX/Flax NNX 中復現了 Dreamer 4 管道，包含訓練程式碼和 Minecraft 演示。 動態模型為 1.6B 引數，30 層，d_model 1920，使用 Muon 訓練 200K 步。 報告的工程資料：B200 上 MFU 為 57-58%，每 GPU 256 幀，模型狀態約 24 GiB。 穩定性是瓶頸，損失曲線掩蓋了大部分生成質量退化的現象。

檢視部落格文章和演示、訓練倉庫、推理倉庫以及 Reactor on X。所有研究歸功於該專案的研究人員。