AI Research 技术调研

Temporal Difference Learning for Model Predictive Control

13分钟阅读

一句话定位

在隐空间里联合训练一个只需预测奖励/回报的任务导向动力学模型（TOLD）和一个 TD-learning 学出的终端价值函数，推理时用 MPPI 在短 horizon（H=5）内做在线轨迹优化、拿价值函数 bootstrap 长期回报，在 DMControl + Meta-World 共 92 个连续控制任务上以远少于既有工作的环境步数超过 SAC、LOOP、Dreamer 等基线，并是文中首个记录在案解出 DMControl Dog locomotion 任务的方法。

背景与定位

  • 属于 model-based RL 中「MPC + 学出的模型与终端价值函数」一支，介于两种极端之间：纯 model-free 的 actor-critic（SAC/DDPG，不建模环境）与用真实模拟器做 MPC（论文自身的 MPC:sim baseline，无价值函数做长期估计、只能做局部最优）。
  • 核心论点：建模整个环境（reconstruction/视频预测）对控制而言是”过度建模”——world-models-ha-schmidhuber、planet、dreamer-v1、dreamer-v2 都通过重建观测（state/pixel prediction）学隐动态，容易把无关的视觉细节（如光影）也建进模型、引入 model bias。TD-MPC 改为只学”预测性 of reward”的隐表示，用 latent state consistency loss（预测的 zt+1 与 target 网络编码的 hθ−(st+1) 做 L2）替代重建，不需要 decoder。
  • 与 muzero / efficientzero 比：两者同样学习 reward-centric 的隐动态模型，但用 MCTS 做离散动作选择，动作空间维度爆炸使其难以扩展到高维连续控制（如 Dog 的 38 维）；TD-MPC 用 MPPI 处理连续动作，可直接扩展到高维 locomotion。
  • 与 LOOP（Sikchi et al., 2022）最相似：都结合 SAC-style 价值学习与 MPC 规划，但 LOOP 学的是 state predictor、且约束规划轨迹贴近 SAC 策略；TD-MPC 用任务导向的 latent dynamics model 完全替代 parameterized policy 做规划主体，策略网络 πθ 只是辅助采样与算 TD-target 的副产物。
  • 是其直接续作 td-mpc2（ICLR 2024 Spotlight）的前身：TD-MPC2 沿用了 latent consistency 的思路，并加入 SimNorm、离散回归、5-Q ensemble 等改动把规模从百万级参数放大到 3 亿+ 参数、支持多任务多具身。
  • 作者 Nicklas Hansen（一作）、Xiaolong Wang、Hao Su（共同通讯），均属 UC San Diego；发表于 ICML 2022。

模型架构

Backbone 是纯 MLP 的确定性组件（无 RNN、无概率分布建模），构成 TOLD（Task-Oriented Latent Dynamics）模型，共五个组件：

  • 表征 hθ: st → zt。状态输入用单隐层 MLP（dim=256）；图像输入用 4 层 CNN（kernel size (7,5,3,3)、stride (2,2,2,2)、每层 32 filters），观测为 3 帧堆叠的 84×84 RGB，配合 ±4 像素随机平移增广。
  • 隐动态 dθ: (zt, at) → zt+1（2 层 MLP，dim=512）。
  • 单步奖励 Rθ: (zt, at) → r̂t（2 层 MLP，dim=512）。
  • 终端价值 Qθ: (zt, at) → q̂t，含双 Q 通路 Q1/Q2（各 2 层 MLP，dim=512，第一层后接 LayerNorm + Tanh）。
  • 策略 πθ: zt → ât（2 层 MLP，dim=512），仅用于引导规划采样、并降低算 TD-target 的开销，不作为最终 actor 使用。

隐维度 Z：Humanoid/Dog 任务用 100，其余任务用 50（更高维 locomotion 需要更大隐空间）。奖励预测器与价值函数最后一层用 zero-init（降低训练早期的模型/价值偏差），其余全连接层用 orthogonal init。以 Walker Run 任务（S=24, A=6）为例，完整 TOLD 模型总参数量约 1,507,000。多任务实验（MT10）中把状态输入拼接 one-hot 任务向量；多模态实验（proprioceptive + 相机）分别编码两种模态再相加融合特征，无跨模态 attention。

规划（MPPI）：推理全部发生在 latent space。从初始 (µ0, σ0)=(0,2) 出发，采样 N=512 条长度 H=5 的 latent rollout（另加 5% 采样自 πθ 作 policy-guided sampling），按 φΓ = γ^H·Qθ(zH,aH) + Σ γ^t·Rθ(zt,at) 估计每条轨迹总回报，取 top-64（12.5%）做 φΓ-normalized 重加权更新 (µ,σ)，迭代 6 次（Humanoid 用 12 次、Dog/图像任务用 8 次），执行首个动作（receding-horizon MPC），下一步用 shift-1 warm-start 复用上一步的 µ。

数据

纯在线 RL，无预存数据集或示教数据：agent 边用 TD-MPC 规划采集环境交互轨迹，边用采到的数据更新 TOLD，replay buffer 容量不限，用优先经验回放（PER, α=0.6, β=0.4，按 value loss 定优先级）采样训练批次。评测覆盖 92 个连续控制任务：

  • 15 个 DMControl 状态任务（其中 6 个为稀疏奖励）；
  • 6 个高维 Humanoid（A∈R²¹）/ Dog（A∈R³⁸）locomotion 任务；
  • 50 个目标条件 Meta-World v2 操作任务，另设 MT10 多任务设置（10 个任务联合训练一个策略）；
  • 2 个多模态（proprioceptive + 第一视角相机）3D locomotion 任务（四足机器人绕障碍导航）；
  • 6 个 DMControl-100k 图像任务（数据效率基准，100k 环境步）+ 12 个 DMControl Dreamer-benchmark 图像任务（最多 3M 步，action repeat=2）。

Action repeat 因任务域而异：Humanoid/Dog/Walker/Finger=2，Cartpole=8，其余 DMControl=4，Meta-World 不使用（=1）；Dreamer-benchmark 全部任务统一用 2。无 sim-to-real、无跨域 co-training，全部为单环境在线交互数据。

训练方法

联合目标（公式 7–10），对 replay buffer 采出的长度 H 轨迹按时间加权 λ^(i-t)（λ=0.5）累加三项 per-step loss：

  • 奖励损失：c1·‖Rθ(zi,ai) − ri‖²（c1=0.5）；
  • 价值损失（fitted Q-iteration 改造）：c2·‖Qθ(zi,ai) − (ri + γ·Qθ−(zi+1, πθ(zi+1)))‖²（c2=0.1），θ− 为 θ 的 EMA（动量系数 ζ=0.99）；
  • 隐一致性损失：c3·‖dθ(zi,ai) − hθ−(si+1)‖²（c3=2）。

所有预测在 latent 空间递归展开（仅首个观测被 hθ 编码），梯度沿时间反传通过全部 H 步，以减轻 model rollout 的 compounding error。策略 πθ 单独用公式 11 训练（对 zi 做 stop-grad，只优化策略参数），目标是最大化 Qθ，形式上近似 DDPG/SAC 的 actor loss 但不含熵项——探索靠对 πθ/Πθ 采样加线性退火高斯噪声（ε 从 0.5 退火到 0.05，超过前 25k 步）。训练早期还把规划 horizon 本身从 1 线性增到 5（同样超过前 25k 步），避免模型初期不准时长 horizon 规划被 model bias 主导。

关键超参（Table 4）：Adam（β1=0.9, β2=0.999），学习率 3e-4（Dog/图像任务）或 1e-3（其余），batch size 2048（Dog）/ 256（图像）/ 512（其余），θ− 每 2 次梯度更新同步一次，discount γ=0.99，seed steps=5,000。全程无蒸馏、无离线 RL、无 action tokenization（连续动作直接由 MPPI 采样优化）。

Infra（训练 / 推理工程）

训练与推理 wall-time 均在单张 NVIDIA RTX3090 GPU上测得，论文未涉及分布式/多卡训练（方法定位为单卡实时 RL）：

  • 训练成本（Table 2）：达到目标回报所需时间——Walker Walk 上 TD-MPC 0.47h（SAC 0.41h、LOOP 7.72h、MPC:sim 0.91h）；Humanoid Stand 上 TD-MPC 9.39h（SAC 9.31h）。每 50 万环境步耗时——Walker Walk 上 TD-MPC 5.60h（SAC 1.41h、LOOP 18.5h）；Humanoid Stand 上 TD-MPC 12.94h（SAC 1.82h）。论文原话：TD-MPC 比 LOOP 快 16× 解出 Walker Walk，同时每 50 万步计算量少 3.3×，且 time-to-solve 追平 SAC（同时样本效率远高于 SAC）。
  • 推理延迟：默认超参（H=5, 6 次 MPPI 迭代）下约 20ms/决策步（约 50Hz），作者称”足以支持多数实时机器人应用（操作、导航、部分 locomotion，假设机载 GPU）“。在 Quadruped Run 上把 horizon 降到 1 可把延迟压到约 12ms/步；仅用学出的策略 πθ（放弃规划）则比完整规划再快约 6×，但性能通常更差。
  • 未披露：FLOPs 拆解、分布式并行策略、混合精度（fp16/bf16）——单卡小规模 MLP，论文未涉及这些维度。

评测 benchmark

  • DMControl-100k（图像输入，6 任务，10 seeds 均值±标准差，Table 1）：TD-MPC 得分 Cartpole Swingup 770±70、Reacher Easy 628±105、Cup Catch 933±24、Finger Spin 943±59、Walker Walk 577±208、Cheetah Run 222±88；与专为图像 RL 调参的 DrQ（759±92/601±213/913±53/901±104/612±164/344±67）相当或更优，全面优于 CURL、PlaNet、Dreamer；MuZero/EfficientZero 在低维动作任务上更强（813±19/952±34/942±17 vs. TD-MPC 770±70/628±105/933±24），但因动作离散化在 Walker Walk / Cheetah Run（A∈R⁶）上无法运行，TD-MPC 可以。
  • 15 个 DMControl 状态任务（400k 步内）：TD-MPC 全面超过 SAC、LOOP、MPC:sim（用真实模拟器规划、horizon 翻倍到 10 仍不敌），在 Quadruped、Acrobot 等复杂动力学任务上领先幅度尤其大；在稀疏奖励的 Finger Turn Hard 上样本效率反而不如 SAC/LOOP。
  • Humanoid/Dog 高维 locomotion（最高 A∈R³⁸）：TD-MPC 在 3M 环境步内解出全部 6 个任务（既有工作通常需 30M 步，即 10 倍），且是文中首个记录在案解出 DMControl Dog 任务的方法。
  • Meta-World 50 个目标条件任务 + MT10 多任务：均超过 SAC baseline，Bin Picking 等复杂操作任务上领先幅度大；MT10 联合训练下 TOLD 仍能从任务间数据共享中受益（尽管其目标只 reward-centric）。
  • 多模态 locomotion：TD-MPC 成功融合 proprioceptive + 相机两个模态解出任务；去掉相机的 blind agent 失败。
  • 消融（latent dynamics objective）：把隐一致性损失替换为 no-regularization / reconstruction / contrastive（SimSiam 式）— reconstruction 和 contrastive 均优于无正则，但本文提出的 latent state consistency 结果最一致；去掉 latent（state predictor 换成恒等映射）在高维任务上明显更差。
  • 基线：SAC（Yarats & Kostrikov 实现）、LOOP（官方实现）、MPC:sim（真实模拟器、无终端价值函数）、CURL、DrQ、DrQ-v2、PlaNet、Dreamer、Dreamer-v2、MuZero、EfficientZero。

创新点与影响

  • 贡献：(1) 提出 Task-Oriented Latent Dynamics（TOLD）模型——只学习预测 reward/value 所需的隐表示，用 latent state consistency loss 替代观测重建，不需要 decoder；(2) 首次把 MPC 的终端价值函数与 latent dynamics model 通过 TD-learning 联合训练（此前 MPC + terminal value 的工作要么用 ground-truth 模型、要么模型学习与价值学习相互解耦）；(3) 用 MPPI + policy-guided sampling 把基于规划的方法扩展到 38 维连续动作空间的高维 locomotion，超越受限于低维/离散动作的 MCTS 类方法（MuZero/EfficientZero）。
  • 改变了什么：证明 model-based 方法经过恰当设计（reward-centric 建模 + 联合 TD 学习）可以同时保持规划的渐进最优性、又不落入 model bias/compounding error 的传统陷阱，缩小了 model-free 方法在 wall-time 效率与 model-based 方法在样本效率之间的差距；latent consistency 的思路被直接续作 td-mpc2 沿用并放大到多任务、多具身、3 亿+ 参数规模。
  • 作者自陈局限：在难探索、稀疏奖励任务（Finger Turn Hard）上样本效率不如 SAC/LOOP，提示需要更精细的探索策略；单独使用策略 πθ（不规划）性能通常弱于完整 MPC 规划；规划预算（horizon/迭代次数）与最终性能强相关，复杂任务上收益更明显，这意味着实时性与性能之间存在需要按任务调节的权衡。

原始链接

一手源存档（sources/）

  • td-mpc—github-readme — GitHub 官方 README（方法概述、安装、训练命令、更新日志，fetched 2026-07-16）
  • arXiv 原文 PDF（2203.04955，arXiv 原文 PDF，不入 git）——正文与附录全部数字均来源于此

关系图谱

反向链接