HarmonyDream: Task Harmonization Inside World Models

一句话定位

HarmonyDream 把 dreamer-v3 一脉的显式世界模型学习重新框定为一个双任务的多任务学习问题（观测建模 vs 奖励建模），发现默认等权（w_o=w_r=w_d=1）会让观测建模系统性地压制奖励建模、诱发虚假相关，进而提出一套只加三个可学习标量、不引入新超参的”harmonizer”自动配平损失尺度，让 DreamerV2/V3/DreamerPro 在视觉机器人任务上获得 10%–69% 的绝对性能提升，并把 DreamerV3 在 Atari 100K 上刷到 136.5% 人类归一化均分的新纪录（ICML 2024）。

背景与定位

显式 model-based RL（dreamer-v3 一脉：Recurrent World Models → SimPLe → Dreamer/DreamerV2/DreamerV3）让世界模型同时学两件事——观测建模（预测未来观测 p(o_{t+1:T}|o_{1:t},a_{1:T})）和奖励建模（预测未来奖励 p(r_{t+1:T}|o_{1:t},a_{1:T})）——并共享同一个隐表示和转移模型，是一种”世界模型内部的多任务学习”。而隐式 MBRL（MuZero、TD-MPC/td-mpc2、RePo）只学奖励/价值一侧，追求 value equivalence。本文的问题意识：默认的等权重损失（w_o=w_r=w_d=1）从未被系统检验过，而两个任务的损失尺度天然悬殊——观测损失是 H×W×C 维求和、奖励损失只是一个标量，前者比后者大两个数量级（Fig 2b），导致显式 MBRL 事实上一直被观测建模主导。

作者先做诊断性实验（Sec 2.3，基于 DreamerV2，在 Meta-world 的 Lever Pull/Handle Pull Side/Hammer 上）：单纯调大奖励损失系数（w_r∈{1,10,100}）就能显著提升样本效率（Finding 1）；离线分析显示默认权重下世界模型会学到”虚假相关”——例如把机器人动作误关联到杠杆运动，从而给出不准确的奖励预测（Finding 2，图 3 开环预测可视化）；但反过来纯粹只学奖励（w_o=0，即隐式 MBRL 精神）效果更差、方差更大，因为奖励信号本身太稀疏、不足以学到有意义的表示（Finding 3）。三个发现共同指向：两个任务都不该被单方面压制，需要动态平衡。

模型架构

HarmonyDream 不改变基座世界模型结构，只在其损失函数上加三个轻量”harmonizer”。以 Dreamer 系为例，世界模型四组件不变：

• Representation model z_t ~ q_θ(z_t|z_{t-1},a_{t-1},o_t)
• Transition model ẑ_t ~ p_θ(ẑ_t|z_{t-1},a_{t-1})
• Observation model ô_t ~ p_θ(ô_t|z_t)，损失 L_o = -log p_θ(o_t|z_t)
• Reward model r̂_t ~ p_θ(r̂_t|z_t)，损失 L_r = -log p_θ(r_t|z_t)
• 加上 dynamics/KL 损失 L_d（DreamerV2 单项 KL；DreamerV3 拆成 L_dyn、L_rep 两项 KL-balancing）

默认目标 L(θ)=w_o L_o + w_r L_r + w_d L_d，三权重固定为 1。HarmonyDream 把每个固定权重换成**可学习标量 σ_i=exp(s_i)（i∈{o,r,d}）**驱动的”harmonious loss”（式 4）：

L(θ,σ_o,σ_r,σ_d) = Σ_i [ (1/σ_i) L_i(θ) + log σ_i ]

Proposition 3.1（附录 B 有证明）：对 σ 求导为零时最优解 σ* = E[L]，即调和后的损失尺度 E[L/σ*]=1——每个任务损失被自动缩放到同一常数量级，不需要手工试权重。但直接用式 4 会出问题：奖励损失这类小尺度损失会导致系数 1/σ≈L⁻¹≫1 爆炸、损害训练稳定性。为此加一个常数做整流（rectification，式 5）：

L̂(θ,σ_o,σ_r,σ_d) = Σ_i [ (1/σ_i) L_i(θ) + log(1+σ_i) ]

整流后解析解 σ = (E[L] + sqrt(E[L]² + 4E[L])) / 2，对应的调和损失尺度为 2/(1+sqrt(1+4/E[L])) < 1——严格小于未整流版本（Prop 3.1）恒等于 1 的调和尺度，即整流项 log(1+σ) 会主动把小尺度损失（如奖励损失）的系数往下压，从而避免其变得极端大（附录 B 给出解析推导）。作者强调这与 Kendall et al. 2018 的 Uncertainty Weighting（UW）神似但本质不同：UW 从最大似然导出、把每个观测像素与标量奖励”同等看待”（仍忽略维度差异悬殊的问题），HarmonyDream 直接对高维观测损失整体做尺度平衡，且不对损失背后的分布做假设（因此也能平衡 KL 损失，UW 做不到）。

跨基座通用性：论文在三个基座上验证方法通用（Sec 4.4）：

• DreamerV2（主实验基座）：直接套用式 5，三个 harmonizer 分别对应 L_o, L_r, L_d。
• DreamerV3：因其把 dynamics 损失拆成 L_dyn 与 L_rep（KL balancing，值相同、系数不同），Harmony DreamerV3 用 DreamerV3 预定义的 KL balancing 系数 α 重组回单个 L_d=αL_dyn+(1-α)L_rep，再套用同一目标（式 15），说明该方法与 KL balancing 正交、可直接叠加。
• DreamerPro（reconstruction-free，用 SwAV 原型表示代替像素重建）：目标改写为对 {SwAV, Temp, R, KL} 四项做整流调和（式 17），因为 L_SwAV + L_Temp 替代了 L_o。

规模配置（DreamerV2 基座，附录 Table 2）：观测尺寸 64×64×3（线性缩放到 [-0.5,0.5]）；action repeat Meta-world=1，RLBench/DMCR/Natural BG DMC=2；replay buffer 容量 10⁶；想象 horizon H=15；折扣 γ=0.99；λ-target 折扣 0.95；熵正则 η=1e-4；batch size Meta-world/RLBench=50，DMCR/Natural BG DMC=16；RSSM hidden size 1024；world model/actor/critic 均用 Adam，lr 分别 3e-4 / 8e-5 / 8e-5；评估 10 个 episode。Minecraft 任务用 DreamerV3 Large 变体（约 77M 参数），观测 128×128×3 RGB-only，每 5 个 episode 强制重置生成新地形，episode 长度 500、无 action repeat。RLBench 上原版 DreamerV2 卷积编解码器不够用，改用 Wu et al. 2023 的 ResNet 风格编解码器（作者强调这是与本文贡献正交的架构改动）。HarmonyDream 本身不引入任何新超参——三个 harmonizer 各对应一个可学习标量，随主网络一起用同一优化器训练。

数据

HarmonyDream 是纯在线 model-based RL，无外部数据集，数据来自 agent 与仿真环境交互后写入 replay buffer 回放训练：

• Meta-world（Yu et al. 2020，50 任务机器人操作基准）：按 Seo et al. 2022a 的难度分级选 6 个任务——easy（Lever Pull、Handle Pull Side）、medium（Hammer、Sweep Into）、hard（Push、Assembly）。训练步数：easy + Hammer 共 250K 环境步，Sweep Into 500K，Push/Assembly 1M；episode 长度 500，无 action repeat。
• RLBench（James et al. 2020）：多数任务需专家示范/专用网络结构（超出本文范围），故仅选两个稠密奖励、相对简单的任务 Push Button 与 Reach Target（Push Button 原生无稠密奖励，按 Seo et al. 2022a 方法人工赋予）；episode 最长 400 步，action repeat 2。
• DMC Remastered（DMCR）（Grigsby & Qi 2020）：DeepMind Control Suite 的视觉扩展版，每个 episode 初始化时随机重置 7 个视觉因子（地板纹理、背景、机器人体色、目标色、反射率、相机位置、光照）；训练 Cheetah Run、Walker Run、Cartpole Balance 三个任务，各 1M 环境步，episode 长度 1000、action repeat 2。
• Atari 100K（Kaiser et al. 2020）：26 个游戏，agent 限交互 100K 步（frameskip 4 → 400K 帧），约合 2 小时真实游戏时长，遵循该基准标准协议报告 human-normalized score。
• Natural Background DMC（Zhang et al. 2018，基于 RePo 代码库实现）：用 Kinetics-400 的 driving_car 视频替换静态背景，训练 Cheetah Run、Walker Run、Cartpole Swingup，episode 长度 1000、action repeat 2。
• Minecraft（MineDojo，Fan et al. 2022）：聚焦 Hunt Cow 基础技能，动作空间裁剪为 5 组共 18 个离散动作（沿用 STG-Transformer 裁剪方案），仅 RGB 观测 128×128×3，每 5 个 episode 重置一次地形保持多样性，episode 长度 500、无 action repeat。
• 离线分析用数据（Sec 2.3/Fig 2c、3 的诊断实验）：单独训练一个 DreamerV2 agent、收集其整个训练过程的轨迹作为固定 replay buffer（250K 环境步，覆盖低回报到高回报的完整分布），在此固定 buffer 上离线训练不同奖励损失系数的 DreamerV2 做对比；state-regression 探针数据集由 1 万段 × 2（原始权重/平衡权重）= 2 万段轨迹片段（每段 50 步）构成，9:1 切分训练/验证，用 4 层 MLP（隐层 400）回归到任务相关物体的真实位置。

训练方法

• 核心目标：把基座世界模型损失中各任务项（观测 L_o、奖励 L_r、动力学 L_d，或 DreamerPro 的 {SwAV,Temp,R,KL}）用整流调和损失（式 5）替换固定权重求和，三个（或四个）harmonizer σ_i 与主网络参数联合用同一优化器反向传播更新，无需额外调度或预热。
• behavior learning 不变：actor-critic 部分完全沿用基座 Dreamer（随机 actor + 确定性 critic，在世界模型想象出的隐轨迹上联合训练，critic 回归 λ-target，actor 通过世界模型反传值梯度并加熵正则），HarmonyDream 只作用于世界模型损失，不触碰策略学习目标。
• 与不稳定性的博弈：直接用未整流的调和损失（式 4）在实践中会给小尺度损失（如奖励损失）算出极端大的系数 1/σ≈L⁻¹≫1，损害训练稳定性；论文附录消融（Fig 17）验证整流项 log(1+σ) 相对未整流版本能提升稳定性与最终性能。
• 与其他多任务加权方法的对比训练实验：额外实现 Uncertainty Weighting（UW）、Dynamic Weight Average（DWA）、NashMTL 三个基线，在 Meta-world Lever Pull、RLBench Push Button、DMCR Cheetah Run 上对比（Fig 10b）。NashMTL 在 Lever Pull 上与 HarmonyDream 相近，但在另外两个任务上因优化不当产生极端（近零）权重而效果差；DWA 假设 replay buffer 数据平稳，但 RL 中 buffer 持续增长非平稳，导致其”学习进度”度量失真；UW 把像素与标量奖励同等视为”同方差高斯输出”仍隐含忽略维度差异。
• 训练精度：Meta-world 与 RLBench 用自动混合精度（AMP），DMCR 用全精度（作者备注：AMP 默认开启但可能触发 nan，可用 --precision 32 关闭）。

Infra（训练 / 推理工程）

• 框架：PyTorch（DreamerV2/DreamerV3 均有 PyTorch 复现，Atari 100K 用官方 DreamerV3 JAX 代码库 danijar/dreamerv3，commit 8fa35f）。
• 单卡训练时长（附录 C.5）：Meta-world（250K 环境步）单次运行约 24 小时；RLBench（500K 环境步）约 24 小时；DMCR（1M 环境步）约 23 小时。HarmonyDream 引入的轻量 harmonizer 不增加训练时长。
• 显存：Meta-world/RLBench 约 10GB，DMCR 约 5GB，作者指出实验可在”典型 12GB 显卡”上完成——即消费级 GPU（如单张 3090/4080 级别）即可复现全部主实验，未披露具体卡型号。
• Atari 100K/DreamerV3 实验的具体训练硬件、GPU 数、GPU-小时：论文未披露（仅说明基于官方 JAX 代码库复现,评测 100 个 episode / 每 2 万环境步）。
• 推理侧（FPS、control-Hz、边缘设备延迟）：论文未披露，本工作聚焦样本效率而非部署性能。

评测 benchmark

• Meta-world（6 任务，与基座 DreamerV2 对比，5 次随机种子、95% CI）：绝对性能提升 10%–69%（abstract 口径）；在最难的 Push 与 Assembly 上，HarmonyDream 分别达到 >75% 与 >90% 成功率，而原版 DreamerV2 在预算内学不到有意义的策略（success rate 接近 0）——对应 abstract 所述”Assembly 上最多 90% 的绝对成功率提升”。
• RLBench（Push Button、Reach Target）：HarmonyDream 相对 DreamerV2 分别带来 +28 个百分点（Push Button）与 +50 个百分点（Reach Target）的绝对最终性能提升。
• DMC Remastered（Cheetah Run、Walker Run、Cartpole Balance）：HarmonyDream 明显强化基座 DreamerV2（Fig 7a）；论文进一步展示 dynamics loss 曲线（Fig 7b）——DreamerV2 在该任务上 dynamics loss 发散，HarmonyDream 通过强调奖励建模绕开了干扰视觉因素、dynamics loss 收敛，为定性论证提供了直接证据（未给出具体数值增益，图示对比）。
• Atari 100K（26 games，基座 DreamerV3，Table 3，人类归一化分数）：

方法	Mean	Median
SimPLe (2020)	33%	13%
TWM (2023)	96%	51%
IRIS (2023)	105%	29%
DreamerV3 (论文原报)	112%	49%
DreamerV3（本文复现）	108%	42%
Harmony DreamerV3（本文）	136.5%	67.1%

Harmony DreamerV3 在 26 个游戏中的 23 个上追平或超过 DreamerV3，且未新增任何超参数或网络结构改动，是”不依赖前瞻搜索（lookahead search）“方法中的新 SOTA（论文指出 EfficientZero 靠在线树搜索 + 优先回放 + 超参调度等额外复杂度维持更高分，比较口径不完全公平）。

• Minecraft Hunt Cow（基座 DreamerV3-Large）：Fig 9 显示 Harmony DreamerV3 相对 DreamerV3 有明显提升（论文文字仅定性描述”exhibits great improvement”，未在正文给出具体成功率数值）。
• 与隐式 MBRL 对比（TD-MPC，Meta-world 三任务）：TD-MPC 因缺乏观测建模引导表示学习而学习效率低，HarmonyDream 表现更优（Fig 10a，定性对比，无量化数值）。
• 与其他多任务加权方法对比（UW/DWA/NashMTL，Fig 10b）：HarmonyDream 综合最优且实现最简单；NashMTL 仅在 Lever Pull 上接近，但在另两个任务上因梯度方向优化产生极端权重而效果差。
• 与 Dreamer 系 task-centric 方法对比（Denoised MDP、RePo，Fig 11，含 DMCR 与 Natural Background DMC）：HarmonyDream 样本效率更高（定性对比）。
• DreamerPro 对比（Fig 8 右，DMCR）：DreamerPro 默认手调 w_r=1000；HarmonyDream 在不手调权重的前提下平均优于 w_r=1 与 w_r=1000 两个手调版本，且免去了手工调参的算力成本。
• 无害性验证：在观测”干净”（无明显干扰因素）的标准 DMC 与 Crafter 上，HarmonyDream 增益边际（因为这些基准本就不存在观测建模主导的问题），但未观察到任何负面性能变化（附录 Fig 16、Fig 22）。

创新点与影响

• 贡献：首次把显式世界模型学习系统地框定为”观测建模 vs 奖励建模”的双任务多任务学习问题，并通过诊断实验证明默认等权重会导致某一任务主导、损害样本效率（此前文献虽有手调奖励权重的先例，如 decoder-free MBRL 中常见 w_r∈{100,1000}，但均是经验性调参，未被系统分析）；提出的 harmonizer 极简（每个损失一个可学习标量、无新超参），且被证明可跨 DreamerV2/DreamerV3/DreamerPro 三个不同基座迁移。
• 改变了什么：把 Atari 100K 上”不用前瞻搜索”方法的 SOTA 从 DreamerV3 的 112% 推到 136.5%（纯靠重新配平损失、不改网络结构），说明现有 model-based RL 方法可能普遍存在未被察觉的任务失衡问题；提供了一个比手调权重（网格搜索 w_r）更省算力、比 UW/DWA/NashMTL 等通用多任务学习方法更简单有效的领域特定方案。
• 作者自陈局限：(1) 方法的提出主要基于经验和直觉观察，缺乏理论刻画（未来方向是理论化解释,或超越”仅考虑损失尺度”去更好地测量/平衡各任务贡献）；(2) 在观测本就”干净”（无强干扰、目标物体显著）的基准（标准 DMC、Crafter）上收益有限，因为这些场景本不存在显著的任务主导问题；(3) RLBench 实验局限于两个可获得稠密奖励的相对简单任务，多数 RLBench 任务因需专家示范和专用网络结构而超出本文范围；(4) RLBench 上为适配任务改用了 ResNet 编解码器，这一架构改动与本文核心贡献正交、非本文方法本身带来。

原始链接

• 论文（arXiv abs，v1 2023-09-30 / v3 2024-06-05）：https://arxiv.org/abs/2310.00344
• 论文 PDF：https://arxiv.org/pdf/2310.00344
• 官方代码（PyTorch + JAX，ICML 2024）：https://github.com/thuml/HarmonyDream

一手源存档（sources/）

• harmonydream—github-readme — GitHub 官方 README 存档（sources/world-model/2023/harmonydream—github-readme.md）
• 论文全文按约定不入库，仅引用 arXiv URL（上）；本页数字均取自抓取的 arXiv v3 PDF 正文与附录（含 Table 2/3、Appendix C.5 计算资源、Proposition 3.1 证明）。