Stabilizing Transformers for Reinforcement Learning
Authors: Emilio Parisotto
Francis Song
Jack Rae
Razvan Pascanu
Caglar Gulcehre
Siddhant Jayakumar
Max Jaderberg
Raphaël Lopez Kaufman
Aidan Clark
Seb Noury
Matt Botvinick
Nicolas Heess
Raia Hadsell
Venue:  (2020)
Keywords: Deep Reinforcement Learning
Transformer
Reinforcement Learning
Self-Attention
Memory
Memory for Reinforcement Learning

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

Under review as a conference paper at ICLR 2020
STABILIZING TRANSFORMERS FOR REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Owing their ability to both effectively integrate information over long time horizons and scale to massive amounts of data, self-attention architectures have re- cently shown breakthrough success in natural language processing (NLP), achiev- ing state-of-the-art results in domains such as language modeling and machine translation. Harnessing the transformer's ability to process long time horizons of information could provide similar performance boost in partially observable re- inforcement learning (RL) domains, but the large-scale transformers used in NLP have yet to be successfully applied to the RL setting, In this work we demon- strate that the standard transformer architecture difficult to optimize, which was previously observed in the supervised learning setting but becomes especially pro- nounced with RL objectives. We propose architectural modifications that substan- tially improve the stability and learning speed of the original Transformer and XL variant. The proposed architecture, the Gated Transformer-2 XL (GTrXL), sur- passes LSTMs on challenging memory environments and achieves -tate-of-the--. results on the multi-task DMLab-30 benchmark suite, exceeding the performance of an external memory architecture. We show that the GTrXL,, trained using the same losses, has stability and performance that consistently matches or exceeds a competitive LSTM baseline, including on more reactive tasks where memory is less critical. GTrXL offers easy-to-train, simple-to-implement but substan- tially more expressive architectural alternative to the standard multi-layer LSTM ubiquitously used for RL agents in partially observable environments.
INTRODUCTION
It has been argued that self-attention architectures VVswaan et al. 2017) deal better with longer tem- poral horizons than recurrent neural networks (RNNs): by construction, they avoid compressing the whole past into fixed-size hidden state and they do not suffer from vanishing exploding gradients the same way as RNNs. Recent work has empirically validated these claims, demonstrating that self-attention architectures recurrent architectures such as the LSTM (Dai et al. 2019 Radford et al.] [2019] Devlin et al.. 2019 Yang et al. 2019). In partieular, the Transformer architecture ( /Vaswani al. (207)) has had break- through success in a wide variety of domains: language modeling (Dai et al. 2019 Radford et al.]
provide significant gains in performance over the more traditional
2019 Yang al. (2019). machine translation (Vaswani al.) 2017 Edunov al. 2018). summa- rization (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 pro- cessing 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 L.TTMs as the main mecha- nism for providing memory to the agent (Espeholt et al. 2018 Kapturowski et al. 2019: Mnih et al. 2016). Despite progress 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 obbervaly- 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 environ- ments where memory is required. In contrast these other memory architectures, the transformer

---

Under review as a conference paper at ICLR 2020
is well-tested in many challenging domains and has seen several open-source implementations in variety of deep learning frameworks
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 well. Typically a complex learning rate schedule is required (e.g., linear warmup or cosine decay) in order to train (Vaswani el al. 2017 Dai et al. 2019), or specialized weight initialization schemes are used 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 normaliza- tion coupled with the addition of new gating mechanism to key points in the submodules of the transformer. Our novel gated architecture, the Gated Transformer-XL. (GTrXL) (shown in Figure [ Right), is able to learn much faster and more reliably and exhibit significantly better final perfor- mance than the canonical transformer. We further demonstrate that the GTrXL achieves state-of- the-art results when compared to the external memory architecture MERLIN (Wayne et al., 2018) on the multitask DMLLab .ab330 suite (Beattie et al. 2016). Additionally, we surpass LSTMM signifi- cantly on memory- based DMLab-30 levels while matching performance on the reactive set, well as significantly outperforming LSTMs on memory-based continuous control and navigation envi- ronments. 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 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 RI..
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 8 Radford et al. 2019). Each layer consists of two submodules: an attention operation followed by a position-wise multi-layer network (see Figure (left)). The input to the transformer block is an embedding from the previous layer E(I-1) E R where T is the number of time steps, D is the hidden dimension, and / E [0, L.] is the layer index with L being the total number of layers. We assume (()) is an arbitrarily-obtaine 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 parallel soft- attention operations for every time step. A residual connection (He et al.) (2166a) and layer normal- ization [Ba ct al. 2016) are then applied to the output (see Appendix C for more details):
Y MultiHeedddAttentioo(((11))))))) ()(0) = E(-1) + +
( = LayerNorm ()),,, (I)
Multi-Layer Perceptron: The Multi-Layer Perceptron (MLP) submodule applies 1 1 temporal convolutional network (0) (i.e., kernel size 1, stride 1) over every step in the sequence, producing new embedding tensor Et)) E IRT'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 (Y)()))
E -0=y"0 ( + (T))
E = LayerNorm (E").
(2)
Relative Position Encodings: The basic MHA operation does not take sequence order into account explicitly because it is permutation invariant. Positional encodings are widely used solution in
e.gg wensorflow/teeserrero wwt.o//eeetttetooonnrr
httpp: httpas///itttuutttt com/
2

---

Under review as a conference paper at ICLR 2020
TrXL
TrXL.11
GTrXL
Lx
Lx
Lx
Galing Lover
Position Wise MLP
Position- Wise MLP
Posttion Wise
Lapdr-Norn
LavenNom
Galling Layer 1
Multi-Head Attention
Mutt-Head Atlention
Multi-Haad
Attantion
Layeehism
Laver-tiom
Figure 1: Transformer variants, showing just a single layer block (there are L layers total). Left:: Canonical Transformer(-XL) block with multi-head attention and position-wise MLP submodules and the standard layer normalization (Ba et al. 2016) placement with respect to the residual con- nection (He al., 2016a). Center: TrXL-I moves the layer normalization the input stream submodules. Coupled with the residual connections, there a gradient path that flows from output to input without any transformations. Right: The GTrXL block, which additionally adds gating layer in place of the residual connection of the TrXL-I.
domains like language where order is important semantic cue, appearing in the original trans- former 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 al. (2019). In this setting, there is additional T-step memory tensor (0 Tx which is treated as constant during weight updates. The MHA submodule then becomes:
Yd RelttreMuubiieccaAttttttoooooooo (U-I), EU-1))
(3)
(1(1) (0-1) +
y = LLayeroormm ())
(4)
where StopGrad is stop- goopggradient function that prevents gradients flowing backwards during back- propagation. We refer to Appendix [ for more detailed description.
GATED TRANSFORMER ARCHITECTURES
While the transformer architecture has achieved breakthrough results modeling sequences for supervised learning tasks (Vaswani et al.. 2017 Liu al. 2018 Dai al. 2019). 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 been used for temporal memory but instead as a mechanism for attention over the input space, they have had success notalyy the challenging multi-agent Starcraft 2 environment (Vinyals al. 2019). Here, the transformer was applied solely across Starcraft units and not over time.
Multiplicative interactions have been successful at stabilizing learning across wide variety of chitectures (Hochreiter & Schmidhuber 1997] Srivastava et al. 2015 Cho et al.) 2014). Motivated by this, we propose the introduction of powerful gating mechanisms place of the residual con- nections within the transformer block, coupled with changes to the order of layer normalization in the submodules. As will 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 to place the layer normalization on only the input stream the submodules, a modification described in several previous works (He al. 2016b Radford al.] 2019) Baevski
3

---

Under review as a conference paper at ICLR 2020
& Auli 2019). The model using this Identity Map Reordering is termed TrXL-I in the following, and is depicted visually in Figure (center) A key benefit to this reordering is that now enables an identity map from the imput of the transformer at the first layer 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 [ equations).
The TrXL-I already exhibits large improvement in stability and performance over TrXL (see Sec- why the Identity Map Reordering improves results is as fol- lows: assuming that the submodules at initialization produce values that are in expectation near zero, the state encoding passed un-transformed to the policy and value heads, enabling the agent to learn Markovian policy at the start of training (i.e., the network is initialized such that T (-18t....,88 .1) R T-|[s;) and / (*r (s 81) V (sg((:-1)) In many environments, reac- tive behaviours need to be learned before memory- based ones can be effectively utilized, i.e., an agent needs to learn how to walk before can learn how to remember where has walked.
tion 4.3.1 One hypothesis
3.2 GATING LAYERS
We further improve performance and optimization stability by replacing the residual connections in Equations 4 and 2 with gating layers. We call the gated architecture with the identity map reordering the Gated Transformer(-XL) (GTrXL). The final GTrXL layer block is written below:
Y Rectiiahhhliieuaaattttetoteeoooe
E(I-1)))
(5)
y() (1) ((lUU(")))LL(())
(6)
E = - (LayerNorm((""))
(7)
E = MY (1), ReLU(E")))
(8)
where g is a gating layer function. A visualization of our final architecture is shown in Figure) (right), with the modifications from the canonical transformer highlighted in red. In our experiments we ablate a variety of gating layers with increasing expressivity:
Input: The gated input connection has sigmoid modulation on the input stream, similar to the short-cut-only gating from He et al. (0016b):
(z. v) (W Ex+y
Output: The gated output connection has a sigmoid modulation on the output stream:
(x, v) + oW b()) @
Highway: The highway connection (Srivastava al. 2015 modulates both streams with sigmoid:
gd v) = oW x b(t) @ x (1 o( x b())) C
Sigmoid-Tanke Tanh: The sigmoid-tante -tanh (Sig Tagh) gate (Van den Oord et al. 2016 similar to the Output gate but with an additional tanh activation on the output stream:
g 1 (a. y)
ow
b) tun g
-nitrtttttntttppeeee gating: The Gated Recurrent Unit (GRU) (Chung et 2014 is recurrent network that performs similarly to an LSTM (Hochreiter & Schmidhuber) 1997) but has fewer parameters. We adapt its powerful gating mechanism as an untied activation function in depth:
r= -=o + U(Tx),
z = (w U.! -- b()), (0) (x. v) = 02
h = tanh(W + U 2 x)),
Gated Identity Initialization: We have claimed that the Identity Map Reordering aids policy op- timization because it initializes the agent close to Markovian policy / value function. If this is indeed the cause of improved stability, we can explicitly initialize the various gating mechanisms to be close to the identity map. This is the purpose of the bias (9) in the applicable gating layers. We later demonstrate in ablation that initially setting (() > 0 can greatly improve learning speed.
4

---

Under review as a conference paper at ICLR 2020
260
DMLab30 (Memary
340
DMLab30
120
220
200-
Mean Human Normalized Scores
90-
100
Mesn Hunnor
80
20-
-20
60
3w30
CMLab30 (Reaclive!
40
GTrXL (GRU)
LSTM
300
20
TrXL
MERLIN@100B
0.0
0.2
0.4
0.6
0.8
1.0
Env. Steps
1e10
-20
0.2
Figure 2: Average return on DMLab-30, re-scaled such that a human has mean 100 score on each level and a random policy has O. Left: Results averaged over the full DMLab-30 suite. Right: DMLab-30 partitioned into "Memory" and "Reactive" split (described in Appendix D). The GTrXL has substantial gain over LSTM in memory-based environments, while even slightly sur- passing performance on the reactive set. We plot 6-8 hyperparameter settings per architecture (see Appendix B). MERLIN scores obtained from personal communication with the authors.
EXPERIMENTS
In this section, we provide experiments on a variety of challenging single and multi-task RL do- mains: DMLab-30 (Beattie al., (016)), Numpad and Memory Maze (see Fig. 8). Crucially we demonstrate that the proposed Gated Transformer-Xe XL (GTrXL) not only shows substantial improve- ments over LSTMs on memory- based enviromments, but suffers no degradation of performance on reactive environments. The GTrXL also exceeds MERLIN (Wayne et al. 2018). an external mem- ory architecture which used a Differentiable Neural Computer (Graves al.] 2016) coupled with auxiliary losses, surpassing its performance on both memory and reactive tasks.
For all transformer architectures except when otherwise stated, we train relatively deep 12-layer networks with embedding size 256 and memory size 512. 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)). We chose to train deep networks in order to demonstrate that our results do not necessarily sacrifice complexity for stability, i.c. we are not making transformers stable for RL simply by making them shallow. Our networks have receptive fields that can potentially span any episode in the environments tested, with an upper bound on the receptive field of 6144 (12 layers x 512 memory (Dai al.) 2019)). Future work will look at scaling transformers in RL even further, e.g. towards the 52-layer network in Radford et al. (2019). See App. B for experimental details.
For all experiments, we used V-MPO (Anonymous Authors 2019), on-policy adaptation of Max- imum a Posteriori Policy Optimization (MPO) (Abdolmaleki et al. 2018a b) that performs approx- imate policy iteration based learned state-value function V(s) instead of the state-action value function used in MPO. Rather than directly updating the parameters in the direction of the policy gradient, V-MPO uses the estimated advantages first construct larget distribution for the policy update subject to sample-based KI. constraint, then calculates the gradient that partially moves the parameters toward that target, again subject to a KL constraint. V-MPO was shown achieve state-of-the-art results for LSTM-basee agents the multi-task DMLab-30 benchmark suite.
4.1 TRANSFORMER AS EFFECTIVE RL MEMORY ARCHITECTURE
We first present results of the best performing GTXXL variant, the GRU-type gating, against a com- petitive LSTM baseline, demonstrating substantial improvement the multi-task DMLabb-38 do-
5

---

Under review as a conference paper at ICLR 2020
Model
Mean Human Norm.
Mean Human Norm., 100-capped
LSTM
99.3 + 1.0
84.0
TrXL
5.0 - 0.2
5.0 3.0
T-XL-I
107.0 = 1.2
87.4
MERLIN@ 100B
115.2
89.4
GTxXL (GRU)
1
89.1 = 0.2
GTrXL (Input)
51.2 13.2
47.6 = 12.1
GTrXL (Output)
112.8 - 0.8
87.8 + 0.3
GTrXL (Highway)
90.9 + 12.9
75.2 - 10.4
GTrXL (SigTanh)
101.0 1.3
83.9 0.7
Final human-normalized return averaged across all 30 DMLab levels for baselines and GTrXL variants. We also include the 100-capped score where the per-level mean score is clipped at 100, providing a metric that is proportional to the percentage of levels that the agent is superhuman. We see that the GTrXL (GRU) surpasses LSTM by a significant gap and exceeds the performance of MERLIN Wayne al. 2018) trained for 100 billion environment steps. The GTrXL (Output) and the proposed reordered TrXL-I also surpass LSTM but perform slightly worse than MERLIN and are not as robust as GTrXL (GRU) (see Sec. 4.3.2). We sample 6-8 hyperparameters per model.
Table I:
We include standard error over runs.
Numpped Sue v.u. Performance
2X2 Numpad
Numpad
CTXKL
30-
25-
20
5-
5-
1.5
220
189
Figure 3: Numpad results demonstrating that the GTrXL has much better memory scaling properties than LSTM. Left: As the Numpad environment's memory requirement increases (because of larger pad size). the GTrXL suffers much less than LSTM. However, because of the combinatorial nature of Numpad, the GTrXI. eventually also starts dropping performance at 4x4, We plot mean and standard error of the last 200 episodes after training each model for O. 0.15B, 1.0B and 2.0B environ- ment steps for Numpad size 2, 3 and 4, respectively. Center, Right: Learning curves for the GTrXL on 2 x 2 and 4 x 4 Numpad. Even when the LSTM is trained for twice as long, the GTrXL still has a substantial improvement over it. We plot hyperparameter settings per model for learning curves.
main (Beattie et al., 2016). DMLab-30 is a large-scale, multitask benchmark comprising 30 first- person 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. [018]] Kapturowski ctt al. 2019; Hessel et al.) 2018). The levels test a wide agent competencies such as language com- prehension, navigation, handling of partial observability, memory, planning, and other forms of long horizon reasoning, with episodes lasting over 4000 environment steps. Figure shows mean return over all levels 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 random policy has a score of O), while Table I has the final performance at 10 billion environment steps. The GTrXL has signifi- cant gap over a 3-layer LSTM baseline trained using the same V-MPO algorithm. Furthermore, we external memory architecture, MERLIN (Wayne) et al. 2018). Because MERLIN was trained for 100 billion environment steps with different algo- rithm, IMPALA (Espeholt 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 final performance of the architecture as a dotted line. Despite the differences, our results demonstrate that the GTrXL can match the previous state-of-the-arr on DMLab-30. An informative split between a set of memory-based levels and more reactive ones (listed in Appendix D) reveals that our model specifically has large improvements in environments where memory plays critical role. Mean- while, GTrXL also shows improvement over LSTMs on the set of reactive levels, as memory can still be effectively utilized in some of these levels.
included the final results of a previously-publice

---

Under review as a conference paper at ICLR 2020
160
DMLab30 (Memory)
(Reactive)
140
Normalzed Scores
120
GT/LL (GRU)
120
GTXXL (Input)
Mormalzed
100
80
(Output)
GT/XL (SigTanh)
60
(Highway)
40
40
TXL-1
20
20
TXXL
MERLIN@100B
-20
-20
0.4
0.6
0.8
1.0
0.2
0.4
Env. Stops
1e10
Env. Staps
1e10
Figure 4: Learning curves for the gating mechanisms, along with MERLIN score at 100 billion frames reference point. We anism on the reactive set of tasks. On the memory environments, the GRU gating has significant gain in learning speed and attains the highest final performance at the fastest rate. We plot both mean (bold) and the individual 6-8 hyperparameter samples per model (light).
that the GRU performs well as any other gating mech-
sensiiiiity anclysis (0.38 cnv stops) Nemory Maze
Sensitivity Arilest 11.48 ent slepst Momory Maxe
Sensstiiiyy Armlysi 2.00 MTV S.RES0 Mimory Mars
Poon vetun
TXXL
Human
12
25
20
25
30
25
20
25
10
15
20
Hark
Kark
tank
Figure 5: Sensitivity analysis of GTrXL variants versus TrXL and LSTM baselines. We sample 25 different hyperparameter sets and seeds and plot the ranked average return at 3 points during training (0.5B, 1.0B and 2.0B environment steps). Higher and flatter lines indicate more robust architectures. The same hyperparameter sampling distributions plot human performance as a dotted line.
used across models (see Appendix [B. We
4.2 SCALING WITH MEMORY HORIZON
We next demonstrate that the GTrXL scales better compared to an LSTM when an environment's 's temporal horizon is increased, using the "Numpad" continuous control task of Humplik et al. (2019) which allows an easy combinatorial increase the temporal horizon. Numpad, a robotic agent is situated a platform resembling the 3x3 number pad of a telephone (generalizable to N N pads). The agent can interact with the pads by colliding with them, causing them to be activated (visualized in the environment state the number pad glowing). The goal of the agent is activate a specific sequence of up to N² numbers, but without knowing this sequence a priori. The only feedback the agent gets by activating numbers: 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. Each correct number in the sequence only provides reward once, i.e. each subsequent activation of that number will no longer provide rewards. Therefore the agent must explicitly develop search strategy determine the correct pad sequence. Once the agent completes the full sequence, all pads are reset and the agent gets chance to repeat the sequence again for more reward. This means higher reward directly translates into how well the pad sequence has been memorized. An image of the scenario is provided in Figure 3 There is the restriction that contiguous pads in the sequence must be contiguous in space, i.c. the next pad in the sequence can only in the Moore neighborhood of the previous pad. Furthermore,
pad can be pressed twice in the sequence.
We present two results in this environment in Figure S The first measures the final performance of the trained models function of the pad size. We can see that LSTM performs badly on all 3 pad sizes, and performs worse as the pad size increases from 2 to 4. The GTrXL performs much better, and almost instantly solves the environment with its much more expressive memory. On the center
7

---

Under review as a conference paper at ICLR 2020
(Memory)
120
DMLab30 (Reactive)
Mexm Human Normaizeed Scores
Thin GTXXL (GRU)
120
(Input)
100
GTXXL (Output)
60
(SigTanh)
60
40
Large
GTXXL (Hiigwway)
40-
TXXL-1
Mean
20
MERLING 1000
-20
-20
0.4
Staps
Erv Steps
Figure 6: Learning curves comparing a thinner GTrXL (GRU) with half the embedding dimension of the other presented gated variants and TrXL baselines. The Thin GTrXL (GRU) has fewer pa- rameters than any other model presented but still matches the performance of the best performing counterpart, the GTrXL (Output), which has over 10 million more parameters. We plot both mean (bold) and 6-8 hyperparameter settings (light) per model.
Mean Human Norm. Score
Param
Model
Millions
L.TM
99.3 + 1.0 103.5 + 0.9
9.25M
Model
Diverged
Large LSTM
51.3M
LSTM
0%
TrXI.
+ 0.2
28.6M
TEXL.
0%
T-XI.-I
107.0 1.2
28.6M
TrXL-I
16%
Thin GT:XL (GRU)
111.5 0.6
22.4M
GTrXI (GRU)
0%
GTrXL (GRU)
0.3
66.4M
GTyXL (Output)
12%
GTrXL (Input)
51.2
34.9M
GTrXL (Output)
112.8 +.8
34.9M
GTrXL (Highway)
90.9 12.9
34.9M
Table Percentage of the 25 parameter
GTrXL (SigTanh)
101.0
41.2M
settings where the training loss diverged within 2 billion env. steps. We do not report numbers for GTrXL gating types that were unstable in DMLab-30. For diverged runs we plot the returns in Fig- ure alas 0 afterwards.
Table 2: Parameter-conttooled ablation. We report standard error of the means of 6-8 runs per model.
and right images, we provide learning curves for the 2 2 and 4 x 4 Numpad environments, and show that even when the LSTM is trained twice as long it does not reach GTrXL's performance.
4.3 GATING VARIANTS + IDENTITY MAP REORDERING
We demonstrated that the GRU-type-gatte gated GTrXL can achieve state-of-the-an results on DMLab- 30, surpassing both 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. However, the ques- tion remains whether the expressive gating mechanisms of the GRU could be replaced by simpler alternatives. In this section, we perform extensive ablations on the gating variants descrihed Sec- tion 3.2 and show that the GTrXL (GRU) has improvements in learning speed, final performance and optimization stability over all other models, even when controlling for the number parameters.
4.3.1 PERFORMANCE A ABLATION
We first report the performance of the gating variants in DMLab-30. Table Figure 4 show the final performance and training curves of the various gating types in both the memory / reactive split, respectively. The canonical TrXL completely fails to learn, while the TrXL-I improves over the LSTM. Of the gating varieties, the GTrXL (Output) can recover large amount of the performance of the GTrXL (GRU), especially the reactive set, but as shown in Sec. 4.3.2 is generally far less stable. The GTrXL (Input) performs worse than the TrXL-I, reinforcing the identity map path hypothesis. 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.
8

---

Under review as a conference paper at ICLR 2020
Gabed identity Inbairaation Adtation .00 steps) Mamory Mace
Gated
inttaa intaliaation Ablation Memory Maze
Hamary Maze
0-
inc.
STTLL tonu) dity
Humer
Figure 7: Ablation of the gated identity initialization Memory Maze by comparing 10 runs of model run with the bias initialization and 10 runs of model without. Every run has independently sampled hyperparameters from a distribution. We plot the ranked mean return of the 10 runs of each each model at 1, 2. and 4 billion environment steps. Each mean return is the average of the past 200 episodes at the point of the model snapshot. We plot human performance as dotted line.
4.3.2 HYPERPARAMETER AND SEED SENSITIVITY
Beyond improved performance, we next demonstrate a significant reduction in hyperparameter and seed sensitivity for the GTrXL (GRU) compared to baselines and other GTrXL variants. We use the "Memory Maze" environment, a memory-basee navigation task in which the agent must discover the location of apple randomly placed in maze of blocks. 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 fixed. The agent can make of landmarks situated around the room return as quickly as possible to the apple for subsequent rewards. Therefore, an effective mapping of the environment results in more frequent returns to the apple and higher reward.
We chose to perform the sensitivity ablation on Memory Maze because (1) it requires the use of long- range memory to be effective and (2) it includes both continuous and discrete action sets (details in Appendix A which makes optimization more difficult. In Figure 5 we sample 25 independent V-MPO hyperparameter settings from a wide range of values and train the networks to 2 billion environment steps (see Appendix B). 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. Models with curves which are both higher and flatter are thus more robust to hyperparameters and random seeds. 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 large margin, thereby offering a substantial reduction in necessary hyperparameter tuning. The values in Table 3 list what percentage of the 25 runs per model had losses that diverged to infinity. 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.
4.3.3 PARAMETER COUNT-CONTROLLED COMPARISONS
For the final gating ablation, 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. 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 the convention that the embedding size the number of heads multiplied by the attention head dimension). The effect is substantial reduction in parameter count, resulting in less parameters than even the canonical TrXI.. Fig. 6 and Tab. 2 compare the different models to the "Thin" GTrXL (GRU), with Tab. listing the parameter counts. We include a parameter- parameter-matched matched LSTM model with 12 layers and 512 hidden size. The Thin GTxXL (GRU) surpasses every other model (within variance) except the GTrXL (GRU), even surpassing the next best best-performing model, the GTrXL (Output), with over 10 million less parameters.
4.3.4 GATED IDENTITY INITIALIZATION ABLATION
All applicable gating variants in the previous sections were trained with the gated identity initializa- tion. We observed in initial Memory Maze results that the gated identity initialization significantly
9

---

Under review as a conference paper at ICLR 2020
improved optimization stability and learning speed. Figure 7 compares otherwise identical 4- layer GTxXL (GRU) trained with and without the gated identity initialization. Similarly to the pre- vious sensitivity plots, plot the ranked mean return of 10 runs at various times during training. As can be seen from Fig. 7 there is significant gap caused by the bias initialization, suggesting that preconditioning the transformer to be close to Markovian results in large learning speed gains.
5 RELATED WORK
Gating has been shown to be effective to address the vanishing gradient problem and thus improve the learnability of recurrent models. LSTM networks (Hochreiter & Schmidhuber, 1997, Graves 2013) rely on an input, forget and output gate that protect the update of the cell. GRU (Chung al, 2014 Cho et al., 2014) is another popular gated recurrent architecture that simplifies the LSTM cell, reducing the number of gates to two. Finding an optimal gating mechanism remains active area of research, with other existing proposals (Krause et al.) 2016 Kalchbrenner et al., 2015; Wu let al.
2016), as well as works trying discover optimal gating by neural architecture search (Zoph 2017 More generally, gating and multiplicative interactions have a long history (Rumelhart al.
& Le
1986). Gating has been investigated previously for improving the representational power of feedforward and recurrent models (Van den Oord et al., [0166] Dauphin et al.] 2017]. as well as learnability (Srivastava et al.) 2015 Zilly et al. 2017). Initialization has also played a crucial role in making deep models trainable (LeCun et al. 1998; Glorot & Bengio 2010 Sutskever et al. 2013).
There has been a wide variety of work looking at improving memory in reinforcement learning agents. External memory approaches typically have a regular feedforward or recurrent policy in- teract with a memory database through read and write operations. Priors are induced through the design of the specific read/write operations, such as those resembling a digital computer (Wayne) al., 2018 Graves et al., 2016) or an environment map (Parisotto & Salakhutdinov 2018: Gupta et al. 2017). An alternative non-parametric perspective to memory stores entire replay buffer
of the agent's past observations, which is made available for the agent to itself reason over either
through fixed rules (Blundell et al., 2016) an attention operation (Pritzel et al., 2017). Others have
looked at improving performance of LSTM agents by extending the architecture with stacked hier-
archical connections/ multiple temporal scales and auxiliary losses (Jaderberg et al. 2019 Stooke
let al., 2019) or allowing an inner-loop update to the RNN weights (Miconi et al. 2018). Other
work has examined self-attention in the context of exploiting relational structure within the input-
space (Zambaldi et al. 2019) or within recurrent memories (Santoro et al. 2018).
6 CONCLUSION
In this paper we provided evidence that confirms previous observations in the literature that standard transformer models are unstable to train in the RL setting and often fail to learn completely (Mishra ctt al. 2018). We presented a new architectural variant of the transformer model, the GTrXL, which has increased performance, more stable optimization, and greater robustness initial seed and hy- perparameters than the canonical architecture. The key contributions of the GTrXL are reordered layer normalization modules and a gating layer instead of the standard residual connection. We performed extensive ablation experiments testing the robustness, ease of optimization and final per- formance of the gating layer variations, as well as the effect of the reordered layer normalization. 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 performance improvement. Furthermore, the GTrXL. (GRU) learns faster, more sta- bly and achieves a higher final performance (even when controlled for parameters) than the other gating variants on the challenging multitask DMLab-30 30 benchmark suite.
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. A core benefit 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. In future work, we hope to test the limits of the GTrXL's ability to scale in the RL setting by providing it with large and varied set of training environments.
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Abhas Abdolmaleki, Jost Tobias Springenberg. Jonas Degrave, Steven Bohez, Yuval Tassa, Dan Belov, Nicolas Heess, and Martin Riedmiller. Relative Entropy Regularized Policy Iteration. arXiv preprint, 20.
pat/ 02256. pdf
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Mar- tin Riedmiller. Maximum Posteriori Policy Optimisation. Int. Conf. Learn. Represent, 2018b.
URL http:::///// httt:https www.xxiw...././///e//./a
06920. 0
Anonymous Authors. V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control. Under review, Int. Conf. Learn. Represent. 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 06450, 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. Int. Conf. Learn. Represent., 2019.
Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq. Simon Green, Víctor Valdés, Amir Sadik, et al. Deepmind lab, arXiv preprint a12.vv161..03381, 2016.
Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z. Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. arXiv preprint arXiv: 190.:66064466 2016.
Peter Buchlovsky, David Budden, Dominik Grewe, Chris Jones, John Aslanides, Frederic Besse, Andy Brock, Aidan Clark, Sergio Gomez Colmenarejo, Aedan Pope, Fabio Viola, and Dan Belov. Te-Replicato:: Distributed Machine Learning for Researchers. arXiv preprint. 2019.
URL http::/tt//ii...../tp/ org. abs 00465
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1078, 2014.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks sequence modeling. arXiv preprint arXiv:1412.3355, 2014.
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime Carbonell, Quoe Le, and Ruslan Salakhutdinov.
Transformer- -e:::ormee--ee:.: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 2978-9988999 -978-,988. Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. N0.18833/v1/P19-111 URL
https: www. acl.aclllle...... wwwpgrgaatoooggoooeoooo
1285
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. Proceedings of the 34th International Conference on Machine Learming- Learning-Volumee Volume 70, pp. 933-941. JMLR. org, 2017.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. Int. Conf. Learn. Represent., 2018.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long and Short Papers), pp. 4171-4186, 2019.
Sergey Edunov, Myle Out, Michael Auli, and David Grangier. Understanding back-translation at scale. In Proceedings the 2018 Conference on Empirical Methods in Natural Language Pro- cessing, pp. 489-500, 2018.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu. Tim Harley, Iain Dunning, Shane Legg, and Koray Kavukcuoglu. Impala: Scalable distributed deep-rl with importance weighted lctor-learner architectures. International Conference Machine Learning, pp. 1406-1415, 2018.

---

Under review as a conference paper at ICLR 2020
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neu- ral networks. In In Proceedings of the International Conference on Artificial Intelligence and Statistics (AISTATSIO). Society for Artificial Intelligence and Statistics, 2010.
Google. Cloud TPU, 2018. URL
Alex Graves. Generating sequences with recurrent neural networks, 2013.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwiñska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538
(7626):471,2))) 2016.
Saurabh Gupta, James Davidson, Sergey Levine, Rahul Sukthankar, and Jitendra Malik. Cognitive mapping and planning for visual navigation. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 2616-2625, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pp. 630-645. Springer, 2016b.
Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt, and Hado van Hasselt. Multi-task Deep Reinforcement Learning with PopArt. arXiv preprint, 2018. URL pdf
https::/arrx: http://tttt..9.0402200
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A. Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint, 2019. URL
Max Jaderberg, Wojeiech M Czarnecki, Iain Dunning, Luke Marris, Guy Lever, Antonio Garcia Castaneda, Charles Beattie, Neil C Rabinowitz, Ari Morcos, Avraham Ruderman, et al. Human- level performance in 3d multiplayer games with population-based reinforcement learning. Sci- ence, 364(6444):859-865, 2019.
Nal Kalchbrenner, Ivo Danihelka, and Alex Graves. Grid long short-term memory. 07 2015.
Steven Kapturowski, Georg Ostrovski. John Quan, Remi Munos, and Will Dabney. Recurrent expe- rience replay in distributed reinforcement learning. Int. Conf. Learn. Represent. 2019.
Ben Krause, Liang Lu, lain Murray, and Steve Renals. Multiplicative Istm for sequence modelling, 2016.
Yann LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient backprop. In Neurul Networks: Tricks of the Trude, This Book an Ouigrowth of 1996 NIPS Workshop. pp. 9-50, London, UK, UK, 1998. Springer- Verlag. ISBN 3-540-65311-2. URL http://dl. acm.
www.iioonn...m..mmmmmmmmmmmm.. efm?id 545754.5555 658382]
Peter Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. Int. Conf. Learn. Represent., 2018.
Yang Liu and Mirella Lapata. Text summarization with pretrained encoders. Conference on Empir- ical Methods in Natural Language Processing.
Thomas Miconi, Kenneth Stanley, and Jeff Clune. Differentiable plasticity: training plastic neural networks with backpropagation. In International Conference on Machine Learning, pp. 3556- 3565, 2018
12

---

Under review as a conference paper at ICLR 2020
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A Simple Neural Attentive Meta- Learner. Int. Conf. Learn. Represent., 2018. URL htttttsss//ppeeeeeeeeee http://wwwwweeepddeeeeeeeee
B B1DmUzZAn
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuogiu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928-1937, 2016.
Michael Neunert, Abbas Abdolmaleki, Markus Wulfmeier, Thomas Lampe, Nicolas Heess, Roland Hafner, and Martin Riedmiller. Continuous- Discreee Deep Reinforcement Learning for Hybrid Control in Robotics. To appear in Conf. on Robot Learn. (CoRI.), 2019.
Emilio Parisotto and Ruslan Salakhutdinov. Neural map: Structured memory for deep reinforcement learning. Int. Conf. Learn. Represent., 2018.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomènech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In International Conference on Machine Learning, pp. 2827-2836, 2017.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. 2019.
D. E. Rumelhart, G. E. Hinton, and J. L. McClelland. Parallel distributed processing: Explorations in the microstructure of cognition, vol. 1. chapter 2: A General Framework for Parallel Distributed Processing, pp. 45-76. MIT Press, Cambridge, MA, USA, 1986. ISBN 0-262-68053-X, URL
http://d:: dl. acm. org/ citatitatta. ion. cfm wiffmiid=100 1tt.:/id-1422220 004779 104286
Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski. Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap. Relational recurrent neural networks. In Advances in Neural Information Processing Systems, pp. 7299-7310, 2018.
Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber: Highway networks. arXiv preprint 105iv:1505.00387, 2015.
Adam Stooke, Valentin Dalibard, Siddhant M Jayakumar, Wojeiech M Czarnecki, and Max Jader- berg. Perception-predicee prediction-reactinn reacttonn agents for deep reinforcement learning. SPIRL Workshop, 2019.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In Sanjoy Dasgupta and David McAllester (eds.), Pro- ceedings of the 30th International Conference on Machine Learning. volume 28 of Proceedings of Machine Learning Research, pp. 1139-1147, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL http:/tttp/:pceeeeeeeeeeeeeeee/e
ngs www.rr.rreeee press w28/ sutskeveri3 html
Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, and Koray Kavukcuoglu. Conditional image generation with pixelenn decoders. In Advances in neural information processing systems, pp. 4790-4798, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz. Kaiser, and Illia Polosukhin. Attention all you need. Advances Neural Informa- tion Processing Systems, pp.
2017.
wwapperrs.nip...o cc/ cc/pa/r//
118l-attention is all you need, pdE
Oriol Vinyals, Igor Babuschkin, Junyoung Chung, Michael Mathicu, Max Jaderberg, Wojtek Czarnecki, Andrew Dudzik, Aja Huang. Petko Georgiev, Richard Powell, Timo Ewalds, Dan Horgan, Manuel Kroiss, Ivo Danihelka, John Agapiou, Junhyuk Oh, Valentin Dal- ibard, David Choi, Laurent Sifre, Yury Sulsky, Sasha Vezhnevets, James Molloy, Trevor Cai, David Budden, Tom Paine, Caglar Gulcehre, Ziyu Wang, Tobias Pfaff, Toby Pohlen, Dani Yogatama, Julia Cohen, Katrina McKinney, Oliver Smith, Tom Schaul, Timothy Lil- licrap. Chris Apps, Koray Kavukcuoglu, Demis Hassabis, and David Silver. AlphaStar: Mastering the Real-Time Strategy Game StarCraft II. /ttp://///deeeedddddd wwt.wad...ccoooo/o// -eeal-ime----e commee-strttee strategy stw.craft-rrr////aoo stt 2019.
13

---

Under review as a conference paper at ICLR 2020
Greg Wayne, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Z. Leibo, Adam Santoro, Meviana Gemici, Malcolm Reynolds, Tim Harley, Josh Abramson, Shakir Mohamed. Danilo Rezende, David Saxton, Adam Cain, Chloe Hillier, David Silver, Koray Kavukcuoglu, Matt Botvinick, Demis Hassabis, and Timothy Lillicrap. Unsupervised predictive memory in goal-directed agent. arXiv preprini arXiv: 10760, 2018.
Yuhuai Wu, Saizheng Zhang, Ying Zhang, Yoshua Bengio, and Ruslan R Salakhutdinov. On multi- plicative integration with recurrent neural networks. Advances in neural information processing systems, pp. 2856-2866 2016.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in Neural Information Processing Systems, 2019.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, Murray Shanahan, Victoria Langston, Razvan Pascanu, Matthew Botvinick, Oriol Vinyals, and Peter Battaglia. Deep re- inforcement learning with relational inductive biases. In International Conference on Learning net/ Forum?]G=-HXaFOCCK)
Representations, 2019. URL httphttssss//peeeeeiee
Julian Georg Zilly, Rupesh Kumar Srivastava, Jan Koutník, and Jürgen Schmidhuber. Recurrent highway networks. In Proceedings of the 34th International Conference Machine Learning- Volume 70, pp. 4189-4198. JMLR. org, 2017.
Barret Zoph and Quoc V. Le. Neural architecture search with reinforcement learning. In ICLR, 2017. URL httpstps: arw.vaxxvvvvv..... http///wwww..op/b///lll1..1
14

---

Under review as a conference paper at ICLR 2020
The Numpad environment, showing the controllable "sphere" robot and full 3x3 pad. Pads are activated when the robot collides with their center. The robot can move on the plane as well as jump to avoid pressing numbers. Right: Top down view of "Memory Maze": (1) Central chamber, (2) blocks among which the apple is placed, (3) landmarks the agent can use to locate the apple. (4) one of the possible location of the apple.
Figure 8: Left:
APPENDIX
A ENVIRONMENT DETAILS
Numpad: Numpad has three actions, two of which move the sphere towards direction in the x.y plane and the third allows the agent to jump in order to get over a pad faster. 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 an incorrect pad is later hit), and the previous action and reward. Episodes last a fixed 500 steps and the agent can repeat the correct sequence any number of times to receive reward. Observations were processed using a simple 2-1ayer MLP with tanh activations to produce the transformer's 's input embedding.
DMLab-30: Ignoring the "jump" and "crouch" actions which we do not use, action in the native DMLab action space consists of integers whose meaning and allowed values are given in Table Following previous work DMLab (Hessel et al., 2018), we used the reduced action set given in Table with an action repeat 4. Observations are 72 96 RGB images. Some levels require language input, and for that all models use an additional 64-dimension LSTM to process the sentence.
In Wayne et al. (2018). the DMLab Arbitrary Visuomotor Mapping task was specifically used to highlight the MERLIN architecture's 's ability to utilize memory. Figure D we show that, given similarly reduced action set used in WWyne et al. (218)), see Table 6 the GTrXL architecture can also reliably attain human-level performance on this task.
ACTION NAME
RANGE
LOOK LEFT RIGHT PIXELS FRAME LOOK LOOK_DOWN, DOWN, PIXELS_PEI PER_FRAME
[-512, 512]
[-512, 512] [-1, 1]
STRAFE RIGHT
MOVE BACK FORWARD FIRE
[-1, 1]
[0. 1]
Table 4: Native action space for DMLab. See https: ptttps:/githhen http//ggtttbtttt///// com/ dewpminddeeeeeeeeeeeeeeeee.eeeee lab/ httb://asterr//oss//eess////
mdl for more details.
An action in the native Memory Maze action space consists of 8 continuous actions and single discrete action whose meaning and allowed values given in Table 7 Unlike for DMLab, we used hybrid continuous-discred distribution (Neunert et al. 2019) to directly output policies in the game's native action space. Observations are 72 x 96 RGB images.
Memory Maze:
15

---

Under review as a conference paper at ICLR 2020
70
Episode reward
50
40
30
20
psychlab arbitrary, visuomotor mapping
10
GTrXL (10)
Environment (4xagent) steps (B)
Figure 9: Learning curves for the DMLab Arbitrary Visuomotor Mapping task using a reduced action set.
ACTION
NATIVE DMLAB ACTION
Forward (FW)
0,
0]
Backward (BW)
0, 0,
C,
-1, 01
Strafe left
0,
Strafe right
0,
O, 1, 0, 01
Small look left (LL)
[-10,
O,
01
Small look right (LR)
10,
O, C, 01 O,
Large look left (LII)
[-60,
C, 0, 01 O, C, G, 0]
Large look right (LR)
60,
Look down
O, 10, C, U, 01
Look up
0, -10, C, 6, 01
FW + small
[-10, 0, C,
01
+ small
[10, 0, C, 1, 01
FW + large LL
[-60,
0, C, 1, 0]
FW large LR
60, 0, C, 1, 0]
Fire
0, 0, C, C, 11
Table 5: Simplified action set for DMLab from Hessel et al. (2018).
ACTION
NATIVE DMLAB ACTION
Small look left (LL) Small look right (LR)
[-10,
0] 0, C, 0, 0]
10,
Look down Look up
0, 10,
C,
0,
01
0, -10,
C, 0,
01
No-op
0,
0]
Table 6: Simplified action set for DMLab Arbitrary Visuomotor Mapping (AVM). This action set is the same as the one used for AVM in Wayne al. (2018) but with an additional no-op, which may also be replaced with the Fire action.
Image Encoder: For DMLab-30 30 and Memory Maze, we used the same image encoder as in (Anony-
mous Authors 2019) for multitask DMLab-30. The ResNet adapted from Hessel et al. (2018)
and each of its layer blocks consists of a (3 x 3, stride 1) convolution, followed by (3 3, stride 2) max-pooling, followed by 2 3 residual blocks with ReLU non-linearities.
16

---

Under review as a conference paper at ICLR 2020
ACTION NAME
RANGE
LEFT_RIGHT
[-1.0. 1.0]
LOOK DOWN_UP
[-1.0, 1.01
STRAFE, LEFT RIGHT
[-1.0, 1.0)
MOVE_ BACK FORWARD
BAACBBARRARRRRRREE RIGHT
[-1.0, 1.0]
HAND HAND_ROTATE, AROUND_RIGHT
[-1.0, 1.0]
HAND ROTATE_AROUND _UP
[-1.0, 1.01
HAND ROTATE AROUND FORWARD
[-1.0, 1.0]
HAND PUSH PULI
[-10.0, 10.0)
HAND GRIP
(0.
Table Hybrid action set for Memory Maze, consisting of continuous actions and single discrete action.
Agent Output: As in (Anonymous Authors 2019). in all cases we use a 256-umit MLP with linear output get the policy logits (for discrete actions), Gaussian distribution parameters (for continuous actions) or value function estimates.
B EXPERIMENTAL DETAILS
For all experiments, beyond sampling independent random seeds, each run also has V-MPO hyper- parameters sampled from a distribution (see Table [). The sampled hyperparameters are kept fixed across all models for specific experiment, meaning that if one of the &a sampled is 0.002, then all models will have 1 run with ca = 0.002 and so on for the rest of the samples. The exception is for the DMLab-30 LSTM, where experiments. Each model had 8 seeds started, but not all runs ran to completion due to compute issues. These hyperparameter settings were dropped randomly and not due to poor environment per- formance. We report how many seeds ran to completion for all models. At least seeds finished for every model tested. We list architecture details by section below. All LSTM models have residual skip connections in depth.
constrained range
found perform better in preliminary
Hyperparameter
Environment Numpad
DMLab-30
Memory Maze
Batch Size
128
128
128
Unroll Length
95
95
95
Discount
0.99
0.99
0.99
Action Repeat
1
4
Pixel Control Cost
10-3
Target Update Period
10
10
10
Initial n
1.0
10.0
1.0
Initial a
5.0
5,0
Initial ag
1.0
1.0
Initial oz
1.0
1.0
0.1
0.1
LSTM [0.001, 0.025)
ca (log-uniform)
[0.001, 0.1)
TrXL Variants [0.001, 0.1)
Eay, (log-uniform)
[0.005, 0.01)
(0.005,0.01)
Eag (log-uniform)
15 10 6 4 x 10-4)
[5 x 10-5,40 x 10-5)
Table V-MPO hyperparameters per environment.
B.1 TRAINING SETUP
All experiments in this work were carried out an actor-learner framework (Espeholt et al. 2018
that utilizes Te-Repcia (Buchlovsky al. 2019) for distributed training on TPUs in the 16-core
17

---

Under review as a conference paper at ICLR 2020
Hidden Dim.
Runs
Model
# Layers
Head Dim.
# Heads
Memory Size
Completed
LSTM
3
256
Large LSTM
12
512
TrXL
12
64
512
512
TrXL-I
12
512
512
GTrXL (GRU)
12
64
512
512
GTrXL (Input)
12
64
512
512
GTrXL (Output)
12
64
512
512
7
GTrXL (Highway)
12
64
512
512
7
GTrXL (SigTanh)
12
64
512
512
Thin GTrXL (GRU)
12
64
256
512
8
Table DMLab-30 Ablation Architecture Details. We report the number of runs per model that ran to completion (i.e. 10 billion environment steps). We follow the standard convention that the hid- den/embedding dimension of transformers is equal to the head dimension multiplied by the number
of heads. (Sec. 4.1 Sec. 4.3).
Hidden
Runs Completed
Model
Layers
Head Dim.
Heads
Memory Size
Dim.
LSTM
256
5
GTrXL (GRU)
12
64
8
256
512
5
Table 10: Numpad Architecture Details. (Sec.
4.2
Hidden Dim.
Model
Layers
Head Dim.
# Heads
Memory Size
LSTM
256
TrXL
12
64
8
256
512
TrXL-I
12
64
8
256
512
GTrXL (GRU)
12
64
8
256
512
GTrXL (Output)
12
64
8
256
512
Table 11: Sensitivity ablation architecture details (Sec. 4.3.2)
Hidden
Runs Completed 8
Model
Layers
Head Dim.
Heads
Memory Size
Dim.
GTrXL (GRU)
64
4
256
512
Table 12: Gated identity initialization ablation architecture details (Sec. 4.3.4
configuration (Google 2018 "Actors" running CPUs performed network inference and interac- tions with the environment, and transmitted the resulting trajectories to the centralised "learner".
C MULTI-HEAD ATTENTION DETAILS
C.I MULTI-HEAD ATTENTION
The Multi-Head Attention (MHA) submodule computes in parallel H soft-attention operations for every time step, producing output tensor (1) Tx MHA operates by first calculating the query Q!)) E 4xTxd keys K(t) E "xTxd and values y( E "xxxx (where d = D/H) through trainable linear projections
w w and w respectively, and then using the combined Q. K. V. tensors to compute the soft attention. A residual connection (He al. 2016a) to the resulting embedding E! is then applied and finally layer normalization (Ba al. 2016
18

---

Under review as a conference paper at ICLR 2020
000 00
CKO
goa 6
Colo. 10000
000
Figure 10: The 25 hyperparameter settings sampled for the sensitivity ablation (Sec. in log scale and values are sampled from the corresponding ranges given in Table
4.3.2
X-axis
MultiHeedAAtennie B(-1))):
Q K(0, V() T E(-1)) w E(-1) w E0-11)
@ QheaKama
(1)
(10)
W(l
MaskedSoftmax axis=m) htm
(11)
(())
htd w htm v!th
(12)
fr(1) F(111) + Linear(y
(13)
( (Y
(14)
where used Einstein summation notation to denote the tensor multiplications, MaskedSoftmax is a causally-maske masked softmax to prevent addressing future information, Linear a linear layer applied per time-step and omit reshaping operations for simplicity.
C.2 RELATIVE MULTI-HEAD ATTENTION
The basic MHA operation does not take sequence order into account explicitly because is permu- tation invariant, positional encodings 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 the relative position encodings and memory scheme described in Dai et al. (2019). In this setting, there an additional T-step memory tensor (!) Txx,, which is treated as constant during weight updates.
widely used solution in domains like language where
19

---

Under review as a conference paper at ICLR 2020
RelativeMulllHHeadAtteennniee (Maiiionnnm))) E(!-1))
D(1-1) [M F(
(15)
Q0, K V -WO E(I-1), w E(1-1) w E(--1)
(16)
R W (0 +
(17)
O htm
Qid Kimdd + Qhed Rhand + Uhed + Uhad Rhmdd
(18)
W
MaskedSoftmax (a axis=m)
(19)
htm
(I) w(() ((())
htd
=
htm
(00)
y() ( + Linear(Y
(21)
y LayerNorm(Y10))
(22)
where the standard sinusoid encoding matrix, u vl0) E Rxdd are trainable parameters, the represents the broadcast operation, and WR is linear projection used to produce the relative location-based keys (see Dai et al. (2019) for a detailed derivation).
C.3 IDENTITY MAP REORDERING
The Identity Map Reordering modifies the standard transformer formulation follows: the layer norm operations are applied only to the input of the sub-module and a non-linear ReLU activation applied to the output stream.
((--
(23)
yl E(-1) + (e
(24)
F 1 (LLayerNorm((
(25)
E = y + (e
(26)
See Figure 0 (Center) for visual depiction of the TrXL.-I.
20

---

Under review as a conference paper at ICLR 2020
D
DMLAB-33 30 MEMORY/REACTIVE PARTITION
Memory
Reactive
nomnsceleeL_oonmachiig...o.oi
voons_collee__goddobjjeer.....
watesawatermaze
womns.exloot..eeffrrd..ccc train
explore_obi obstructed goals small explore.gov -gwllgoallloooao locations small explore.object.ree rew.rdd..esddds..cec few explore obstructed goals large explore_goal loppaloocation_laat Jarge
rooms rooms_keys.
cetisskeesdooooss.s.eeeoo doors doorsspuzzzle
gegguagesselet.... www.crebeeddbppleocc _object
language execute waadommanoommmame task
execuue
language_answer www.aanswweauu www.qquiaaiieeequeeeeea. -ueestion
explore.object many many
lasertag_onee .ww.oopponent.argg
lasertag three opponents large
lasertag opponent small
htr.iaaggtrreeeoooeetttttt smmall
natlab_fixed Laxee_large_mee map
variabg..rygg..... map regrooth
natlab.varying varbbig.yyyngggpppp www.:mappaaaa man_randdmized skymaze. irreversible path hard skymaze
skymaze inreversible path varied
syychab_rbiirary.vvsoooooo....m
syychaab.conimuoussrcoogggi
psychlab sequential comparison
psychlab_ Lis.uallssuaaeeiiiil
cotoreoobjec__lccation small
Table 13: Partition of DMLab-30 levels into memory- and reactive split of levels.
21

---

Under review as a conference paper at ICLR 2020
Model
Median Human Normalized Score
LSTM
136.6 + 3,4
GTrXL
137.1 5.0
Table 14: Final numan-normallen median return across all 57 Atari levels for LSTM and GTrXL at 11.4 billion environment steps (equivalent to 200 million per individual game). Both models are 256 dimensions in width. We include standard error over runs.
Atari-57
160
Median Human Normalized Scores
140
120
100
80
60
40
GTrXL
LSTM
0
0.0
0.2
0.4
0.6
0.8
1.0
Env. Steps
1e10
Figure 11: Median human-normalized returns training progresses for both GTrXL and LSTM models. We run 8 hyperparameter settings per model.
E ATARI-57 RESULTS
In this section, we run the GTrXL on the multitask Atari-57 benchmark. Although Atari-57 was not designed specifically to test an agent's 's memory capabilities, we include these results here to demonstrate that we suffer no performance regression on popular environment suite, providing further evidence that GTrXL can be used as an architectural replacement to the LSTM.
The LSTM and GTrXL are matched in width at 256 dimensions. The GTrXL is 12 layers deep to show our model's learning stability even at large capacity. The LSTM architecture matches the one reported in Anonymous Authors (2019). We train for 11.4 billion environment steps, equivalent to 200 million per environment. We run 8 hyperparameter settings per model.
22


Reviews:
[{"review_id": "rklAUbhRFr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\nMy decision is Weak Reject, considering the following aspects.\n\nPositive points: (1) The experiments seem solid. The authors have evaluated the overall performance, as well as hyperparameters, seeds, and ablations. (2) Moving the layer normalization layer to the input stream seems to be surprisingly effective. This could be an interesting finding. (3) The paper is well organized.\n\nNegative points: (1) Lack of experiments on benchmark and large environments. The authors did not evaluate their model on the widely used benchmark Atari-57. Also, it is unclear whether the proposed transformer can scale to large environments. (2) Lack of understanding of the layer normalization. The authors provide some explanations about why the reordering works, but they seem not intuitive. More analysis about why the reordering works would significantly enhance this paper. \n\nSpecific questions: (1) Have you tried simply removing the layer normalization layer? (2) TrXL-I moves two layer-normalization layers together. Have you tried only moving one of them? Which modification contributes more? (3) Could you provide more explanations about why the modification of the layer normalization layer works? (4) Have you experimentally validated the proposed hypothesis as to why the Identity Map Reordering, such as recording the evolution of the produced values in the submodules?\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "HkeMR-1QiB", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "* Summary\nThis paper introduces architecture modifications for self-attention to stabilize transformers in reinforcement learning.\nThe new architecture, Gated Transformer-XL, replaces the order of the layer norm blocks to preserve an identity mapping.\nMultiple existing gating layers are proposed to replace the residual connections of transformer.\nThe new architecture is compared against MERLIN and LSTM on DMlab-30, and further ablation studies are done on Numpad and Memory Maze.\nIt is noted that they use the recent V-MPO objective to train LSTM and transformer.\nTheir results show that transformers are able to learn in memory-intensive environments, with some gating combinations surpassing LSTM.\n\n* Decision\nThis paper presents promising empirical results, however the experiments are limited, making it difficult to place in the broader work.\nIn addition, the contribution is incremental and not well-motivated.\nI would recommend a weak rejection.\nStill, I think the paper is well written and could be improved upon.\n\n* Reasons\nWhile the empirical results are impressive, they are not put into context.\nDMlab-30 is still a relatively new environment suite and it is difficult to place the result of this paper in the context of broader work.\nIn addition, the comparison is against LSTM on a new objective.\nWhile Transformer is able to beat LSTM on the same objective, it is unclear whether that is a success of the objective or the architecture.\nIn Numpad, the transformer architecture shows an improvement over LSTM, but no comparisons are made to any other memory-based agents.\nThe hyperparameter studies on Memory Maze also show improvements in memory-related tasks, but do not help in understanding of the proposed work.\n\nThe choice of architecture modification is also not well motivated.\nAs the paper mentions, initializing near identity has been shown to be important in the supervised learning literature.\nFor the topic of this paper however, I do not think this adequetly explains the instability of transformers in reinforcement learning.\nIn the related work section for example, the paper notes that gating mechanisms have been used to handle the vanishing gradients problem.\nThe paper also notes that vanishing gradients is not an issue in transformers.\nHence, it is unclear why gating would stabilize transformers for reinforcement learning.\n\nThe paper is overall well written and the ideas developed are clear.\nUnfortunately, the impressive results on DMlab are not sufficient for both the lack of deeper empirical study and better theoretical motivation for the architecture modifications.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "BJerZG59YS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper is motivated by the unstable performance of Transformer in reinforcement learning, and tried several variants of Transformer to see whether some of them can stabilize the Transformer. The experimental results look good, however, I have problems in understanding the motivation, the intuition of the proposed methods, the experimental design, and the general implication to the research community that is using the Transformer in their day-to-day research.\nFirst, the paper was based on the hypothesis of the authors that the Transformer is not stable, however, there is no comprehensive study on the unstability, and deep understanding on the root cause of it. It would be much more convincing to give a form definition of unstability and to add experimental study and theoretical analysis to the motivation part, instead of just based on a hypothesis.\n\nSecond, the proposal of the new structures (e.g., reordering the layer normalization, adding the gating layer) are quite ad hoc. There is not very solid motivation and theoretical analysis on why they could solve the unstable problem of the Transformer.  For example, by changing the order of layer normalization, there are direct identity mapping from the first layer to the last layer, making the information flow smoother. However, why this will make the Transformer more stable? The hypothesis and intuitive analysis are not very convincing. For another example, why replacing the residual connection with the gating layer can make the Transformer more stable? It seems to me that these are mostly heuristics, but not verified or strongly motivated solutions.\n\nThird, if the proposal is sound, it should not be effective only for reinforcement learning. It should be able to improve the performance or stability of the Transformer in general (e.g., in NLP tasks). However, there is no experiments and discussions regarding this.\n\nFourth, the experiments on the reinforcement learning is a little narrow, and many famous RL benchmarks and environments were not tested. This makes it unclear whether the proposed approach is generally effective.\n\n\n**I read the author responses, however, they do not really change my assessment on the paper.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes  architectural modifications to transformers, which are promising for sequential tasks requiring memory but can be unstable to optimize, and applies the resulting method to the RL setting, evaluated in the DMLab-30 benchmark.

While I thought the approach was interesting and the results promising, the reviewers unanimously felt that the experimental evaluation could be more thorough, and were concerned with the motivation behind of some of the proposed changes.

Decision: Reject