TD-MPC2：可扩展、稳健的连续控制世界模型

24年ICLR来自UCSD的论文“TD-MPC2: Scalable, Robust World Models for Continuous Control”。

TD-MPC 是一种基于模型的强化学习 (RL) 算法，它在学习的隐式（无解码器）世界模型潜空间中执行局部轨迹优化。这项工作提出 TD-MPC2：对 TD-MPC 算法的一系列改进。TD-MPC2 在 104 个在线 RL 任务中显著改善了基线，这些任务涵盖 4 个不同的任务领域，使用一组超参即可获得持续强劲的结果。智体能力随模型和数据大小的增加而增加，并成功训练单个 317M 参数智体执行跨多个任务域、具身和动作空间的 80 项任务。最后，总结与大型 TD-MPC2 智体相关的经验教训、机遇和风险。

在互联网规模的数据集上训练大型模型已经产生了能够执行各种语言和视觉任务的通用模型（Brown，2020；He，2022；Kirillov，2023）。这些模型的成功很大程度上可以归因于庞大数据集的可用性，以及精心设计的架构，这些架构可以随模型和数据大小的增加而可靠地扩展。虽然研究人员最近将这一范式扩展到机器人技术（Reed，2022；Brohan，2023），但一个从大型未经整理（即混合质量）数据集中、跨多个具身、通过低层动作执行多样控制任务的通才具身智体，仍然是一个难以实现的目标。当前对通才具身智体的方法存在以下问题：(a) 为行为克隆那些近乎专家的轨迹假设，严重限制可用数据的数量 (Reed，2022；Lee，2022；Kumar，2022；Schubert，2023；Driess，2023；Brohan，2023)，以及 (b) 缺乏消费大量未经整理的数据集、可扩展的连续控制算法。

强化学习 (RL) 是从未经整理的数据集中提取专家行为的理想框架。然而，大多数现有的 RL 算法（Lillicrap，2016；Haarnoja，2018）都是为单任务学习而设计的，依赖于每个任务的超参，没有选择这些超参的原则性方法（Zhang，2021）。能够处理大型多任务数据集的算法，必然需要对不同任务之间的变化具有鲁棒性（例如，动作空间维度、探索难度和奖励分配）。

这项工作提出 TD-MPC2：朝着实现这一目标迈出了重要一步。TD-MPC2 是一种基于模型的 RL 算法，旨在在由多个任务域、具身和动作空间组成的大型未经整理数据集上学习通用世界模型，数据来源于涵盖广泛技能水平的行为策略，无需进行超参调整。

该工作以 TD-MPC（Hansen，2022）为基础，可在学习的隐世界模型潜空间中执行局部轨迹优化（规划）。TD-MPC2 是一种用于训练大规模多任务世界模型的实用算法。具体来说，提出一系列对 TD-MPC 算法的改进，这些改进对于实现强大的算法鲁棒性（可以在所有任务中使用相同的超参）并将其世界模型的参数扩展到比以前多 300 倍至关重要。

学习一个隐世界模型

使用重建（解码器）目标来学习环境的生成模型很诱人，因为它具有丰富的学习信号。然而，准确预测长时范围的原始将来观察结果（例如图像或本体感受特征）是一个难题，而且不一定能实现有效的控制（Lambert，2020）。TD-MPC2 不是使用重建明确地对动态进行建模，而是旨在学习一个最有用的模型：一个能够准确预测以一系列动作为条件的结果（回报）模型。具体来说，TD-MPC2 结合使用联合嵌入预测（Grill，2020）、奖励预测和 TD 学习（Sutton，1998），从环境交互中学习一个隐式的、以控制为中心的世界模型，而无需解码观察结果。这种基于模型 RL 的替代公式，是用适中的模型大小对大型数据集进行建模的关键。随后，世界模型可以按照 MPC 框架执行局部轨迹优化（规划）进行决策。

组件。TD-MPC2 架构如图所示，由五个组件组成：
编码器 z=h(s, e)
潜在动态 z’=d(z, a, e)
奖励 r^=R(z, a, e)
终值 q^=Q(z, a, e)
策略先验 a^=p(z, e)

其中 s 和 a 是状态和动作，z 是潜表示，e 是可学习的任务嵌入（可忽略），用于多任务世界模型。

策略先验 p 用于指导基于样本的轨迹优化器 (规划器)，并降低 TD 学习的计算成本。在线交互期间，TD-MPC2 维护带有轨迹的重放缓冲区 B，并迭代地 (i) 使用从 B 采样的数据更新世界模型，以及 (ii) 通过使用学习模型进行规划来收集新的环境数据。

模型目标。h、d、R、Q 函数联合优化，最小化以下目标：

其中sg是stop-grad算子，而qt是t时刻的TD-目标，即

其中Q-是Q的EMA。由于不同任务的奖励幅度可能存在巨大差异，TD-MPC2 将奖励和价值预测公式化为log-变换空间中的离散回归（多类分类）问题，通过最小化 r、q 作为软目标的交叉熵来进行优化。

策略目标。策略先验 p 是一个随机最大熵（Ziebart，2008；Haarnoja，2018）策略，其学习最大化目标

由于价值估计 Q(zt,p(zt)) 和熵 H 的大小在数据集和训练的不同阶段之间可能有很大差异，因此有必要平衡这两个损失以防止过早出现熵崩溃 (Yarats et al., 2021)。自动调整 α、β 的常见选择，是保持其中一个不变，并根据熵目标 (Haarnoja, 2018) 或移动统计 (Hafner, 2023) 调整另一个。在实践中，选择通过移动统计来调整 α，但从经验上看，这两个选项之间的结果没有显著差异。

架构。TD-MPC2 的所有组件都实现为 MLP，中间有线性层，后面跟着 LayerNorm（Ba，2016）和 Mish（Misra，2019）激活。为了缓解梯度爆炸，用 softmax 操作将 z 投影到 L 个固定维度的单纯形（simplices）中，从而对潜表示进行归一化（Lavoie，2022）。将 z 嵌入为 simplices（而不是离散表示或压缩）的一个关键好处是，它自然地使表示偏向稀疏性，而无需强制执行硬约束。此归一化方案称为 SimNorm。让 V 成为每个 simplex g 的维数，其由 z 的 L 个分区（组）构建。然后，SimNorm 应用以下转换：

带一个策略先验的MPC

TD-MPC2 用学习的世界模型进行规划来得出其闭环控制策略。具体来说，该方法利用 MPC 框架进行局部轨迹优化，使用模型预测路径积分 (MPPI) (Williams，2015) 作为无导数优化器，使用长度为 H 的采样动作序列 (at , at+1 , . . . , at+H ) 用该模型推出潜轨迹进行评估。在每个决策步中，估计具有对角协方差的时间相关多元高斯的参数 μ∗, σ∗，使预期回报最大化，即

方程的求解方法是从 N (μ, σ2) 中迭代采样动作序列，评估其预期回报，并根据加权平均值更新 μ、σ。值得注意的是，方程用超出范围 H 的学习终值函数进行引导，估计完整的 RL 目标。TD-MPC2 重复此迭代规划过程固定次数，并在环境中的 at 处执行第一个动作。为了加速规划的收敛，一部分动作序列源自策略先验 p，初始化 (μ, σ) 作为前一个决策步的解，移位一步，进行热启动规划。

训练一个通才TD-MPC2智体

TD-MPC2 在各种单任务问题中的成功可以归因于上述算法。然而，学习一个大型通才 TD-MPC2 智体，可以在多个任务域、具身和动作空间中执行各种任务，这带来了几个独特的挑战：（i）如何学习和表示任务语义？（ii）如何在没有特定域知识的情况下适应多个观察和动作空间？（iii）如何利用学习的模型进行新任务的少样本学习？

可学习的任务嵌入。为了在多任务设置中取得成功，智体需要学习一种利用任务相似性的通用表示，同时仍保留在测试时区分任务的能力。当任务或域知识可用时，例如以自然语言指令的形式，任务嵌入 e 可以编码此类信息。但是，在一般情况下，当无法假设域知识时，可以选择从数据中学习任务嵌入（以及隐任务关系）。TD-MPC2 使用可学习的固定维度任务嵌入 e 来调节其所有五个组件，该嵌入与模型的其他组件一起进行联合训练。为了提高训练稳定性，将 e 的 l2 范数限制为 ≤ 1；这也会在实验中产生语义上更一致的任务嵌入。多任务 TD-MPC2 智体微调到新任务时，可以选择将 e 初始化为语义上相似的任务嵌入，或者简单地初始化为随机向量。

动作掩码。TD-MPC2 学习使用各种观察和动作空间执行任务，而无需任何域知识。为此，将所有模型输入和输出零填充到它们各自的最大维度，并在训练和推理期间屏蔽策略先验 p 所做预测中的无效动作维。这可确保无效维中的预测误差不会影响 TD 目标估计，并防止 p 在具有较小动作空间的任务中错误地增加其熵。同样，在规划期间仅沿有效维采样动作。

如下是TD-MPC2的改进点。

架构设计。TD-MPC2 的所有组件都是 MLP，每层后都有 LayerNorm（Ba，2016）和 Mish（Misra，2019）激活。对潜状态 z 应用 SimNorm 正则化，使表示偏向稀疏性并保持较小的 l2 范数。训练一组 Q 函数（默认为 5 个），并在每个 Q 函数的第一个线性层后额外应用 1% Dropout（Srivastava，2014）。TD-目标计算为两个随机子采样 Q 函数的最小值（Chen，2021）。相比之下，TD-MPC 实现没有 LayerNorm 的 MLP，而是使用 ELU（Clevert，2015）激活。 TD-MPC 根本不限制潜状态，这在某些情况下会导致梯度爆炸。最后，TD-MPC 仅学习 2 个 Q 函数并且不使用 Dropout。与 TD-MPC 的 1M 参数相比，TD-MPC2 中的架构差异导致默认的单任务模型大小的可学习参数净增加 4M（总共 5M）。但是，单纯增加 TD-MPC 的模型大小并不能带来持续更好的性能，而 TD-MPC2 则可以。

策略先验。TD-MPC2 的策略先验使用最大熵 RL 进行训练（Ziebart，2008；Haarnoja，2018），而 TD-MPC 的策略先验是作为确定性策略进行训练的，并将高斯噪声应用于动作。精心调整的高斯噪声调度与使用最大熵训练的策略先验相当。但是，最大熵 RL 可以更轻松地应用于与任务无关的超参。仅在具有多个动作空间的多任务学习中计算有效动作维度上的策略熵。

规划。TD-MPC2 的规划过程与 TD-MPC 的规划过程非常接近。但是，不利用迭代之间的动量略微简化了规划，因为它可以产生可比的结果。还通过一系列代码级优化将规划的吞吐量提高约 2 倍。

模型目标。重新审视 TD-MPC 的训练目标，并提高其对任务变化（例如奖励幅度）的鲁棒性。TD-MPC2 在log-变换空间中使用奖励和价值的离散回归（软交叉熵），这使得两个损失项的幅度与任务奖励的幅度无关。 TD-MPC 使用连续回归，这会导致在奖励较大的任务中训练不稳定。虽然可以通过基于移动统计对任务奖励进行规范化等方式缓解此问题，但在单任务情况下，很难为多任务学习设计稳健的奖励规范化方案。TD-MPC2 保留了用于联合嵌入预测的连续回归项，因为潜表示已由 SimNorm 规范化，并且离散回归对于高维空间的计算成本很高（z 的每个维度都需要 N 个 bin）。

多任务模型。TD-MPC2 引入一个用于跨多个域、具身和动作空间学习多任务世界模型的框架。引入一个规范化的可学习任务嵌入空间，TD-MPC 的所有组件都以此为条件，并且通过在训练和推理期间应用零填充和动作掩码来适应多个观察和动作空间。在大量任务上训练多任务模型，并使用在线 RL 将模型微调到保留任务（跨具身）。TD-MPC 仅考虑在具有共享观察和动作空间的少数任务上进行多任务学习，而不考虑对学习的多任务模型进行微调。

简化算法和实施。TD-MPC2 消除 MPPI 中的动量（Williams，2015），并用均匀采样取代重放缓冲区中的优先经验重放采样，这两者都简化实施，而实验结果没有显着变化。最后，还使用更快的重放缓冲区实现，使用多个工作器进行采样，并通过代码级优化（例如 Q 函数集合矢量化）来提高训练和规划吞吐量，这使得 TD-MPC2 的挂钟时间与 TD-MPC 相当，尽管架构更大（5M vs. 1M）。

对 TD-MPC2 进行评估，涉及 4 个任务域，总共 104 个不同的连续控制任务：DMControl（Tassa，2018）、Meta-World（Yu，2019）、ManiSkill2（Gu，2023）和 MyoSuite（Caggiano，2022）。任务包括高维状态和动作空间（最高 A）、稀疏奖励、多目标操纵、生理上准确的肌肉骨骼运动控制、复杂的运动（例如狗和人形机器人的实现），并涵盖了广泛的任务难度。还包括 10 个带有视觉观察的 DMControl 任务。公开发布 300 多个模型检查点、数据集和代码，用于训练和评估 TD-MPC2 智体。

TD-MPC2 的所有组件都以 MLP 的形式实现。编码器 h 包含可变数量的层 (2 − 5)，具体取决于架构大小；所有其他组件都是 3 层 MLP。中间层由线性层、LayerNorm 和 Mish 激活函数组成。潜表示被归一化为单纯嵌入。Q 函数还使用 Dropout。