Stabilizing Transformers for Reinforcement Learning

By Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury and 3 more. 2019

486 citations recorded by the source corpus.

Owing to their ability to both effectively integrate information over long time horizons and scale to massive amounts of data, self-attention architectures have recently shown breakthrough success in natural language processing (NLP), achieving state-of-the-art results in domains such as language modeling and machine translation. Harnessing the transformer's ability to process long time horizons of information could provide a similar performance boost in partially observable reinforcement learning (RL) domains, but the large-scale transformers used in NLP have yet to be successfully applied to the RL setting. In this work we demonstrate that the standard transformer architecture is difficult to optimize, which was previously observed in the supervised learning setting but becomes especially pronounced with RL objectives. We propose architectural modifications that substantially improve the stability and learning speed of the original Transformer and XL variant. The proposed architecture, the Gated Transformer-XL (GTrXL), surpasses LSTMs on challenging memory environments and achieves state-of-the-art results on the multi-task DMLab-30 benchmark suite, exceeding the performance of an external memory architecture. We show that the GTrXL, trained using the same losses, has stability and performance that consistently matches or exceeds a competitive LSTM baseline, including on more reactive tasks where memory is less critical. GTrXL offers an easy-to-train, simple-to-implement but substantially more expressive architectural alternative to the standard multi-layer LSTM ubiquitously used for RL agents in partially observable environments.

Subfields: Artificial Intelligence, Cognitive Neuroscience, Reinforcement Learning, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Cognitive Neuroscience

In Reinforcement Learning

In Model Architectures

Citations in the corpus

Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers (2021, background): “Scaling Op Description NL Num. layers EL Num enc. layers DL Num. dec. layers DM dmodel KV dKV NH Num. of heads FF dff SH Shared heads SKV Tied key-values learning (Parisotto et al., 2020) and computational biology (Senior et al., 2020).”
Enhancing the Transformer with Explicit Relational Encoding for Math Problem Solving (2019, neutral): “tention layer of the regular Transformer (Parisotto et al., 2019).”
Decision Transformer: Reinforcement Learning via Sequence Modeling (2021, methodology): “[4] discussed design decisions for more stable training of transformers in the high-variance RL setting.”
Offline Reinforcement Learning as One Big Sequence Modeling Problem (2021, methodology): “…works have studied applying sequence models of various types to represent components in standard RL algorithms, such as policies, value functions, and models (Bakker, 2002; Heess et al., 2015a; Chiappa et al., 2017; Parisotto et al., 2020; Parisotto & Salakhutdinov, 2021; Kumar et al., 2020b).”
Pretrained Transformers as Universal Computation Engines (2021, background): “, 2021), reinforcement learning (Parisotto et al., 2020), and imitation learning (Abramson et al.”
Random Feature Attention (2021, background): “Transformer architectures (Vaswani et al., 2017) have achieved tremendous success on a variety of sequence modeling tasks (Ott et al., 2018; Radford et al., 2018; Parmar et al., 2018; Devlin et al., 2019; Parisotto et al., 2020, inter alia ).”
Towards mental time travel: a hierarchical memory for reinforcement learning agents (2021, methodology): “We find that gating each layer [43] is not beneficial to HCAM on our tasks (App.”
Learning Natural Language Generation with Truncated Reinforcement Learning (2022, background): “Yet, they still face optimization and computational challenges (Parisotto et al., 2020).”
Data Movement Is All You Need: A Case Study on Optimizing Transformers (2020, background): “…have also begun to be applied to domains beyond NLP, including vision (Doso-vitskiy et al., 2020), speech recognition (Yeh et al., 2019), reinforcement learning (Parisotto et al., 2019), molecular property prediction (Maziarka et al., 2020), and symbolic mathematics (Lample & Charton, 2019).”
Transformers are Meta-Reinforcement Learners (2022, background): “Parisotto et al. (2019) then proposed some architectural changes in the vanilla transformer, reordering layer normalization modules and replacing residual connections with expressing gating mechanisms, improving state-of-the-art performance for a set of memory environments.”
Transformers are Sample Efficient World Models (2022, background): “Parisotto et al. (2020) make the observation that the standard Transformer architecture is difﬁcult to optimize with RL objectives.”
Episodic Transformer for Vision-and-Language Navigation (2021, background): “In [48], the authors show that training transformers for long time horizon planning with RL is challenging and propose a solution.”
S4RL: Surprisingly Simple Self-Supervision for Offline Reinforcement Learning (2021, methodology): “ferent network types (Hausknecht & Stone, 2015; Parisotto et al., 2020) and exploiting data augmentations using selfsupervised learning (Laskin et al.”
Few-shot Sequence Learning with Transformers (2020, background): “rested in performing few-shot adaptation to a given new task. Transformers (Vaswani et al., 2017) have been very successful at modeling discrete sequences (Barrault et al., 2019; Devlin et al., 2018; Parisotto et al., 2019). Further, they have been shown to use context tokens appended to an input to”
History Aware Multimodal Transformer for Vision-and-Language Navigation (2021, methodology): “Reinforcement Learning (RL) [12, 26] is one of the most effective approach among them, but it is considered unstable to directly train large-scale transformers via RL [27].”
Reinforcement Learning as One Big Sequence Modeling Problem (2021, background): “Indeed, a number of prior works have studied applying sequence models of various types to represent components in standard RL algorithms, such as policies, value functions, and models [4, 21, 9, 45, 44, 34].”
Adaptive Transformers in RL (2020, methodology): “Transformers in RL literature has been quite scarce. [6] attempted to use the original Transformer on simple tabular MDPs but was unable to achieve better than random performance. Stable Transformers [1] was the ﬁrst to successfully apply Transformers to RL tasks. They again found that the standard ”
Tell me why! Explanations support learning relational and causal structure (2021, background): “The agent memory is a 4-layer Gated TransformerXL Memory (Parisotto et al., 2020), which receives the visual encoder output and previous reward as inputs.”