4. 世界模型中的多任务损失
世界模型（World Model）的核心目标是通过学习环境的内在动态规律，使智能体能够在“想象”中推演未来。与单一任务的监督学习不同，世界模型通常需要同时预测多个维度的信息：下一帧观测（图像或点云）、隐状态向量、即时奖励、动作分布乃至终止信号。这些预测任务共享同一个骨干网络，却各自拥有不同的量纲、分布特性和学习难度。如果简单地将各任务损失相加，往往会导致梯度竞争、训练不稳定甚至某一任务被完全忽略。本章将系统讲解多任务损失组合的核心原理、组合策略、动态调权方法以及实战中的常见陷阱。
上图展示了世界模型多任务学习的整体架构。共享骨干网络从历史观测与动作序列中提取特征，分别送入四个任务头：观测预测头（重建像素）、状态预测头（特征空间回归）、奖励预测头（标量回归或分类）、动作与终止预测头。底部列出了四种主流的损失组合策略。下面逐一展开。
4.1 核心原理：为什么不能简单相加
多任务损失组合的本质是求解一个多目标优化问题。设总损失为：
L_total = λ_obs · L_obs + λ_state · L_state + λ_rew · L_rew + λ_act · L_act其中每个 L 代表对应任务的损失值，λ 为权重系数。看似简单，但实际训练中会遇到三个核心矛盾：
- 量纲差异：观测重建损失（如 MSE）通常在 0.01~0.1 量级，而奖励预测的交叉熵可能在 1~5 量级。若直接相加，大数值任务会主导梯度方向。
- 收敛速度差异：奖励预测和动作预测通常比像素级重建更容易收敛。当简单任务已经收敛时，其梯度噪声会干扰尚未收敛的困难任务。
- 梯度冲突：不同任务对共享特征的梯度方向可能不一致，甚至相反。例如，观测重建希望保留高频纹理细节，而状态预测希望提取抽象语义，两者在浅层特征上可能产生对抗。
因此，组合策略的核心目标不是找到“最优权重”，而是平衡各任务的梯度贡献，使共享表征能够同时满足所有任务的需求。
4.2 组合策略与操作流程
根据实践经验和文献，以下四种策略被广泛使用，各有适用场景。
| 策略名称 | 核心公式/操作 | 适用场景 | 优点 | 缺点 |
|---|---|---|---|---|
| 固定加权 | L = Σ λᵢLᵢ，λ 为超参数 | 各任务量纲相近、难度均衡 | 简单、稳定、易复现 | 需要手动调参，对任务差异敏感 |
| 不确定性加权 | L = Σ (1/2σᵢ²)Lᵢ + log σᵢ | 任务噪声水平不同，异方差场景 | 自动学习权重，理论优雅 | σ 可能发散，需额外正则 |
| GradNorm | 调整 λ 使各任务梯度范数接近 | 梯度冲突严重、任务难度差异大 | 直接平衡梯度，收敛更稳 | 计算开销大，需额外优化器 |
| 课程学习/渐进式 | 先训练观测重建，再逐步加入其他任务 | 从易到难，稳定预训练 | 降低初期训练难度 | 训练周期长，需要调度策略 |
操作流程建议：
- 第一步：基线设定。先使用固定加权，将所有 λ 设为 1，观察各任务损失的量级和收敛曲线。
- 第二步：量纲归一化。根据第一步的损失量级，将 λ 初始化为各任务损失倒数的量级，使初始总损失中各任务贡献大致相当。
- 第三步：动态监控。每 1000 步记录各任务损失值和梯度范数，绘制曲线。
- 第四步：策略选择。若梯度范数差异超过 10 倍，切换至 GradNorm；若任务噪声明显，使用不确定性加权。
- 第五步：消融验证。每次调整后，在验证集上评估各任务单独指标，确保没有任务被牺牲。
4.3 案例拆解：自动驾驶世界模型的损失调优
我们以某自动驾驶仿真环境中的世界模型训练为例。该模型需要同时预测：下一帧鸟瞰图（BEV）分割图（观测）、自车未来 3 秒的轨迹点（状态）、碰撞风险概率（奖励）、以及方向盘转角（动作）。
初始现象：使用固定加权 λ 全为 1 训练 50k 步后，观测分割的 IoU 仅为 0.32，轨迹预测误差 MAE 为 2.1 米，碰撞预测 AUC 为 0.68，动作预测相关系数为 0.55。训练曲线显示，奖励预测损失在 10k 步后就不再下降，而观测损失仍在缓慢改善。
分析过程：通过梯度日志发现，奖励预测的梯度范数约为 0.8，而观测重建的梯度范数仅为 0.05，相差 16 倍。这导致共享骨干的梯度被奖励任务主导，观测任务几乎“饿死”。同时，奖励任务本身已经收敛，其梯度噪声开始干扰其他任务。
实际操作：
- 将奖励预测的 λ 从 1.0 降至 0.1，观测重建的 λ 提升至 5.0。
- 引入 GradNorm，每 500 步调整一次权重，目标梯度范数设为所有任务的平均值。
- 对动作预测使用 Huber Loss 替代 MSE，降低离群点影响。
结果对比：调整后训练 50k 步，观测 IoU 提升至 0.51，轨迹 MAE 降至 1.2 米，碰撞 AUC 提升至 0.79，动作相关系数提升至 0.72。训练稳定性显著改善，损失曲线无剧烈震荡。
风险与代价：GradNorm 引入了额外的优化器（学习率 0.025），增加了约 15% 的训练时间。且权重调整初期，总损失出现短暂上升，需要耐心等待 2k 步左右恢复。
复盘经验：多任务调优的核心不是追求每个任务都最优，而是找到“帕累托最优”平衡点。在本案例中，牺牲少量奖励预测精度（AUC 从 0.68 到 0.79 其实是提升），换来了观测和轨迹的显著改善，整体驾驶安全性更高。
4.4 常见误区与总结复盘
误区一：认为权重越大任务越重要。实际上，权重需要与任务损失的量级和梯度尺度配合。一个损失值很小的任务可能需要很大的权重才能获得足够的梯度贡献。
误区二：忽略梯度方向冲突。即使两个任务的梯度范数相同，方向也可能相反。此时应使用梯度投影或 PCGrad 等方法，而不是单纯调整权重。
误区三：过早使用复杂策略。在数据量不足或模型较小时，固定加权往往已经足够。GradNorm 和不确定性加权在小模型上可能引入额外方差。
经验建议：
- 始终先做“量纲对齐”，再谈“动态调权”。
- 每 5k 步保存一次各任务损失和梯度范数日志，便于事后分析。
- 如果某个任务长期不收敛，先检查数据预处理和标签质量，再怀疑权重问题。
- 对于像素级观测重建，考虑使用感知损失（LPIPS）替代 MSE，能更好对齐人类感知。
总结：世界模型的多任务损失组合是一个动态平衡过程。固定加权是基线，不确定性加权适合噪声异方差场景，GradNorm 适合梯度冲突严重的情况，课程学习适合从易到难的训练策略。实际工程中，建议以“梯度范数监控”为核心手段，结合任务优先级和收敛状态，定期调整权重。没有一劳永逸的公式，只有持续监控、实验对比和迭代调优。