Open Dreamer：Dreamer 4 管道的 JAX/Flax 开源实现及完整训练配方

研究人员 Francesco Sacco、Diego Martí 和 Edward Hu 在 Reactor 的资助下，发布了 Open Dreamer，这是 Dreamer 4 在 JAX/Flax NNX 中的开源实现。项目发布了两个代码库：一个用于训练（包括分词器、动态模型、rollout 和 FVD 评分），另一个用于本地推理。浏览器演示可流式传输生成的 Minecraft 世界，并支持在真实游戏与世界模型之间切换。其架构对分词器和动态模型均采用块因果 Transformer 主干网络。分词器是基于 Transformer 的掩码自编码器，在不使用 KL 散度或对抗损失的情况下实现了约 100 倍的压缩。动态模型（参数量为 1.6B）通过扩散强制、流匹配和快捷模型进行下一帧预测，同时还能预测后续动作。训练过程使用 Muon 优化器、WSD 调度策略、峰值学习率 3e-4 和指数移动平均（EMA）衰减系数 0.999，共进行 20 万步。团队在 B200 GPU 上实现了 57-58% 的模型 FLOPs 利用率，模型状态仅占用约 24 GiB 内存。激活值曾是主要的内存开销，通过数据并行和激活检查点解决了这一问题。稳定性问题是最大的挑战，团队采取了六项修复措施：使用 Muon 优化器、强制使用 EMA 权重、谨慎的混合精度训练、v 空间损失加权、小批量重心最优传输，以及弱化 μ 参数化。该代码库未包含行为克隆或强化学习（RL）训练循环，FVD 评分也尚未公布。