认识 Open Dreamer：基于 JAX/Flax 复现 Dreamer 4 世界模型管道，完整训练配方已发布

一群 AI 研究人员（Reactor）发布了 Open Dreamer，这是 Dreamer 4 世界模型管道的开源实现，使用 JAX 和 Flax NNX 编写。

实际发布的内容

发布了两个仓库。next-state/open-dreamer 包含训练管道：因果视频分词器、动作条件的潜动态模型、 rollout 生成和 FVD 评分。reactor-team/open-dreamer 包含一个最小的本地 rollout 工具，可从 MP4 和匹配的动作文件生成帧。

第三个成果是托管在 Reactor 运行时上的浏览器演示。它实时流式传输生成的 Minecraft 世界，并公开一个 Game ⟷ Dream 切换开关，可将流从真实游戏逐帧切换到世界模型。

明确的目标是复现 Dreamer 4 研究。研究团队特意避免使用该研究论文之外的方法，以保持搜索空间狭窄。他们从 CoinRun 开始，这是一个可以在单个 GPU 上训练的程序化生成的 2D 平台游戏，然后将工作管道扩展到 Minecraft/VPT 风格的游戏视频。

架构：一个骨干，两个模型

分词器和动态模型都使用相同的块因果变换器骨干。该骨干交替使用两种注意力类型。空间层在同一帧的元素之间传播信息。因果时间层在帧之间传播信息。

分词器是基于变换器的掩码自编码器，而不是 VAE。团队报告大约 100 倍的压缩率，并指出该设计不需要 KL 或对抗性损失。他们认为掩码使潜在空间更易于扩散。

动态模型执行下一帧预测，并使用扩散强制、流匹配和 shortcut 模型进行训练。它还预测下一个动作。rollout 被折叠到每个时间步的块（前一个动作、状态、策略）中，而不是在单独的转换模块和策略之间交替。空间注意力在每个块内部运行；因果时间注意力连接跨时间的块。

关键地，世界模型 token 不能读取智能体 token。因此，任务和策略信息只能通过下一个动作影响未来状态。

训练配方配置

发布的 Minecraft 配置使配方具体化。

动态模型为 1.6B 参数：30 个块因果层，d_model 1920，30 个注意力头，3 个 KV 头用于分组查询注意力。每四层中有一层是时间注意力层。每个时间步携带 32 个学习到的寄存器 token，packing_factor: 2 将相邻的分词器潜在变量打包到每个动态空间 token 中。时间注意力使用 192 步滑动窗口。

训练运行 200,000 步，使用 Muon、WSD 调度和峰值学习率 3e-4。Shortcut/bootstrap 样本在第 100,000 步以 0.25 的批次分数开启。EMA 衰减为 0.999。

分词器配置每帧输出 512 个潜在 token，瓶颈宽度为 16。原始 360×640 帧被填充到 368×640，以便两个维度都能被 16×16 块整除。编码器深度为 12，d_model 1536；解码器深度为 8，d_model 1024。MAE 掩码概率最高为 0.9，LPIPS 以权重 0.2 应用在半个时间步上。

VPT 动作被解析为 27 个二进制动作通道加上 121 个分类鼠标类，没有连续通道。

计算最大化与内存墙

研究团队报告模型 FLOPs 利用率为 57-58%，而健康变换器训练的基准为 60%。原因是 roofline 论点。在 B200 上，带宽受限和计算受限之间的交叉点为 292 FLOP/字节。每 GPU 提供 256 帧将工作负载推过该峰值点。

分片方式与预期相反。1.6B 参数的模型状态——参数、梯度、优化器状态和 EMA——大约占用 24 GiB，适合 B200。激活才是真正的成本。研究团队尝试了数据并行、FSDP、张量并行和序列并行，最终选择了普通数据并行加激活检查点。

数据加载通过将整个数据集预分词为 .arrayrecord 文件，然后使用 Grain 和 GPU 端预取缓冲区解决。ffmpeg 解码速度不足以保持 GPU 满载。

稳定性部分才是真正的重点

研究团队明确指出稳定性消耗了他们大部分时间。他们的关键观察：大多数稳定性问题发生在损失下降的情况下。MSE 平滑改善而生成质量下降。

记录了六种修复方法。Muon 取代了 LaProp，后者随机且频繁地尖峰，在两次约 400 B200 小时的运行中。EMA 权重被视为扩散推断的必要条件。混合精度对边界敏感：参数保持 float32，BF16 覆盖大多数 matmul 激活和注意力输入，而 float32 保留用于归一化和动态流输出头。

在损失加权方面，他们使用 x-prediction 和 v-space 损失，这减少到与 Dreamer 4 类似的加权项，但分母是平方的。他们报告了微小但显著的改进。噪声和潜在序列之间的 Minibatch barycentric optimal transport 使 rollout 生成更加稳定。μ-参数化经过测试并被判断为不必要，部分原因是 Muon 在不同模型大小上保持超参数更稳定。

CoinRun 阶段的一个进一步结果：iso-FLOPs 扫描将计算最优缩放置于大约 N∝C^0.56 和 D∝C^0.44。

不包含的内容

仓库不包括行为克隆或 RL 训练循环；完整的 Dreamer 4 BC/RL 智能体循环被列为开放路线图项目。帖子中描述的 CoinRun 策略工作未用于 Minecraft 且未发布。

该帖子也未发布 FVD 分数，尽管 scripts/eval_fvd.py 附带了一个基于 I3D 的工具，配置了 4 个上下文帧和 240 帧的视野。

关键要点

Open Dreamer 在 JAX/Flax NNX 中复现了 Dreamer 4 管道，包含训练代码和 Minecraft 演示。 动态模型为 1.6B 参数，30 层，d_model 1920，使用 Muon 训练 200K 步。 报告的工程数据：B200 上 MFU 为 57-58%，每 GPU 256 帧，模型状态约 24 GiB。 稳定性是瓶颈，损失曲线掩盖了大部分生成质量退化的现象。

查看博客文章和演示、训练仓库、推理仓库以及 Reactor on X。所有研究归功于该项目的研究人员。