Stabilizing Transformers for Reinforcement Learning

1 INTRODUCTION . It has been argued that self-attention architectures ( Vaswani et al. , 2017 ) deal better with longer temporal horizons than recurrent neural networks ( RNNs ) : by construction , they avoid compressing the whole past into a fixed-size hidden state and they do not suffer from vanishing or exploding gradients in the same way as RNNs . Recent work has empirically validated these claims , demonstrating that self-attention architectures can provide significant gains in performance over the more traditional recurrent architectures such as the LSTM ( Dai et al. , 2019 ; Radford et al. , 2019 ; Devlin et al. , 2019 ; Yang et al. , 2019 ) . In particular , the Transformer architecture ( Vaswani et al. , 2017 ) has had breakthrough success in a wide variety of domains : language modeling ( Dai et al. , 2019 ; Radford et al. , 2019 ; Yang et al. , 2019 ) , machine translation ( Vaswani et al. , 2017 ; Edunov et al. , 2018 ) , summarization ( Liu & Lapata ) , question answering ( Dehghani et al. , 2018 ; Yang et al. , 2019 ) , multi-task representation learning for NLP ( Devlin et al. , 2019 ; Radford et al. , 2019 ; Yang et al. , 2019 ) , and algorithmic tasks ( Dehghani et al. , 2018 ) . The repeated success of the transformer architecture in domains where sequential information processing is critical to performance makes it an ideal candidate for partially observable RL problems , where episodes can extend to thousands of steps and the critical observations for any decision often span the entire episode . Yet , the RL literature is dominated by the use of LSTMs as the main mechanism for providing memory to the agent ( Espeholt et al. , 2018 ; Kapturowski et al. , 2019 ; Mnih et al. , 2016 ) . Despite progress at designing more expressive memory architectures ( Graves et al. , 2016 ; Wayne et al. , 2018 ; Santoro et al. , 2018 ) that perform better than LSTMs in memory-based tasks and partially-observable environments , they have not seen widespread adoption in RL agents perhaps due to their complex implementation , with the LSTM being seen as the go-to solution for environments where memory is required . In contrast to these other memory architectures , the transformer is well-tested in many challenging domains and has seen several open-source implementations in a variety of deep learning frameworks 1 . Motivated by the transformer ’ s superior performance over LSTMs and the widespread availability of implementations , in this work we investigate the transformer architecture in the RL setting . In particular , we find that the canonical transformer is significantly difficult to optimize , often resulting in performance comparable to a random policy . This difficulty in training transformers exists in the supervised case as well . Typically a complex learning rate schedule is required ( e.g. , linear warmup or cosine decay ) in order to train ( Vaswani et al. , 2017 ; Dai et al. , 2019 ) , or specialized weight initialization schemes are used to improve performance ( Radford et al. , 2019 ) . These measures do not seem to be sufficient for RL . In Mishra et al . ( 2018 ) , for example , transformers could not solve even simple bandit tasks and tabular Markov Decision Processes ( MDPs ) , leading the authors to hypothesize that the transformer architecture was not suitable for processing sequential information . However in this work we succeed in stabilizing training with a reordering of the layer normalization coupled with the addition of a new gating mechanism to key points in the submodules of the transformer . Our novel gated architecture , the Gated Transformer-XL ( GTrXL ) ( shown in Figure 1 , Right ) , is able to learn much faster and more reliably and exhibit significantly better final performance than the canonical transformer . We further demonstrate that the GTrXL achieves state-ofthe-art results when compared to the external memory architecture MERLIN ( Wayne et al. , 2018 ) on the multitask DMLab-30 suite ( Beattie et al. , 2016 ) . Additionally , we surpass LSTMs significantly on memory-based DMLab-30 levels while matching performance on the reactive set , as well as significantly outperforming LSTMs on memory-based continuous control and navigation environments . We perform extensive ablations on the GTrXL in challenging environments with both continuous actions and high-dimensional observations , testing the final performance of the various components as well as the GTrXL ’ s robustness to seed and hyperparameter sensitivity compared to LSTMs and the canonical transformer . We demonstrate a consistent superior performance while matching the stability of LSTMs , providing evidence that the GTrXL architecture can function as a drop-in replacement to the LSTM networks ubiquitously used in RL . 2 TRANSFORMER ARCHITECTURE AND VARIANTS . The transformer network consists of several stacked blocks that repeatedly apply self-attention to the input sequence . The transformer layer block itself has remained relatively constant since its original introduction ( Vaswani et al. , 2017 ; Liu et al. , 2018 ; Radford et al. , 2019 ) . Each layer consists of two submodules : an attention operation followed by a position-wise multi-layer network ( see Figure 1 ( left ) ) . The input to the transformer block is an embedding from the previous layerE ( l−1 ) ∈ RT×D , where T is the number of time steps , D is the hidden dimension , and l ∈ [ 0 , L ] is the layer index with L being the total number of layers . We assume E ( 0 ) is an arbitrarily-obtained input embedding of dimension [ T , D ] , e.g . a word embedding in the case of language modeling or an embedding of the per-timestep observations in an RL environment . Multi-Head Attention : The Multi-Head Attention ( MHA ) submodule computes in parallel H softattention operations for every time step . A residual connection ( He et al. , 2016a ) and layer normalization ( Ba et al. , 2016 ) are then applied to the output ( see Appendix C for more details ) : Y ( l ) = MultiHeadAttention ( E ( l−1 ) ) , Ŷ ( l ) = E ( l−1 ) + Y ( l ) , Y ( l ) = LayerNorm ( Ŷ ( l ) ) , ( 1 ) Multi-Layer Perceptron : The Multi-Layer Perceptron ( MLP ) submodule applies a 1× 1 temporal convolutional network f ( l ) ( i.e. , kernel size 1 , stride 1 ) over every step in the sequence , producing a new embedding tensor E ( l ) ∈ RT×D . As in Dai et al . ( 2019 ) , the network output does not include an activation function . After the MLP , there is a residual update followed by layer normalization : E ( l ) = f ( l ) ( Y ( l ) ) , Ê ( l ) = Y ( l ) + E ( l ) , E ( l ) = LayerNorm ( Ê ( l ) ) . ( 2 ) Relative Position Encodings : The basic MHA operation does not take sequence order into account explicitly because it is permutation invariant . Positional encodings are a widely used solution in 1e.g . https : //github.com/kimiyoung/transformer-xl , https : //github.com/ tensorflow/tensor2tensor domains like language where order is an important semantic cue , appearing in the original transformer architecture ( Vaswani et al. , 2017 ) . To enable a much larger contextual horizon than would otherwise be possible , we use the relative position encodings and memory scheme used in Dai et al . ( 2019 ) . In this setting , there is an additional T -step memory tensor M ( l ) ∈ RT ×D , which is treated as constant during weight updates . The MHA submodule then becomes : Y ( l ) = RelativeMultiHeadAttention ( StopGrad ( M ( l−1 ) ) , E ( l−1 ) ) ( 3 ) Ŷ ( l ) = E ( l−1 ) + Y ( l ) , Y ( l ) = LayerNorm ( Ŷ ( l ) ) ( 4 ) where StopGrad is a stop-gradient function that prevents gradients flowing backwards during backpropagation . We refer to Appendix C for a more detailed description . 3 GATED TRANSFORMER ARCHITECTURES . While the transformer architecture has achieved breakthrough results in modeling sequences for supervised learning tasks ( Vaswani et al. , 2017 ; Liu et al. , 2018 ; Dai et al. , 2019 ) , a demonstration of the transformer as a useful RL memory has been notably absent . Previous work has highlighted training difficulties and poor performance ( Mishra et al. , 2018 ) . When transformers have not been used for temporal memory but instead as a mechanism for attention over the input space , they have had success—notably in the challenging multi-agent Starcraft 2 environment ( Vinyals et al. , 2019 ) . Here , the transformer was applied solely across Starcraft units and not over time . Multiplicative interactions have been successful at stabilizing learning across a wide variety of architectures ( Hochreiter & Schmidhuber , 1997 ; Srivastava et al. , 2015 ; Cho et al. , 2014 ) . Motivated by this , we propose the introduction of powerful gating mechanisms in place of the residual connections within the transformer block , coupled with changes to the order of layer normalization in the submodules . As will be empirically demonstrated , the “ Identity Map Reordering ” and gating mechanisms are critical for stabilizing learning and improving performance . 3.1 IDENTITY MAP REORDERING . Our first change is to place the layer normalization on only the input stream of the submodules , a modification described in several previous works ( He et al. , 2016b ; Radford et al. , 2019 ; Baevski & Auli , 2019 ) . The model using this Identity Map Reordering is termed TrXL-I in the following , and is depicted visually in Figure 1 ( center ) . A key benefit to this reordering is that it now enables an identity map from the input of the transformer at the first layer to the output of the transformer after the last layer . This is in contrast to the canonical transformer , where there are a series of layer normalization operations that non-linearly transform the state encoding . Because the layer norm reordering causes a path where two linear layers are applied in sequence , we apply a ReLU activation to each sub-module output before the residual connection ( see Appendix C for equations ) . The TrXL-I already exhibits a large improvement in stability and performance over TrXL ( see Section 4.3.1 ) . One hypothesis as to why the Identity Map Reordering improves results is as follows : assuming that the submodules at initialization produce values that are in expectation near zero , the state encoding is passed un-transformed to the policy and value heads , enabling the agent to learn a Markovian policy at the start of training ( i.e. , the network is initialized such that π ( ·|st , . . . , s1 ) ≈ π ( ·|st ) and V π ( st|st−1 , . . . , s1 ) ≈ V π ( st|st−1 ) ) . In many environments , reactive behaviours need to be learned before memory-based ones can be effectively utilized , i.e. , an agent needs to learn how to walk before it can learn how to remember where it has walked .

Summary: The paper explores a transformer for reinforcement learning. The authors demonstrate that Canonical Transformer is unstable. The authors introduce two modifications to the Canonical Transformer. The first is to move the layer normalization layer to the input stream. The second is to replace residual connections with gating layers. The experimental results show that (1) the first modification, i.e., moving the layer normalization layer to the input stream significantly stabilizes the training; (2) Gated Recurrent Unit (GRU) gating seems to be most effective gating mechanism.