Transformer + RL(1)：是强化学习魔高一尺，还是 Transformer 道高一丈

引言：为什么 Transformer 在 RL 上不 work

Transformer 最早发源于自然语言处理领域（NLP），以一个非常惊人的标题 “Attention is All You Need” 催生了无数后续工作，而 Transformer 最近的一些变体已经在 NLP 领域证明了其可以在大多数任务中取代并超越经典模型 LSTM/GRU。 另一方面，其他领域的研究者也在尝试 Transformer 是否可以将其优势和潜力推广出来，最近在计算机视觉领域的一些例子便是很好的印证。但是在强化学习（RL）领域， Transformer 架构仍未被广泛采用，普通的 MLP 网络结构已经可以解决很多常见的决策问题，带有部分可观测性质（POMDP）的决策也依然以 RNN（比如GRU，LSTM）为主流。究竟是什么制约了 Transformer 在强化学习领域的应用呢？是其固有的一些弱点：比如训练调参比较复杂，需要大量的数据和计算资源才能收敛；还是其他方面的原因？本文将尝试反思和探索 Transformer + RL 的最佳方案。
通常情况下，我们需要去使用一些技巧来稳定 Transformer的训练或者提高其性能，例如
  • 复杂的学习率调整策略（如线性预热或余弦衰减）
  • 复杂的梯度裁剪技巧
  • dropout 的各种变体
  • 特定的权重初始化方案
  • 特殊的 normalization 网络层和使用位置（pre-norm 和 post-norm）
  • 位置编码的各种技巧
  • 序列数据处理和特征工程的各种技巧和等等。
  1. 但除了上述在监督学习中训练 Transformer 遇到的困难外，强化学习还具有监督学习所没有的一些挑战，其中包括
  • 长期的时间依赖性
  • 环境奖励的稀疏性
  • 平衡强化学习的探索和利用
早在2018年，Mishra 等人就进行了一些探索在强化学习中使用 Transformer 的可行性的初步实验，这些环境包括简单的 bandit 任务和表格版马尔科夫决策过程（MDP）。然而，他们的结论是，Transformer 在这两种环境中学习到的策略往往仅能跟随机策略的性能相当，进而得出了 Transformer 这样的架构并不适合去处理 RL 中的时序信息的结论，具体训练中不收敛的可能原因有：
  • Sample efficiency: 在 RL 的一大分类—— actor-critic 算法中，actor和critic是交替训练的。这就要求他们各自的模型需要在较少的样本训练中迅速适应对方的变化。然而，众所周知，Transformer 需要大量的样本来学习，因此自然而然面临 sample efficiency 的问题
  • Implicit sequential information: 经典强化学习任务的主要的目标是去学习一个 MDP 环境中的最优策略，这表示着时序信息在预测下一个行动中起着非常重要的作用。在常规的 RL 场景中，由于智能体根据其自然的时间顺序逐一接收单帧的 observation，并不会产生什么问题。然而，Transformer在每一步都会收到一连串的 observations，而它们的顺序信息通常是由位置编码给出的，即顺序信息不是明确给出的，而是必须由 Transformer 自己来学习的。
 

初探：第一个成功案例——GTrXL

直到2019年，随着 Parisotto 等人的工作，Transformer 在RL方面有了显著的突破。 他们引入了一种被称为 GTrXL（Gated Transformer-XL[2]）的新架构，其核心改进主要有以下几点：
  • Transformer-XL：Transformer-XL[1] 提出了一种特殊的架构，相比常规 Transformer 能够在不破坏时间连贯性的情况下，使其能够学习超过固定的长度的依赖，这使得它可以利用当前的输入轨迹加上过去的轨迹来进行预测。不仅如此，Transformer-XL 还引入了一种新的相对位置编码方案，不仅能够学习到更长期的依赖，而且还能解决上下文碎片化的问题，其核心逻辑如下图所示（绿色的连接便是其优势）。
  • Identity Map Reordering(IMR)[2]：下图左边是常规 Transformer 的一个 block，而中间部分则是采用了 IMR 的 block，两者的差别就在于调整了 Layer-Norm 的位置。为什么要这么做呢，让我们来首先分析智能体需要学习到什么：从时序角度来说，我们可以把智能体需要学习到的动作分为两类——反应型动作（reaction action）和记忆型动作（memory action），其中前者满足标准的马尔科夫性所以又可称产生这种动作的策略为马尔科夫策略。但一般来说，智能体需要先学会反应，然后基于记忆的行为才能够被有效学习到。例如，一个智能体在学习如何记住它所走过的地方之前，需要先学习如何行走。 因此，从左边调整到中间这种设计，即重新排序 Layer-Norm的位置，这种改动的关键好处是：可以实现从第一层的 Transformer 到最后一层的恒等映射（Identity Map）。在训练早期，主干网络的输出值接近于零，当前帧的状态编码就会直接通过残差连接传递到决策输出端，使得智能体在训练开始时就能很方便的学习反应型动作，而忽略来自 Transformer-XL 的对于过去 observation 的记忆（因为这时记忆经常还不精确）。当反应型动作学习的比较完全后，可以再去转向更新长期记忆相关的网络模块。
  • Gating Transformer (GT)[2]：进一步，为了使得模型能够学会控制来自残差连接和主体网络的信息流动，从而提升残差分支和主干分支的组合表示能力（即反应型动作和记忆型动作和灵活运用），GTrXL 中用 Gating Layer 取代了默认的残差链接（关于 Gating Layer 的设计问题，可以参考原论文的对比实验分析，最终 GRU Gate 综合表现最好）。
  • Gated Identity Initialization (GII)[2]：Identity Map Reordering 有助于网络初期的训练，Gating Layer可以提升网络表示能力的上限，那如何充分发挥两者的优势呢？其实可以使用对 Gate 的一种特殊初始化来处理，将 Gate 中的 bias 项初始化为一个较大的负数，这样可以明确地将门控机制初始化为接近恒等映射的状态，从而得而兼之。
通过上述对 Transformer 架构的修改，GTrXL 不仅在 DMLab-30 等需要长期记忆的环境中对比 LSTM有了非常大的提升，而且在反应型动作占主导地位的环境中也能保持很高的性能，成功综合了各类长短期记忆的需求。进一步地，这种网络结构可以与许多策略相结合，比如V-MPO（Actor-Critic），R2D2（Value-Based），充分让 Transformer 更适应强化学习的优化过程，并利用 Transformer 在宽度和深度易于拓展的优点，在更大的数据集，更大型的，复杂的环境中发挥作用。

进阶：更强大的 Transformer 拓展

GTrXL 为强化学习中的 Transformer 使用设定了一个基线，后续又出现了一些改进它的近期工作，其中一些比较重要的工作包括：
  • Adaptive Transformer in RL[3]: Kumar等人在2020年发表了这一篇文章，大大增加了 GTrXL 的memory size，并通过增加 Adaptive Attention Span改进了attention的效率。他们的算法为每个模块学习一个optimal attention span，从而使 attention 能够有选择地使用更大的内存块来关注过去的时间段。具体来说，lower layers 学习较短的 attention span 来捕捉局部信息，而 upper layers 则利用大的跨度来捕捉全局信息，即只有在需要时才关注长期的记忆，减少了不必要的 attention 计算开销，最终可以在大幅提高训练速度和减少资源消耗的前提下，保持跟原版 GTrXL 相近的性能。
  • Contrastive BERT (CoBERL)[4]: 受双向 Transformer 编码器 BERT 中通过掩码预测设计自监督训练任务想法的启发，CoBERL 进一步将这种设计思路推向了极致。它利用 Trainable Gate 将 GTrXL 和 LSTM 结合起来，并结合 BERT 中的双向掩码预测自监督任务和对比学习方法，让这个组合版模型可以学到更好更一致的 RL Transformer representations。这种设计在处理高维输入时优势尤为明显，大大提高了强化学习的数据利用效率，在各种类型的任务中取得了很好的效果。

番外：Transformer + 离线强化学习

在 Transformer 架构与在线强化学习策略相结合取得的杰出成果之后，Transformer 也开始挑战另一热点话题——离线强化学习。在离线强化学习的设定中，智能体不能与环境交互，只能从一个有限数据集（包含各种水平和类型的数据）中进行学习，而传统的强化学习方法，由于难以处理 Out Of Distribution（OOD）数据的价值估计问题，训练效果往往很差，需要针对离线强化学习设计新的算法模块。但从2021年来，研究者们发现 Transformer 在这类问题上有着巨大的潜力：
  • Decision Transformer[5]: 离线强化学习与 Transformer 结合的里程碑之作，启发了后续通用决策模型的研究思路，具体分析可以参考我们之前的博客—— Decision Transformer 前沿追踪 https://zhuanlan.zhihu.com/p/551812218
  • Trajectory Transformer[6]: 作为 Decision Transformer 的同期工作，Trajectory Transformer 也将离线强化学习问题看作一个序列建模问题，但它的不同之处在于训练方式，即专注于轨迹上的模型分布建模。首先将状态和动作沿着每个维度独立地离散化，以便于后续的处理。整个模型的输入由状态、动作和奖励这样的轨迹数据组成，并在 beam search 技术的支持下自回归性地进行预测。Trajectory Transformer 还可以通过在每一组之间进行微小的修改，是得模型可以用于更多场景，例如模仿学习、目标条件强化学习。

总结：Transformer is not all you need but has much potential

从上述工作我们可以发现，合理应用 Transformer 是可以在强化学习领域取得惊人效果的：
  • 在反应型决策环境中，Transformer + RL可以与标准策略梯度或值函数算法的性能相匹配。
  • 在那些需要长期记忆的环境中，使用 Transformer 网络的 RL 智能体可以大幅优于常规的 RNN 模型，利用状态和动作在其轨迹中的时序依赖性，捕捉各种长期时序依赖，学习一个更好的表征来去预测下一步的决策行为。
  • 在离线强化学习问题中，Transformer 强大的序列建模能力可以很好地挖掘有限数据集中的最优策略。
但另一方面，使用基于 Transformer 架构的 RL 智能体仍然有一些典型的、众所周知的缺点：
  • Transformer 目前的训练速度较慢，资源消耗较大。
  • 当时序上下文窗口过大时，需要相对较长的时间去计算动作，实时控制场景中很难部署应用。
  • 当需要处理纯粹的反应型环境，且数据维度较低时，使用 Transformer 架构的优势并不明显，因为它在性能和效率方面都不会带来任何好处。

Reference

  1. Zihang Dai, et al. "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context." arXiv:1901.02860v3, 2019.
  1. Emilio Parisotto, et al."Stabilizing Transformers for Reinforcement Learning." Proceedings of the 37th International Conference on Machine Learning, PMLR 119:7487-7498, 2020.
  1. Shakti Kumar, Jerrod Parker, Panteha Naderian. "Adaptive Transformers in RL." arXiv:2004.03761v1, 2020.
  1. Andrea Banino, et al. "CoBERL: Contrastive BERT for Reinforcement Learning." arXiv:2107.05431v2, 2021
  1. Chen, Lili, et al. "Decision transformer: Reinforcement learning via sequence modeling." Advances in neural information processing systems 34, 2021.
  1. Michael Janner, Qiyang Li, Sergey Levine. "Offline Reinforcement Learning as One Big Sequence Modeling Problem." Advances in neural information processing systems 34, 2021.
posted @ 2022-08-30 21:49  OpenDILab  阅读(615)  评论(0)    收藏  举报