Reactorという名で活動するAI研究者グループが、JAXおよびFlax NNXで構築されたDreamer 4ワールドモデルパイプラインのオープンソース実装であるOpen Dreamerを公開しました。ワールドモデルは観測データから環境がどのように変化するかを学習し、エージェントが実際のシステムと相互作用することなく計画・学習・動画生成を行えるようにするものです。Danijar Hafnerによって開発されたDreamerシリーズは、モデルベース強化学習において最も研究されている研究方向の一つであり、Dreamer 4は大規模動画ワールドモデルへとアプローチを拡張しています。本プロジェクトは、Dreamer 4の研究手法を忠実に再現すると同時に、安定化策やコンピュート設定を含む完全なトレーニングパイプラインを公開することを目的としており、このレベルの詳細は大規模ワールドモデルのリリースでは通常省略されてきました。

公開された成果物

2つのコードリポジトリが公開されました。1つ目のnext-state/open-dreamerには、完全なトレーニングパイプラインが含まれています。因果ビデオトークナイザー、アクション条件付き潜在ダイナミクスモデル、ロールアウト生成ユーティリティ、およびFVD（Fréchet Video Distance）スコアリングで構成されています。2つ目のreactor-team/open-dreamerは、入力MP4ファイルと対応するアクションファイルのペアから動画フレームを生成できる最小限のローカルロールアウトハーネスを提供します。

3つ目の成果物は、Reactorランタイム上で動作するブラウザベースのデモです。生成されたMinecraft環境をリアルタイムでストリーミングし、実際のゲームとワールドモデル出力の間で動画ストリームをフレーム単位で切り替えるGame ⟷ Dreamトグルを備えています。

チームは、検索空間を絞るために元の論文外の手法を意図的に避け、Dreamer 4の研究を再現することを目的としたと述べています。開発は、1台のGPUでトレーニング可能な手続き的生成2DプラットフォーマーであるCoinRunから始まり、その後、稼働するパイプラインをMinecraft/VPTスタイルのゲームプレイ動画へとスケールアップしました。Minecraftは、OpenAIのVideo PreTraining（VPT）プロジェクト以来、ゲームプレイ実演からエージェントを学習するためのベンチマークドメインとして使用されてきました。VPTプロジェクトでは、請負業者がラベル付けしたYouTube映像で逆ダイナミクスモデルをトレーニングし、大規模なアクションラベル付きゲームプレイデータを生成しました。

追加の詳細はプロジェクトブログ、arXiv論文、およびReactorのXでの発表で確認できます。

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

トークナイザーとダイナミクスモデルはどちらも同じブロック因果トランスフォーマーバックボーンを共有し、2種類のアテンションを交互に使用します。Space layersは単一フレーム内の要素間で情報を伝達し、因果時間layersはフレーム間で情報を伝達します。

トークナイザーは、従来の変分オートエンコーダー（VAE）ではなく、トランスフォーマーベースのMasked Autoencoder（MAE）として設計されています。チームは約100倍の圧縮率を報告しており、この設計によりKLダイバージェンスや敵対的損失が不要になると指摘しています。マスキング手法により、潜在空間が拡散ベースの生成により適しやすくなると主張しています。

ダイナミクスモデルは、拡散フォーシング、フローマッチング、およびショートカットモデルを使用して次フレーム予測を行います。また、次のアクションも予測します。別々の遷移モジュールとポリシモジュールを交互に使用するのではなく、ロールアウトは（前のアクション、状態、ポリシー）の構造を持つタイムステップごとのブロックに統合されています。空間アテンションは各ブロック内で動作し、因果時間アテンションがブロック間を時間軸で結びます。

重要な設計制約：ワールドモデルトークンはエージェントトークンを読み取ることができません。したがって、タスクおよびポリシー情報は、予測された次のアクションを通じてのみ将来の状態に影響を与えることができます。

トレーニング設定

公開されたMinecraft設定ファイルにより、トレーニングレシピが具体的に指定されています。

ダイナミクスモデルは、30層のブロック因果layersで構成される16億パラメータであり、d_model 1920、30のアテンションヘッド、grouped-query attentionを使用する3つのKVヘッドを備えています。4層ごとにtime-attention layerとして機能します。各タイムステップは32個の学習済みレジスタートークンを持ち、packing_factor: 2が隣接するトークナイザー潜在表現を各ダイナミクス空間トークンに統合します。Time attentionは192ステップのスライディングウィンドウで動作します。

トレーニングは200,000ステップにわたり、MuonオプティマイザーとWSD（warmup-stable-decay）スケジュール、ピーク学習率3e-4で実行されます。ショートカットおよびブートストラップサンプリングはステップ100,000で0.25のバッチフラクションで有効化されます。指数移動平均（EMA）の減衰率は0.999に設定されています。

トークナイザー設定は、ボトルネック幅16でフレームあたり512個の潜在トークンを生成します。生の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）を報告しています。これは、健全なトランスフォーマートレーニングのベンチマークとして一般的に引用される60%と比較した値です。彼らの分析はルーフライン議論を使用しています。NVIDIA B200では、帯域幅律速とコンピュート律速の境界は292 FLOP/byteにあります。GPUあたり256フレームを処理すると、ワークロードはその境界点を超えてコンピュート律速領域に入ります。

シャーディングの決定は当初の予想とは異なる結果となりました。16億パラメータの場合、完全なモデル状態（パラメータ、勾配、オプティマイザー状態、EMA）は約24 GiBを占有し、単一のB200に収まります。モデル状態ではなく、アクティベーションが主なメモリコストであることが判明しました。チームはデータ並列化、完全シャーディングデータ並列化（FSDP）、テンソル並列化、シーケンス並列化を試験し、最終的にプレーンなデータ並列化とアクティベーションチェックポイントの組み合わせを採用しました。

データローディングについては、チームはデータセット全体を.arrayrecordファイルに事前トークナイズし、GPU側プリフェッチバッファ付きのGrainを使用しました。標準的なffmpegベースのデコーディングではGPUを飽和させるのに十分な速度が得られませんでした。

安定性エンジニアリング

研究チームは、安定性の問題が開発時間の最大の割合を占めたと明確に述べています。彼らの中心的な観察は、ほとんどの安定性問題は損失が減少し続けているにもかかわらず発生するという点です。平均二乗誤差（MSE）は順調に改善される一方で、生成品質が同時に劣化するという現象があり、これにより拡散ベースのワールドモデルに対する従来の損失ベースの監視が信頼できなくなります。

6つの具体的な修正が文書化されています：

1. オプティマイザーの変更：MuonがLaPropに代わって使用されました。LaPropは約400 B200時間の2回の個別トレーニングランにわたり、ランダムかつ頻度を増すスパイクを示していました。

2. 必須としてのEMA：EMA重みは拡散推論においてオプションではなく不可欠として扱われます。

3. 混合精度の境界：パラメータはfloat32のままで、BF16がほとんどの行列積アクティベーションとアテンション入力をカバーし、float32は正規化layersとダイナミクスフロー出力ヘッドに保持されます。

4. 損失重み付け：チームはv-space損失付きのx-predictionを使用し、これはDreamer 4の定式化に類似するが二乗分母を持つ重み付け項に還元されます。わずかだが顕著な改善を報告しています。

5. 最適輸送：ノイズと潜在シーケンス間に適用されたミニバッチ重心最適輸送が、ロールアウト生成の安定性を向上させました。

6. μ-パラメータ化：テストされましたが不要と判断されました。これは部分的に、Muonがモデルサイズ全体にわたりハイパーパラメータの安定性をより効果的に維持するためです。

CoinRun開発フェーズの追加の発見：iso-FLOPsスイープにより、コンピュート最適スケーリングがおよそN ∝ C^0.56、D ∝ C^0.44と推定されました。

含まれていないもの

リポジトリには、行動クローン（BC）または強化学習（RL）のトレーニングループは含まれていません。完全なDreamer 4 BC/RLエージェントループはオープンロードマップ項目としてリストされています。プロジェクト文書で説明されているCoinRunポリシーの成果は、Minecraft実装には使用されず、公開もされていません。

プロジェクトはFVDスコアを公開していませんが、リポジトリには4コンテキストフレームと240フレームの地平線に設定されたI3Dベースの評価ハーネスであるscripts/eval_fvd.pyが同梱されています。

主要なエンジニアリング数値のまとめ

• ダイナミクスモデル：16億パラメータ、30層、d_model 1920、Muonで200,000ステップのトレーニング
• ハードウェア効率：NVIDIA B200 GPUで57〜58%のMFU、GPUあたり256フレーム、約24 GiBのモデル状態
• 主な課題：スループットではなく安定性 — 損失曲線が生成品質の劣化の大部分を隠蔽