Reactor开源Open Dreamer 实现JAX/Flax 重现Dreamer 4世界模型
项目概述
Reactor团队在GitHub上同步发布了两套代码库：next-state/open-dreamer 包含完整的训练管线——因果视频分词器、动作条件潜在动力模型、Rollout生成与FVD评分；reactor-team/open-dreamer 则提供本地演示工具，可将MP4视频及对应动作文件喂入模型，实时生成Minecraft世界画面。项目的核心目标是忠实复现《Dreamer 4》论文中的世界模型架构，避免引入任何外部技巧，以保持搜索空间的纯粹性。
关键技术实现
- 统一Transformer骨干：分词器与动力模型共享同一块因果Transformer骨干，交替使用空间注意力层（帧内信息）和时间注意力层（帧间信息）。
- 分词器采用MAE：与传统VAE不同，使用基于Mask的自编码器实现约100倍压缩，无需KL或对抗损失，掩码机制提升潜在空间的可扩散性。
- 动力模型结构：1.6B参数、30层块因果结构，d_model 1920，30个注意力头并配备3个KV头的分组查询注意力；每四层插入一次时间注意力。每个时间步携带32个学习型寄存器token，采用2:1的packing_factor将相邻的分词token打包进动力模型空间token。
- 动作预测融合：模型在每个时间步同时预测下一个帧和动作，采用块内空间注意力和跨块因果时间注意力，确保世界模型token无法直接读取agent token，动作信息只能通过显式动作传递影响后续状态。
训练细节与硬件表现
训练在NVIDIA B200 GPU上进行，整体计算利用率（MFU）达到57%‑58%，接近健康Transformer训练的60%基准。每GPU一次性处理256帧，实现了带宽受限向计算受限的转折点（292 FLOP/byte）。模型状态约24 GiB，能够完整装入单卡显存；激活占用是主要显存瓶颈。
训练时长200 k步，采用Muon调度器配合WSD学习率计划，峰值学习率3e-4。在第100 k步引入shortcut/bootstrap采样（占批次的0.25），EMA衰减0.999。分词器输出每帧512个latent token，使用12层深度的编码器（d_model 1536）和8层解码器（d_model 1024），掩码概率最高0.9，LPIPS权重0.2作用于一半时间步。动作空间采用27维二进制 + 121类离散鼠标动作的离散表示。
稳定性挑战与解决方案
团队指出，训练过程中的稳定性问题占用了最多调试时间。尽管MSE持续下降，生成质量却出现倒退。六项关键修复包括：
- 将LaProp替换为Muon，解决随机梯度波动；
- 强制在扩散推理阶段使用EMA权重；
- 采用混合精度策略：参数保持float32，主要矩阵乘与注意力使用BF16，归一化层和动力流输出保持float32；
- 使用x‑prediction + v‑space损失，使权重分配更贴近Dreamer 4；
- 引入基于最小二乘最优传输的噪声‑latent序列对齐，提高Rollout生成的平稳性；
- 通过μ‑参数化实验验证其对大模型并无显著收益。
未公开的内容与未来路线
当前仓库未包含行为克隆或完整RL训练回路，完整的Dreamer 4 BC/RL代理仍在路线图上。CoinRun阶段的策略工作未迁移至Minecraft，相关代码亦未公开。评估脚本提供FVD计算（基于I3D），但官方并未公布具体FVD分数。
行业意义
Open Dreamer的发布为研究社区提供了一个高保真、可复现的世界模型基准，特别是对JAX/Flax生态的支持有助于推动大规模并行训练与硬件效率的进一步探索。随着生成式模型向交互式、时序任务扩展，能够在单卡上训练并实时推理的Dreamer 4实现，为游戏AI、仿真和机器人学习提供了可落地的技术路径。未来若补全完整RL回路，预计将在强化学习与生成式视频领域产生更大影响。