Transformers are Meta-Reinforcement Learners

1 INTRODUCTION . In recent years , the Transformer architecture ( Vaswani et al. , 2017 ) achieved exceptional performance on many machine learning applications , especially for text ( Devlin et al. , 2019 ; Raffel et al. , 2020 ) and image processing ( Dosovitskiy et al. , 2021b ; Caron et al. , 2021 ; Yuan et al. , 2021 ) . This intrinsically relates to its few-shot learning nature Brown et al . ( 2020a ) : the attention weights work as context-dependent parameters , inducing better generalization . Furthermore , this architecture parallelizes token processing by design . This capability avoids the vanishing gradients problem , very common for recurrent models . As a result , they can handle longer sequences more efficiently . This work argues that these two capabilities are essential for a Meta-Reinforcement Learning ( metaRL ) agent . We propose TrMRL ( Transformers for Meta-Reinforcement Learning ) , a memorybased meta-Reinforcement Learner which uses the transformer architecture to formulate the learning process . It works as a memory reinstatement mechanism ( Rovee-Collier , 2012 ) during learning , associating recent working memories to create an episodic memory which is used to contextualize the policy . Figure 1 illustrates the process . We formulated each task as a distribution over working memories . TrMRL associates these memories using self-attention blocks to create a task representation in each head . These task representations are combined in the position-wise MLP to create an episodic output ( which we identify as episodic memory ) . We recursively apply this procedure through layers to refine the episodic memory . In the end , we select the memory associated with the current timestep and feed it into the policy head . Nonetheless , transformer optimization is often unstable , especially in the RL setting . Past attempts either fail to stabilize ( Mishra et al. , 2018 ) or required architectural additions ( Parisotto et al. , 2019 ) or restrictions on the observations space ( Loynd et al. , 2020 ) . We argue that this challenge can be mitigated through a proper weight initialization scheme . For this matter , we applied T-Fixup initialization ( Huang et al. , 2020 ) . We conducted a series of experiments to evaluate meta-training , fast adaptation , and out-ofdistribution generalization in continuous control environments for locomotion and robotic manipulation . Results show that TrMRL consistently achieves or surpasses the current state-of-the-art meta-RL agents in performance and sample efficiency . It presents online adaptation , requiring as few as 20 timesteps to identify and achieve the desired performance on test tasks . We also conducted an ablation study to show the effectiveness of the T-Fixup initialization , and the sensibility to network depth , sequence size , and the number of attention heads . 2 RELATED WORK . Meta-Learning is an established Machine Learning ( ML ) principle to learn inductive biases from the distribution of tasks to produce a data-efficient learning system ( Bengio et al. , 1991 ; Schmidhuber et al. , 1996 ; Thrun & Pratt , 1998 ) . This principle spanned in a variety of methods in recent years , learning different components of an ML system , such as the optimizer ( Andrychowicz et al. , 2016 ; Li & Malik , 2016 ; Chen et al. , 2017 ) , neural architectures ( Hutter et al. , 2019 ; Zoph & Le , 2017 ) , metric spaces ( Vinyals et al. , 2016 ) , weight initializations ( Finn et al. , 2017 ; Nichol et al. , 2018 ; Finn et al. , 2018 ) , and conditional distributions ( Zintgraf et al. , 2019 ; Melo et al. , 2019 ) . Another branch of methods learns the entire system using memory-based architectures ( Ortega et al. , 2019 ; Wang et al. , 2017 ; Duan et al. , 2016 ; Ritter et al. , 2018a ) or generating update rules by discovery ( Oh et al. , 2020 ) or evolution ( Co-Reyes et al. , 2021 ) . Memory-Based Meta-Learning is the particular class of methods where we focus on in this work . In this context , Wang et al . ( 2017 ) ; Duan et al . ( 2016 ) concurrently proposed the RL2 framework , which formulates the learning process as a Recurrent Neural Network ( RNN ) where the hidden state works as the memory mechanism . Given the recent rise of attention-based architectures , one natural idea is to use it as a replacement for RNNs . Mishra et al . ( 2018 ) proposed an architecture composed of causal convolutions ( to aggregate information from past experience ) and soft-attention ( to pinpoint specific pieces of information ) . In contrast , our work applies causal , multi-head self-attention by stabilizing the complete transformer architecture with an arbitrarily large context window . Finally , Ritter et al . ( 2021 ) also applied multi-head self-attention for rapid task solving for RL environments . However , in a different dynamic : their work applied the attention mechanism iteratively in a pre-defined episodic memory , while ours applies it recursively through transformer layers to build an episodic memory from the association of recent working memories . Our work has intersections with Cognitive Neuroscience research on memory for learning systems ( Hoskin et al. , 2018 ; Rovee-Collier , 2012 ; Wang et al. , 2018 ) . In this context , Ritter et al . ( 2018c ) extended the RL2 framework incorporating a differentiable neural dictionary as the inductive bias for episodic memory recall . In the same line , Ritter et al . ( 2018b ) also extended RL2 but integrating a different episodic memory system inspired by the reinstatement mechanism . In our work , we also mimic reinstatement to retrieve episodic memories from working memories but using self-attention . Lastly , Fortunato et al . ( 2019 ) studied the association between working and episodic memories for RL agents , specifically for memory tasks , proposing separated inductive biases for these memories based on LSTMs and auxiliary unsupervised losses . In contrast , our work studies this association for the Meta-RL problem , using memory as a task proxy implemented by the transformer architecture . Meta-Reinforcement Learning is a branch of Meta-Learning for RL agents . Some of the algorithms described in past paragraphs extend to the Meta-RL setting by design ( Finn et al. , 2017 ; Mishra et al. , 2018 ; Wang et al. , 2017 ; Duan et al. , 2016 ) . Others were explicitly designed for RL and often aimed to create a task representation to condition the policy . PEARL ( Rakelly et al. , 2019 ) is an off-policy method that learns a latent representation of the task and explores via posterior sampling . MAESN ( Gupta et al. , 2018 ) also creates task variables but optimizes them with on-policy gradient descent and explores by sampling from the prior . MQL ( Fakoor et al. , 2020 ) is also an off-policy method , but it uses a deterministic context that is not permutation invariant implemented by an RNN . Lastly , VariBAD Zintgraf et al . ( 2020 ) formulates the problem as a Bayes-Adaptive MDP and extends the RL2 framework by incorporating a stochastic latent representation of the task trained with a VAE objective . Our work contrasts all the previous methods in this task representation : we condition the policy in the episodic memory generated by the transformer architecture from the association of past working memories . We show that this episodic memory works as a proxy to the task representation . Transformers for RL . The application of the transformer architecture in the RL setting is still an open challenge . Mishra et al . ( 2018 ) tried to apply this architecture for simple bandit tasks and tabular MDPs and reported unstable train and random performance . Parisotto et al . ( 2019 ) then proposed some architectural changes in the vanilla transformer , reordering layer normalization modules and replacing residual connections with expressing gating mechanisms , improving state-of-the-art performance for a set of memory environments . Loynd et al . ( 2020 ) also studied how transformerbased models can improve the performance of sequential decision-making agents . It stabilized the architecture using factored observations and an intense hyperparameter tuning procedure , resulting in improved sample efficiency . In contrast to these methods , our work stabilizes the transformer model by improving optimization through a better weight initialization . In this way , we could use the vanilla transformer without architectural additions or imposing restrictions on the observations . Finally , recent work studied how to replace RL algorithms with transformer-based language models ( Janner et al. , 2021 ; Chen et al. , 2021 ) . Using a supervised prediction loss in the offline RL setting , they modeled the agent as a sequence problem . Our work , on the other hand , considers the standard RL formulation in the meta-RL setting . 3 PRELIMINARIES . We define a Markov decision process ( MDP ) by a tuple M = ( S , A , P , R , P0 , γ , H ) , where S is a state space , A is an action space , P : S × A × S → [ 0 , ∞ ) is a transition dynamics , R : S × A → [ −Rmax , Rmax ] is a bounded reward function , P0 : S → [ 0 , ∞ ) is an initial state distribution , γ ∈ [ 0 , 1 ] is a discount factor , and H is the horizon . The standard RL objective is to maximize the cumulative reward , i.e. , maxE [ ∑T t=0 γ tR ( st , at ) ] , with at ∼ πθ ( at | st ) and st ∼ P ( st | st−1 , at−1 ) , where πθ : S ×A → [ 0 , ∞ ) is a policy parameterized by θ . 3.1 PROBLEM SETUP : META-REINFORCEMENT LEARNING . In the meta-RL setting , we define p ( M ) : M → [ 0 , ∞ ) a distribution over a set of MDPs M. During meta-training , we sample Mi ∼ p ( M ) from this distribution , where Mi = ( S , A , Pi , Ri , P0 , i , γ , H ) . Therefore , the tasks1 share a similar structure in this setting , but reward function and transition dynamics vary . The goal is to learn a policy that , during meta-testing , can adapt to a new task sampled from the same distribution p ( M ) . In this context , adaptation means 1We use the terms task and MDP interchangeably . maximizing the reward under the task in the most efficient way . To achieve this , the meta-RL agent should learn the prior knowledge shared across the distribution of tasks . Simultaneously , it should learn how to differentiate and identify these tasks using only a few episodes . 3.2 TRANSFORMER ARCHITECTURE . The transformer architecture ( Vaswani et al. , 2017 ) was first proposed as an encoder-decoder architecture for neural machine translation . Since then , many variants have emerged , proposing simplifications or architectural changes across many ML problems ( Dosovitskiy et al. , 2021a ; Brown et al. , 2020b ; Parisotto et al. , 2019 ) . Here , we describe the encoder architecture as it composes our memory-based meta-learner . The transformer encoder is a stack of multiple equivalent layers . There are two main components in each layer : a multi-head self-attention block , followed by a position-wise feed-forward network . Each component contains a residual connection ( He et al. , 2015 ) around them , followed by layer normalization ( Ba et al. , 2016 ) . The multi-head self-attention ( MHSA ) block computes the selfattention operation across many different heads , whose outputs are concatenated to serve as input to a linear projection module , as in Equation 1 : MHSA ( K , Q , V ) = Concat ( h1 , h2 , . . . , hω ) Wo , hi = softmax ( QKT√ d ·M ) V , ( 1 ) where K , Q , V are the keys , queries , and values for the sequence input , respectively . Additionally , d represents the dimension size of keys and queries representation and ω the number of attention heads . M represents the attention masking operation . Wo represents a linear projection operation . The position-wise feed-forward block is a 2-layer dense network with a ReLU activation between these layers . All positions in the sequence input share the parameters of this network , equivalently to a 1× 1 temporal convolution over every step in the sequence . Finally , we describe the positional encoding . It injects the relative position information among the elements in the sequence input since the transformer architecture fully parallelizes the input processing . The standard positional encoding is a sinusoidal function added to the sequence input ( Vaswani et al. , 2017 ) .

Summary: This work studies the use of the transformer architecture to enable memory-based meta RL. The authors argue that transformer's attention is  inherently suitable to facilitate episodic memory which can be then inputed to the policy. Furthermore, the authors point out that transformer demonstrates better properties at  handling longer transition sequences. Finally, the authors suggest that training a transformer in RL is often unstable, and to this end they propose to apply T-Fixup initialization that evidently enables stable training. The paper provides an empirical study on a several high-dimensional continuous control tasks that feature locomotion and dexterous manipulation to support the aforementioned claims.