新聞宏觀經濟Open Dreamer:Reactor 釋出 Dreamer 4 世界模型管線的開源 JAX/Flax 重現版本,包含完整訓練配方

Open Dreamer:Reactor 釋出 Dreamer 4 世界模型管線的開源 JAX/Flax 重現版本,包含完整訓練配方

作者: MarkTechPost·

重點速覽

  • Open Dreamer 提供了 Dreamer 4 管線的開源實作,包含因果影片分詞器、動作條件潛在動態模型、展開工具和 FVD 評估程式碼。
  • Minecraft 動態模型使用 16 億參數,分布在 30 個區塊因果層中,配置為以 Muon 優化器訓練 200,000 步。
  • Reactor 報告在 NVIDIA B200 GPU 上達到 57% 至 58% 的模型 FLOPs 使用率,激活值被識別為主要的記憶體成本而非模型狀態。
  • 團隊記錄了多項穩定性措施,包括 EMA 使用、混合精度邊界、優化器變更、損失加權和小批量最佳傳輸。
  • 此次發布不包含行為克隆或強化學習訓練迴圈,且專案尚未發布 FVD 分數。
Open Dreamer:Reactor 釋出 Dreamer 4 世界模型管線的開源 JAX/Flax 重現版本,包含完整訓練配方

一個以 Reactor 為名運作的研究團隊釋出了 Open Dreamer,這是一個以 JAX 和 Flax NNX 建構的 Dreamer 4 世界模型管線開源實作。世界模型學習從觀測資料預測環境如何演進,使代理人能夠在無需與真實系統互動的情況下進行規劃、學習和生成影片。由 Danijar Hafner 開發的 Dreamer 系列一直是基於模型的強化學習中最受關注的研究方向之一,Dreamer 4 則將此方法延伸至大規模影片世界模型。該專案旨在忠實重現 Dreamer 4 的研究方法,同時將完整的訓練管線——包括穩定性修復和運算配置——公開釋出,這是大規模世界模型發布中通常省略的細節層級。

已發布的成品

兩個程式碼庫已發布。第一個是 next-state/open-dreamer,包含完整的訓練管線:因果影片分詞器、動作條件潛在動態模型、展開生成工具,以及 FVD(Fréchet 影片距離)評分。第二個是 reactor-team/open-dreamer,提供一個最小化的本地展開工具,能夠從輸入 MP4 檔案搭配對應的動作檔案生成影片幀。

第三個成品是一個在 Reactor 執行環境上運作的瀏覽器展示。它即時串流生成的 Minecraft 環境,並包含一個「遊戲 ⟷ 夢境」切換鈕,可在實際遊戲與世界模型輸出之間逐幀轉換影片串流。

團隊表示,他們的目標是刻意重現 Dreamer 4 的研究,避免使用原始論文以外的技術,以維持窄化的搜尋空間。開發工作始於 CoinRun,這是一個可程式生成的 2D 平台遊戲,可在單一 GPU 上訓練,之後再將可運作的管線擴展至 Minecraft/VPT 風格的遊戲影片。自從 OpenAI 的 Video PreTraining(VPT)專案以來,Minecraft 一直作為從遊戲示範中學習代理人的基準領域,該專案在承包商標註的 YouTube 影片上訓練逆動態模型,以產生大規模動作標註的遊戲資料。

更多詳情可參閱專案部落格、arXiv 論文,以及 Reactor 在 X 上的公告

架構:一個骨幹,兩個模型

分詞器和動態模型都共用同一個區塊因果 Transformer 骨幹,該骨幹在兩種注意力類型之間交替。空間層在單一幀的元素之間傳播資訊,而因果時間層則在幀之間傳播資訊。

分詞器被設計為基於 Transformer 的遮罩自編碼器(MAE),而非傳統的變分自編碼器(VAE)。團隊報告約 100 倍的壓縮率,並指出此設計無需 KL 散度或對抗損失。他們認為遮罩方法使潛在空間更適合基於擴散的生成。

動態模型使用擴散強制、流匹配和捷徑模型進行下一幀預測。它還預測下一個動作。展開不是在獨立的轉移模組和策略模組之間交替,而是整合到每個時間步的區塊中,結構為(前一動作、狀態、策略)。空間注意力在每個區塊內運作,而因果時間注意力跨時間連接各區塊。

一個關鍵的設計約束:世界模型 Token 無法讀取代理人 Token。因此,任務和策略資訊只能透過預測的下一個動作影響未來狀態。

訓練配置

隨附的 Minecraft 配置檔案詳細說明了訓練配方。

動態模型包含 16 億個參數,分布在 30 個區塊因果層中,d_model 為 1920,30 個注意力頭,使用群組查詢注意力的 3 個 KV 頭。每隔四層作為一個時間注意力層。每個時間步攜帶 32 個學習式暫存器 Token,封裝因子為 2,將相鄰的分詞器潛在表示整合到每個動態空間 Token 中。時間注意力在 192 步的滑動視窗上運作。

訓練涵蓋 200,000 步,使用 Muon 優化器搭配 WSD(預熱-穩定-衰減)排程,峰值學習率為 3e-4。捷徑和自助抽樣在第 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 個類別滑鼠類別,無連續通道。

運算效能與記憶體策略

團隊報告 57–58% 的模型 FLOPs 使用率(MFU),相較於健康 Transformer 訓練常被引用的 60% 基準。他們的分析使用屋頂線論證:在 NVIDIA B200 上,頻寬限制與運算限制的交叉點出現在 292 FLOP/byte。每 GPU 處理 256 幀使工作負載超越該脊線進入運算限制區域。

分片決策與最初預期相反。在 16 億參數下,完整的模型狀態——參數、梯度、優化器狀態和 EMA——佔用約 24 GiB,可放入單一 B200。證實激活值而非模型狀態才是主要的記憶體成本。團隊嘗試了資料平行、全分片資料平行(FSDP)、張量平行和序列平行,最終選擇了簡單的資料平行結合激活檢查點。

在資料載入方面,團隊將整個資料集預先分詞為 .arrayrecord 檔案,並使用 Grain 搭配 GPU 端預取緩衝區。標準的基於 ffmpeg 的解碼不足以讓 GPU 維持滿載。

穩定性工程

研究團隊明確表示,穩定性問題消耗了他們最大的開發時間。他們的核心觀察:大多數穩定性問題在損失持續下降的同時出現。均方誤差(MSE)平穩改善,而生成品質卻同時惡化——這種現象使傳統的基於損失的監控對基於擴散的世界模型不可靠。

六個具體修復已記錄如下:

  1. 優化器變更:Muon 取代了 LaProp,後者在兩次各約 400 B200 小時的獨立訓練中出現隨機且頻率漸增的尖峰。

  2. EMA 為必要項:EMA 權重被視為擴散推論的必要條件,而非可選項。

  3. 混合精度邊界:參數維持 float32,BF16 涵蓋大多數矩陣乘法激活和注意力輸入,float32 則保留用於正規化層和動態流輸出頭。

  4. 損失加權:團隊使用 x 預測搭配 v 空間損失,這簡化為類似 Dreamer 4 公式但分母為平方的加權項。他們報告了小幅但可察覺的改善。

  5. 最佳傳輸:在雜訊和潛在序列之間應用小批量重心最佳傳輸改善了展開生成的穩定性。

  6. μ 參數化:經過測試但被認為不必要,部分原因是 Muon 在不同模型規模間更有效地維持超參數穩定性。

CoinRun 開發階段的另一個發現:一項等 FLOPs 掃描估計運算最佳擴展約為 N ∝ C^0.56 且 D ∝ C^0.44。

未包含的內容

程式庫不包含行為克隆(BC)或強化學習(RL)訓練迴圈。完整的 Dreamer 4 BC/RL 代理人迴圈被列為開放的路線圖項目。專案文件中描述的 CoinRun 策略工作並未用於 Minecraft 實作,也未發布。

該專案未發布 FVD 分數,但程式庫附帶 scripts/eval_fvd.py,這是一個基於 I3D 的評估工具,配置為 4 個上下文幀和 240 幀的範圍。

關鍵工程數據摘要

  • 動態模型:16 億參數,30 層,d_model 1920,以 Muon 訓練 200,000 步
  • 硬體效率:NVIDIA B200 GPU 上 57–58% MFU,每 GPU 256 幀,約 24 GiB 模型狀態
  • 主要挑戰:穩定性,而非吞吐量——損失曲線隱藏了大多數生成品質的回退