用于表征学习的 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 训练。
感兴趣的读者可以阅读论文原文,了解更多研究细节。