Open Dreamer の紹介：Dreamer 4 ワールドモデルパイプラインの JAX/Flax 再現、完全なトレーニングレシピを公開

AI 研究者の小グループ（Reactor）が Open Dreamer をリリースしました。これは、Dreamer 4 ワールドモデルパイプラインのオープン実装で、JAX と Flax NNX で書かれています。

実際にリリースされたもの

2 つのリポジトリがリリースされました。next-state/open-dreamer にはトレーニングパイプライン（因果ビデオトークナイザー、行動条件付き潜在ダイナミクスモデル、ロールアウト生成、FVD スコアリング）が含まれています。reactor-team/open-dreamer には、MP4 と対応するアクションファイルからフレームを生成する最小限のローカルロールアウトハーネスが含まれています。

3 つ目の成果物は、Reactor ランタイム上でホストされているブラウザデモです。生成された Minecraft ワールドをリアルタイムでストリーミングし、Game ⟷ Dream トグルを公開して、実際のゲームからワールドモデルにフレーム単位でストリームを渡します。

明確な目的は Dreamer 4 研究を再現することでした。研究チームは、検索空間を狭く保つために、その研究論文以外の方法を意図的に避けました。彼らは CoinRun（単一 GPU でトレーニング可能な手続き型生成 2D プラットフォーマー）から始め、その後、動作パイプラインを Minecraft/VPT スタイルのゲームプレイビデオにスケールアップしました。

アーキテクチャ：1 つのバックボーン、2 つのモデル

トークナイザーとダイナミクスモデルは両方とも同じブロック因果トランスフォーマーバックボーンを使用しています。このバックボーンは 2 種類のアテンションを交互に使用します。空間レイヤーはフレーム内の要素間で情報を伝播します。因果時間レイヤーはフレーム間で情報を伝播します。

トークナイザーは VAE ではなく、トランスフォーマーベースの Masked Autoencoder です。チームは約 100 倍の圧縮率を報告し、設計に KL ロスや敵対的ロスが必要ないと述べています。マスキングにより、潜在空間がより拡散しやすくなると主張しています。

ダイナミクスモデルは次フレーム予測を実行し、拡散フォーシング、フローマッチング、ショートカットモデルでトレーニングされます。また、次の行動も予測します。ロールアウトは、別個の遷移モジュールとポリシーを交互に使用する代わりに、タイムステップごとのブロック（前の行動、状態、ポリシー）に折りたたまれています。空間アテンションは各ブロック内で実行されます。因果時間アテンションはブロックを時間的に接続します。

重要なのは、ワールドモデルトークンはエージェントトークンを読み取れないことです。したがって、タスクとポリシー情報は次の行動を通じてのみ将来の状態に影響を与えることができます。

設定されたトレーニングレシピ

リリースされた Minecraft 設定により、レシピが具体化されています。

ダイナミクスモデルは 1.6B パラメータ：30 ブロック因果層、d_model 1920、30 アテンションヘッド、グループ化クエリアテンション用の 3 KV ヘッド。4 層ごとに時間アテンション層があります。各タイムステップは 32 の学習済みレジスタトークンを運び、packing_factor: 2 は隣接するトークナイザー潜在変数を各ダイナミクス空間トークンにパックします。時間アテンションは 192 ステップのスライディングウィンドウを使用します。

トレーニングは Muon、WSD スケジュール、ピーク学習率 3e-4 で 200,000 ステップ実行されます。Shortcut/bootstrap サンプルはステップ 100,000 で 0.25 のバッチ分数でオンになります。EMA 減衰は 0.999 です。

トークナイザー設定は、ボトルネック幅 16 でフレームあたり 512 の潜在トークンを出力します。生の 360×640 フレームは 368×640 にパディングされ、両方の次元が 16×16 パッチで割り切れるようになります。エンコーダ深度は d_model 1536 で 12、デコーダ深度は d_model 1024 で 8。MAE マスキング確率は最大 0.9、LPIPS は半分のタイムステップで重み 0.2 で適用されます。

VPT アクションは、連続チャネルなしで、27 のバイナリアクションチャネルと 121 の分類マウスクラスに解析されます。

Computemaxxing とメモリウォール

研究チームは、モデル FLOPs 利用率が 57-58% であると報告しています。これは、健全なトランスフォーマートレーニングのベンチマークである 60% に対してのものです。理由はルーフラインの議論です。B200 では、帯域幅制限と計算制限の間のクロスオーバーは 292 FLOP/バイトにあります。GPU あたり 256 フレームを供給すると、ワークロードがそのピークポイントを超えます。

シャーディングは期待とは逆の方向に進みました。1.6B パラメータのモデル状態（パラメータ、勾配、オプティマイザ状態、EMA）は約 24 GiB を占め、B200 に収まります。アクティベーションが本当のコストでした。研究チームはデータ並列、FSDP、テンソル並列、シーケンス並列を試しましたが、最終的にプレーンデータ並列とアクティベーションチェックポイントに落ち着きました。

データローディングは、データセット全体を .arrayrecord ファイルに事前トークナイズし、Grain と GPU 側のプリフェッチバッファを使用することで解決されました。ffmpeg デコードは GPU を供給するのに十分な速度ではありませんでした。

安定性のセクションが本当のペイロード

研究チームは、安定性が最大の時間を費やしたことを明確に述べています。彼らの重要な観察：ほとんどの安定性問題は損失が下がっているにもかかわらず発生します。MSE は滑らかに改善する一方で、生成品質は低下します。

6 つの修正が文書化されています。Muon が LaProp に取って代わりました。LaProp はランダムに、そしてますます頻繁にスパイクし、約 400 B200 時間の 2 回の実行にわたって。EMA ウェイトは拡散推論に必須として扱われます。混合精度は境界に敏感です：パラメータは float32 のまま、BF16 はほとんどの matmul アクティベーションとアテンション入力をカバーし、float32 は正規化とダイナミクスフロー出力ヘッドに保持されます。

損失重み付けに関しては、x-prediction と v-space 損失を使用しています。これは Dreamer 4 のものと似た重み付け項に還元されますが、分母が二乗されています。小さく顕著な改善を報告しています。ノイズと潜在シーケンス間のミニバッチ重心最適輸送により、ロールアウト生成がより安定しました。μ-パラメータ化はテストされ、不要と判断されました。これは、Muon がモデルサイズ全体でハイパーパラメータをより安定に保つためです。

CoinRun フェーズからのさらなる結果：iso-FLOPs スイープは、計算最適スケーリングをおよそ N∝C^0.56 および D∝C^0.44 としました。

含まれていないもの

リポジトリには行動クローニングや RL トレーニングループは含まれていません。完全な Dreamer 4 BC/RL エージェントループはオープンなロードマップ項目としてリストされています。記事で説明されている CoinRun ポリシーの作業は Minecraft には使用されず、リリースされませんでした。

また、記事は FVD スコアを公開していませんが、scripts/eval_fvd.py には、4 つのコンテキストフレームと 240 フレームの地平線で構成された I3D ベースのハーネスが含まれています。

主なポイント

Open Dreamer は JAX/Flax NNX で Dreamer 4 パイプラインを再現し、トレーニングコードと Minecraft デモを提供。 ダイナミクスモデルは 1.6B パラメータ、30 層、d_model 1920、Muon で 200K ステップ学習。 報告されたエンジニアリング数値：B200 で 57-58% MFU、GPU あたり 256 フレーム、モデル状態約 24 GiB。 安定性がボトルネックであり、損失曲線は生成品質の低下を隠していた。

ブログ記事とデモ、トレーニングリポジトリ、推論リポジトリ、および Reactor の X をチェックしてください。この研究のすべての功績は、このプロジェクトの研究者に帰属します。