Original: S TABILIZING T RANSFORMERS FOR R EINFORCEMENT L EARNING
Revised: S TABILIZING T RANSFORMERS FOR R EINFORCEMENT L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In this work we demonstrate that the standard transformer architecture is difﬁcult to optimize, which was previously observed in the supervised learning setting but becomes especially pronounced with RL objectives.
Revised: In this work we demonstrate that the standard transformer architecture is difﬁcult to optimize, which was previously observed in the supervised learning setting but becomes especially pronounced with RL objectives.

Original: We propose architectural modiﬁcations that substantially improve the stability and learning speed of the original Transformer and
Revised: We propose architectural modiﬁcations that substantially improve the stability and learning speed of the original Transformer and

Original: XL variant.
Revised: XL variant.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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 ﬁxed-size hidden state and they do not suffer from vanishing or exploding gradients in the same way as RNNs.
Revised: 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 ﬁxed-size hidden state and they do not suffer from vanishing or exploding gradients in the same way as RNNs.

Original: Recent work has empirically validated these claims, demonstrating that self-attention architectures can provide signiﬁcant gains in performance over the more traditional recurrent architectures such as the LSTM (Dai et al., 2019;
Revised: Recent work has empirically validated these claims, demonstrating that self-attention architectures can provide signiﬁcant gains in performance over the more traditional recurrent architectures such as the LSTM (Dai et al., 2019;

Original: Radford et al., 2019;
Revised: Radford et al., 2019;

Original: Devlin et al., 2019;
Revised: Devlin et al., 2019;

Original: Yang et al., 2019).
Revised: Yang et al., 2019).

Original: 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;
Revised: 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;

Original: Radford et al., 2019;
Revised: Radford et al., 2019;

Original: Yang et al., 2019), machine translation (Vaswani et al., 2017; Edunov et al., 2018), summarization (Liu & Lapata), question answering (Dehghani et al., 2018;
Revised: Yang et al., 2019), machine translation (Vaswani et al., 2017; Edunov et al., 2018), summarization (Liu & Lapata), question answering (Dehghani et al., 2018;

Original: Yang et al., 2019), multi-task representation learning for NLP (Devlin et al., 2019; Radford et al., 2019;
Revised: Yang et al., 2019), multi-task representation learning for NLP (Devlin et al., 2019; Radford et al., 2019;

Original: Yang et al., 2019), and algorithmic tasks (Dehghani et al., 2018).
Revised: Yang et al., 2019), and algorithmic tasks (Dehghani et al., 2018).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Despite progress at designing more expressive memory architectures (Graves et al., 2016;
Revised: Despite progress at designing more expressive memory architectures (Graves et al., 2016;

Original: Wayne et al., 2018;
Revised: Wayne et al., 2018;

Original: 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.
Revised: 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.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: In particular, we ﬁnd that the canonical transformer is signiﬁcantly difﬁcult to optimize, often resulting in performance comparable to a random policy.
Revised: In particular, we ﬁnd that the canonical transformer is signiﬁcantly difﬁcult to optimize, often resulting in performance comparable to a random policy.

Original: This difﬁculty in training transformers exists in the supervised case as well.
Revised: This difﬁculty in training transformers exists in the supervised case as well.

Original: 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).
Revised: 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).

Original: These measures do not seem to be sufﬁcient for RL.
Revised: These measures do not seem to be sufﬁcient for RL.

Original: In Mishra et al.
Revised: In Mishra et al.

Original: (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.
Revised: (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.

Original: 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.
Revised: 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.

Original: Our novel gated architecture, the Gated Transformer-XL (GTrXL) (shown in Figure 1,
Revised: Our novel gated architecture, the Gated Transformer-XL (GTrXL) (shown in Figure 1,

Original: Right), is able to learn much faster and more reliably and exhibit signiﬁcantly better ﬁnal performance than the canonical transformer.
Revised: Right), is able to learn much faster and more reliably and exhibit signiﬁcantly better ﬁnal performance than the canonical transformer.

Original: 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).
Revised: 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).

Original: Additionally, we surpass LSTMs signiﬁcantly on memory-based DMLab-30 levels while matching performance on the reactive set, as well as signiﬁcantly outperforming LSTMs on memory-based continuous control and navigation environments.
Revised: Additionally, we surpass LSTMs signiﬁcantly on memory-based DMLab-30 levels while matching performance on the reactive set, as well as signiﬁcantly outperforming LSTMs on memory-based continuous control and navigation environments.

Original: We perform extensive ablations on the GTrXL in challenging environments with both continuous actions and high-dimensional observations, testing the ﬁnal performance of the various components as well as the GTrXL’s robustness to seed and hyperparameter sensitivity compared to
Revised: We perform extensive ablations on the GTrXL in challenging environments with both continuous actions and high-dimensional observations, testing the ﬁnal performance of the various components as well as the GTrXL’s robustness to seed and hyperparameter sensitivity compared to

Original: LSTMs and the canonical transformer.
Revised: LSTMs and the canonical transformer.

Original: 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.
Revised: 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.

Original:  T RANSFORMER A RCHITECTURE AND V ARIANTS
Revised:  T RANSFORMER A RCHITECTURE AND V ARIANTS

Original: The transformer network consists of several stacked blocks that repeatedly apply self-attention to the input sequence.
Revised: The transformer network consists of several stacked blocks that repeatedly apply self-attention to the input sequence.

Original: The transformer layer block itself has remained relatively constant since its original introduction (Vaswani et al., 2017; Liu et al., 2018;
Revised: The transformer layer block itself has remained relatively constant since its original introduction (Vaswani et al., 2017; Liu et al., 2018;

Original: Radford et al., 2019).
Revised: Radford et al., 2019).

Original: Each layer consists of two submodules: an attention operation followed by a position-wise multi-layer network (see Figure 1, Left).
Revised: Each layer consists of two submodules: an attention operation followed by a position-wise multi-layer network (see Figure(left)).

Original: The input to the transformer block is an embedding from the previous layer E ( l − 1)
Revised: The input to the transformer block is an embedding from the previous layer E ( l − 1)

Original: ∈ R T × D , where T is the number of time steps
Revised: ∈ R T × D , where T is the number of time steps

Original: , D is the hidden dimension, and l ∈ [0 , L ] is the layer index with L being the total number of layers.
Revised: , D is the hidden dimension, and l ∈ [0 , L ] is the layer index with L being the total number of layers.

Original: 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.
Revised: 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.

Original: Multi-Head Attention: The Multi-Head Attention (MHA) submodule computes in parallel H softattention operations for every time step.
Revised: Multi-Head Attention: The Multi-Head Attention (MHA) submodule computes in parallel H softattention operations for every time step.

Original: 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): 
Revised: 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): [Equation]

Revised: MultiHeadAttention

Original: [Equation]
Revised: [Equation]

Original: Multi Layer Perceptron
Revised: Multi-Layer Perceptron

Original: : The
Revised: : The

Original: Multi-
Revised: Multi

Original: 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 ) ∈ R T × D .
Revised: -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 ) ∈ R T × D .

Original: As in Dai et al.
Revised: As in Dai et al.

Original: (2019), the network output does not include an activation function.
Revised: (2019), the network output does not include an activation function.

Original: After the MLP, there is a residual update followed by layer normalization:
Revised: After the MLP, there is a residual update followed by layer normalization:

Original: [Equation]
Revised: [Equation]

Original: Relative Position Encodings: The basic MHA operation does not take sequence order into account explicitly because it is permutation invariant.
Revised: Relative Position Encodings: The basic MHA operation does not take sequence order into account explicitly because it is permutation invariant.

Original: Positional encodings are a widely used solution in domains like language where order is an important semantic cue, appearing in the original transformer architecture (Vaswani et al., 2017).
Revised: Positional encodings are a widely used solution in domains like language where order is an important semantic cue, appearing in the original transformer architecture (Vaswani et al., 2017).

Original: 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.
Revised: 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.

Original: In this setting, there is an additional  memory tensor M ( l ) ∈ R L × D , which is treated as constant during weight updates.
Revised: In this setting, there is an additional T -step memory tensor M ( l ) ∈ R T × D , which is treated as constant during weight updates.

Original: The MHA submodule then becomes:
Revised: The MHA submodule then becomes:

Original: [Equation]
Revised: [Equation]

Original: where StopGrad is a stop-gradient function that prevents gradients ﬂowing backwards during backpropagation.
Revised: where StopGrad is a stop-gradient function that prevents gradients ﬂowing backwards during backpropagation.

Original: We refer to Appendix C for a more detailed description.
Revised: We refer to Appendix C for a more detailed description.

Original:  G ATED T RANSFORMER A RCHITECTURES
Revised:  G ATED T RANSFORMER A RCHITECTURES

Original: 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.
Revised: 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.

Original: Previous work has highlighted training difﬁculties and poor performance (Mishra et al., 2018).
Revised: Previous work has highlighted training difﬁculties and poor performance (Mishra et al., 2018).

Original: 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).
Revised: 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).

Original: Here, the transformer was applied solely across Starcraft units and not over time.
Revised: Here, the transformer was applied solely across Starcraft units and not over time.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: As will be empirically demonstrated, the “Identity Map Reordering” and gating mechanisms are critical for stabilizing learning and improving performance.
Revised: As will be empirically demonstrated, the “Identity Map Reordering” and gating mechanisms are critical for stabilizing learning and improving performance.

Original: 3. I DENTITY M AP R EORDERING
Revised: 3. I DENTITY M AP R EORDERING

Original: Our ﬁrst change is to place the layer normalization on only the input stream of the submodules, a modiﬁcation described in several previous works (He et al., 2016b;
Revised: Our ﬁrst change is to place the layer normalization on only the input stream of the submodules, a modiﬁcation described in several previous works (He et al., 2016b;

Original: Radford et al., 2019; Baevski
Revised: Radford et al., 2019; Baevski

Original: & Auli, 2019).
Revised: & Auli, 2019).

Original: The model using this Identity Map Reordering is termed TrXL-I in the following, and is depicted visually in Figure 1 (center).
Revised: The model using this Identity Map Reordering is termed TrXL-I in the following, and is depicted visually in Figure 1 (center).

Original: A key beneﬁt to this reordering is that it now enables an identity map from the input of the transformer at the ﬁrst layer to the output of the transformer after the last layer.
Revised: A key beneﬁt to this reordering is that it now enables an identity map from the input of the transformer at the ﬁrst layer to the output of the transformer after the last layer.

Original: This is in contrast to the canonical transformer, where there are a series of layer normalization operations that non-linearly transform the state encoding.
Revised: This is in contrast to the canonical transformer, where there are a series of layer normalization operations that non-linearly transform the state encoding.

Original: 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).
Revised: 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).

Original: The TrXL-I already exhibits a large improvement in stability and performance over TrXL (see Section 4.3.1).
Revised: The TrXL-I already exhibits a large improvement in stability and performance over TrXL (see Section 4.3.1).

Original: 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 π ( ·| s t , . . .
Revised: 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 π ( ·| s t , . . .

Original: , s 1 )
Revised: , s 1 )

Original: ≈ π ( ·| s t ) and V π ( s t | s t − 1 , . . .
Revised: ≈ π ( ·| s t ) and V π ( s t | s t − 1 , . . .

Original: , s 1 )
Revised: , s 1 )

Original: ≈ V π
Revised: ≈ V π

Original: ( s t | s t − 1 ) ).
Revised: ( s t | s t − 1 ) ).

Original: 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.
Revised: 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.

Original: 3. G ATING L AYERS
Revised: 3. G ATING L AYERS

Original: We further improve performance and optimization stability by replacing the residual connections in
Revised: We further improve performance and optimization stability by replacing the residual connections in

Original: Equations 4 and 2 with gating layers.
Revised: Equations 4 and 2 with gating layers.

Original: We call the gated architecture with the identity map reordering the Gated Transformer(-XL) (GTrXL).
Revised: We call the gated architecture with the identity map reordering the Gated Transformer(-XL) (GTrXL).

Original: The ﬁnal GTrXL transformer block is written below:
Revised: The ﬁnal GTrXL layer block is written below:

Original: [Equation]
Revised: [Equation]

Original: where g is a gating layer function.
Revised: where g is a gating layer function.

Original: A visualization of our ﬁnal architecture is shown in Figure(right), with the modiﬁcations from the canonical transformer highlighted in red.
Revised: A visualization of our ﬁnal architecture is shown in Figure(right), with the modiﬁcations from the canonical transformer highlighted in red.

Original: In our experiments we ablate a variety of gating layers with increasing expressivity:
Revised: In our experiments we ablate a variety of gating layers with increasing expressivity:

Original: Input: The gated input connection has a sigmoid modulation on the input stream, similar to the short-cut-only gating from He et al.
Revised: Input: The gated input connection has a sigmoid modulation on the input stream, similar to the short-cut-only gating from He et al.

Original: (2016b):
Revised: (2016b):

Original: [Equation]
Revised: [Equation]

Original: Output: The gated output connection has a sigmoid modulation on the output stream:
Revised: Output: The gated output connection has a sigmoid modulation on the output stream:

Original: [Equation]
Revised: [Equation]

Original: Highway: The highway connection (Srivastava et al., 2015) modulates both streams with a sigmoid:
Revised: Highway: The highway connection (Srivastava et al., 2015) modulates both streams with a sigmoid:

Original: [Equation]
Revised: [Equation]

Original: Sigmoid-Tanh: The sigmoid-tanh (SigTanh) gate (Van den Oord et al., 2016) is similar to the Output gate but with an additional tanh activation on the output stream:
Revised: Sigmoid-Tanh: The sigmoid-tanh (SigTanh) gate (Van den Oord et al., 2016) is similar to the Output gate but with an additional tanh activation on the output stream:

Original: [Equation]
Revised: [Equation]

Original: Gated
Revised: Gated

Revised: -Recurrent-Unit-type

Revised: gating

Revised: The Gated Recurrent

Original: Recurrent-Unit-type gating: The Gated Recurrent Unit (GRU) (Chung et al., 2014) is a recurrent network that performs similarly to an LSTM (Hochreiter & Schmidhuber, 1997) but has fewer parameters.
Revised: Unit (GRU) (Chung et al., 2014) is a recurrent network that performs similarly to an LSTM (Hochreiter & Schmidhuber, 1997) but has fewer parameters.

Original: We adapt its powerful gating mechanism as an untied activation function in depth:
Revised: We adapt its powerful gating mechanism as an untied activation function in depth:

Original: [Equation]
Revised: [Equation]

Original: Gated Identity Initialization: We have claimed that the Identity Map Reordering aids policy optimization because it initializes the agent close to a Markovian policy / value function.
Revised: Gated Identity Initialization: We have claimed that the Identity Map Reordering aids policy optimization because it initializes the agent close to a Markovian policy / value function.

Original: If this is indeed the cause of improved stability, we can explicitly initialize the various gating mechanisms to be close to the identity map.
Revised: If this is indeed the cause of improved stability, we can explicitly initialize the various gating mechanisms to be close to the identity map.

Original: This is the purpose of the bias b ( l ) g in the applicable gating layers.
Revised: This is the purpose of the bias b ( l ) g in the applicable gating layers.

Original: We later demonstrate in an ablation that initially setting b ( l ) g > 0 can greatly improve learning speed.
Revised: We later demonstrate in an ablation that initially setting b ( l ) g > 0 can greatly improve learning speed.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we provide experiments on a variety of challenging single and multi-task RL domains: DMLab-30 (Beattie et al., 2016), Numpad and Memory Maze (see Fig.
Revised: In this section, we provide experiments on a variety of challenging single and multi-task RL domains: DMLab-30 (Beattie et al., 2016), Numpad and Memory Maze (see Fig.

Original: Crucially we show that the proposed Gated Transformer-XL (GTrXL) not only shows substantial improvements over LSTMs on memory-based environments, but suffers no degradation of performance on reactive environments.
Revised: Crucially we demonstrate that the proposed Gated Transformer-XL (GTrXL) not only shows substantial improvements over LSTMs on memory-based environments, but suffers no degradation of performance on reactive environments.

Original: The GTrXL also exceeds MERLIN (Wayne et al., 2018), an external memory architecture which used a Differentiable Neural Computer (Graves et al., 2016) coupled with auxiliary losses, surpassing its performance on both memory and reactive tasks.
Revised: The GTrXL also exceeds MERLIN (Wayne et al., 2018), an external memory architecture which used a Differentiable Neural Computer (Graves et al., 2016) coupled with auxiliary losses, surpassing its performance on both memory and reactive tasks.

Original: For all transformer architectures except when otherwise stated, we train relatively deep 12-layer networks with embedding size 256 and memory size 512.
Revised: For all transformer architectures except when otherwise stated, we train relatively deep 12-layer networks with embedding size 256 and memory size 512.

Original: These networks are comparable to the state-of-the-art networks in use for small language modeling datasets (see enwik8 results in (Dai et al., 2019)).
Revised: These networks are comparable to the state-of-the-art networks in use for small language modeling datasets (see enwik8 results in (Dai et al., 2019)).

Original: We chose to train deep networks in order to demonstrate that our results do not necessarily sacriﬁce complexity for stability, i.e. we are not making transformers stable for RL simply by making them shallow.
Revised: We chose to train deep networks in order to demonstrate that our results do not necessarily sacriﬁce complexity for stability, i.e. we are not making transformers stable for RL simply by making them shallow.

Original: Our networks have receptive ﬁelds that can potentially span any episode in the environments tested, with an upper bound on the receptive ﬁeld of 6144 ( 12 layers × 512 memory (Dai et al., 2019)).
Revised: Our networks have receptive ﬁelds that can potentially span any episode in the environments tested, with an upper bound on the receptive ﬁeld of 6144 ( 12 layers × 512 memory (Dai et al., 2019)).

Original: Future work will look at scaling transformers in RL even further, e.g. towards the 52-layer network in Radford et al.
Revised: Future work will look at scaling transformers in RL even further, e.g. towards the 52-layer network in Radford et al.

Original: See App.
Revised: See App.

Original: B for experimental details.
Revised: B for experimental details.

Original: For all experiments, we used V-MPO (Anonymous Authors, 2019), an on-policy adaptation of Maximum a Posteriori Policy Optimization (MPO) (Abdolmaleki et al., 2018a;b) that performs approximate policy iteration based on a learned state-value function V ( s ) instead of the state-action value function used in MPO.
Revised: For all experiments, we used V-MPO (Anonymous Authors, 2019), an on-policy adaptation of Maximum a Posteriori Policy Optimization (MPO) (Abdolmaleki et al., 2018a;b) that performs approximate policy iteration based on a learned state-value function V ( s ) instead of the state-action value function used in MPO.

Original: Rather than directly updating the parameters in the direction of the policy gradient, V-MPO uses the estimated advantages to ﬁrst construct a target distribution for the policy update subject to a sample-based KL constraint, then calculates the gradient that partially moves the parameters toward that target, again subject to a KL constraint.
Revised: Rather than directly updating the parameters in the direction of the policy gradient, V-MPO uses the estimated advantages to ﬁrst construct a target distribution for the policy update subject to a sample-based KL constraint, then calculates the gradient that partially moves the parameters toward that target, again subject to a KL constraint.

Original: V-MPO was shown to achieve state-of-the-art results for LSTM-based agents on the multi-task DMLab-30 benchmark suite.
Revised: V-MPO was shown to achieve state-of-the-art results for LSTM-based agents on the multi-task DMLab-30 benchmark suite.

Original: 4. T RANSFORMER AS E FFECTIVE RL M EMORY A RCHITECTURE
Revised: 4. T RANSFORMER AS E FFECTIVE RL M EMORY A RCHITECTURE

Original: We ﬁrst present results of the best performing GTrXL variant, the GRU-type gating, against a competitive LSTM baseline, demonstrating a substantial improvement on the multi-task DMLab-30 do- main (Beattie et al., 2016).
Revised: We ﬁrst present results of the best performing GTrXL variant, the GRU-type gating, against a competitive LSTM baseline, demonstrating a substantial improvement on the multi-task DMLab-30 do- [Table] surpasses LSTM by a signiﬁcant gap and exceeds the performance of MERLIN (Wayne et al

Original: Figure 2 shows mean return over all levels as training progresses, where the return is human normalized as done in previous work (meaning a human has a per-level mean score of 100and a random policy has a score of 0), while Table 1 has the ﬁnal performance at 10 billion environment steps. 
Revised: , 2018) trained for 100 billion environment steps. The GTrXL ( and the proposed reordered TrXL-I also surpass LSTM but perform slightly worse than MERLIN and are not as robust as GTrXL (GRU) ( Sec. 4.3.2). We sample 6-8 hyperparameters per model.

Revised: We include standard error over runs.

Revised: main (Beattie et al., 2016).

Revised: DMLab-30 is a large-scale, multitask benchmark comprising 30 ﬁrstperson 3D environments with image observations and has been widely used as a benchmark for architectural and algorithmic improvements (Wayne et al., 2018; Espeholt et al., 2018; Kapturowski et al., 2019; Hessel et al., 2018).

Revised: The levels test a wide agent competencies such as language comprehension, navigation, handling of partial observability, memory, planning, and other forms of long horizon reasoning, with episodes lasting over 4000 environment steps.

Revised: Figure 2 shows mean return over all levels as training progresses, where the return is human normalized as done in previous work (meaning a human has a per-level mean score of 100 and a random policy has a score of 0), while Table 1 has the ﬁnal performance at 10 billion environment steps.

Original: The GTrXL has a signiﬁcant gap over a 3-layer LSTM baseline trained using the same V-MPO algorithm.
Revised: The GTrXL has a signiﬁcant gap over a 3-layer LSTM baseline trained using the same V-MPO algorithm.

Original: Furthermore, we included the ﬁnal results of a previously-published external memory architecture, MERLIN (Wayne et al., 2018).
Revised: Furthermore, we included the ﬁnal results of a previously-published external memory architecture, MERLIN (Wayne et al., 2018).

Original: Because MERLIN was trained for 100 billion environment steps with a different algorithm, IMPALA (Espeholt et al., 2018), and also involved an auxiliary loss critical for the memory component to function, the learning curves are not directly comparable and we only report the ﬁnal performance of the architecture as a dotted line.
Revised: Because MERLIN was trained for 100 billion environment steps with a different algorithm, IMPALA (Espeholt et al., 2018), and also involved an auxiliary loss critical for the memory component to function, the learning curves are not directly comparable and we only report the ﬁnal performance of the architecture as a dotted line.

Original: Despite the differences, our results demonstrate that the GTrXL can match a state-of-the-art result.
Revised: Despite the differences, our results demonstrate that the GTrXL can match the previous state-of-the-art on DMLab-30.

Original: An informative split between a set of memory-based levels and more reactive ones (listed in Appendix D) reveals that our model speciﬁcally has large improvements in environments where memory plays a critical role.
Revised: An informative split between a set of memory-based levels and more reactive ones (listed in Appendix D) reveals that our model speciﬁcally has large improvements in environments where memory plays a critical role.

Original: Meanwhile, GTrXL also shows improvement over LSTMs on the set of reactive levels, as memory can still be effectively utilized in some of these levels.
Revised: Meanwhile, GTrXL also shows improvement over LSTMs on the set of reactive levels, as memory can still be effectively utilized in some of these levels.

Original: 4. S CALING WITH M EMORY H ORIZON
Revised: 4. S CALING WITH M EMORY H ORIZON

Original: We next demonstrate that the GTrXL scales better compared to an LSTM when an environment’s temporal horizon is increased, using the “Numpad” continuous control task of Humplik et al.
Revised: We next demonstrate that the GTrXL scales better compared to an LSTM when an environment’s temporal horizon is increased, using the “Numpad” continuous control task of Humplik et al.

Original: (2019) which allows an easy combinatorial increase in the temporal horizon.
Revised: (2019) which allows an easy combinatorial increase in the temporal horizon.

Original: In Numpad, a robotic agent is situated on a platform resembling the 3x3 number pad of a telephone (generalizable to N × N pads).
Revised: In Numpad, a robotic agent is situated on a platform resembling the 3x3 number pad of a telephone (generalizable to N × N pads).

Original: The agent can interact with the pads by colliding with them, causing them to be activated (visualized in the environment state as the number pad glowing).
Revised: The agent can interact with the pads by colliding with them, causing them to be activated (visualized in the environment state as the number pad glowing).

Original: The goal of the agent is to activate a speciﬁc sequence of up to N 2 numbers, but without knowing this sequence a priori.
Revised: The goal of the agent is to activate a speciﬁc sequence of up to N 2 numbers, but without knowing this sequence a priori.

Original: The only feedback the agent gets is by activating numbers: if the pad is the next one in the sequence, the agent gains a reward of +1, otherwise all activated pads are cleared and the agent must restart the sequence.
Revised: The only feedback the agent gets is by activating numbers: if the pad is the next one in the sequence, the agent gains a reward of +1, otherwise all activated pads are cleared and the agent must restart the sequence.

Original: Each correct number in the sequence only provides reward once, i.e. each subsequent activation of that number will no longer provide rewards.
Revised: Each correct number in the sequence only provides reward once, i e. each subsequent activation of that number will no longer provide rewards.

Original: Therefore the agent must explicitly develop a search strategy to determine the correct pad sequence.
Revised: Therefore the agent must explicitly develop a search strategy to determine the correct pad sequence.

Original: Once the agent completes the full sequence, all pads are reset and the agent gets a chance to repeat the sequence again for more reward.
Revised: Once the agent completes the full sequence, all pads are reset and the agent gets a chance to repeat the sequence again for more reward.

Original: This means higher reward directly translates into how well the pad sequence has been memorized.
Revised: This means higher reward directly translates into how well the pad sequence has been memorized.

Original: An image of the scenario is provided in Figure 3.
Revised: An image of the scenario is provided in Figure 3.

Original: There is the restriction that contiguous pads in the sequence must be contiguous in space, i.e. the next pad in the sequence can only be in the Moore neighborhood of the previous pad.
Revised: There is the restriction that contiguous pads in the sequence must be contiguous in space, i.e. the next pad in the sequence can only be in the Moore neighborhood of the previous pad.

Original: Furthermore, no pad can be pressed twice in the sequence.
Revised: Furthermore, no pad can be pressed twice in the sequence.

Original: We present two results in this environment in Figure 3.
Revised: We present two results in this environment in Figure 3.

Original: The ﬁrst measures the ﬁnal performance of the trained models as a function of the pad size.
Revised: The ﬁrst measures the ﬁnal performance of the trained models as a function of the pad size.

Original: We can see that LSTM performs badly on all 3 pad sizes, and performs worse as the pad size increases from 2 to 4.
Revised: We can see that LSTM performs badly on all 3 pad sizes, and performs worse as the pad size increases from 2 to 4.

Original: The GTrXL performs much better, and almost instantly solves the environment with its much more expressive memory.
Revised: The GTrXL performs much better, and almost instantly solves the environment with its much more expressive memory.

Original: On the center[Table] and the individual 6-8 hyperparameter samples per model (light). and right images, we provide learning curves for the 2 × 2 and 4 × 4 Numpad environments, and show that even when the LSTM is trained twice as long it does not reach GTrXL’s performance.
Revised: On the center and right images, we provide learning curves for the 2 × 2 and 4 × 4 Numpad environments, and show that even when the LSTM is trained twice as long it does not reach GTrXL’s performance.

Original: 4. G ATING V ARIANTS + I DENTITY M AP R EORDERING
Revised: 4. G ATING V ARIANTS + I DENTITY M AP R EORDERING

Original: We demonstrated that the GRU-type-gated GTrXL can achieve state-of-the-art results on DMLab30, surpassing both a deep LSTM and an external memory architecture, and also that the GTrXL has a memory which scales better with the memory horizon of the environment.
Revised: We demonstrated that the GRU-type-gated GTrXL can achieve state-of-the-art results on DMLab30, surpassing both a deep LSTM and an external memory architecture, and also that the GTrXL has a memory which scales better with the memory horizon of the environment.

Original: However, the question remains whether the expressive gating mechanisms of the GRU could be replaced by simpler alternatives.
Revised: However, the question remains whether the expressive gating mechanisms of the GRU could be replaced by simpler alternatives.

Original: In this section, we perform extensive ablations on the gating variants described in Section 3.2, and show that the GTrXL (GRU) has improvements in learning speed, ﬁnal performance and optimization stability over all other models, even when controlling for the number of parameters.
Revised: In this section, we perform extensive ablations on the gating variants described in Section 3.2, and show that the GTrXL (GRU) has improvements in learning speed, ﬁnal performance and optimization stability over all other models, even when controlling for the number of parameters.

Original: 4.3. P ERFORMANCE A BLATION
Revised: 4.3. P ERFORMANCE A BLATION

Original: We ﬁrst report the performance of the gating variants in DMLab-30.
Revised: We ﬁrst report the performance of the gating variants in DMLab-30.

Original: Table 1 and Figure 4 show the ﬁnal performance and training curves of the various gating types in both the memory / reactive split, respectively.
Revised: Table 1 and Figure 4 show the ﬁnal performance and training curves of the various gating types in both the memory / reactive split, respectively.

Original: The canonical TrXL completely fails to learn, while the TrXL-I improves over the 
Revised: The canonical TrXL completely fails to learn, while the TrXL-I improves over the [Table] [Table]

Revised: we plot the returns in Figure 5 as 0 af-

Original: LSTM.
Revised: LSTM.

Original: Of the gating varieties, the GTrXL (Output) can recover a large amount of the performance of the GTrXL (GRU), especially in the reactive set, but as shown below is generally far less stable.
Revised: Of the gating varieties, the GTrXL (Output) can recover a large amount of the performance of the GTrXL (GRU), especially in the reactive set, but as shown in Sec.

Original: The
Revised: 4.3.2 is generally far less stable.

Original:  GTrXL (Input) performs worse than even the TrXL-I, reinforcing the identity map path hypothesis.[Table] [Table]
Revised: The GTrXL (Input) performs worse than even the TrXL-I, reinforcing the identity map path hypothesis.

Original: plot Finally, the GTrXL (Highway) and GTrXL (SigTanh) are more sensitive to the hyperparameter choice compared to the alternatives, with some choices doing worse than TrXL-I.
Revised: Finally, the GTrXL (Highway) and GTrXL (SigTanh) are more sensitive to the hyperparameter settings compared to the alternatives, with some settings doing worse than TrXL-I.

Original: 4.3. H YPERPARAMETER AND S EED S ENSITIVITY
Revised: 4.3. H YPERPARAMETER AND S EED S ENSITIVITY

Original: Beyond improved performance, we next demonstrate a signiﬁcant reduction in hyperparameter and seed sensitivity for the GTrXL (GRU) compared to baselines and other GTrXL variants.
Revised: Beyond improved performance, we next demonstrate a signiﬁcant reduction in hyperparameter and seed sensitivity for the GTrXL (GRU) compared to baselines and other GTrXL variants.

Original: We use the “Memory Maze” environment, a memory-based navigation task in which the agent must discover the location of an apple randomly placed in a maze of blocks.
Revised: We use the “Memory Maze” environment, a memory-based navigation task in which the agent must discover the location of an apple randomly placed in a maze of blocks.

Original: The agent receives a positive reward for collecting the apple and is then teleported to a random location in the maze, with the apple’s position held ﬁxed.
Revised: The agent receives a positive reward for collecting the apple and is then teleported to a random location in the maze, with the apple’s position held ﬁxed.

Original: The agent can make use of landmarks situated around the room to return as quickly as possible to the apple for subsequent rewards.
Revised: The agent can make use of landmarks situated around the room to return as quickly as possible to the apple for subsequent rewards.

Original: Therefore, an effective mapping of the environment results in more frequent returns to the apple and higher reward.
Revised: Therefore, an effective mapping of the environment results in more frequent returns to the apple and higher reward.

Original: We chose to perform the sensitivity ablation on Memory Maze because (1) it requires the use of longrange memory to be effective and (2) it includes both continuous and discrete action sets (details in
Revised: We chose to perform the sensitivity ablation on Memory Maze because (1) it requires the use of longrange memory to be effective and (2) it includes both continuous and discrete action sets (details in

Original: Appendix A) which makes optimization more difﬁcult.
Revised: Appendix A) which makes optimization more difﬁcult.

Original: In Figure 5, we sample 25 independent
Revised: In Figure 5, we sample 25 independent

Original: V-MPO hyperparameter settings from a wide range of values and train the networks to 2 billion environment steps (see Appendix B).
Revised: V-MPO hyperparameter settings from a wide range of values and train the networks to 2 billion environment steps (see Appendix B).

Original: Then, at various points in training (0.5B, 1.0B and 2.0B), we rank all runs by their mean return and plot this ranking.
Revised: Then, at various points in training (0.5B, 1.0B and 2.0B), we rank all runs by their mean return and plot this ranking.

Original: Models with curves which are both higher and ﬂatter are thus more robust to hyperparameters and random seeds.
Revised: Models with curves which are both higher and ﬂatter are thus more robust to hyperparameters and random seeds.

Original: Our results demonstrate that (1) the GTrXL (GRU) can learn this challenging memory environment in much fewer environment steps than LSTM, and (2) that GTrXL (GRU) beats the other gating variants in stability by a large margin, thereby offering a substantial reduction in necessary hyperparameter tuning.
Revised: Our results demonstrate that (1) the GTrXL (GRU) can learn this challenging memory environment in much fewer environment steps than LSTM, and (2) that GTrXL (GRU) beats the other gating variants in stability by a large margin, thereby offering a substantial reduction in necessary hyperparameter tuning.

Original: The values in
Revised: The values in

Original: Table 3 list what percentage of the 25 runs per model had losses that diverged to inﬁnity.
Revised: Table 3 list what percentage of the 25 runs per model had losses that diverged to inﬁnity.

Original: We can see that the only model reaching human performance in 2 billion environment steps is the GTrXL (GRU), with 10 runs having a mean score 8 and above.
Revised: We can see that the only model reaching human performance in 2 billion environment steps is the GTrXL (GRU), with 10 runs having a mean score 8 and above.

Original: 4.3. P ARAMETER C OUNT -C ONTROLLED C OMPARISONS
Revised: 4.3. P ARAMETER C OUNT -C ONTROLLED C OMPARISONS

Original: For the ﬁnal gating ablation, we compare transformer variants while tracking their total parameter count to control for the increase in capacity caused by the introduction of additional parameters in the gating mechanisms.
Revised: For the ﬁnal gating ablation, we compare transformer variants while tracking their total parameter count to control for the increase in capacity caused by the introduction of additional parameters in the gating mechanisms.

Original: To demonstrate that the advantages of the GTrXL (GRU) are notdue solely  to an increase in parameter count, we halve the number of attention heads (which also effectively halves the embedding dimension due to the convention that the embedding size is the number of heads multiplied by the attention head dimension).
Revised: To demonstrate that the advantages of the GTrXL (GRU) are not solely due to an increase in parameter count, we halve the number of attention heads (which also effectively halves the embedding dimension due to the convention that the embedding size is the number of heads multiplied by the attention head dimension).

Original: The effect is a substantial reduction in parameter count, resulting in less parameters than even the canonical TrXL.
Revised: The effect is a substantial reduction in parameter count, resulting in less parameters than even the canonical TrXL.

Original: Table 2 compares the different models to the “Thin” GTrXL (GRU) while listing their parameter counts.
Revised: Fig. 6 and Tab. 2 compare the different models to the “Thin” GTrXL (GRU), with Tab. 2 listing the parameter counts.

Original: We can see that the

Original:  Thin GTrXL (GRU) matches every other model (within variance) except the GTrXL (GRU). It even matches the next best-performing model, the GTrXL (Output), with over 10 million less parameters.
Revised: The Thin GTrXL (GRU) surpasses every other model (within variance) except the GTrXL (GRU), even surpassing the next best-performing model, the GTrXL (Output), with over 10 million less parameters.

Original: 4.3. G ATED I DENTITY I NITIALIZATION A BLATION
Revised: 4.3. G ATED I DENTITY I NITIALIZATION A BLATION

Original: All applicable gating variants in the previous sections were trained with the gated identity initialization (initial experiments found values b ( l ) g = 2 for GRU-type gating and b ( l ) g = 1 for other gating types to work well).
Revised: All applicable gating variants in the previous sections were trained with the gated identity initialization.

Original: We observed in initial Memory Maze results that the gated identity initialization signiﬁcantly improved optimization stability and learning speed.
Revised: We observed in initial Memory Maze results that the gated identity initialization signiﬁcantly improved optimization stability and learning speed.

Original: Figure 7 compares an otherwise identical 4-layer GTrXL (GRU) trained with( b ( l ) g = 2 ) and without ( b ( l ) g = 0 )
Revised: Figure 7 compares an otherwise identical 4- layer GTrXL (GRU) trained with and without the gated identity initialization.

Original: the gated identity initialization, with 10 hyperparameter samples per initial bias setting.

Original: Similarly to the previous sensitivity plots, we plot the ranked mean return ofall 10 runs at various times during training.
Revised: Similarly to the previous sensitivity plots, we plot the ranked mean return of 10 runs at various times during training.

Original: As can be seen from Fig.
Revised: As can be seen from Fig.

Original: 7, there is a signiﬁcant gap caused by the bias initialization, suggesting that preconditioning the transformer to be close to Markovian results in large learning speed gains.
Revised: 7, there is a signiﬁcant gap caused by the bias initialization, suggesting that preconditioning the transformer to be close to Markovian results in large learning speed gains.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Gating has been shown to be effective to address the vanishing gradient problem and thus improve the learnability of recurrent models.
Revised: Gating has been shown to be effective to address the vanishing gradient problem and thus improve the learnability of recurrent models.

Original: LSTM networks (Hochreiter & Schmidhuber, 1997; Graves, 2013) rely on an input, forget and output gate that protect the update of the cell.
Revised: LSTM networks (Hochreiter & Schmidhuber, 1997; Graves, 2013) rely on an input, forget and output gate that protect the update of the cell.

Original: GRU (Chung et al., 2014; Cho et al., 2014) is another popular gated recurrent architecture that simpliﬁes the LSTM cell, reducing the number of gates to two.
Revised: GRU (Chung et al., 2014; Cho et al., 2014) is another popular gated recurrent architecture that simpliﬁes the LSTM cell, reducing the number of gates to two.

Original: Finding an optimal gating mechanism remains an active area of research, with other existing proposals (Krause et al., 2016; Kalchbrenner et al., 2015; Wu et al., 2016), as well as works trying to discover optimal gating by neural architecture search (Zoph & Le, 2017)
Revised: Finding an optimal gating mechanism remains an active area of research, with other existing proposals (Krause et al., 2016; Kalchbrenner et al., 2015; Wu et al., 2016), as well as works trying to discover optimal gating by neural architecture search (Zoph & Le, 2017)

Original: More generally, gating and multiplicative interactions have a long history (Rumelhart et al., 1986).
Revised: More generally, gating and multiplicative interactions have a long history (Rumelhart et al., 1986).

Original: Gating has been investigated previously for improving the representational power of feedforward and recurrent models (Van den Oord et al., 2016;
Revised: Gating has been investigated previously for improving the representational power of feedforward and recurrent models (Van den Oord et al., 2016;

Original: Dauphin et al., 2017), as well as learnability (Srivastava et al., 2015; Zilly et al., 2017).
Revised: Dauphin et al., 2017), as well as learnability (Srivastava et al., 2015; Zilly et al., 2017).

Original: Initialization has also played a crucial role in making deep models trainable (LeCun et al., 1998; Glorot & Bengio, 2010; Sutskever et al., 2013).
Revised: Initialization has also played a crucial role in making deep models trainable (LeCun et al., 1998; Glorot & Bengio, 2010; Sutskever et al., 2013).

Original: There has been a wide variety of work looking at improving memory in reinforcement learning agents.
Revised: There has been a wide variety of work looking at improving memory in reinforcement learning agents.

Original: External memory approaches typically have a regular feedforward or recurrent policy interact with a memory database through read and write operations.
Revised: External memory approaches typically have a regular feedforward or recurrent policy interact with a memory database through read and write operations.

Original: Priors are induced through the design of the speciﬁc read/write operations, such as those resembling a digital computer (Wayne et al., 2018; Graves et al., 2016) or an environment map (Parisotto & Salakhutdinov, 2018; Gupta et al., 2017).
Revised: Priors are induced through the design of the speciﬁc read/write operations, such as those resembling a digital computer (Wayne et al., 2018; Graves et al., 2016) or an environment map (Parisotto & Salakhutdinov, 2018; Gupta et al., 2017).

Original: An alternative non-parametric perspective to memory stores an entire replay buffer of the agent’s past observations, which is made available for the agent to itself reason over either through ﬁxed rules (Blundell et al., 2016) or an attention operation (Pritzel et al., 2017).
Revised: An alternative non-parametric perspective to memory stores an entire replay buffer of the agent’s past observations, which is made available for the agent to itself reason over either through ﬁxed rules (Blundell et al., 2016) or an attention operation (Pritzel et al., 2017).

Original: Others have looked at improving performance of LSTM agents by extending the architecture with stacked hierarchical connections / multiple temporal scales and auxiliary losses (Jaderberg et al., 2019; Stooke et al., 2019) or allowing an inner-loop update to the RNN weights (Miconi et al., 2018).
Revised: Others have looked at improving performance of LSTM agents by extending the architecture with stacked hierarchical connections / multiple temporal scales and auxiliary losses (Jaderberg et al., 2019; Stooke et al., 2019) or allowing an inner-loop update to the RNN weights (Miconi et al., 2018).

Revised: Other work has examined self-attention in the context of exploiting relational structure within the inputspace (Zambaldi et al., 2019) or within recurrent memories (Santoro et al., 2018).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper we provided evidence that conﬁrms previous observations in the literature that standard transformer models, despite the recent successes in supervised learning (Devlin et al., 2019; Dai et al., 2019; Yang et al., 2019; Radford et al., 2019), are too unstable to train in the RL setting and often fail to learn completely (Mishra et al., 2018).
Revised: In this paper we provided evidence that conﬁrms previous observations in the literature that standard transformer models, despite the recent successes in supervised learning (Devlin et al., 2019; Dai et al., 2019; Yang et al., 2019; Radford et al., 2019), are too unstable to train in the RL setting and often fail to learn completely (Mishra et al., 2018).

Original: We presented a new architectural variant of the transformer model, the GTrXL, which has increased performance, more stable optimization, and greater robustness to initial seed and hyperparameters than the canonical architecture.
Revised: We presented a new architectural variant of the transformer model, the GTrXL, which has increased performance, more stable optimization, and greater robustness to initial seed and hyperparameters than the canonical architecture.

Original: The key contributions of the GTrXL are reordered layer normalization modules, enabling an initially Markov regime of training, and a gating layer instead of the standard residual connections.
Revised: The key contributions of the GTrXL are reordered layer normalization modules, enabling an initially Markov regime of training, and a gating layer instead of the standard residual connections.

Original: We performed extensive ablation experiments testing the robustness, ease of optimization and ﬁnal performance of the gating layer variations, as well as the effect of the reordered layer normalization.
Revised: We performed extensive ablation experiments testing the robustness, ease of optimization and ﬁnal performance of the gating layer variations, as well as the effect of the reordered layer normalization.

Original: These results empirically demonstrate that the GRU-type gating performs best across all metrics, exhibiting comparable robustness to hyperparameters and random seeds as an LSTM while still maintaining a performance improvement.
Revised: These results empirically demonstrate that the GRU-type gating performs best across all metrics, exhibiting comparable robustness to hyperparameters and random seeds as an LSTM while still maintaining a performance improvement.

Original: Furthermore, the GTrXL (GRU) learns faster, more stably and achieves a higher ﬁnal performance (even when controlled for parameters) than the other gating variants on the challenging multitask DMLab-30 benchmark suite.
Revised: Furthermore, the GTrXL (GRU) learns faster, more stably and achieves a higher ﬁnal performance (even when controlled for parameters) than the other gating variants on the challenging multitask DMLab-30 benchmark suite.

Original: Having demonstrated substantial and consistent improvement in DMLab-30, Numpad and Memory Maze over the ubiquitous LSTM architectures currently in use, the GTrXL makes the case for wider adoption of transformers in RL.
Revised: Having demonstrated substantial and consistent improvement in DMLab-30, Numpad and Memory Maze over the ubiquitous LSTM architectures currently in use, the GTrXL makes the case for wider adoption of transformers in RL.

Original: A core beneﬁt of the transformer architecture is its ability to scale to very large and deep models, and to effectively utilize this additional capacity in larger datasets.
Revised: A core beneﬁt of the transformer architecture is its ability to scale to very large and deep models, and to effectively utilize this additional capacity in larger datasets.

Original: In future work, we hope to test the limits of the GTrXL’s ability to scale in the RL setting by providing it with a large and varied set of training environments.
Revised: In future work, we hope to test the limits of the GTrXL’s ability to scale in the RL setting by providing it with a large and varied set of training environments.

Original: A PPENDIX A E NVIRONMENT D ETAILS
Revised: A PPENDIX A E NVIRONMENT D ETAILS

Original: Numpad: Numpad has three actions, two of which move the sphere towards some direction in the x,y plane and the third allows the agent to jump in order to get over a pad faster.
Revised: Numpad: Numpad has three actions, two of which move the sphere towards some direction in the x,y plane and the third allows the agent to jump in order to get over a pad faster.

Original: The observation consists of a variety of proprioceptive information (e.g. position, velocity, acceleration) as well as which pads in the sequence have been correctly activated (these will shut off if an incorrect pad is later hit), and the previous action and reward.
Revised: The observation consists of a variety of proprioceptive information (e.g. position, velocity, acceleration) as well as which pads in the sequence have been correctly activated (these will shut off if an incorrect pad is later hit), and the previous action and reward.

Original: Episodes last a ﬁxed 500 steps and the agent can repeat the correct sequence any number of times to receive reward.
Revised: Episodes last a ﬁxed 500 steps and the agent can repeat the correct sequence any number of times to receive reward.

Original: Observations were processed using a simple 2-layer MLP with tanh activations to produce the transformer’s input embedding.
Revised: Observations were processed using a simple 2-layer MLP with tanh activations to produce the transformer’s input embedding.

Original: DMLab-30:
Revised: DMLab-30:

Original: Ignoring the “jump” and “crouch” actions which we do not use, an action in the native
Revised: Ignoring the “jump” and “crouch” actions which we do not use, an action in the native

Original: DMLab action space consists of 5 integers whose meaning and allowed values are given in Table 4.
Revised: DMLab action space consists of 5 integers whose meaning and allowed values are given in Table 4.

Original: Following previous work on DMLab (Hessel et al., 2018), we used the reduced action set given in
Revised: Following previous work on DMLab (Hessel et al., 2018), we used the reduced action set given in

Original: Table 5 with an action repeat of 4.
Revised: Table 5 with an action repeat of 4.

Original: Observations are 72 × 96 RGB images.
Revised: Observations are 72 × 96 RGB images.

Original: Some levels require a language input, and for that all models use an additional 64-dimension LSTM to process the sentence.
Revised: Some levels require a language input, and for that all models use an additional 64-dimension LSTM to process the sentence.

Original: In Wayne et al.
Revised: In Wayne et al.

Original: (2018), the DMLab Arbitrary Visuomotor Mapping task was speciﬁcally used to highlight the MERLIN architecture’s ability to utilize memory.
Revised: (2018), the DMLab Arbitrary Visuomotor Mapping task was speciﬁcally used to highlight the MERLIN architecture’s ability to utilize memory.

Original: In Figure 9 we show that, given a similarly reduced action set as used in Wayne et al.
Revised: In Figure 9 we show that, given a similarly reduced action set as used in Wayne et al.

Original: (2018), see Table 6, the GTrXL architecture can also reliably attain human-level performance on this task.
Revised: (2018), see Table 6, the GTrXL architecture can also reliably attain human-level performance on this task.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: Memory Maze: An action in the native Memory Maze action space consists of 8 continuous actions and a single discrete action whose meaning and allowed values are given in Table 7.
Revised: Memory Maze: An action in the native Memory Maze action space consists of 8 continuous actions and a single discrete action whose meaning and allowed values are given in Table 7.

Original: Unlike for
Revised: Unlike for

Original: DMLab, we used a hybrid continuous-discrete distribution (Neunert et al., 2019) to directly output policies in the game’s native action space.
Revised: DMLab, we used a hybrid continuous-discrete distribution (Neunert et al., 2019) to directly output policies in the game’s native action space.

Original: Observations are 72 × 96 RGB images.
Revised: Observations are 72 × 96 RGB images.

Original: A CTION
Revised: A CTION

Original: N ATIVE DML AB ACTION
Revised: N ATIVE DML AB ACTION

Original: Forward (FW) [ 0, 0, 0, 1, 0]
Revised: Forward (FW) [ 0, 0, 0, 1, 0]

Original: Backward (BW) [ 0, 0, 0, -1, 0]
Revised: Backward (BW) [ 0, 0, 0, -1, 0]

Original: Strafe left [ 0, 0, -1, 0, 0]
Revised: Strafe left [ 0, 0, -1, 0, 0]

Original: Strafe right [ 0, 0, 1, 0, 0]
Revised: Strafe right [ 0, 0, 1, 0, 0]

Original: Small look left (LL)
Revised: Small look left (LL)

Original: Small look right (LR) [ 10, 0, 0, 0, 0]
Revised: Small look right (LR) [ 10, 0, 0, 0, 0]

Original: Large look left (LL )
Revised: Large look left (LL )

Original: Large look right (LR) [ 60, 0, 0, 0, 0]
Revised: Large look right (LR) [ 60, 0, 0, 0, 0]

Original: Look down [ 0, 10, 0, 0, 0]
Revised: Look down [ 0, 10, 0, 0, 0]

Original: Look up [ 0, -10, 0, 0, 0]
Revised: Look up [ 0, -10, 0, 0, 0]

Original: FW + small LL [-10, 0, 0, 1, 0]
Revised: FW + small LL [-10, 0, 0, 1, 0]

Original: FW + small LR [ 10, 0, 0, 1, 0]
Revised: FW + small LR [ 10, 0, 0, 1, 0]

Original: FW + large LL [-60, 0, 0, 1, 0]
Revised: FW + large LL [-60, 0, 0, 1, 0]

Original: FW + large LR [ 60, 0, 0, 1, 0]
Revised: FW + large LR [ 60, 0, 0, 1, 0]

Original: Fire [ 0, 0, 0, 0, 1]
Revised: Fire [ 0, 0, 0, 0, 1]

Original: [Table]
Revised: [Table]

Original: ACTION
Revised: ACTION

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: Look down [ 0, 10, 0, 0, 0] Look up [ 0, -10, 0, 0, 0]
Revised: Look down [ 0, 10, 0, 0, 0] Look up [ 0, -10, 0, 0, 0]

Original: No-op [ 0, 0, 0, 0, 0]
Revised: No-op [ 0, 0, 0, 0, 0]

Original: Image Encoder:
Revised: Image Encoder:

Original: For DMLab-30 and Memory Maze, we used the same image encoder as in (Anonymous Authors, 2019) for multitask DMLab-30.
Revised: For DMLab-30 and Memory Maze, we used the same image encoder as in (Anonymous Authors, 2019) for multitask DMLab-30.

Original: The ResNet was adapted from Hessel et al.
Revised: The ResNet was adapted from Hessel et al.

Original: (2018) and each of its layer blocks consists of a ( 3 × 3 , stride 1) convolution, followed by ( 3 × 3 , stride 2) max-pooling, followed by 2 3 × 3 residual blocks with ReLU non-linearities.
Revised: (2018) and each of its layer blocks consists of a ( 3 × 3 , stride 1) convolution, followed by ( 3 × 3 , stride 2) max-pooling, followed by 2 3 × 3 residual blocks with ReLU non-linearities.

Original: [Table]
Revised: [Table]

Original: Agent Output: As in (Anonymous Authors, 2019), in all cases we use a 256-unit MLP with a linear output to get the policy logits (for discrete actions), Gaussian distribution parameters (for continuous actions) or value function estimates.
Revised: Agent Output: As in (Anonymous Authors, 2019), in all cases we use a 256-unit MLP with a linear output to get the policy logits (for discrete actions), Gaussian distribution parameters (for continuous actions) or value function estimates.

Original: B E XPERIMENTAL DETAILS
Revised: B E XPERIMENTAL DETAILS

Original: For all experiments, beyond sampling independent random seeds, each run also has V-MPO hyperparameters sampled from a distribution (see Table 8).
Revised: For all experiments, beyond sampling independent random seeds, each run also has V-MPO hyperparameters sampled from a distribution (see Table 8).

Original: The sampled hyperparameters are kept ﬁxed across all models for a speciﬁc experiment, meaning that if one of the (cid:15) α sampled is 0.002, then all models will have 1 run with (cid:15) α = 0 .
Revised: The sampled hyperparameters are kept ﬁxed across all models for a speciﬁc experiment, meaning that if one of the (cid:15) α sampled is 0.002, then all models will have 1 run with (cid:15) α = 0 .

Original: 002 and so on for the rest of the samples.
Revised: 002 and so on for the rest of the samples.

Original: The exception is for the DMLab-30 LSTM, where a more constrained range was found to perform better in preliminary experiments.
Revised: The exception is for the DMLab-30 LSTM, where a more constrained range was found to perform better in preliminary experiments.

Original: Each model had 8 seeds started, but not all runs ran to completion due to compute issues.
Revised: Each model had 8 seeds started, but not all runs ran to completion due to compute issues.

Original: These hyperparameter settings were dropped randomly and not due to poor environment performance.
Revised: These hyperparameter settings were dropped randomly and not due to poor environment performance.

Original: We report how many seeds ran to completion for all models.
Revised: We report how many seeds ran to completion for all models.

Original: At least 6 seeds ﬁnished for every model tested.
Revised: At least 6 seeds ﬁnished for every model tested.

Original: We list architecture details by section below.
Revised: We list architecture details by section below.

Revised: All LSTM models have residual skip connections in depth.

Original: [Table]
Revised: [Table]

Original: B. T RAINING SETUP
Revised: B. T RAINING SETUP

Original: All experiments in this work were carried out in an actor-learner framework (Espeholt et al., 2018) that utilizes TF-Replicator (Buchlovsky et al., 2019) for distributed training on TPUs in the 16-coreconﬁguration (Google, 2018).
Revised: All experiments in this work were carried out in an actor-learner framework (Espeholt et al., 2018) that utilizes TF-Replicator (Buchlovsky et al., 2019) for distributed training on TPUs in the 16-core

Original: “Actors” running on CPUs performed network inference and interactions with the environment

Original: [Table]
Revised: [Table]

Original: and transmitted the resulting trajectories to the centralised “learner“
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table] [Table]
Revised: conﬁguration (Google, 2018). “Actors” running on CPUs performed network inference and interactions with the environment, and transmitted the resulting trajectories to the centralised “learner“.

Original: C M ULTI -H EAD A TTENTION D ETAILS C. M ULTI -H EAD A TTENTION
Revised: C M ULTI -H EAD A TTENTION D ETAILS C. M ULTI -H EAD A TTENTION

Original: The Multi-Head Attention (MHA) submodule computes in parallel H soft-attention operations for every time step, producing an output tensor Y ( l ) ∈ R T × D .
Revised: The Multi-Head Attention (MHA) submodule computes in parallel H soft-attention operations for every time step, producing an output tensor Y ( l ) ∈ R T × D .

Original: MHA operates by ﬁrst calculating
Revised: MHA operates by ﬁrst calculating

Original: the query Q ( l ) ∈
Revised: the query Q ( l ) ∈

Original: R H × T × d , keys K ( l ) ∈ R H
Revised: R H × T × d , keys K ( l ) ∈ R H

Original: × T × d , and values V ( l ) ∈
Revised: × T × d , and values V ( l ) ∈

Original: R H × T × d (where d = D/H ) through trainable linear projections W ( l ) Q , W ( l ) K , and W ( l ) V , respectively, and then using the combined
Revised: R H × T × d (where d = D/H ) through trainable linear projections W ( l ) Q , W ( l ) K , and W ( l ) V , respectively, and then using the combined

Original: Q , K , V , tensors to compute the soft attention.
Revised: Q , K , V , tensors to compute the soft attention.

Original: A residual connection (He et al., 2016a) to the resulting embedding E ( l ) is then applied and ﬁnally layer normalization (Ba et al., 2016).
Revised: A residual connection (He et al., 2016a) to the resulting embedding E ( l ) is then applied and ﬁnally layer normalization (Ba et al., 2016).

Original: MultiHeadAttention ( E ( l − 1) ) :
Revised: MultiHeadAttention ( E ( l − 1) ) :

Original: [Equation]
Revised: [Equation]

Original: where we used Einstein summation notation to denote the tensor multiplications, MaskedSoftmax is a causally-masked softmax to prevent addressing future information, Linear is a linear layer applied per time-step and we omit reshaping operations for simplicity.
Revised: where we used Einstein summation notation to denote the tensor multiplications, MaskedSoftmax is a causally-masked softmax to prevent addressing future information, Linear is a linear layer applied per time-step and we omit reshaping operations for simplicity.

Original: C. R ELATIVE M ULTI -H EAD A TTENTION
Revised: C. R ELATIVE M ULTI -H EAD A TTENTION

Original: The basic MHA operation does not take sequence order into account explicitly because it is permutation invariant, so positional encodings are a widely used solution in domains like language where order is an important semantic cue, appearing in the original transformer architecture (Vaswani et al., 2017).
Revised: The basic MHA operation does not take sequence order into account explicitly because it is permutation invariant, so positional encodings are a widely used solution in domains like language where order is an important semantic cue, appearing in the original transformer architecture (Vaswani et al., 2017).

Original: To enable a much larger contextual horizon than would otherwise be possible, we use the relative position encodings and memory scheme described in Dai et al.
Revised: To enable a much larger contextual horizon than would otherwise be possible, we use the relative position encodings and memory scheme described in Dai et al.

Original: In this setting, there is an additional  memory tensor M ( l ) ∈ R L × D , which is treated as constant during weight updates.
Revised: In this setting, there is an additional T -step memory tensor M ( l ) ∈ R T × D , which is treated as constant during weight updates.

Original: RelativeMultiHeadAttention( M (
Revised: RelativeMultiHeadAttention

Original: [Equation]
Revised: [Equation]

Original: , E [Equation] [Equation]
Revised: where Φ is the standard sinusoid encoding matrix

Original: [Equation]
Revised: [Equation]

Original: where Φ is the standard sinusoid encoding matrix, u ( l )

Original:  , v ( l ) ∈
Revised: u ( l ) , v ( l ) ∈

Original: R H
Revised: R H

Original: × d are trainable parameters, the ∗ represents the broadcast operation, and W R is a linear projection used to produce the relative location-based keys (see Dai et al.
Revised: × d are trainable parameters, the ∗ represents the broadcast operation [Equation] and W R is a linear projection used to produce the relative location-based keys (see Dai et al.

Original: for a detailed derivation).
Revised: for a detailed derivation).

Original: C. I DENTITY M AP R EORDERING
Revised: C. I DENTITY M AP R EORDERING

Original: The Identity Map Reordering modiﬁes the standard transformer formulation as follows: the layer norm operations are applied only to the input of the sub-module and a non-linear ReLU activation is applied to the output stream.
Revised: The Identity Map Reordering modiﬁes the standard transformer formulation as follows: the layer norm operations are applied only to the input of the sub-module and a non-linear ReLU activation is applied to the output stream.

Original: [Equation]
Revised: [Equation]

Original: See Figure 1 (Center) for a visual depiction of the TrXL-I.
Revised: See Figure 1 (Center) for a visual depiction of the TrXL-I.

Original: D DML AB -30 M EMORY /R EACTIVE P ARTITION
Revised: D DML AB -30 M EMORY /R EACTIVE P ARTITION

Original: Memory
Revised: Memory

Original: Reactive rooms select nonmatching object rooms watermaze explore obstructed goals small explore goal locations small explore object rewards few explore obstructed goals large explore goal locations large explore object rewards many rooms collect good objects train rooms exploit deferred effects train rooms keys doors puzzle language select described object language select located object language execute random task language answer quantitative question lasertag one opponent large lasertag three opponents large lasertag one opponent small lasertag three opponents small natlab ﬁxed large map natlab varying map regrowth natlab varying map randomized skymaze irreversible path hard skymaze irreversible path varied psychlab arbitrary visuomotor mapping psychlab continuous recognition psychlab sequential comparison psychlab visual search explore object locations small explore object locations large based and reactive split of levels.
Revised: Reactive rooms select nonmatching object rooms watermaze explore obstructed goals small explore goal locations small explore object rewards few explore obstructed goals large explore goal locations large explore object rewards many rooms collect good objects train rooms exploit deferred effects train rooms keys doors puzzle language select described object language select located object language execute random task language answer quantitative question lasertag one opponent large lasertag three opponents large lasertag one opponent small lasertag three opponents small natlab ﬁxed large map natlab varying map regrowth natlab varying map randomized skymaze irreversible path hard skymaze irreversible path varied psychlab arbitrary visuomotor mapping psychlab continuous recognition psychlab sequential comparison psychlab visual search explore object locations small explore object locations large [Table] E A TARI -57 R ESULTS In this section, we run the GTrXL on the multitask Atari-57 benchmark.

Revised: Although Atari-57 was not designed speciﬁcally to test an agent’s memory capabilities, we include these results here to demonstrate that we suffer no performance regression on a popular environment suite, providing further evidence that GTrXL can be used as an architectural replacement to the LSTM.

Revised: The LSTM and GTrXL are matched in width at 256 dimensions.

Revised: The GTrXL is 12 layers deep to show our model’s learning stability even at large capacity.

Revised: The LSTM architecture matches the one reported in Anonymous Authors (2019).

Revised: We train for 11.4 billion environment steps, equivalent to 200 million per environment.

Revised: We run 8 hyperparameter settings per model.