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 分片（包含视频字节、动作等的 pickle 字典）。
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               # UV 是此处使用的快速 Python 包管理器
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 行为克隆 / 强化学习循环。

许可与引用

• 仓库目前带有“保留所有权利”的占位符通知；正式开源许可证计划在后续版本中发布。
• 若在研究中使用代码，请引用 README 中提供的 Zenodo 条目和原始 Dreamer 4 论文。

谁可能用得上？

• 探索可扩展世界模型学习的研究人员，尤其是在 Minecraft 等视频丰富的环境中。
• 需要快速 JAX 实现的实时生成代理的工程师。
• 寻找完整视频分词器 + 潜在动力学管道具体示例的学生。

---

以上所有信息均直接取自仓库的 README；未推断任何额外功能。