Valdi: 价值扩散世界模型
世界模型可以启用模型预测控制(MPC),但这要求动态预测既足够快以用于在线使用,又足够表达力以表示不确定的未来。扩散模型提供了一种对不确定动态进行建模的自然机制,但其迭代推理过程使其难以用于低延迟的潜在规划。我们通过价值扩散世界模型(Valdi)弥合了这一差距,将用于MPC的端到端在线训练与潜在扩散动态模型相结合。在CarRacing环境的初步实验中,我们展示了Valdi在训练和推理时均使用单步扩散,与确定性MLP基线性能相当。我们的实验揭示了在此设置下预测多模态性与控制性能之间的权衡。代码已开源。
论文精读
TL;DR Valdi 将潜在扩散动力学模型融入 TD-MPC 式在线 MPC 训练,单步扩散即可在 CarRacing 中匹配确定性 MLP 基线,同时支持不确定性的多模态轨迹预测,揭示了预测多样性与控制性能之间的权衡。
问题
问题背景
在 Model Predictive Control (MPC) 与基于模型的强化学习中,世界模型需要在保证预测精度的同时满足在线推理的低延迟要求,以便在实时控制循环中生成可靠的未来轨迹。
现有方法局限
当前主流方法如 TD-MPC 使用确定性 MLP 作为动态模型,虽推理快速,但无法表达环境中的多模态不确定性(例如交通场景中其他智能体的多种可能行为)。扩散模型天然擅长建模多模态分布,但常应用于像素空间,其迭代去噪过程导致单次预测耗时过久,难以集成到需要毫秒级决策的 MPC 在线规划中。已往工作往往牺牲表达力换取速度,或牺牲速度换取不确定性建模能力,两者难以兼得。
为什么这个问题难且重要
核心挑战在于表达力与推理延迟的根本性权衡:扩散模型用多步采样逼近真实分布,每增加一步采样都会累积计算开销,而 MPC 的实时约束要求动力学预测必须在极短的规划窗口内完成。业界对不确定感知的快速规划需求迫切,例如在自动驾驶、机器人抓取等场景中,准确预测多种可能的未来对安全决策至关重要。若能在保持低延迟的前提下引入多模态预测,将显著提升控制策略的鲁棒性与泛化性。
工程启示:在资源受限的在线场景中部署扩散世界模型,必须通过架构简化(如单步扩散)或专用推理优化来达到延迟要求,这直接决定了算法能否落地到实际嵌入式系统。
行业类比
类似自动驾驶预测模块需在 100ms 内输出多种可能轨迹,既不能因计算超时而错过控制时机,也不能因简化模型而丢失关键多模态信息。
核心洞察
- 潜在扩散模型仅需单步去噪即可达到确定性MLP的在线控制性能。Valdi 首次在 TD-MPC 框架内用端到端训练的潜在扩散模型替代 MLP 动力学模型,发现训练和推理均使用一步扩散就足以匹配基线,打破了扩散模型必须多步迭代的固有认知。这与现有像素空间扩散世界模型(如 Alonso et al., Hafner et al.)的高延迟推理形成鲜明对比,证明在压缩至低维潜在空间后,单步去噪即可有效捕捉过渡动态,为扩散模型在低延迟 MPC 中的实用化提供了简洁范式。
- 增加扩散步数会增强预测的视觉多样性,但导致控制性能下降。实验表明,多步扩散产生更逼真的多模态未来轨迹,却在 CarRacing 任务上回报降低,而单步扩散的确定性预测反而更优。这揭示了不确定性建模中“表达力-控制性能”的权衡:高随机性的世界模型可能引入价值估计的较大方差,干扰 MPC 的动作排序。该发现为在不同复杂度和不确定性需求的环境中选择扩散调度提供了经验指导,区别于想当然追求丰富预测的做法。
- 扩散动力学模型在长时域规划中的价值估计比 MLP 基线更可靠,即使短期误差更大。尽管扩散 rollout 的即时预测误差较高,但其产生的价值估计在更深的规划 horizon 上比 MLP 保持更好的一致性,这表明扩散模型能更合理地捕捉未来不确定性,辅助规划器规避短视决策。这一特性挑战了“精确短期预测是规划基础”的传统假设,提示在基于搜索的规划中,模型对长期后果的合理预期可能比短期精度更关键,为设计面向规划的动力学模型提供了新视角。
方法
核心架构:端到端在线潜在扩散世界模型
Valdi 在 TD-MPC 框架基础上,将确定性 MLP 动力学模型替换为 潜在扩散模型(latent diffusion dynamics model),实现面向 模型预测控制(MPC) 的快速、不确定性感知规划。
输入与编码
系统接收高维观测(如图像),通过一个 编码器 将其映射到紧凑的潜在状态 z。编码器与动力学、奖励、价值网络联合训练,形成端到端的表示。
扩散动力学模块
核心动力学预测采用 条件扩散过程。给定当前潜在状态 z_t 和动作 a_t,扩散模型学习生成下一个潜在状态 z_{t+1} 的分布。为满足在线推理的低延迟要求,训练和推理均仅使用 单步扩散(即直接预测噪声,再由状态估计还原)。该设计使得动力学预测速度接近 MLP,同时保留了扩散模型对不确定性的建模能力。
训练目标
Valdi 联合优化多个损失:
- 表示学习:通过时间差分(TD)目标促使潜在空间保留与控制和回报相关的信息。
- 奖励与价值预测:在潜在空间中预测即时奖励和状态价值,服务于规划器。
- 扩散损失:标准去噪分数匹配损失,使生成的状态分布贴近真实环境动态。 训练全程在线进行,采用经验回放池平衡样本效率与在线自适应。
规划与动作输出
推理时,规划器利用扩散动力学模型在潜在空间内展开候选轨迹,估计累积回报,并通过 MPC 选出最优动作序列。与 TD-MPC 的确定性 rollout 不同,扩散模型的随机性可让规划器在不确定区域探索更多可能性。
实验发现,在 CarRacing 环境中,单步扩散已足以匹配 MLP 基线的控制性能;增加扩散步数虽能生成视觉多样、多模态的预测,却导致任务性能下降,揭示了预测多样性与控制性能的 trade-off。
与同类方法的差异:相比现有基于像素空间扩散的世界模型(如 DreamerV3 的联想),Valdi 将扩散生成置于紧凑的潜在空间,并与 TD 学习、MPC 规划无缝融合,避免了像素重建带来的高昂计算开销,更适合低延迟在线控制场景。
实验
实验设计
Valdi 在 CarRacing 环境中进行端到端在线训练,与确定性 MLP 动力学基线(类似 TD-MPC)对比。训练和推理时默认使用单个扩散步骤(T=1),并考察增加推理扩散步数对性能与多模态性的影响。评估指标为任务累计回报,同时分析价值函数 self-consistency 与 rollout 预测的 grounding 质量。
关键发现
- 单步扩散匹配基线:Valdi 使用单扩散步在控制性能上与 MLP 基线持平,表明扩散动力学在低延迟潜在规划中的可行性。
- 多步推理导致性能下降:增加推理扩散步数(如
T>1)产生视觉上更多样的多模态轨迹,但任务回报不升反降,揭示预测多模态性与控制性能在此环境中存在 trade-off。 - 长期价值估计更优:尽管扩散 rollout 的短期预测误差更大,在规划所用的长 rollout 深度上,其价值估计比 MLP 基线更 grounded(更贴近真实环境反馈),暗示扩散模型有助于缓解价值函数崩溃。
与基线对比的深度解读
MLP 基线的确定性预测在简单环境中快速且准确,而扩散模型的多模态能力在此未转化为控制收益,可能因为 CarRacing 环境本身缺乏需要多样预判的不确定性。然而,扩散 rollout 在长程上的价值估计优势意味着,在需要应对不确定未来的复杂任务中,潜在扩散动力学 可能提供更鲁棒的规划依据。当前单步训练足以匹配基线,验证了 端到端在线训练扩散世界模型 的工程可行性,未来工作可探索更高维、随机性更强的环境来释放多步扩散的潜力。
行业影响
落地场景
Valdi 将潜在扩散模型融入模型预测控制(MPC),单步推理即可匹配确定性基线的性能,使其适用于对 时延敏感 且 动态不确定性强 的在线决策场景。典型应用包括:
- 自动驾驶与机器人:在局部轨迹规划中,利用扩散模型生成多模态未来状态,提升对行人、车辆等不确定行为的应对能力。
- 工业过程控制:如化工反应器或能源系统优化,需要快速模拟多种可能演化路径,以做出鲁棒控制决策。
- 交互式游戏 AI:角色行为规划中,Valdi 能同时兼顾实时性与未来分支多样性,增强非玩家角色(NPC)的拟人化水平。
商业价值
- 降本:单步扩散推理在训练与部署阶段均保持低计算开销,避免了传统扩散模型的多步迭代成本,可在边缘设备上运行,减少云端依赖。
- 增收 / 体验提升:模型除确定性预测外,天然携带 价值函数自一致性 约束,使规划更关注长期回报,减少短视错误。在自动驾驶领域,这意味着更低的事故率与更平稳的行驶体验,提升用户信任;在工业控制中,则可提高良品率或能效。
- 相比确定性 MLP 世界模型,Valdi 在长时序推演中价值估计误差更小(论文显示其 rollout 的 value 更“grounded”),有助于避免规划失效,间接降低安全风险与维护成本。
与现有产品 / 工作流的接口
Valdi 的架构建立在 TD-MPC 系列框架之上,可直接替换其中的确定性动力学模型,集成路径清晰:
- 现有基于 MPC 的解决方案(如自动驾驶中的轨迹优化模块、机器人 ROS 2 控制节点)只需将动力学预测组件切换为 Valdi 的单步扩散模型,无需改动 planner 逻辑。
- 训练流程支持端到端在线学习,可与标准 RL 环境(如 Gymnasium)对接,便于迁移至企业自研仿真平台。
- 模型输出为潜在空间编码,与现有 价值网络、策略网络 解耦,允许灵活替换决策组件,适合模块化部署。
具体落地 Use Case
- 自动驾驶仿真与规划:在路口无保护左转等交互密集型场景中,利用 Valdi 对周围车辆未来轨迹进行 单步多模态预测,规划器可在毫秒级内评估多种可能性并选择安全动作,相比纯确定性模型显著降低探索性风险。
- 量化交易中的订单执行:市场微观结构充满随机性,Valdi 可对极短时间窗内的价格跳动与流动性变化进行概率建模,以 MPC 方式动态优化拆单策略,在低延迟约束下捕获更多价格改善机会。
局限
- 实验仅限 **CarRacing** 这一简单环境,其动力学相对确定、多模态需求低,难以充分体现扩散模型的表征优势。在更复杂的视觉控制任务(如 DMControl/Atari)中,扩散模型是否能稳定提升规划质量且不损害效率,仍缺乏证据。因此,方法的泛化性和实用性尚待更广泛验证。
- 研究中发现增加推理扩散步数反而使控制性能下降,与扩散模型通常依靠多步去噪生成高保真样本的直觉相悖。这可能由于短时域内的不准确预测累积,或值函数估计在长时域中虽更准确,但短时域误差影响了最终决策。该现象暴露了将扩散模型直接融入在线 MPC 面临的挑战,限制了其在需要精确多模态预测的场景中的应用。
- 虽然 **Valdi** 采用单步扩散匹配了 **MLP** 基线,但扩散模型自身的参数量和训练代价远高于简单 MLP。文中未与现有扩散世界模型(如 **Diamond**)对比推理延迟与资源消耗,实际部署效率存疑。若单步即可,则可能没必要引入扩散模型,除非在更复杂环境中扩散能带来显著性能提升,才可与额外成本权衡。