認識 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。所有研究歸功於該項目的研究人員。