TD-MPC2: Scalable, Robust World Models for Continuous Control

一句话定位

TD-MPC2 把「在隐式（无解码器）世界模型的隐空间里做局部轨迹优化（MPC 规划）」这条路线做到了可扩展、可复用：用同一套超参在 104 个连续控制任务上超过无模型/有模型基线，并把单个 317M 参数世界模型放大到跨 80 个任务、多具身、多动作空间上训练，证明 model-based RL 也能随模型/数据规模稳定涨点。

背景与定位

• 本文是 TD-MPC（Hansen et al., ICML 2022）的直接续作，属于 model-based RL + 隐式世界模型 + MPC 规划 范式：不做像素/观测重建，而是学一个「control-centric」模型——只要求它准确预测由动作序列决定的回报（reward + value），用 joint-embedding prediction（BYOL 式隐一致性）+ 奖励预测 + TD-learning 三项联合训练，再用采样式 MPC（MPPI）在隐空间做规划。
• 与生成式世界模型路线对比：dreamer-v3 是本文最强的有模型基线，走「学生成式世界模型 + 在想象 rollout 上优化无模型策略」；planet/dreamer-v1 早期靠重建学隐动态；muzero 用 MCTS 在离散动作上规划学得的模型。TD-MPC2 与它们的关键分歧：decoder-free（不重建观测，作者引 Lambert 2020「objective mismatch」论证重建不利于控制），且用连续动作 MPC 规划而非无模型策略或 MCTS。
• joint-embedding 的隐一致性目标与 i-jepa / lecun-path-autonomous-machine-intelligence 的 JEPA 思想同源（预测隐表示而非像素），但 TD-MPC2 把它嵌进 RL 值学习闭环。
• 论文明确批评当代「通才具身智能体」两条路线的缺陷：(a) Gato/RT-1 式行为克隆假设近专家轨迹、可用数据被严重限制；(b) 缺少能吃下大规模「混质量（uncurated）」数据的可扩展连续控制算法。TD-MPC2 定位为后者的解法——从覆盖 random→expert 全谱系行为策略的数据里，用 RL 抽取专家行为。
• 会议归属：ICLR 2024 Spotlight（OpenReview id Oxh5CstDJU）。三位作者 Nicklas Hansen、Hao Su*、Xiaolong Wang*（*equal advising）均属 UC San Diego。

模型架构

Backbone 是 MLP + 隐式世界模型（非 RSSM、非 DiT、非重建），五个组件全部条件于观测/动作与一个可学习任务嵌入 e：

• Encoder z = h(s, e)：观测映射到隐表示；层数随规模变（2–5 层）。
• Latent dynamics z' = d(z, a, e)：隐空间前向动力学（3 层）。
• Reward r̂ = R(z, a, e)：预测单步奖励（3 层）。
• Terminal value q̂ = Q(z, a, e)：预测折扣回报，学 5 个 Q 的 ensemble，TD-target 取其中随机抽 2 个 EMA 目标 Q̄ 的最小值（RED-Q 式）。
• Policy prior â = p(z, e)：最大熵随机策略，作用是引导采样式规划器、并降低 TD-learning 计算成本。

关键设计（相对 TD-MPC 的改动）：

• SimNorm 隐归一化：把 z 切成 L 个固定维度单纯形（simplex），每个 simplex 内做 softmax（维度 V=8，温度 τ=1），使各分区和为 1。这是 VQ-VAE「向量-离散码」的软松弛版，天然把表示偏向稀疏但不施加硬约束。作者实验证明 SimNorm 是训练稳定性的关键（消融见下）。
• NormedLinear：所有中间层 = Linear + LayerNorm + Mish 激活（TD-MPC 原版用 ELU、无 LayerNorm，会在部分任务上梯度爆炸）。
• Q 函数第一层后加 1% Dropout。
• 奖励/值用离散回归：在 log-transformed 空间做多分类（101 个 bin）、以软 cross-entropy 拟合，使损失量级与任务奖励量级解耦（C51/DreamerV3 式 two-hot）；joint-embedding 项仍用连续 L2 回归（因为 z 已被 SimNorm 归一化）。
• 多具身/多动作空间：不靠领域知识，用**零填充（zero-pad）所有模型输入输出到最大维度 + 动作掩码（action masking）屏蔽无效动作维；任务语义靠可学习任务嵌入 e（约束 ℓ2-norm ≤ 1，即 max_norm=1）**编码，联合训练。

规划（MPC with policy prior）：用 MPPI（Model Predictive Path Integral）做无导数采样优化，在隐空间 rollout 候选动作序列估回报，超出 horizon H 用学得的终端值函数 bootstrap 补全完整 RL 目标；一部分候选序列来自 policy prior，且用上一步解 shift-1 做 warm-start。相比 TD-MPC 去掉了 MPPI 的 momentum，规划吞吐经代码级优化提升约 2×。

5M 基础模型精确参数分解（S=39, T=80, A=6）：Encoder 167,936 / Dynamics 843,264 / Reward 631,397 / Policy 582,668 / Q 3,156,985 / Task 7,680，合计 5,389,930。隐状态维 512、任务嵌入维 96、MLP 维 512、Encoder 维 256。

五档规模配置（Encoder dim / MLP dim / Latent dim / enc-layers / Q）：

规模	Enc dim	MLP dim	Latent	enc层	Q
1M	256	384	128	2	2
5M*	256	512	512	2	5
19M	1024	1024	768	3	5
48M	1792	1792	768	4	5
317M	4096	4096	1376	5	8

（*5M 为单任务默认；任务嵌入维恒为 96；跨规模不改其它超参。）视觉任务把 MLP encoder 换成 4 层浅 CNN、64×64 输入 + random shift 增广（TD-MPC 曾支持到 9×224×224）。

数据

• 单任务在线 RL：无固定数据集，边规划采数据边更新（replay buffer 容量 1,000,000，均匀采样，取代 TD-MPC 的优先经验回放）。
• 多任务离线 RL 数据集（本文 scaling 主战场，均来自 240 个单任务 TD-MPC2 智能体的 replay buffer，因此天然覆盖 random→expert 全谱系行为）：
  • 80-task 集：DMControl + Meta-World，80 个任务、12 种具身、2.69M 条 episode、545M transitions、34GB。任务构成 = 全部 50 个 Meta-World 任务 + 30 个 DMControl 任务（19 个原版 + 11 个新造）。
  • 30-task 集：纯 DMControl，30 个任务、11 种具身、690k episode、345M transitions、20GB。
• 104 任务全景（单任务基准，同一套超参跑完）：DMControl 39 个（19 原版 + 11 自造，含 10 个视觉版）、Meta-World 50 个、ManiSkill2 5 个、MyoSuite 10 个。任务覆盖高维状态/动作（最高 A ∈ ℝ³⁹，MyoSuite 骨骼肌手）、稀疏奖励、多物体操作、复杂 locomotion（Dog A∈ℝ³⁸、Humanoid A∈ℝ²¹）；Pick YCB 涉及 YCB 数据集全部 74 个物体。
• Episode 设置（无终止条件、固定长度、infinite-horizon）：DMControl episode 1000 步 / action-repeat 2 / 有效 500 步 / 总步数 4M–14M；Meta-World 200 / 2 / 100 / 2M；ManiSkill2 200 / 2 / 100 / 4M–14M；MyoSuite 100 / 1 / 100 / 2M。成功判定用更严格口径：仅当 episode 最后一步成功才算成功（拿起又掉落不算）。

训练方法

• 模型目标 L(θ)：对 buffer 采样的 H 步轨迹，按时间加权 λ^t 累加三项——① joint-embedding prediction（‖z'_t − sg(h(s'_t))‖²，stop-grad 目标）② 奖励预测 CE ③ 值预测 CE（TD-target q_t = r_t + γ Q̄(z'_t, p(z'_t))，Q̄ 为 Q 的 EMA）。奖励/值以 soft cross-entropy 在 log 空间做离散回归。
• 策略目标 Lp(θ)：最大熵 RL，最大化 Σ λ^t [α Q(z_t, p(z_t)) − β H(p)]，梯度只对 p。α 用 moving (5%,95%) 百分位统计自动调，避免熵早塌；多动作空间下只对有效动作维算熵。
• 多任务/离线处理：直接把多任务离线训练当普通训练做，不加任何保守正则、不调超参——目标是让「大规模多任务离线预训练 → 单任务在线微调」无缝衔接、超参不变。附录 J 另给一个可选的测试时正则（只在规划时惩罚 Q 方差大的轨迹、系数 c 任务无关自动缩放），19M/80-task 上 c=0.01 把分数从 56.54 提到 62.01（c 过大反降，c=0.1→44.13）。
• 少样本微调：19M 智能体在 70 个任务上预训练，朴素全量微调到 10 个 held-out 任务（在线 RL、空 buffer、超参不变），新任务嵌入初始化为语义相近源任务的嵌入。20k 步（低数据）下比 from-scratch 提升约 2×。
• 关键超参（全任务统一）：UTD=1、batch 256（多任务放大到 1024）、Adam、lr 3e-4（encoder 1e-4）、grad-clip norm 20、λ=0.5、策略熵系数 β=1e-4、Q EMA 系数 0.99。discount 用启发式 γ = clip((T/5−1)/(T/5), [0.95, 0.995])（DMControl T=500 → 0.99），seed steps S = max(5T, 1000)。
• 规划超参：horizon H=3、迭代 6（若 ‖A‖≥20 再 +2）、population 512、policy-prior 采样 24、elites 64、min/max std 0.05/2、温度 0.5、无 momentum。

Infra（训练 / 推理工程）

• 训练硬件与成本：论文 Table 1 报 80-task 数据集训练成本（单张 NVIDIA GeForce RTX 3090 GPU-days）——1M：3.7 天 / 5M：4.2 / 19M：5.3 / 48M：12 / 317M：33 GPU-days。作者强调「317M 模型可在有限算力下训练」，是刻意「democratize RL」的低门槛路线（对比需大团队大算力的 RL）。
• 显存/内存要求（README）：单任务在线 RL 需 GPU ≥8GB、机器 ≥12GB RAM；80-task 多任务离线 RL 需 128GB RAM；训 317M 需 GPU ≥24GB。
• 工程改动：Q-function ensemble 向量化、更快的多 worker replay buffer、规划代码级优化（吞吐 ~2×），使 5M 模型 wall-time 与 1M 的 TD-MPC 相当。相比基线，TD-MPC2 用 batch 256（SAC/TD-MPC 用 512）、UTD=1（DreamerV3 用 UTD=512）即可稳定学习。
• 推理/控制频率：论文未给具体 FPS / control-Hz / 边缘部署延迟数字（未披露）；推理即每决策步跑 6 次 MPPI 迭代的采样规划。
• 并行策略、精度（fp16/bf16/fp8）等分布式细节论文未披露（单卡训练，无需多卡并行）。

评测 benchmark

一手结果（均 3 seeds、95% CI）：

• 单任务对比（104 任务 4 域）：TD-MPC2 用同一套超参在全部 4 个域上数据效率与渐进性能均超 SAC / DreamerV3 / TD-MPC。高维 locomotion（Dog A∈ℝ³⁸、Humanoid）与多物体操作（Pick YCB）上大幅领先：Pick YCB 在 14M 步达 >60% 成功率，其它方法在预算内学不会。TD-MPC 在 Walker Stand/Walk 等任务因梯度爆炸发散，DreamerV3 在 Dog 上数值不稳、且难做精细操作（lift/pick/stack）。MyoSuite 结果尤其亮眼——作者称报告前没在该基准调过任何实验。视觉 RL（10 个图像版 DMControl）上与最强基线 DrQ-v2、DreamerV3 相当。
• 多任务 scaling（归一化分数，任务成功率 + 回报归一到 [0,100] 平均）：
  • 80-task（DMControl+Meta-World）：1M→16.0，5M→49.5，19M→57.1，48M→68.0，317M→70.6。
  • 30-task（DMControl）：1M→18.9，5M→54.2，19M→59.4，48M→71.4，317M→70.6（图中 30 任务 5M=54.2）。
  • 分数随 log(模型参数) 近似线性上升，317M 处未饱和，预期继续放大仍涨。对比 TD-MPC 同样放大反而掉点（80-task 1M=16.0 但更大规模不涨）。
• 消融（在 Dog Run / Humanoid Walk / Pick YCB 三个最难任务 + 19M 80-task 多任务上，红色为默认）：
  • 归一化：No Norm vs SimNorm vs LN+SimNorm，多任务 80 上 SimNorm 默认 54.2、No Norm 掉到 46.8——SimNorm 对稳定性关键。
  • Q 数量：2/5/10，多任务上分别 54.2 / 53.5 / 57.0（单任务默认 5）。
  • 回归目标：Continuous vs Discrete，离散回归更稳（大奖励任务上连续回归不稳）。
  • Actor：Policy-only / Planning-only / Planning+policy，规划+策略最好（证明规划贡献显著）。
  • 激活函数：ELU/ReLU/Mish 渐进性能相当，但 Mish 梯度更平滑。
  • 任务嵌入归一化：ℓ2-norm=1 的归一化嵌入把 80-task 从 46.6 提到 54.2，且 t-SNE 上语义更连贯（Door Open 与 Door Close 相邻；相似度更贴近动力学/具身而非目标）。
• 基线：SAC、DreamerV3（官方复现「S」20M、UTD 512）、TD-MPC（官方 1M、逐任务调参）、CURL、DrQ-v2。

创新点与影响

• 贡献：(1) 算法鲁棒性——重访 TD-MPC 核心设计（SimNorm、LayerNorm+Mish、离散回归、5-Q ensemble、最大熵 policy prior），做到单套超参吃 104 任务；(2) 多具身/多动作空间架构——零填充 + 动作掩码 + 归一化可学习任务嵌入，无需领域知识；两者共同让 model-based RL 首次随规模稳定涨点（放大到比原 TD-MPC 多 300× 参数）。
• 改变了什么：证明「decoder-free 隐式世界模型 + 采样式 MPC 规划」是可扩展的通才世界模型路线，且能从混质量离线数据里抽专家行为——不依赖近专家 BC 数据（区别于 Gato/RT-1），也不依赖离散动作 token 化。开源 324 个 checkpoint（含 12 个多任务模型）+ 545M/345M transition 两套数据集 + 训练评测代码，成为 model-based RL 社区基础设施。
• 作者自陈局限：(i) 尚未有「开箱即通吃一切」的算法——TD-MPC2 强在难连续控制，但扩展到离散动作空间仍是 open problem（MPC 面向连续动作，或需换 MCTS/采样式离散规划，见附录 I）；(ii) 通才级泛化预计需比现有任务多数个数量级的任务量；(iii) 奖励误设、把无约束自主权交给学得模型可能致灾难性失败、某些应用数据获取成本高致算力/数据集中。作者建议采用广义奖励（成功标签/人类偏好/目标嵌入距离），但大规模用其预训练仍是开放问题。

原始链接

• arXiv：https://arxiv.org/abs/2310.16828 （PDF：https://arxiv.org/pdf/2310.16828）
• OpenReview（ICLR 2024 Spotlight）：https://openreview.net/forum?id=Oxh5CstDJU
• 项目主页：https://www.tdmpc2.com/ （Models：/models，Dataset：/dataset）
• 代码：https://github.com/nicklashansen/tdmpc2
• 前作 TD-MPC（ICML 2022）：https://arxiv.org/abs/2203.04955

一手源存档（sources/）

• td-mpc2—github-readme — GitHub 官方 README（安装/任务/用法/开源资源，fetched 2026-07-16）
• td-mpc2—project-page — tdmpc2.com 项目页快照（数据集规模表、324 checkpoint 披露，fetched 2026-07-16）
• arXiv 原文 PDF（2310.16828，arXiv 原文 PDF，不入 git）——正文/附录全部数字来源