强化学习中的Transformer发展到哪一步了？清北联合发布TransformRL综述

2023-02-12 12:02

©作者 | Wenzhe Li等

来源 | 机器之心

强化学习（RL）为顺序决策提供了一种数学形式，深度强化学习（DRL）近年来也取得巨大进展。然而，样本效率问题阻碍了在现实世界中广泛应用深度强化学习方法。为了解决这个问题，一种有效的机制是在 DRL 框架中引入归纳偏置。

在深度强化学习中，函数逼近器是非常重要的。然而，与监督学习（SL）中的架构设计相比，DRL 中的架构设计问题仍然很少被研究。大多数关于 RL 架构的现有工作都是由监督学习 / 半监督学习社区推动的。

例如，在 DRL 中处理基于高维图像的输入，常见的做法是引入卷积神经网络（CNN）[LeCun et al., 1998; Mnih et al., 2015]；处理部分可观测性（partial observability）图像的常见做法则是引入递归神经网络（RNN） [Hochreiter and Schmidhuber, 1997; Hausknecht and Stone, 2015]。

近年来，Transformer 架构 [Vaswani et al., 2017] 展现出优于 CNN 和 RNN 的性能，成为越来越多 SL 任务中的学习范式 [Devlin et al., 2018; Dosovitskiy et al., 2020; Dong et al., 2018]。Transformer 架构支持对长程（long-range）依赖关系进行建模，并具有优异的可扩展性 [Khan et al., 2022]。受 SL 成功的启发，人们对将 Transformer 应用于强化学习产生了浓厚的兴趣，希望将 Transformer 的优势应用于 RL 领域。

Transformer 在 RL 中的使用可以追溯到 Zambaldi 等人 2018 年的一项研究，其中自注意力（self-attention）机制被用于结构化状态表征的关系推理。随后，许多研究人员寻求将自注意力应用于表征学习，以提取实体之间的关系，从而更好地进行策略学习 [Vinyals et al., 2019; Baker et al., 2019]。

除了利用 Transformer 进行表征学习，之前的工作还使用 Transformer 捕获多时序依赖，以处理部分可观测性问题 [Parisotto et al., 2020; Parisotto and Salakhutdinov, 2021]。离线 RL [Levine et al., 2020] 因其使用离线大规模数据集的能力而受到关注。受离线 RL 的启发，最近的研究表明，Transformer 结构可以直接作为顺序决策的模型 [Chen et al., 2021; Janner et al., 2021] ，并推广到多个任务和领域 [Lee et al., 2022; Carroll et al., 2022]。

实际上，在强化学习中使用 Transformer 做函数逼近器面临一些特殊的挑战，包括：

强化学习智能体（agent）的训练数据通常是当前策略的函数，这在学习 Transformer 的时候会导致不平稳性（non-stationarity）；
现有的 RL 算法通常对训练过程中的设计选择高度敏感，包括模型架构和模型容量 [Henderson et al., 2018]；
基于 Transformer 的架构经常受制于高性能计算和内存成本，这使得 RL 学习过程中的训练和推理都很昂贵。

例如，在用于视频游戏的 AI 中，样本生成的效率（在很大程度上影响训练性能）取决于 RL 策略网络和估值网络（value network）的计算成本 [Ye et al., 2020a; Berner et al., 2019]。

为了更好地推动强化学习领域发展，来自清华大学、北京大学、智源人工智能研究院和腾讯公司的研究者联合发表了一篇关于强化学习中 Transformer（即 TransformRL）的综述论文，归纳总结了当前的已有方法和面临的挑战，并讨论了未来的发展方向，作者认为 TransformRL 将在激发强化学习潜力方面发挥重要作用。

论文标题：

A Survey on Transformers in Reinforcement Learning

论文地址：

https://arxiv.org/pdf/2301.03044.pdf

论文的总体结构如下：

第 2 章介绍了 RL 和 Transformer 的背景知识，然后简要介绍了这两者是如何结合在一起的；
第 3 章描述了 RL 中网络架构的演变，以及长期以来 RL 中阻碍广泛探索 Transformer 架构的挑战；
第 4 章论文作者对 RL 中的 Transformer 进行了分类，并讨论了目前具有代表性的方法；
第 5 章总结并指出了未来潜在的研究方向。

核心内容从论文第 3 章开始，下面我们来看一下论文的主要内容。

RL中的网络架构

在介绍 TransformRL 的分类方法之前，论文回顾了 RL 中网络架构设计的早期进展，并总结了其存在的挑战。作者认为 Transformer 是一种先进的神经网络架构，将有助于深度强化学习（DRL）的发展。

函数逼近器的架构

自 Deep Q-Network [Mnih et al., 2015] 的开创性工作以来，人们为 DRL 智能体的网络架构做了许多努力。强化学习中网络架构的改进主要分为两类。

一类是设计新的结构，结合 RL 归纳偏置来降低训练策略或价值函数的难度。例如 [Wang et al. 2016] 提出决斗（dueling）网络架构，其中一个网络用于状态价值函数，另一个用于状态相关的行动优势函数（action advantage function），这种架构设计结合了归纳偏置。

另一类是研究常用的神经网络技术（如正则化、残差连接（skip connection）、批归一化）是否可以应用于 RL。例如，[Ota et al. 2020] 发现在使用在线特征提取器增强状态表征的同时增加输入维度，会有助于提高 DRL 算法的性能和样本效率。[Sinha et al. 2020] 为 DRL 智能体提出了一种深度密集架构，使用残差连接进行有效学习，并使用归纳偏置来缓解数据处理不平等问题。

[Ota et al. 2021] 使用 DenseNet [Huang et al., 2017] 和解耦表征学习来改善大型网络的信息流和梯度。最近，由于 Transformer 的优越性能，研究人员尝试将 Transformer 架构应用于策略优化算法，但发现普通的 Transformer 设计无法在 RL 任务中达到理想的性能 [Parisotto et al., 2020]。

面临的挑战

虽然过去几年基于 Transformer 的架构在 SL 领域取得了诸多进展，但将 Transformer 应用于 RL 并不简单。实际上，这存在多个特有的挑战。

从 RL 的角度看，许多研究指出现有的 RL 算法对深度神经网络的架构非常敏感 [Henderson et al., 2018; Engstrom et al., 2019; Andrychowicz et al., 2020]。首先，RL 中数据收集和策略优化之间的范式交替导致训练的不平稳。其次，RL 算法通常对训练过程中的设计选择高度敏感。[Emmons et al. 2021] 证明仔细选择模型架构和正则化对于 DRL 智能体的性能至关重要。

从 Transformer 的角度看，基于 Transformer 的架构存在内存占用大、延迟高的问题，这阻碍了它们的高效部署和推理。最近，许多研究围绕原始 Transformer 架构对计算和内存效率进行改进，但其中大部分工作都集中在 SL 领域。

在 RL 领域，Parisotto 和 Salakhutdinov 提出把基于大容量 Transformer 的学习器模型转变为小容量的 actor 模型，以避免 Transformer 的高推理延迟。然而，这种方法在内存和计算方面仍然很昂贵。目前，RL 社区还未充分探索高效或轻量级的 Transformer。

强化学习中的Transformer

尽管 Transformer 已成为大多数监督学习研究的基础模型，但由于前述挑战，它在 RL 社区长期未得到广泛应用。实际上，TransformRL 的大多数早期尝试都将 Transformer 用于状态表征学习或提供记忆信息（memory information），同时仍然将标准 RL 算法用于智能体学习，例如时序差分学习和策略优化。

因此，尽管引入 Transformer 作为函数逼近器，这些方法仍然受到传统 RL 框架的挑战。直到最近，离线 RL 使得从大规模离线数据中学习最优策略成为可能。受离线 RL 的启发，最近的工作进一步将 RL 问题视作固定经验的条件序列建模问题。这样做有助于绕过传统 RL 中的 bootstrapping error 挑战，从而使 Transformer 架构释放其强大的顺序建模能力。

论文回顾了 TransformRL 的进展，并按分类展示现有方法。作者将现有方法分为四类：表征学习、模型学习、顺序决策和通用智能体。图 2 显示相关分类的概览。

用于表征学习的Transformer

考虑到 RL 任务的顺序性，使用 Transformer 编码器模块是合理的。事实上，RL 任务中的各种序列都需要处理，例如局部 per-timestep 序列（多实体序列 [Vinyals et al., 2019; Baker et al., 2019]、多智能体序列 [Wen et al., 2022]）、时序序列（[Parisotto et al., 2020; Banino et al., 2021]）等。

局部per-timestep序列的编码器

这种方法早期显著的成功体现在使用 Transformer 处理智能体观察到的可变数量实体的复杂信息。[Zambaldi et al. 2018a] 首先提出用多头点积注意力捕获结构化观察的关系推理，随后 AlphaStar [Vinyals et al., 2019] 实现了在具有挑战性的多智能体环境（星际争霸 II）中处理多实体观察。在这种称为实体 Transformer 的机制中，观察结果以如下形式编码：

其中 e_i 代表智能体对实体 i 的观察，要么直接从整个观察中切片，要么由实体 tokenizer 给出。

一些后续工作丰富了实体 Transformer 机制。[Hu et al. 2020] 提出了一种兼容的解耦策略，以明确地将行动与各种实体相关联，并利用注意力机制进行策略解释。为了实现具有挑战性的 one-shot 视觉模仿，Dasari 和 Gupta [2021] 使用 Transformer 来学习专注于特定任务元素的表征。

类似于分散在观察中的实体，一些研究利用 Transformer 来处理其他局部的 per-timestep 序列。Tang 和 Ha [2021] 利用 Transformer 的注意力机制来处理感知序列并构建一个置换不变性输入策略。在不兼容的多任务 RL 中，[Kurin et al., 2020] 提出使用 Transformer 来提取形态域知识。

时序编码器

同时，用 Transformer 处理时序序列也是合理的。时序编码器被用作存储架构，

其中 o_t 表示智能体在时间 t 的观察，Emb_0:t 表示从初始观察到当前观察的历史观察的嵌入。

在早期的工作中，[Mishra et al. 2018] 无法使用 vanilla Transformer 处理时序序列，并且发现它在某些特定任务中甚至比随机策略表现更差。Gated Transformer-XL (GTrXL) [Parisotto et al., 2020] 是第一个使用 Transformer 作为存储架构来处理轨迹的有效方案。

GTrXL 通过 Identity Map Reordering 修改 Transformer-XL 架构 [Dai et al., 2019]，以提供从时序输入到 Transformer 输出的「skip」路径，这可能有助于形成从一开始就稳定的训练过程。

[Loynd et al. 2020] 提出了一种用于长期依赖的记忆向量快捷机制，[Irie et al. 2021] 将线性 Transformer 与快速加权编程器（Fast Weight Programmer）相结合以获得更好的性能。[Melo 2022] 提出使用自注意力机制来模拟基于存储的元 RL 的存储恢复。

虽然随着存储的增长和参数规模的扩大，Transformer 的性能优于 LSTM/RNN，但它在 RL 上的数据效率不佳。后续工作利用一些辅助自监督任务来促进学习 [Banino et al., 2021] 或使用预训练的 Transformer 架构作为时序编码器 [Li et al., 2022; Fan et al.，2022]。

用于模型学习的Transformer

除了使用 Transformer 作为序列嵌入的编码器，Transformer 架构还在基于模型的算法中作为环境模型的 backbone。与以单步观察和行动为条件的预测不同，Transformer 使环境模型能够预测以一定长度的历史信息为条件的变换（transition）。

实际上，Dreamer 及其后续算法的成功 [Hafner et al., 2020, 2021; Seo et al., 2022] 已经在一些部分可观察的环境或需要记忆机制的任务中证明了基于历史信息的世界模型的优点。以历史信息为条件的世界模型由一个捕获抽象信息的观察编码器和一个学习潜在空间中变换的变换模型组成。

已有研究使用 Transformer 架构而不是 RNN 来构建基于历史的世界模型。[Chen et al. 2022] 用基于 Transformer 的模型 TSSM（Transformer State-Space Model）替换 Dreamer 中基于 RNN 的循环状态空间模型（RSSM）。

IRIS（Imagination with autoRegression over an Inner Speech）[Micheli et al., 2022] 通过对 rollout 经验的自回归学习来学习基于 Transformer 的世界模型，而没有像 Dreamer 那样的 KL 平衡，并在 Atari [Bellemare et al., 2013] 上取得了很好的结果。

此外，还有研究尝试用基于 Transformer 的世界模型做规划。[Ozair et al. 2021] 验证了使用 Transformer 变换模型进行规划来完成随机任务的有效性。[Sun et al. 2022] 提出了一种以目标为条件的 Transformer 变换模型，该模型在程序任务的视觉基础规划中是很有效的。

RNN 和 Transformer 都适合学习基于历史信息的世界模型。然而，[Micheli et al. 2022] 发现与 Dreamer 相比，Transformer 架构是数据效率更高的世界模型。TSSM 的实验结果表明，Transformer 架构在需要长期记忆的任务中表现出众。

用于顺序决策的Transformer

除了融入到传统 RL 算法中作为高性能架构以外，Transformer 还可以直接用作顺序决策模型。这是因为可以把 RL 看作一个条件序列建模问题：生成可以产生高回报的行动序列。

鉴于 Transformer 在序列预测方面的优异准确性，Bootstrapped Transformer (BooT) [Wang et al., 2022] 提出通过 bootstrap Transformer 来生成数据，同时优化数据以进行顺序决策。Bootstrapping Transformer 用于数据增强可以扩大离线数据集的数量和覆盖范围，从而提升性能。具体地说，BooT 比较了不同的数据生成方案和 bootstraping 方案，以分析 BooT 如何助力策略学习。结果表明，它可以生成与底层 MDP 一致的数据，而无需额外的约束。

用于通用智能体的Transformer

Decision Transformer 已经在离线数据的各种任务中发挥巨大作用，有研究者开始考虑 Transformer 是否可以像 CV 和 NLP 领域那样让通用智能体解决多个不同任务或问题。

泛化到多个任务

一些研究者借鉴了 CV 和 NLP 中对大规模数据集进行预训练的思想，并尝试从大规模多任务数据集中抽象出通用策略。Multi-Game Decision Transformer (MGDT) [Lee et al., 2022] 是 DT 的一个变体，该模型在由专家和非专家数据组成的多样化数据集上学习 DT，并使用一组参数在多个 Atari 游戏上实现接近人类的水平。

为了在非专家级数据集上获得专家级的表现，MGDT 设计了专家行动推理机制，从 return-to-go 的先验分布计算专家级的 return-to-go 后验分布并根据贝叶斯公式预设专家级 return-to-go 的概率。

同样，Switch Trajectory Transformer (SwitchTT) [Lin et al., 2022] 是 TT 的多任务扩展，利用稀疏激活模型，将 FFN 层替换为混合专家层，以实现高效的多任务离线学习。此外，SwitchTT 还采用分布式 trajectory 值估计器对值估计的不确定性进行建模。依靠这两个增强功能，SwitchTT 在性能和训练速度方面都比 TT 提升了很多。MGDT 和 SwitchTT 利用从多个任务和各种性能级策略中收集的经验来学习通用策略。然而，构建大规模的多任务数据集并非易事。

与 CV 和 NLP 中的大规模数据集通常使用来自互联网的海量数据和简单的人工标记不同，RL 中的顺序决策数据总是缺少行动信息，并且不易标记。因此，[Baker et al. 2022] 提出了一种半监督方案，利用没有行动信息的大规模在线数据，学习基于 Transformer 的逆向动态模型（IDM）。该模型利用对过去和未来的观察来预测行动信息，能够标记大量在线视频数据。IDM 是在包含手动标记行动的小型数据集上学习的，并且足够准确。

NLP 的许多已有工作证明了 prompt 在适应新任务方面的有效性，一些工作利用基于 DT 方法的 prompt 技术来实现快速适应。Prompt-based Decision Transformer (Prompt-DT) [Xu et al., 2022] 从少样本（few-shot）演示数据集中采样一系列变换作为 prompt，并将少样本策略泛化到离线元 RL 任务上。

[Reed et al. 2022] 进一步利用基于 prompt 的架构，通过在涵盖自然语言、图像、时间决策和多模态数据的超大规模数据集上进行自回归序列建模来学习通用智能体（Gato）。Gato 能够执行来自不同领域的一系列任务，包括文本生成和决策。

[Laskin et al. 2022] 提出了算法蒸馏 (AD)，以在单任务 RL 算法学习过程的 across-episode 序列上训练 Transformer。因此，即使在新任务中，Transformer 也可以学会在自回归生成过程中逐步改进其策略。

泛化到更广泛领域

除了泛化到多个任务，Transformer 还是一个强大的「通用」模型，可以用于与顺序决策相关的一系列领域。受 NLP 中掩码语言建模（masked language modeling）[Devlin et al., 2018] 技术的启发，[Carroll et al. 2022] 提出了 Uni [MASK]，它将各种常用研究领域统一为 mask 推理问题，包括行为克隆、离线 RL、GCRL、过去 / 未来推理和动态预测。

Uni [MASK] 比较了不同的 mask 方案，包括任务特定 mask、随机 mask 和微调变体。结果表明，用随机 mask 训练的单个 Transformer 可以解决任意推理任务。

此外，[Reid et al. 2022] 发现，使用在语言数据集或包含语言模态的多模态数据集上预训练的 Transformer 对 DT 进行微调是有益的。这表明，即使是来自非 RL 领域的知识也可以通过 transformer 进行 RL 训练。

感兴趣的读者可以阅读论文原文，了解更多研究细节。

更多阅读