In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought

📄 arXiv: 2405.20692v1 📥 PDF

作者: Sili Huang, Jifeng Hu, Hechang Chen, Lichao Sun, Bo Yang

分类: cs.LG, cs.AI

发布日期: 2024-05-31

💡 一句话要点

提出In-context Decision Transformer，通过分层思维链加速离线强化学习。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 上下文强化学习 分层决策 Transformer 长时程任务 离线强化学习

📋 核心要点

  1. 现有上下文强化学习方法在处理长时程任务时，由于序列过长导致计算成本高昂。
  2. IDT通过模仿人类分层决策，将动作序列替换为高层决策序列，从而缩短序列长度。
  3. 实验表明，IDT在长时程任务上取得了显著的性能提升，并大幅降低了在线评估时间。

📝 摘要（中文）

本文提出了一种In-context Decision Transformer (IDT)，旨在通过提供任务提示，在离线强化学习中实现上下文学习，从而处理在线任务。现有方法将强化学习任务视为跨回合的序列预测问题，虽然能在试错中实现自我改进，但当跨回合序列随着任务范围增加时，计算成本很高。IDT受到人类决策分层结构的启发，重构序列，使其由高层决策而非与环境交互的底层动作组成。由于一个高层决策可以指导多步底层动作，IDT自然避免了过长的序列，更有效地解决在线任务。实验结果表明，IDT在长时程任务上优于现有的上下文强化学习方法。特别是在D4RL基准测试中，IDT的在线评估时间比基线快36倍，在Grid World基准测试中快27倍。

🔬 方法详解

问题定义：现有基于Transformer的上下文强化学习方法，例如Decision Transformer，在处理长时程任务时，需要处理非常长的序列，导致计算复杂度极高，难以进行在线部署和应用。这些方法将每一个时间步的动作都视为序列中的一个token，序列长度与任务的horizon直接相关。

核心思路：本文的核心思路是借鉴人类决策的分层结构，将低层次的动作序列抽象成高层次的决策序列。一个高层次的决策可以指导多个低层次动作的执行，从而有效地缩短序列长度，降低计算复杂度。通过这种方式，模型可以更快地进行推理和决策，更适用于在线任务。

技术框架：IDT的整体框架仍然基于Transformer架构，但其输入序列不再是原始的动作序列，而是高层次的决策序列。具体来说，IDT首先使用一个策略网络将状态映射到高层次的决策空间。然后，Transformer模型基于上下文信息（例如，历史状态、奖励等）预测下一个高层次决策。最后，一个低层次的控制器将高层次决策转化为具体的动作序列，与环境进行交互。

关键创新：IDT的关键创新在于引入了分层决策结构，将原始的动作序列替换为高层次的决策序列。这种分层结构有效地缩短了序列长度，降低了计算复杂度，使得模型能够更快地进行推理和决策。此外，IDT还能够更好地捕捉任务的长期依赖关系，从而提高策略的性能。

关键设计：IDT的关键设计包括：1) 高层次决策空间的定义：如何将状态映射到高层次的决策空间？论文中可能使用了聚类、编码器-解码器等方法。2) 低层次控制器的设计：如何将高层次决策转化为具体的动作序列？论文中可能使用了PID控制器、强化学习策略等方法。3) Transformer模型的结构和参数设置：例如，Transformer的层数、注意力头的数量、嵌入维度等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，IDT在D4RL和Grid World基准测试中取得了显著的性能提升。在D4RL基准测试中，IDT的在线评估时间比基线快36倍，在Grid World基准测试中快27倍。这表明IDT能够有效地降低计算复杂度，提高决策效率。此外，IDT在长时程任务上的性能也优于现有的上下文强化学习方法，表明IDT能够更好地捕捉任务的长期依赖关系。

🎯 应用场景

IDT具有广泛的应用前景，例如机器人控制、游戏AI、自动驾驶等领域。它可以应用于需要处理长时程任务的场景，例如，机器人需要在复杂环境中完成一系列动作才能达到目标，自动驾驶汽车需要在复杂的交通环境中进行导航和决策。IDT能够有效地降低计算复杂度，提高决策效率，从而使得这些应用成为可能。

📄 摘要（原文）

In-context learning is a promising approach for offline reinforcement learning (RL) to handle online tasks, which can be achieved by providing task prompts. Recent works demonstrated that in-context RL could emerge with self-improvement in a trial-and-error manner when treating RL tasks as an across-episodic sequential prediction problem. Despite the self-improvement not requiring gradient updates, current works still suffer from high computational costs when the across-episodic sequence increases with task horizons. To this end, we propose an In-context Decision Transformer (IDT) to achieve self-improvement in a high-level trial-and-error manner. Specifically, IDT is inspired by the efficient hierarchical structure of human decision-making and thus reconstructs the sequence to consist of high-level decisions instead of low-level actions that interact with environments. As one high-level decision can guide multi-step low-level actions, IDT naturally avoids excessively long sequences and solves online tasks more efficiently. Experimental results show that IDT achieves state-of-the-art in long-horizon tasks over current in-context RL methods. In particular, the online evaluation time of our IDT is \textbf{36$\times$} times faster than baselines in the D4RL benchmark and \textbf{27$\times$} times faster in the Grid World benchmark.