Planning Transformer: Long-Horizon Offline Reinforcement Learning with Planning Tokens

📄 arXiv: 2409.09513v1 📥 PDF

作者: Joseph Clinton, Robert Lieck

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-09-14

备注: 11 pages, 5 figures, Submitted to AAAI

💡 一句话要点

提出Planning Transformer，利用规划令牌解决离线强化学习中的长时程任务难题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 离线强化学习 长时程规划 Transformer 规划令牌 决策Transformer

📋 核心要点

  1. 现有离线强化学习方法，特别是基于决策Transformer的方法，在长时程任务中面临累积误差增大的挑战。
  2. 论文提出Planning Transformer，通过引入规划令牌，使模型能够进行隐式规划，从而指导低级策略并减少误差。
  3. 实验结果表明，Planning Transformer在复杂D4RL环境中取得了新的state-of-the-art，并提高了模型策略的可解释性。

📝 摘要（中文）

监督学习方法在离线强化学习中表现出色，特别是使用决策Transformer的方法，在连续环境和稀疏奖励方面表现良好。然而，由于自回归模型的高累积误差，它们通常在长时程任务中表现不佳。为了克服这个限制，我们超越了下一个token预测，引入了规划令牌（Planning Tokens），其中包含关于agent未来行为的高级、长时间尺度信息。以规则的时间间隔预测双时间尺度token，使我们的模型能够使用这些长时程规划令牌作为一种隐式规划形式来指导其低级策略并减少累积误差。这种架构修改显著提高了长时程任务的性能，在复杂的D4RL环境中建立了新的state-of-the-art。此外，我们证明了规划令牌通过可解释的计划可视化和注意力图提高了模型策略的可解释性。

🔬 方法详解

问题定义：现有的基于Transformer的离线强化学习方法，如Decision Transformer，在处理长时程任务时，由于自回归模型的特性，会产生严重的累积误差。这意味着模型在预测未来状态和动作时，每一步的误差都会累积，导致最终策略的性能下降。因此，如何减少长时程任务中的累积误差是亟待解决的问题。

核心思路：论文的核心思路是引入“规划令牌”（Planning Tokens），这些令牌包含agent未来行为的高级、长时间尺度信息。通过在常规时间间隔预测这些规划令牌，模型可以利用它们作为一种隐式规划的形式，从而指导其低级策略。这种方式类似于人类在做决策时，会先制定一个大致的计划，然后再逐步执行。

技术框架：Planning Transformer的整体架构基于标准的Transformer结构，但增加了规划令牌的预测分支。模型首先接收历史的状态、动作和奖励序列作为输入，然后通过Transformer编码器提取特征。除了预测下一个动作之外，模型还预测一系列规划令牌，这些令牌代表了未来一段时间内的期望状态或目标。在解码阶段，模型同时利用历史信息和规划令牌来生成最终的动作序列。

关键创新：最重要的技术创新点是规划令牌的引入。与传统的只预测下一个动作的方法不同，Planning Transformer能够预测未来一段时间内的规划，从而使模型具有更强的全局意识和规划能力。这种方法有效地减少了长时程任务中的累积误差，并提高了模型的性能。

关键设计：规划令牌的具体内容可以根据任务的特点进行设计，例如，可以是未来几个时间步的状态、奖励或目标。论文中使用了双时间尺度的token预测，即同时预测低级的动作token和高级的规划token。损失函数包括动作预测损失和规划令牌预测损失。网络结构方面，主要采用标准的Transformer结构，并针对规划令牌的预测分支进行了一些调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Planning Transformer在D4RL benchmark中的多个长时程任务上取得了显著的性能提升，建立了新的state-of-the-art。与Decision Transformer相比，Planning Transformer在某些任务上的性能提升超过了50%。此外，通过可视化规划令牌和注意力图，论文展示了Planning Transformer具有更强的可解释性。

🎯 应用场景

Planning Transformer具有广泛的应用前景，例如机器人导航、游戏AI、自动驾驶等需要长时程规划的任务。通过学习离线数据，该模型可以有效地解决复杂环境中的决策问题，并为实际应用提供更可靠的策略。此外，规划令牌的引入也为模型的可解释性提供了新的途径，有助于理解模型的决策过程。

📄 摘要（原文）

Supervised learning approaches to offline reinforcement learning, particularly those utilizing the Decision Transformer, have shown effectiveness in continuous environments and for sparse rewards. However, they often struggle with long-horizon tasks due to the high compounding error of auto-regressive models. To overcome this limitation, we go beyond next-token prediction and introduce Planning Tokens, which contain high-level, long time-scale information about the agent's future. Predicting dual time-scale tokens at regular intervals enables our model to use these long-horizon Planning Tokens as a form of implicit planning to guide its low-level policy and reduce compounding error. This architectural modification significantly enhances performance on long-horizon tasks, establishing a new state-of-the-art in complex D4RL environments. Additionally, we demonstrate that Planning Tokens improve the interpretability of the model's policy through the interpretable plan visualisations and attention map.