Open Dreamer: JAX/FlaxによるDreamer 4パイプラインのオープンソース実装、完全なトレーニングレシピ付き

Reactorがスポンサーを務める研究者フランチェスコ・サッコ、ディエゴ・マルティ、エドワード・フーは、Dreamer 4 のオープン実装であるOpen DreamerをJAX/Flax NNXで公開しました。学習用（トークナイザー、ダイナミクスモデル、ロールアウト、FVDスコアリング）とローカル推論用の2つのリポジトリが公開されています。ブラウザデモでは、生成されたMinecraftの世界がストリーミング表示され、実際のゲームとワールドモデルを切り替えることができます。アーキテクチャは、トークナイザーとダイナミクスモデルの両方にブロック因果トランスフォーマーバックボーンを使用しています。トークナイザーはトランスフォーマーベースのマスクオートエンコーダーで、KL損失や敵対的損失なしに約100倍の圧縮を実現しています。ダイナミクスモデル（16億パラメータ）は、拡散フォーシング、フローマッチング、ショートカットモデルを用いて次のフレームを予測し、次のアクションも予測します。学習は20万ステップで、Muonオプティマイザー、WSDスケジュール、ピーク学習率3e-4、EMA減衰0.999を使用します。チームはB200 GPU上で57-58%のモデルFLOPs効率を達成し、モデル状態は約24 GiBに収まりました。アクティベーションが主なメモリコストであり、データ並列処理とアクティベーションチェックポインティングで解決しました。安定性の問題が最大の課題であり、6つの修正（Muonオプティマイザー、必須のEMA重み、注意深い混合精度、v-space損失重み付け、ミニバッチ重心最適輸送、μパラメータ化の軽視）が施されています。リポジトリには行動クローニングや強化学習のループは含まれておらず、FVDスコアも公開されていません。