论文理解【Offline RL】——【DT】Decision Transformer: Reinforcement Learning via Sequence Modeling

• 标题：Decision Transformer: Reinforcement Learning via Sequence Modeling
• 文章链接：Decision Transformer: Reinforcement Learning via Sequence Modeling
• 代码：kzl/decision-transformer
• 官方主页：Decision Transformer: Reinforcement Learning via Sequence Modeling
• 发表：NIPS 2021
• 领域：离线强化学习（offline/batch RL）—— Transformer Based / Hindsight 监督思想

---

• 摘要：我们提出了一个将强化学习（RL）抽象为序列建模问题的框架。这使得我们能够利用 transformer 模型的简单性（simplicity）和可扩展性（scalability），以及 GPT-x 和 BERT 等语言建模方面的相关进展。特别地，我们提出了 Decision Transformer（DT），一个将 RL 问题转换为条件序列建模的架构。与之前基于值函数或计算策略梯度的 RL 方法不同，DT 只是通过 causally masked Transformer 来输出最优操作。通过训练一个以期望 return、过去的状态和动作作为条件的自回归模型，DT 可以产生实现期望回报的未来行动。尽管 DT 很简单，但它在 Atari、OpenAI gym 和 Key-to-Door 任务上媲美或超过了 SOTA 的 model-free offline RL Baseline 的性能。

1.背景

• Offline RL 是这样一种问题设定：Learner 可以获取由一批 episodes 或 transitions 构成的固定交互数据集，要求 Learner 直接利用它训练得到一个好的策略，而且禁止 Learner 和环境进行任何交互，示意图如下
  
  关于 Offline RL 的详细介绍，请参考 Offline/Batch RL简介
• Offline RL 是近年来很火的一个方向，下图显示了 2019 年以来该领域的重要工作，本文出现在 21 年，和同期的 TT 一样是最先纯粹使用 Transformer 模型解 Offline RL 问题的文章，可能也是最先把 Offline RL 当做序列建模问题来解的文章
  
• 本文使用的 Transformer 模型是 GPT，这是 Transformer 的 Decoder 部分，可以作为标准语言模型使用。所谓 “标准语言模型”，就是它会吃进去句尾的若干个词（token）然后预测下一个词是什么，这个过程可以反复进行从而实现文本生成的效果。GPT 相对传统语言模型的优势在于，GPT 模型结构中的 masked self attention 层利用 $Q,K,V$ 矩阵在各个 token 间建立起了显式连接，从而解决了 RNN/LSTM 等传统序列模型的长跨度信息的遗忘问题。详细说明可以参考 快速串联 RNN / LSTM / Attention / transformer / BERT / GPT

2. 本文方法

2.1 思想

• Transformer 作为近年非常火的模型，有很多人想要把它融合到 Offline RL 领域。过去的一些方法把将 Transformer 融入到传统 RL 的网络结构中，本文则是想要验证是否能单纯使用 Transformer 模型来解决 Offline RL 问题，也就是说要完全丢掉基于 Bellman 等式的 TD Learning 那一套，只把 Offline RL 作为序列预测问题来解。这有若干好处

  1. 抛弃了 Bootstrap，不会出现传统 RL 的致命三要素问题。关于 “致命三要素问题”，可以参考 Why is there a Deadly Triad issue and how to handle it ?
  2. 不再需要引入折扣系数 $\gamma$，这种随时间跨度增加而增大的损耗可能导致不良的短视行为
  3. Transformer 结构可通过 self attention 执行大跨度的 credit assignment，而传统 RL 只能利用 Bellman 迭代缓慢地传播奖励，还容易产生 “干扰” 信号。这一特性使 Transformer 模型在稀疏奖励或存在误导奖励的情况下仍然有效工作
  4. 相比传统 RL 通过仔细评估价值函数拟合一个非常窄的策略分布，借助 Transformer 这类监督学习框架可以建模更多样的行为分布，有利于更好地进行泛化
  5. 从问题形式上看，Offline RL 数据中的轨迹序列恰好也是序列数据，正好适用于 Transformer 类先进的序列建模模型

• 标准 GPT 只是一个模仿数据集序列进行生成的模型，如果直接把它用到离线数据集上，其实就是做了个 “轨迹级” 的 BC 模仿学习，而我们都知道在 Offline 数据集上直接这样做会由于次优轨迹的存在而导致性能不佳。作者借鉴了 UDRL 论文的处理方法，将预期得到的一些目标也作为附加输入来做监督训练，这样训练完后就能泛化到 “面对近期的历史轨迹，为了实现目标，应该采取什么 action”，从而允许我们通过给定目标信息来控制策略质量。UDRL 中目标信息设计为 “在一定 horizon 内收到一定 return”，如下图所示
  
  此想法也很像 Hindsight Experience Replay（HER）论文的 Hindsight 思想，这种根据 Offline 轨迹重实际实现的目标进行重新标注的思路其实和 “自监督学习” 非常类似，可以无缝衔接到 GPT 这种自监督预训练模型上。最近的文章 论文理解【Offline RL】——【RvS】What is Essential for Offline RL via Supervised Learning? 把这类方法定义为 “RvS” 类方法，并进行了一些分析

2.2 方法

• 标准设定下，offline 数据集中的一条轨迹形如
  $$\tau =(s_1,a_1,r_2,s_2,a_2,r_3,...,r_T,s_T,a_T)$$ 作者首先将 $\tau$ “重标记” 为以下形式
  $$\tau =({\hat{R}}_1,s_1,a_1,{\hat{R}}_2,s_2,a_2,...,{\hat{R}}_T,s_T,a_T)$$ 其中 $s,a$ 都没变，return-to-go $\hat{R}$ 定义为轨迹在 $t$ 时刻之后的所有 reward 之和（从 $t$ 时刻开始的无折扣 return）
  $$R_t=\sum _{t^{\prime }=t}^T{r}_{t^{\prime }}$$

  下图给出了一个 return-to-go 的示意
  
  这个 MDP 中每走一步会得到 -1 的 reward，agent 从随机位置出发按随机策略行动直到到达 goal 结束轨迹。每个圈边上的数字即是该轨迹中对应状态下的 return-to-go。在训练完成后的 generation 阶段，给定起点位置后，只须在每一步决策时以训练数据中此位置收到的最大 return-to-go 作为条件选择动作，就能实现 offline 数据集中次优轨迹的拼接。这里只是一个示意，作者提出的方法中对 return-to-go 条件的选取细节有所不同

• 本文的设计的 DT 模型结构如下
  
  从下往上看

  1. 处理后的轨迹按顺序输入，任意 $\hat{R},s,a$ 都是一个 token
  2. 蓝绿黄的小方块代表嵌入层，这里 $\hat{R},s,a$ 先分别使用三个嵌入层进行嵌入，然后再直接加上位置编码完成嵌入。

     这里有几个细节

     1. 同一时刻的一组 ${\hat{R}}_t,s_t,a_t$ 使用相同的位置编码
     2. 另外这里嵌入层不是像传统 GPT 那样使用 nn.Embedding 配合 vocab_size 实现，而是直接用 nn.Linear 实现，这可能是因为很多 RL 环境状态、动作空间过大，或者直接就是连续的环境。另外，如果是图像形式的状态输入（如 Atria 环境），则用 nn.Conv2d 进行嵌入
     3. state token 嵌入之前首先要经过状态归一化，即先用数据集中所有 state 计算均值和标准差，然后将每个 state 减去均值除以标准差，使得所有 state 宏观上呈现 0 均值 1 标准差的分布

  3. 接下来所谓的 “causal transformer” 其实就是 GPT 这种带有 mask self attention 结构的 Transformer 结构，不过 DT 里这个 mask 是按时刻 $t$ 一组一组地遮盖的
  4. 再上面的橘黄色块是一个就是线性层，把 GPT 结构得到的动态 token 嵌入向量转变为实际的 action。具体而言，如果环境的动作空间是离散的，这里就用线性层调整下维度然后加 softmax 优化交叉熵损失；如果动作空间是连续的，这里就用线性层直接输出动作然后优化 L2 损失

• 上图中的红色虚线显示了推断阶段的 autoregress 过程，而训练阶段使用的是 teacher-forcing。具体而言

  1. 训练时首先从 offline 数据集中采样一段连续 $k$ step 的轨迹，设这组轨迹的最后时刻为 $t$，将 $t$ 时刻之前的轨迹序列
     $$\tau =({\hat{R}}_{t-k},s_{t-k},a_{t-k},...,{\hat{R}}_{t-1},s_{t-1},a_{t-1})$$ 送入 DT 模型，依次预测出 ${\hat{R}}_t,s_t,a_t$，然后只拿出其中预测的动作 $a_t$ 和轨迹中真实的 $a_t$ 计算损失来优化 DT
  2. 测试阶段，人为选定一个目标的 return-to-go（比如 offline 数据集中最高的轨迹 return），从初始状态开始不断预测 action，配合环境交互实现自回归。注意每一步交互得到真实 reward 后，就从初始设定的 return-to-go 中减去这个值，这样就自然地得到下一时刻的 return-to-go，不断这样操作，DT 就能作为策略使用了

• 下面给出伪代码
  

3. 实验

• 作者主要在 Atria 的离散环境和 gym 的经典连续控制环境中进行了测试，前者要求 long-term credit assignment 的能力，后者要求精细控制能力。另外还在一个 Key-To-Door 环境进行附加测试，该网格世界环境要求 agent 在必须第一个房间拿到钥匙并在第三个房间走到门才能得到 reward，奖励非常稀疏，是对 long-term credit assignment 能力要求非常高的特化环境。对比方法主要是一些 RL-based 的 TD-Learning 方法以及 BC，总体性能如下
  
  可见 DT 媲美或超过过去的最好方法

3.1 Atari

• Atari 环境的难点主要在于高维的视觉输入以及延迟 reward 导致的信用分配困难。具体而言，作者在各个测试环境首先训练一个 DQN，再从其 replay-buffer 中随机采样 1% 的轨迹（约50万 transition）作为 offline 数据集来训练。我们知道 GPT 这类 Transformer 模型由于 self attention 的计算量无法输入太长的序列，因此这里 DT 输入的长度依环境不同被限制为 50 或 30，四个环境下实验效果如下
  
  这里的 BC 就是 DT 不加 return-to-go 条件，直接像普通 GPT 那样工作实现的。四个环境中三个表现不错

3.2 gym

• 这里除了 Hopper、walker 几个经典环境，还加了一个 goal-condition 的 reacher环境，使用 D4RL 数据。对比方法都是 TD learning based 方法，CQL 作为 SOTA，结果如下
  
  可见 DT 基本都是最好的

4. 讨论

• 原文做了很多讨论，这里简单搬过来
  1. DT 是仅仅在一个子集上做 BC 吗：不是，DT 可以通过数据集中的所有轨迹来改进泛化性，即使这些轨迹的真实 return 和目标不同。具体而言，如果数据充足，DT 和用最优的部分数据做 BC（称为 %BC）性能差不多；而数据不足时，DT 远远好于这种 %BC 的表现
  2. DT 对 return 分布的建模效果如何：下图显示了对于不同的目标 return，agent 在评估过程中实现的平均抽样 return
     
     可见每个任务上期望的目标 return 和真实观察到的 return 都是高度相关的。此外，在一些 Atria 任务上，我们可以用比数据集中存在的最大 return 更高的 return 条件来提示 DT，这表明 DT 有时能够外推
  3. DT 是否执行了有效的长期信用分配：是的。作者在 Key-To-Door 这个环境进行测试来验证这一点，该网格世界环境要求 agent 在必须第一个房间拿到钥匙，并在带着钥匙走到第三个房间的门才能得到 binary reward，奖励非常稀疏。这个问题对于信用分配来说是困难的，因为信用必须从轨迹的开始传播到结束，跳过在中间采取的行动。作者使用随机策略生成轨迹组成 offline 数据集进行训练，并使用整个轨迹的长度作为上下文，而不是像在其他环境中那样有一个固定长度的上下文窗口。效果如下
     
     可见 DT 和 %BC 效果不错，基于 Bellman 等式的 CQL 无法很好的进行长期信用分配
  4. DT 能在稀疏的奖励设置中成为准确的 critic 吗：可以。这里作者还是用 Key-To-Door 这个环境来进行说明，这里 DT 经过修改，可以同时输出预测的 action 和 return-to-go。
     
     如左图所示，作者发现 DT 可以根据发生的事件不断更新获得 reward 的概率（因为 binary reward，所以预测的 return 可以理解为成功概率）。此外如右图，作者发现 DT 关注到了关键事件（拿起钥匙或到达门）这种将状态和奖励关联的能力允许 DT 形成准确的价值预测
  5. DT 在稀疏奖励环境表现良好吗：表现良好。由于 DT 对 reward 密度做出了最小假设，它不像 TD 类方法那样必须要稠密奖励才能良好工作。作者使用 D4RL 的 delay-reward 版本来验证这一点，这时轨迹中的所有 reward 都仅在轨迹结束时给出。结果如下
     可见 CQL 这种 TD Learning 方法性能严重下降，DT 则没有太多影响，说明 DT 对于延迟奖励更鲁棒