next-state/open-dreamer

Open Dreamer – JAX/Flax によるリアルタイム Dreamer 4 世界モデル

何であるか

• JAX/Flax（NNX ライブラリ）で書かれた Dreamer 4 世界モデルパイプラインのオープンソース実装。
• 動画トークナイザー、アクション条件付き潜在動力学モデルの完全な トレーニングスタック と評価ツール（ロールアウト、Fréchet Video Distance）を提供。
• モデルは Minecraft/VPTスタイルのプレイ動画 で訓練され、クラウドホスティングされたデモによりリアルタイムで実行可能。

主要コンポーネント

コンポーネント	役割
dreamer/models.py	因果的動画トークナイザーと潜在動力学モデルの定義。
dreamer/training.py	トレーニングループ、損失関数、チェックポイント処理（Orbax経由）。
dreamer/generation.py	ノイズ除去スケジュール、ロールアウト生成、動画レンダリングのユーティリティ。
dreamer/fvd/	評価のための特徴抽出と Fréchet Video Distance 計算。
scripts/	実行可能なエントリポイント：train_tokenizer.py、tokenize_minecraft_dataset.py、train_dynamics.py、eval_fvd.py。
configs/	データセット、モデルハイパーパラメータ、評価設定を記述する Hydra/OmegaConf YAML ファイル。
site/	インタラクティブデモとドキュメントをホストする Next.js ウェブサイト。

一般的なワークフロー

1. データ準備 – ローカル Minecraft MP4録画を ArrayRecord シャード（ビデオバイト、アクションなどを持つピクル化された辞書）に変換。
2. トークナイザーのトレーニング – 動画フレームの潜在表現を学習（scripts/train_tokenizer.py）。
3. データセットのトークナイズ – すべての原始エピソードを潜在トークンにエンコード（scripts/tokenize_minecraft_dataset.py）。
4. 動力学のトレーニング – トークン化されたデータ上でアクション条件付き潜在動力学モデルを適合（scripts/train_dynamics.py）。
5. 生成と評価 – モデルをロールアウトして新しい動画フレームを合成し、FVD を計算（scripts/eval_fvd.py）。

試してみる方法

• ライブデモ – インストール不要。ウェブサイトがブラウザ内デモをホストし、Minecraft世界をストリーミングし、本物のゲームとモデルの予測の切り替えが可能。
• ローカル推論 – コンパニオンリポジトリ reactor-team/open-dreamer には、トレーニング済みチェックポイントを自前の動画/アクションシーケンスで実行する軽量推論スクリプトが含まれている。

インストール（トレーニング側）

# Python 3.11 と CUDA-12 対応の JAX ビルドが必要
pip install uv               # ここでは高速な Python パッケージマネージャーとして UV を使用
uv sync                       # ピン止めされた依存関係を仮想環境にインストール
source .venv/bin/activate     # 環境をアクティベート

異なる JAX ウィール（例：異なる GPU や CPU）が必要な場合は、uv sync ステップ後にインストールしてください。

トレーニングスクリプトの実行（トークナイザーの例）

# まず configs/tokenizer.yaml と configs/dataset/minecraft_vpt.yaml を編集
uv run scripts/train_tokenizer.py

トークナイズ、動力学トレーニング、評価に対応する同様のコマンドが存在し、すべて configs/ 下の YAML 設定ファイルによって駆動される。

ロードマップ

• 現在は世界モデルトレーニングパイプラインのみをサポート。
• 将来的には、完全な Dreamer 4 行動クローン / 強化学習ループの追加を計画。

ライセンスと引用

• リポジトリは現在「All rights reserved」というプレースホルダーメッセージを保持しているが、後続リリースで正式なオープンソースライセンスを計画。
• 研究でコードを使用する場合は、READMEに記載の Zenodo エントリと元の Dreamer 4 論文を引用すること。

誰にとって役立つか？

• Minecraft などのビデオ豊富な環境でスケーラブルな世界モデル学習を探索する研究者。
• 高速な JAX 実装が必要なリアルタイム生成エージェントを開発するエンジニア。
• 動画トークナイザー＋潜在動力学パイプラインの具体的な例を探している学生。

---

上記のすべての情報はリポジトリの README から直接取得；追加機能は推測されていない。