WHALE：面向具身决策的可泛化和可扩展的世界模型

2024年11月来自南京大学和南栖仙策公司的论文“WHALE: Towards Generalizable And Scalable World Models For Embodied Decision-making”。

世界模型在具身环境中的决策中起着至关重要的作用，使在现实世界中成本高昂的探索成为可能。为了促进有效的决策，世界模型必须具备强大的泛化能力，支持分布外 (OOD) 的忠实想象，并提供可靠的不确定性估计来评估模拟经验的可信度，这两者都对之前的可扩展方法提出了重大挑战。 WHALE ，是一个用于学习可泛化世界模型的框架，由两项关键技术组成：行为调节和回溯展开。行为调节解决了策略分布偏移，这是世界模型泛化误差的主要来源之一，而回溯扩展则无需模型集成即可实现有效的不确定性估计。这些技术是通用的，可以与任何神经网络架构相结合进行世界模型学习。结合这两种技术， Whale-ST ，一种可扩展的时空变换式世界模型，具有增强的泛化能力。

评估价值估计准确性和视频生成保真度，证明 Whale-ST 在模拟任务中的优势。此外，还检验不确定性估计技术的有效性，可增强完全离线场景中基于模型的策略优化。此外， Whale-X ，一个 414M 参数的世界模型，在 Open X-Embodiment 数据集的 970K 条轨迹上进行训练。通过最少的演示表明，Whale-X 在现实世界的操作场景中表现出良好的可扩展性和强大的泛化性。

人类可以在脑海中想象一个想象的世界，预测不同的行为可能导致不同的结果 [1, 2]。受人类智能这一方面的启发，世界模型 [3] 旨在抽象现实世界的动态并提供这种“假设”预测。因此，具身智体可以与世界模型而不是现实世界环境交互以生成模拟数据，这些数据可用于各种下游任务，包括反事实预测 [4]、离策略评估 [5] 和离线强化学习 [6]。促进更有效决策的要求对世界模型的普遍性提出了重大挑战，这是以前的方法尚未充分解决的问题 [7]。此外，实现对想象视觉体验的可靠不确定性估计仍然是一项重大挑战，影响离线策略优化中合成数据的可信利用 [8]。这两个未解决的问题阻碍了世界模型在支持决策方面的进一步成功。

世界模型 的研究历史悠久，最近才开始引起广泛关注。世界模型最初在简单任务中以“动作模型”的名义引入，作为决策系统的组成部分 [14, 15]，在文献中也被称为“环境模型” [16]、“动态模型” [17, 18]，或简称为“模型” [19]。自神经网络时代到来以来，动态模型已更广泛地应用于深度强化学习算法中，以提高学习效率，从而产生了一系列基于模型的强化算法 [19, 20, 8, 21, 18]，而这些算法主要侧重于低维本体感受状态空间中的环境建模。 [3] 首次提出了具有高维视觉观察的现代世界模型通用框架，其中视觉模块将观察的图像编码为紧凑的潜向量以提取当前时间步的视觉信息，记忆序列模型集成历史代码以创建可以预测未来状态的表示。这种通用的世界模型架构很快在复杂的决策任务中取得了一系列显著的成功 [22、23、24、25、26]。

尽管在世界模型学习中取得了这些成功，但 分布外泛化 仍然是世界模型面临的一个根本挑战，尚未得到充分解决。与没有或只有轻微 OOD 假设的传统监督学习设置（其中目标分布与训练数据分布相似）相比，世界模型回答“假设”问题：“如果智体做出任何可能的决策，环境中会发生什么？”，这必须是高度不分布的。解决这一泛化问题的一个潜在方法是收集大量数据来训练大型世界模型。最近，先进的方法利用现代动作条件视频预测模型 [27, 28] 来建模视觉动态并从大规模视频体验数据中进行预训练 [11, 29, 30]。这些方法采用了各种复杂的模型架构，包括 RNN [31, 22, 32]、扩散模型 [33, 34] 和 Transformer [35, 36, 11]。然而，可用于模型学习的训练数据通常由专家或近乎专家的策略收集，导致状态动作空间中的数据覆盖率低，这对推理学习世界模型中次优策略的决策结果提出了挑战[7]。

另一项工作研究了学习方法对世界模型普遍性的影响。对于单步转换的标准最大似然目标，自回归推出误差或价值差距与分布误差和策略发散相关，并被推出范围二次放大，这种现象称为 复合误差 [19, 37, 38]。为了克服标准 MLE 学习中的局限性，已经进行一系列改进，包括训练多步模型以减少推出误差 [39, 40]、使用控制目标训练转换模型 [41, 42]、对抗性训练反事实目标策略的模型 [4]、学习动态奖励以提高模型泛化能力 [43]、对比学习能量转换函数 [44]，并将策略信息纳入模型输入以使测试时模型适应目标策略 [9]。尽管在低维本体感受观察任务中取得了成功，但将这些方法扩展到大量高维视觉数据仍然缺乏。引入适合大规模数据训练大型世界模型的高级学习方法，提高模型的泛化能力具有前所未有的重要性。

除了泛化之外，用于控制的世界模型的另一个关键主题是 不确定性估计 。研究表明，从部分覆盖数据中学习的世界模型中的离线决策容易受到模型预测误差的利用[8，45]，这需要量化世界模型的预测不确定性并提醒智体对模型不确定性保持悲观态度[8，46，47，18]。这些典型的算法通过并行学习的多个模型的集成来估计模型不确定性[48]，这非常耗费计算资源，尤其是对于大规模任务和模型。最近的研究[49，50]采用分类标记预测分布的熵作为大语言模型（LLM）的不确定性指标，尽管对其在基于离线模型的控制中有效性的支持有限。

序贯决策任务的典型公式是马尔可夫决策过程 (MDP) [51, 52]，由元组 M = (S, A, r, T^∗, γ, H, ρ0) 指定，其中 S 是状态空间，A 是动作空间，r(s, a) 是奖励函数，T^∗ (s′ |s, a) 是实际转移概率，γ ∈ (0, 1] 是折扣因子，H 是决策范围，ρ0 (s) 是初始状态分布。仅考虑 γ = 1 和 H < ∞ 的情况。

在强化学习 [53] 中，目标是学习一种最大化 MDP 中预期回报的策略，这涉及估计不同策略的价值。

一种常见的情况是，预先收集了大量的经验数据，但直接与环境交互要么被禁止，要么成本高昂，因此需要离线进行价值估计。

可以从离线数据中明确学习环境模型T，以模仿真实状态转换T^∗，合成虚构经验来模拟真实环境交互。对于视觉观察任务，智体无法直接观察状态；相反，它会接收观察空间O内的高维图像，这通常会引入冗余信息和部分可观测性。 视觉世界模型 [3]通常学习视觉模块E/θ：O→S，从当前帧观察ot中提取压缩表示z/t=E/θ(o/t)，并使用序列模型整合过去帧的潜表示z/1：t以及动作a/1：t，以克服未来预测的部分可观测性。这种世界模型架构为任何给定的策略提供了自回归想象，从而可以在没有现实世界交互的情况下进行策略评估和改进。

假设 V^π/T 是模型 T 中估计的值，环境模型误差会导致策略 π 的价值差距 |V^π/T∗ − V^π/T |。如果模型全局准确，则任何策略的价值差距都会缩小。然而，离线经验通常由一小部分策略收集（例如，接近专家的策略）。因此，学习的环境模型可能不熟悉新决策模式的结果，并且有望超越训练经验进行反事实推理以评估各种策略。

世界模型的常见学习方法将转换学习视为标准监督学习问题，以教师-强制的方式最小化预收集的轨迹上单步转换概率的负对数似然 (NLL)。NLL等价于最小化KL发散。

学习世界模型通常用于通过以自回归方式模拟轨迹来评估任何目标策略 π。 在句子生成、翻译等经典的序列建模任务中，随着模型准确率的提高，从教师-强制训练到自回归生成的分布迁移逐渐减小，因此不会导致显著的负面影响。然而，对于世界模型的学习，分布漂移既来自模型预测的不准确性，也来自目标策略和行为策略之间的分歧，从而加剧了评估的不准确性。

即使有大量专家训练数据可用，世界模型学习的泛化问题也被视为一项严峻的挑战 [7]。此前，已经提出一些解决方案。

1. • 用分布匹配代替教师-强制目标解决复合误差 [37]，从而将 H^2 项简化为 H。

2. • 基于能量的模型可以帮助减少分布误差 [44]，从而免除神经网络拟合病态转换函数的需要。

3. • 对抗性反因子学习（adversarial counterfactor learning） [4] 缓解策略发散问题。

4. • 学习环境奖励来约束世界模型，可以令人惊讶地帮助从数据分布中泛化 [43]。

同时，上述方法都是基于对抗性学习的，目前很难扩展。在 [9] 中，发现了一种增强泛化能力的新方法，即通过强制执行策略条件来学习在训练过程中进行泛化。

世界模型中的序贯决策通常需要智体探索超出训练数据集的分布外 (OOD) 。这要求世界模型表现出强大的泛化能力，使其能够做出与现实世界动态密切相关的准确预测。同时，可靠地量化预测不确定性（特别是对于世界模型有效泛化范围之外的转换）对于稳健的决策至关重要，这可以防止离线策略优化利用错误的模型预测。考虑到这些问题，提出 WHALE，一个用于学习可泛化世界模型的框架，具有增强的泛化性和有效的不确定性估计。

泛化的行为条件

根据分布漂移的误差分解，世界模型的泛化误差主要来自策略发散引起的误差复合。为了减少这种误差，解决此策略泛化问题的一个可能方法是将行为信息嵌入世界模型中，使模型能够主动识别策略的行为模式并适应策略引起的分布转变 [9]。这种适应效应已被证明可以减少由策略发散引起的模型泛化误差。基于行为条件的概念，引入一个学习目标，从训练轨迹中获取行为嵌入并集成学习的嵌入。

希望将训练轨迹 τ/H 中的决策模式提取到行为嵌入中，这让人想起以历史 τ/h 为条件轨迹似然的 证据下限 (ELBO) 最大化 [54, 55, 56]。

信息瓶颈要求学习变量 z 能够有效地捕捉轨迹中的决策模式，从而嵌入有关相应行为策略的信息。根据这一论点，通过最大化 H 个决策步骤上的 ELBO 并调整类似于 β-VAE [57] 的 KL 约束量来学习行为嵌入。这确保这个表征保持策略一致，这意味着由相同策略生成的轨迹表现出相似的行为模式，从而具有相似的表征。

然后使用学习的先验预测器 p/ψ 从历史 τ/h 中获得行为嵌入 z/h，以便在世界模型学习期间进行行为调节，其中行为嵌入被接受为未来预测的附加协变量。

在学习的世界模型中推出目标策略或执行动作序列时，先前的预测器会从交互历史中推断出潜行为意图，从而使自回归生成过程能够动态调整目标分布，以便自适应地想象未来。

不确定性估计的回溯-扩展

智体可以在世界模型中采取任何行动来生成想象的未来，这可能与离线数据集有很大不同。因此，世界模型不可避免地会产生不准确和不可靠的样本。先前的研究从理论和实验上都证明，如果无限制地使用模型生成的数据，策略的性能可能会受到严重损害 [8, 45]。因此，不确定性估计对于世界模型至关重要，因为它可以指示何时信任模型。

以前基于集成的世界模型不确定性估计方法通常需要训练多个模型 [48, 18]，这使得它们在计算上非常昂贵，尤其是对于大规模任务和复杂模型。除此之外，任意-步不确定性估计 [58] 已成为一种计算效率高的替代方案，它利用了不同长度历史信息之间的预测差异，而不依赖于集成。然而，任意-步是专门为循环神经网络设计的，不能直接应用于 transformer，这限制了它在实际任务中的可扩展性和适用性。为了克服这一限制，引入一种不确定性估计方法，即 回溯-扩展 ，其灵感来自如意-步中可变长度历史的概念。回溯-扩展的核心创新在于引入了 回溯-动作 ，它利用了具身控制中动作空间的语义结构，为基于 Transformer 的世界模型提供更准确、更高效的不确定性估计。

具体来说， 回溯-动作 可以等效替代任何给定的动作序列。对于动作序列 a/1:k，其对应的回溯-动作定义为 a^r/k，目标是机器人从任何观察 o/1 执行动作序列 a/1:k 产生的结果 o/k+1 与执行回溯-动作 a^r/k 时产生的结果 o^r/k+1 几乎相同。得益于具身控制中动作空间的语义结构，回溯-动作对于末端执行器姿势控制在计算上是可行的。例如，在 Open X-embodied 数据集中，动作空间由控制末端执行器的 7 维向量定义。前三个维度表示夹持器位置的变化（∆x、∆y、∆z），接下来的三个维度表示腕部方向的变化（∆roll、∆pitch、∆yaw），最后一个维度决定夹持器是打开还是关闭。因此，可以回溯动作 a^r/k，其中输入 a^(j)/i 表示动作 a/i 第 j 维的值。

下一个概念是 回溯-展开 。给定一个回溯步 k，该过程首先回溯 o/t−k 作为推出的起始帧。从 o/t−k 执行动作序列 a/t−k:k 的相应回溯动作方舟，产生相应的结果 o/k+1。在实践中，为了防止多步累积回溯动作 a^r/k 超出动作空间范围，a^r/k 被分为 k 个步骤。在每个步骤中，前六个维度设置为 1/k * a^r(i)/k，而最后一个维度 a^r(6)/k 保持不变，通过多步推出实现预期结果。

最后，提出 回溯-展开 不确定性量化器，如图所示。为了估计 (o/t, a/t) 的不确定性，设置各种回溯步来生成相应的回溯-展开预测。然后，计算特征级差异，以感知损失表示，即使用回溯-推出生成的输出与不使用回溯-展开输出之间的差异。此外，结合动态模型熵 (T (·|ot , at )) 的预测熵，该熵通常用于衡量模型本身的置信度。通过将平均差异与预测熵相乘，得出最终的不确定性估计结果。值得注意的是，回溯-展开在训练阶段不需要任何修改，与基于集成的方法相比，这显著降低了计算开销。

使用 ST-transformer 架构进行实例化：Whale-ST

如图展示 Whale-ST 的整体架构。具体来说，Whale-ST 包含三个主要组件：行为调节模型、视频token化器和动态模型。受之前研究 [11] 的启发，这些模块采用时空 ST-transformer 架构。在这个框架内，每个token都设计为只关注当前帧中其他tokens以及前几帧中相应位置的tokens。此外，Whale-ST 可以一次并行生成下一帧的所有tokens。这些设计显著简化了计算需求，从相对于序列长度的二次依赖关系简化为线性依赖关系，从而降低了模型训练的内存使用量和计算成本，同时提高了模型推理速度。