在现实世界中微调离线的世界模型

23年10月来自UCSD和清华的论文“Finetuning Offline World Models in the Real World”。

强化学习 (RL) 以数据效率低而闻名，这使得在真实机器人上进行训练变得困难。虽然基于模型的 RL 算法（世界模型）在一定程度上提高了数据效率，但它们仍然需要数小时或数天的交互才能学习技能。最近，离线 RL 已被提议作为一种框架，用于在预先存在的数据集上训练 RL 策略，而无需任何在线交互。然而，将算法限制在固定数据集上会导致训练和推理之间的状态动作分布偏移，并限制其对新任务的适用性。

这项工作考虑使用在真实机器人上收集的离线数据对世界模型进行预训练的问题，然后用学习的模型进行规划，根据收集的在线数据对模型进行微调。为了减轻在线交互过程中的外推误差，在测试时通过平衡估计回报和（认知）模型不确定性来正则化规划器。在模拟和真实机器人的各种视觉运动控制任务上评估该方法，即使在离线数据有限的情况下，其也可以对见过和未见过的任务进行少量微调。

代码和数据开源：https://yunhaifeng.com/FOWM

如图所示是**FOWM（Finetuning Offline World Models）**方法。其提出一个框架，用于直接在现实世界中对世界模型进行离线预训练和在线微调，而无需依赖模拟器或合成数据。该方法用学习的模型进行规划来迭代收集新数据，并根据预先存在的数据和新收集的数据的组合对模型进行微调。在规划过程中利用新的测试-时间正则化，可以在≤20次试验中对未见过的任务变化进行少量微调。

TD-MPC [26] 是一种 MBRL 算法，它使用模型预测控制 (MPC) 进行规划，其中的世界模型和终值函数是通过时间差分 (TD) 学习共同学习的。TD-MPC 有两个有趣的特性：(i) 它使用规划，能够在测试时正则化动作选择，(ii) 与其他 MBRL 算法相比，很轻量，能够实时运行。如图总结该架构：世界模型将观察值 s0 编码为其潜表示 z0 ，然后循环预测未来潜状态 z1:h 以及最佳动作 aˆ0:h、奖励 rˆ0:h 和价值 qˆ0:h 。未来状态 s1:h 为学习提供监督，但不是规划所必需的。

具体来说，TD-MPC 学习了五个组件：（1）表示 z = hθ(s)，将高维输入 s 映射到紧凑的潜表示 z，（2）潜在动态模型 z′ = dθ (z, a)，用于预测下一个时间步的潜表示，以及三个预测头：（3）奖励预测器 rˆ = Rθ (z, a)，用于预测瞬时奖励，（4）终端价值函数 qˆ = Qθ (z, a)，以及（5）潜策略指南 aˆ = πθ(z)，用作规划的行为先验。用 z′ , s′ 表示子序列中 z, s 的后继（潜）状态，并使用 aˆ, rˆ, qˆ 区分预测与观测的（真值）a，r，q。在原始定义中， TD-MPC 是一种在线异-策略 RL 算法，最小化目标来联合优化所有组件。

在推理过程中，TD-MPC 使用基于采样的规划器 (MPPI) [28] 来规划动作，该规划器在动作序列空间上迭代拟合具有对角协方差的时间相关多元高斯分布，使回报（通过学习模型模拟动作来评估）最大化。

本文提出一个从离线到在线的世界模型微调框架，通过规划过程中测试-时间正则化减轻模型中的外推误差。框架包括两个阶段：(1) 离线阶段，其中世界模型在预存在的离线数据上进行预训练；(2) 在线阶段，其中学习的模型随后在有限的在线交互数据上进行微调。虽然用 TD-MPC [26] 作为主干世界模型和规划器，但该方法广泛应用于任何规划的 MBRL 算法。

由于两个数据集之间的状态-动作分布偏移，所有方法在离线数据上训练并在未见过的数据评估时，都会受到外推误差的影响。在这种情况下，无模型 Q- 学习方法中的价值高估，是最容易理解的错误类型 [14、16、17、19]。然而，像 TD-MPC 这样的 MBRL 算法在离线环境中面临着独特的挑战：状态-动作分布偏移不仅存在于价值估计中，而且在估计采样轨迹的回报时也存在于（潜）动态和奖励预测中。

虽然仅对样本内动作应用 TD-backups 可以有效缓解离线训练期间的价值高估，但在规划期间，仍可能在未见过的状态-动作对查询世界模型（包括动态、奖励预测器和价值函数）。尽管价值函数是经过谨慎学习的，但这仍可能导致严重的外推误差。为了解决这个额外的错误源，提出一种测试-时间行为正则化技术，平衡估计回报和（认知）模型不确定性。通过正则化估计的回报，尽管状态-动作覆盖不完善，仍保留了使用世界模型进行规划的表现力。在无法收集额外数据的离线 RL 设置中，避免结果高度不确定的动作，但保守策略同样会限制在线 RL 环境中的探索。与之前主要学习明确且一致保守的价值函数和/或策略的离线强化学习方法不同，基于模型不确定性的正则化规划具有一个有趣的特性：随着规划继续谨慎探索，模型在新数据上进行微调，认知不确定性自然会降低。这使得基于模型不确定性的测试时，正则化非常适合少量微调和多条轨迹的持续微调。

真实的机器人设置。如图（右）所示。智体使用位置控制来控制带有钳口夹持器的 xArm 7 机器人，静态第三人称英特尔实感摄像头捕获 224 × 224 RGB 图像观测值（厨房使用额外的俯视摄像头）；智体还可以访问机器人本体感受信息。 该设置不需要进一步的仪器。考虑三个任务：伸手、拾取和厨房，以及从它们派生出的几种任务变体。任务如图所示。伸手的目标是用末端执行器伸手到达目标，拾取的目标是拾起目标物体并将其举到高度阈值以上，厨房的目标是抓起锅并将其放入水槽。用手动设计的检测器来确定任务成功率，并为离线和在线 RL 自动提供稀疏奖励（尽管有噪声）。对伸手使用 120 条离线轨迹，对拾取使用 200 条，对厨房使用 216 条。

模拟任务和数据集。考虑一组多样化的任务和数据集，包括来自 D4RL [38] 基准的四个任务，如Hopper（中等）、Hopper（中等重播）、AntMaze（中等游戏）和 AntMaze（中等多样化）、来自 xArm [40] 基准的两个视觉运动控制任务（推和挑）和一个四足动物运动任务（行走）；这两个 xArm 任务与真实世界任务相似，只是使用较低的图像分辨率（84×84）和密集奖励。有关任务可视化，请参阅上图（左）。还为每个 xArm 任务考虑两种数据集变体：中等，包含从次优智体采样的 40k 个转换（800 条轨迹）；中等重播，包含从头开始训练 TD-MPC 代理的重放缓冲区前 40k 个转换（800 条轨迹）。