Open Dreamer:Reactor发布基于JAX/Flax的Dreamer 4世界模型管道开源复现版,附带完整训练方案
要点速览
- •Open Dreamer提供了Dreamer 4管道的开源实现,包含因果视频分词器、动作条件潜在动力学模型、rollout工具和FVD评估代码。
- •Minecraft动力学模型使用16亿参数,分布在30个块因果层中,配置为使用Muon优化器训练200,000步。
- •Reactor报告在NVIDIA B200 GPU上达到57%至58%的模型FLOPs利用率,激活值而非模型状态被识别为主要内存开销。
- •团队记录了多项稳定性措施,包括EMA的使用、混合精度边界、优化器更换、损失加权和小批量最优传输。
- •该发布不包含行为克隆或强化学习训练循环,项目也尚未公布FVD分数。

一个名为Reactor的研究团队发布了Open Dreamer,这是一个基于JAX和Flax NNX构建的Dreamer 4世界模型管道开源实现。世界模型通过从观测数据中学习预测环境的演化方式,使智能体能够在不与真实系统交互的情况下进行规划、学习和生成视频。由Danijar Hafner开发的Dreamer系列一直是最受关注的基于模型的强化学习研究路线之一,Dreamer 4将该方法进一步拓展至大规模视频世界模型。该项目旨在忠实复现Dreamer 4的研究方法,同时将完整的训练管道——包括稳定性修复和计算配置——公开提供,这种详细程度在此前大规模世界模型发布中通常被省略。
发布内容
项目共发布了两个代码仓库。第一个仓库next-state/open-dreamer包含完整的训练管道:因果视频分词器、动作条件潜在动力学模型、rollout生成工具以及FVD(Fréchet视频距离)评分。第二个仓库reactor-team/open-dreamer提供了一个最小化的本地rollout工具,能够根据输入的MP4文件和对应的动作文件生成视频帧。
第三个交付物是运行在Reactor运行时上的浏览器演示。它可以实时流式传输生成的Minecraft环境,并包含一个Game ⟷ Dream切换功能,可在逐帧基础上在真实游戏画面和世界模型输出之间切换视频流。
团队表示,他们的目标是有意复现Dreamer 4的研究,避免使用原论文之外的技术以保持较窄的搜索空间。开发工作首先在CoinRun上进行——这是一个可程序化生成的2D平台游戏,能够在单个GPU上训练——随后将可用的管道扩展至Minecraft/VPT风格的游戏视频。自OpenAI的Video PreTraining(VPT)项目以来,Minecraft一直是基于游戏演示学习智能体的基准领域,该项目在承包商标注的YouTube视频上训练了一个逆动力学模型,以产生大规模动作标注的游戏数据。
更多详细信息可在项目博客、arXiv论文以及Reactor在X上的公告中查看。
架构:一个骨干网络,两个模型
分词器和动力学模型共享同一个块因果Transformer骨干网络,该骨干在两种注意力类型之间交替。空间层在同一帧的元素之间传播信息,而因果时间层则在帧之间传播信息。
分词器被设计为基于Transformer的掩码自编码器(MAE),而非传统的变分自编码器(VAE)。团队报告约100倍的压缩率,并指出该设计无需KL散度或对抗损失。他们认为,掩码方法使潜在空间更适合基于扩散的生成。
动力学模型使用扩散强制、流匹配和快捷模型进行下一帧预测,同时还会预测下一个动作。rollout不再在独立的转移模块和策略模块之间交替,而是被整合到按时间步组织的块结构中,结构为(前一动作,状态,策略)。空间注意力在每个块内部运行,因果时间注意力则跨时间连接各块。
一个关键的设计约束:世界模型token无法读取智能体token。因此,任务和策略信息只能通过预测的下一个动作来影响未来状态。
训练配置
发布的Minecraft配置文件详细规定了训练方案。
动力学模型包含16亿参数,分布在30个块因果层中,d_model为1920,30个注意力头,以及使用分组查询注意力的3个KV头。每隔四层为一个时间注意力层。每个时间步携带32个学习到的寄存器token,打包因子为2,将相邻的分词器潜变量整合到每个动力学空间token中。时间注意力在192步的滑动窗口上运行。
训练共进行200,000步,使用Muon优化器和WSD(预热-稳定-衰减)调度,峰值学习率为3e-4。快捷/引导采样在第100,000步以0.25的批次比例激活。指数移动平均(EMA)衰减设为0.999。
分词器配置每帧产生512个潜在token,瓶颈宽度为16。原始360×640帧被填充至368×640,使两个空间维度都能被16×16的patch整除。编码器深度为12,d_model为1536;解码器深度为8,d_model为1024。MAE掩码概率最高达到0.9,LPIPS损失在半数时间步上以0.2的权重应用。
VPT动作被解析为27个二值动作通道和121个分类鼠标类别,不包含连续通道。
计算性能与内存策略
团队报告了57–58%的模型FLOPs利用率(MFU),相比之下,健康的Transformer训练通常以60%作为基准。他们的分析使用了roofline论证:在NVIDIA B200上,带宽受限与计算受限之间的交叉点位于292 FLOP/byte。每GPU处理256帧将工作负载推过该临界点,进入计算受限区域。
分片策略出乎最初的预期。在16亿参数规模下,完整的模型状态——参数、梯度、优化器状态和EMA——约占24 GiB,可容纳在单个B200中。激活值而非模型状态才是主要的内存开销。团队先后尝试了数据并行、完全分片数据并行(FSDP)、张量并行和序列并行,最终选择了普通数据并行结合激活检查点方案。
在数据加载方面,团队将整个数据集预分词为.arrayrecord文件,并使用Grain配合GPU端预取缓冲区。标准的基于ffmpeg的解码不足以保持GPU满载。
稳定性工程
研究团队明确表示,稳定性问题消耗了他们最多的开发时间。他们的核心观察:大多数稳定性问题在损失持续下降的情况下仍然出现。均方误差(MSE)平滑改善,而生成质量却同时退化——这一现象使得传统的基于损失的监控对于基于扩散的世界模型变得不可靠。
文档记录了六项具体修复措施:
-
优化器更换:Muon替代了LaProp,后者在两次各约400 B200小时的独立训练运行中表现出随机且愈发频繁的尖峰。
-
EMA作为必需项:EMA权重被视为扩散推理的必需组件,而非可选项。
-
混合精度边界:参数保持float32,BF16覆盖大多数矩阵乘法激活和注意力输入,float32保留用于归一化层和动力学流输出头。
-
损失加权:团队使用x预测配合v空间损失,该损失简化为类似Dreamer 4公式的加权项,但分母为平方形式。他们报告了小幅但可察觉的改进。
-
最优传输:在噪声和潜在序列之间应用的小批量重心最优传输改善了rollout生成的稳定性。
-
μ参数化:经过测试被认为不必要,部分原因是Muon在不同模型规模下能更好地保持超参数稳定性。
CoinRun开发阶段的另一项发现:iso-FLOPs扫描估计计算最优缩放约为 N ∝ C^0.56 和 D ∝ C^0.44。
未包含的内容
该仓库不包含行为克隆(BC)或强化学习(RL)训练循环。完整的Dreamer 4 BC/RL智能体循环被列为开放路线图项目。项目文档中描述的CoinRun策略工作未用于Minecraft实现,也未发布。
该项目未公布FVD分数,尽管仓库附带了scripts/eval_fvd.py,这是一个基于I3D的评估工具,配置为4个上下文帧和240帧的评估范围。
关键工程数据摘要
- 动力学模型:16亿参数,30层,d_model 1920,使用Muon训练200,000步
- 硬件效率:在NVIDIA B200 GPU上达到57–58% MFU,每GPU 256帧,模型状态约24 GiB
- 主要挑战:稳定性而非吞吐量——损失曲线掩盖了大多数生成质量退化