M$^3$PC: Test-time Model Predictive Control for Pretrained Masked Trajectory Model
作者: Kehan Wen, Yutong Hu, Yao Mu, Lei Ke
分类: cs.LG, cs.RO, eess.SY
发布日期: 2024-12-07 (更新: 2025-02-06)
备注: ICLR 2025
🔗 代码/项目: GITHUB
💡 一句话要点
提出M$^3$PC，利用模型预测控制提升预训练Masked轨迹模型在离线强化学习中的决策性能。
🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)
关键词: 离线强化学习 模型预测控制 Masked轨迹模型 Transformer 决策优化 预训练模型 离线到在线强化学习
📋 核心要点
- 现有离线强化学习方法难以充分利用预训练轨迹模型在不同模态间学习到的丰富关系，尤其是在推理阶段生成最优策略时。
- M$^3$PC的核心思想是利用模型预测控制（MPC），在测试时利用预训练轨迹模型自身的预测能力来指导动作选择，无需额外训练。
- 实验表明，M$^3$PC在D4RL和RoboMimic数据集上显著提升了预训练模型的决策性能，并在离线到在线和目标导向强化学习中表现出更强的优势。
📝 摘要（中文）
本文提出了一种名为M$^3$PC的测试时模型预测控制方法，用于提升预训练Masked轨迹模型的决策性能。该模型通过Masked自编码目标训练，能够有效捕捉轨迹数据集中不同模态（如状态、动作、奖励）之间的关系。然而，在推理阶段，如何充分利用这些信息来生成最优策略是一个挑战。M$^3$PC利用预训练轨迹模型作为策略模型和世界模型，通过模型预测控制来指导动作选择。在D4RL和RoboMimic数据集上的实验结果表明，M$^3$PC在不进行额外参数训练的情况下，显著提高了预训练轨迹模型的决策性能。此外，该框架可以应用于离线到在线强化学习和目标导向强化学习，在提供额外的在线交互预算或指定不同的任务目标时，能够获得更显著的性能提升和更好的泛化能力。
🔬 方法详解
问题定义：现有离线强化学习方法通常难以在推理阶段充分利用预训练轨迹模型所学习到的状态、动作和奖励之间的复杂关系。预训练模型虽然能够重建被Masked的部分轨迹，但如何利用其预测能力来指导策略生成，从而做出最优决策，是一个亟待解决的问题。现有方法缺乏一种有效的机制来利用预训练模型的预测能力进行决策。
核心思路：M$^3$PC的核心思路是在测试阶段使用模型预测控制（MPC），将预训练的轨迹模型同时作为策略模型和世界模型。通过MPC，模型可以预测未来多个时间步的状态和奖励，并选择能够最大化累积奖励的动作序列。这种方法充分利用了预训练模型自身的预测能力，无需额外的参数训练。
技术框架：M$^3$PC的整体框架包括以下几个主要步骤：1) 使用Masked自编码目标预训练轨迹模型；2) 在测试阶段，给定当前状态，使用MPC算法；3) MPC算法通过轨迹模型预测未来多个时间步的状态和奖励；4) 基于预测结果，选择能够最大化累积奖励的动作；5) 执行选定的动作，并重复步骤2-4。
关键创新：M$^3$PC的关键创新在于将预训练的轨迹模型与模型预测控制相结合，从而在测试阶段充分利用模型的预测能力。与传统的离线强化学习方法相比，M$^3$PC无需额外的策略学习或价值函数估计，而是直接利用预训练模型的预测能力进行决策。这种方法简化了训练流程，并提高了决策性能。
关键设计：M$^3$PC的关键设计包括：1) 合适的Masked策略，用于预训练轨迹模型，使其能够学习到状态、动作和奖励之间的关系；2) MPC算法的优化目标，通常是最大化未来一段时间内的累积奖励；3) 预测步长的选择，需要在计算复杂度和预测精度之间进行权衡；4) 动作选择策略，例如使用CEM（Cross-Entropy Method）或随机采样等方法来搜索最优动作序列。具体的损失函数和网络结构与预训练的轨迹模型有关，论文中可能使用了Transformer结构。
🖼️ 关键图片
📊 实验亮点
实验结果表明，M$^3$PC在D4RL和RoboMimic数据集上显著提升了预训练轨迹模型的决策性能。例如，在D4RL数据集上，M$^3$PC在多个任务上取得了超越现有离线强化学习算法的性能。此外，在离线到在线强化学习中，M$^3$PC通过少量在线交互，获得了更显著的性能提升。这些结果表明，M$^3$PC是一种有效的利用预训练模型进行决策的方法。
🎯 应用场景
M$^3$PC具有广泛的应用前景，例如机器人控制、自动驾驶、游戏AI等领域。它可以应用于离线数据丰富的场景，通过预训练模型和测试时MPC，实现高效的决策。此外，M$^3$PC还可以应用于离线到在线强化学习，通过少量在线交互进一步提升性能。该方法有望降低强化学习的训练成本，并提高其在实际应用中的可行性。
📄 摘要（原文）
Recent work in Offline Reinforcement Learning (RL) has shown that a unified Transformer trained under a masked auto-encoding objective can effectively capture the relationships between different modalities (e.g., states, actions, rewards) within given trajectory datasets. However, this information has not been fully exploited during the inference phase, where the agent needs to generate an optimal policy instead of just reconstructing masked components from unmasked ones. Given that a pretrained trajectory model can act as both a Policy Model and a World Model with appropriate mask patterns, we propose using Model Predictive Control (MPC) at test time to leverage the model's own predictive capability to guide its action selection. Empirical results on D4RL and RoboMimic show that our inference-phase MPC significantly improves the decision-making performance of a pretrained trajectory model without any additional parameter training. Furthermore, our framework can be adapted to Offline to Online (O2O) RL and Goal Reaching RL, resulting in more substantial performance gains when an additional online interaction budget is provided, and better generalization capabilities when different task targets are specified. Code is available: https://github.com/wkh923/m3pc.