next-state/open-dreamer

Open Dreamer 是一个简单、高性能且易于使用的 JAX/Flax NNX Dreamer 4 世界模型实现——在 Minecraft/VPT 风格的游戏数据上进行训练，并支持实时游玩。本仓库包含**训练 pipeline**：一个因果视频 tokenizer、一个基于动作的 latent 动力学模型，以及用于生成 rollout 和评估质量的工具。 本仓库目前支持： - 训练因果视频 tokenizer - 对 Minecraft/VPT 风格的 MP4 数据集进行 tokenization - 训练基于动作的 latent 动力学模型 - 生成 rollout 并计算 FVD ## 🎮 立即体验 体验 Open Dreamer 最快的方式是使用**浏览器内的实时演示**——进入生成的 Minecraft 世界并进行实时游玩，通过“游戏 ⟷ Dream”切换功能，将真实游戏的流逐帧交给世界模型处理。无需任何环境配置；实时模型运行在 [Reactor](https://reactor.inc) 云端。 👉 **[打开演示](https://next-state.github.io/open-dreamer/)** [网站](https://next-state.github.io/open-dreamer/) 还详细介绍了该模型的工作原理，并提供了交互式图表和结果展示。 ## ⚡ 自行运行 要在本地使用您自己的视频和动作部署训练好的模型——即从 MP4 和匹配的动作序列中生成新帧——请使用推理仓库： 👉 **[reactor-team/open-dreamer](https://github.com/reactor-team/open-dreamer)** 本 README 的其余部分涵盖了从零开始**训练**模型的过程。 ## 📋 路线图 - [ ] 完整的 Dreamer 4 行为克隆 / RL agent 训练循环 ## ⚙️ 环境要求 - Python 3.11 - [`uv`](https://docs.astral.sh/uv/) - 兼容 CUDA 12 的 JAX 环境 - Minecraft/VPT 风格的 ArrayRecord 数据；请参见 [dreamer/data/README.md](dreamer/data/README.md) ## 安装 ``` pip install uv uv sync source .venv/bin/activate ``` 依赖锁定目标为 CUDA 12 JAX。如果您的机器需要不同的 JAX 构建，请在同步后为您的加速器环境安装正确的 wheel。 ## 工作流程 1. 准备原始的 MP4 ArrayRecord shard。 2. 在原始视频片段上训练 tokenizer。 3. 将完整 episode tokenization 为 latent ArrayRecords。 4. 将生成的 latent 统计信息复制到 latent 数据集配置中。 5. 在 latent episode 和动作上训练动力学模型。 6. 生成视频并计算 FVD。 以下命令假设使用固定 256 帧的原始记录。诸如 `index_max`、`short_T`、`long_T`、`horizon` 和 `fvd_chunk_size` 等示例值在 [dreamer/data/README.md](dreamer/data/README.md) 中有详细说明。 ## 仓库结构 ``` . ├── dreamer/ # Core models, training helpers, data, sampling, FVD │ ├── data/ # Grain/ArrayRecord pipelines and serialization │ ├── fvd/ # FVD feature extraction and scoring │ ├── models.py # Tokenizer and dynamics model definitions │ ├── training.py # Training and evaluation helpers │ ├── generation.py # Denoising schedules and rollout utilities │ └── checkpointing.py # Orbax checkpoint bundles ├── scripts/ │ ├── train_tokenizer.py │ ├── tokenize_minecraft_dataset.py │ ├── train_dynamics.py │ └── eval_fvd.py ├── configs/ │ ├── dataset/ # Raw-video and latent dataset configs │ ├── tokenizer.yaml │ ├── tokenize.yaml │ ├── dynamics.yaml │ └── eval_fvd.yaml └── site/ # Website, blog post, and the live demo (Next.js) ``` ## 数据集 主要工作流程要求使用原始的 Minecraft/VPT 风格的 shard，命名为 `shard-*.array_record`。每条原始记录都是一个序列化的 Python 字典，包含 MP4 字节流、视频形状、动作以及可选的来源元数据。Tokenized 的记录将被写入为包含 tokenizer latent 和动作的 msgpack ArrayRecords。 原始视频的设置位于 [configs/dataset/minecraft_vpt.yaml](configs/dataset/minecraft_vpt.yaml)。请至少更新以下内容： ``` array_record_path: /path/to/mp4-arrayrecords index_max: 1500 dataset_mean: [0.2241, 0.2348, 0.2086] dataset_std: [0.1809, 0.1874, 0.2282] ``` `dataset_mean` 和 `dataset_std` 是用于视频值归一化的像素统计信息。当更改原始数据集时，请重新计算这些值。 ## 训练 tokenizer Tokenizer 用于学习动力学模型所使用的 latent 表示。默认配置文件为 [configs/tokenizer.yaml](configs/tokenizer.yaml)。在运行之前，请先编辑该文件以及 [configs/dataset/minecraft_vpt.yaml](configs/dataset/minecraft_vpt.yaml)： ``` uv run scripts/train_tokenizer.py ``` 输出默认存放到 `logs//`： - `checkpoints/` 包含 tokenizer 和优化器的 checkpoint。 - `vis/` 在 `visualize_every > 0` 时包含重建图像。 训练 tokenizer 时可以使用比完整原始 episode 长度更短的窗口。例如，固定 256 帧的记录可以使用 16 帧的窗口进行训练，随后再将其作为完整的 256 帧 latent episode 进行 tokenization。 ## 对数据集进行 Tokenization 在 tokenizer 训练完成后，将每个 MP4 episode 编码为 tokenizer latent。请先编辑 [configs/tokenize.yaml](configs/tokenize.yaml)： ``` uv run scripts/tokenize_minecraft_dataset.py ``` 这将输出 latent shard 和统计信息： ``` /path/to/tokenized_data/ shard-00000.array_record shard-00001.array_record metadata/latent_stats.npz ``` 打印 latent 统计信息： ``` python - <<'PY' import numpy as np stats = np.load("/path/to/tokenized_data/metadata/latent_stats.npz") print("latent_mean:", stats["mean"].tolist()) print("latent_std:", stats["std"].tolist()) print("num_samples:", int(stats["num_samples"])) print("num_videos:", int(stats["num_videos"])) PY ``` 将 `latent_mean` 和 `latent_std` 复制到 [configs/dataset/minecraft_vpt_latent.yaml](configs/dataset/minecraft_vpt_latent.yaml) 中， 然后将 `array_record_path` 和 `index_max` 指向 tokenization 后的输出结果。 ## 训练动力学模型 动力学模型基于 latent ArrayRecords 和经过偏移处理的动作进行训练。默认配置文件为 [configs/dynamics.yaml](configs/dynamics.yaml)，它会导入 [configs/dataset/minecraft_vpt_latent.yaml](configs/dataset/minecraft_vpt_latent.yaml)。在运行前请编辑这两个文件： ``` uv run scripts/train_dynamics.py ``` ## 评估 基于动力学 checkpoint 生成视频并计算 FVD。请先编辑 [configs/eval_fvd.yaml](configs/eval_fvd.yaml)： ``` uv run scripts/eval_fvd.py ``` 使用 `mode=generate` 可仅保存 MP4，使用 `mode=evaluate` 可基于先前生成的视频计算 FVD。 ## 配置说明 配置文件是位于 `configs/` 目录下的 Hydra/OmegaConf YAML 文件。主脚本配置中对每个字段的用途和限制都进行了注释说明。 常用的配置文件： - [configs/tokenizer.yaml](configs/tokenizer.yaml) - tokenizer 架构与训练。 - [configs/tokenize.yaml](configs/tokenize.yaml) - 离线 tokenization。 - [configs/dynamics.yaml](configs/dynamics.yaml) - 动力学架构与训练。 - [configs/eval_fvd.yaml](configs/eval_fvd.yaml) - rollout 和 FVD 设置。 - [configs/dataset/minecraft_vpt.yaml](configs/dataset/minecraft_vpt.yaml) - 原始 MP4 数据集设置。 - [configs/dataset/minecraft_vpt_latent.yaml](configs/dataset/minecraft_vpt_latent.yaml) - tokenization 后的 latent 数据集设置。 ## 📚 参考文献 - Dreamer 4: [Training Agents Inside of Scalable World Models](https://danijar.com/project/dreamer4/) - Jasmine: [A simple, performant and scalable JAX-based world modeling codebase](https://github.com/p-doom/jasmine) ## 📄 许可证 **保留所有权利。** 请参阅 [LICENSE](LICENSE)。这是一份临时声明；正式的许可证预计将在未来的版本中发布。 ## 📄 引用 如果您的研究中使用了 Open Dreamer，请引用： ``` @misc{marti2026opendreamer, title={How to Train a Frontier-level World Model}, author={Marti Monso, Diego and Sacco, Francesco and Hu, Edward}, month={jul}, year={2026}, publisher={Zenodo}, doi={10.5281/zenodo.21475232}, url={https://next-state.github.io/open-dreamer/}, } ``` 同时请引用原版的 Dreamer 4 论文： ``` @misc{hafner2025trainingagentsinsidescalable, title={Training Agents Inside of Scalable World Models}, author={Danijar Hafner and Wilson Yan and Timothy Lillicrap}, year={2025}, eprint={2509.24527}, archivePrefix={arXiv}, primaryClass={cs.AI}, url={https://arxiv.org/abs/2509.24527}, } ```

标签：JAX, 世界模型, 人工智能, 强化学习, 游戏AI, 用户模式Hook绕过, 视频生成, 逆向工具