Review of FOWM (Finetuning Offline World Models)

离线强化学习（offline RL），作为一种无需在线交互就能在预先存在的数据集上训练策略的框架。然而，将算法限制在固定数据集上会导致状态-动作分布偏移（state-action distribution shift），即训练和推理阶段的分布不一致，同时也限制了算法在新任务上的适用性。

在强化学习中，外推误差（extrapolation errors）是指模型在推理阶段（尤其是在线交互期间）遇到训练时未见过的状态或动作分布时，模型预测的结果可能会出现较大的偏差。这种误差会导致规划器（planner）生成不准确甚至错误的策略，从而影响强化学习的效果。

为了应对这一问题，FOWM提出在测试阶段对规划器进行正则化，通过以下方式减轻外推误差：

平衡估计的回报（returns）和模型的不确定性（epistemic uncertainty）：
- 估计的回报：这是强化学习模型通常优化的目标，即根据规划得到的策略，模型预测在未来会获得的累计奖励（如通过 Q 值或奖励函数预测）。
- 模型的不确定性：这是指模型在预测当前状态或动作的潜在结果时的不确定性程度。特别是在未见过的状态或动作上，不确定性往往较高（即模型缺乏足够的信心）。
正则化的目标：
- 规划器需要同时考虑“选择那些能够带来高回报的策略”和“避免选择模型不确定性高的策略”之间的权衡。
- FOWM通过引入正则化项，将模型的不确定性作为一种惩罚项加入规划过程中，从而限制规划器在不确定性高的区域过度探索。
作用：
- 正则化能有效减少模型由于外推误差导致的错误预测，从而提高模型在在线交互中的鲁棒性。
- 在测试阶段，尤其是当模型面对新任务或未见状态时，这种正则化有助于维持模型的性能并减少策略失败的风险。

Reinforcement Learning and the TD-MPC Algorithm

2.1 强化学习（Reinforcement Learning, RL）

描述了一个强化学习（RL）问题，通过交互学习视觉-运动控制策略，在无限时域部分可观测马尔可夫决策过程（POMDP）框架下进行建模。

目标：

学习一个策略，它是一个条件概率分布，输出在某个状态下执行动作的概率：策略通过最大化期望的累计奖励（expected return, cumulative reward）实现最优性：其中：

：在时间步执行动作时获得的奖励。
：折扣因子，用于权衡长期和短期奖励。
：可学习的模型参数。

模型和规划

采用基于模型的强化学习算法（Model-Based Reinforcement Learning, MBRL），其中：

策略被分解为多个可学习的组件（如世界模型，world model）。
学习到的世界模型用于执行规划（planning）。

为了简化讨论，使用下标表示所有可学习的参数。

POMDP 中的环境与状态

在 POMDP 中，环境遵循一个未知的状态转移函数：

状态是部分可观测的，无法直接获得其真实值。

可以通过传感器观测值（例如摄像头或机器人本体感知信息）定义环境的近似状态：

2.2 TD-MPC

TD-MPC 是一种基于模型的强化学习（Model-Based Reinforcement Learning, MBRL）算法，它结合了模型预测控制（Model Predictive Control, MPC）和时间差分学习（Temporal Difference, TD Learning）。该算法特别适合解决以下问题：

规划（Planning）：能够在测试阶段通过正则化动作选择，提高策略的稳健性。
轻量化：相比其他基于模型的强化学习算法，其计算效率更高，可以实时运行。

核心架构

TD-MPC 的架构包含以下五个主要组件（如上图所示）：

潜在表示（Latent Representation）：
- 映射高维输入到紧凑的潜在表示。
潜在动态模型（Latent Dynamics Model）：
- 预测下一时间步的潜在表示。
奖励预测器（Reward Predictor）：
- 预测即时奖励。
终端值函数（Terminal Value Function）：
- 预测终端值。
策略网络（Policy Network）：
- 生成行为策略，用于规划和动作选择。

目标函数

TD-MPC 的训练通过以下目标函数优化上述所有组件：其中：

：从重放缓冲区中采样的序列长度。
：停止梯度操作符（stop-gradient），用于对动态模型的学习。
：TD 目标，表示累积奖励的估计值。
损失项包括：
- 潜在动态模型误差：，对潜在表示的预测与真实值进行对齐。
- 奖励预测误差：。
- 价值函数误差：。
- 策略网络的正则化：限制预测的 Q 值。

推理阶段的规划过程

在推理阶段，TD-MPC 使用基于采样的规划器（MPPI，Model Predictive Path Integral）进行动作选择：

动作采样：
- 从策略网络中生成动作序列，并评估这些动作的回报。
回报估计：
- 通过以下公式计算动作序列的估计回报：其中：
  - ：终端状态的价值估计。
  - ：即时奖励。
潜在状态的动态更新：
- 使用动态模型递归预测潜在状态：

A Test-Time Regularized World Model and Planner

FOWM提出了一种从离线数据到在线微调的框架，用于训练世界模型（world model）和规划器（planner），目标是：

通过测试阶段正则化（test-time regularization）缓解在线交互过程中的外推误差（extrapolation errors）。
在有限的在线数据和交互下实现世界模型的高效微调。

该框架分为两个阶段：

离线阶段（Offline Stage）：
- 在预先收集的离线数据上对世界模型进行预训练。
在线阶段（Online Stage）：
- 利用规划器从真实环境中收集在线数据，进一步微调预训练模型。

3.1 离线强化学习中的外推误差

离线数据训练并在未见数据上评估时，状态-动作分布偏移会导致外推误差。特别是 Q 学习可能出现值高估（value overestimation）问题。基于模型的强化学习算法（如 TD-MPC）在离线环境中还会遇到潜在动态和奖励预测上的分布偏移。

解决方案：引入状态条件值估计器

仅针对样本内动作进行 TD 备份，为避免 TD 目标中包含样本外动作，FOWM引入了一个状态条件值估计器，TD 目标被重新定义为：

为缓解值高估问题，使用不对称的损失（expectile regression）：是常数，通过调整控制值的保守程度。

替代动作项

为了避免样本外动作，将策略学习中的动作项替换为优势加权回归（AWR）：

3.2 测试阶段行为正则化中的不确定性估计

仅对样本内动作进行 TD 备份能够缓解训练阶段的值高估问题，但在规划阶段，世界模型可能会查询未见状态-动作对，从而导致外推误差。这会影响值函数的准确性。

为了解决这一问题，FOWM提出了一种测试阶段行为正则化方法，在规划时平衡估计的回报和模型的不确定性（epistemic uncertainty）。这种方法可以：

避免选择模型预测高度不确定的动作。
保留规划器的表达能力，即便世界模型的状态-动作覆盖范围不完美。

方法

为了估计模型的不确定性，使用了一小组值函数的集合，计算状态-动作对的 Q 值的标准差作为不确定性指标：

修改后的回报函数包括一个基于不确定性的惩罚项：其中：

：不确定性正则项。
：控制正则强度的常数。

优化与实现

为快速传播在线交互数据中的信息，FOWM同时维护两个重放缓冲区：

：存储离线数据。
：存储在线交互数据。

通过平衡采样离线和在线数据，优化目标函数，提升了模型的微调效果。研究表明，平衡采样可以显著提高世界模型微调的性能。

Results

FOWM在离线到在线强化学习中表现优异，真实机器人和模拟环境测试中均超越 TD-MPC 和 IQL，对数据和任务变化的适应性更强，渐近性能更高。在 5 个真实任务和 11 个模拟任务中成功适应未见场景（如从 "reach" 到 "push"），显著优于 TD-MPC，即使初始模型未获得奖励。消融实验表明，样本内训练、值函数集合、正则化和平衡采样均对性能提升有贡献，尤其是平衡采样显著提高了数据利用效率。离线数据的数量和多样性对零样本任务和微调效果至关重要，多样化数据集将微调任务成功率提升至 89%。

Code

[Future work]

xiangyu fu's blog