Transformers are Meta-Reinforcement Learners
Authors: Luckeciano Carvalho Melo
Venue: ICLR 2022 Submitted (2022)
Keywords: Reinforcement Learning
Meta-Reinforcement Learning
Transformers

Abstract: The transformer architecture and variants presented a remarkable success across many machine learning tasks in recent years. This success is intrinsically related to the capability of handling long sequences and the presence of context-dependent weights from the attention mechanism. We argue that these capabilities suit the central role of a Meta-Reinforcement Learning algorithm. Indeed, a meta-RL agent needs to infer the task from a sequence of trajectories. Furthermore, it requires a fast adaptation strategy to adapt its policy for a new task - which can be achieved using the self-attention mechanism. In this work, we present TrMRL (Transformers for Meta-Reinforcement Learning), a meta-RL agent that mimics the memory reinstatement mechanism using the transformer architecture. It associates the recent past of working memories to build an episodic memory recursively through the transformer layers. This memory works as a proxy to the current task, and we condition a policy head on it. We conducted experiments in high-dimensional continuous control environments for locomotion and dexterous manipulation. Results show that TrMRL achieves or surpasses state-of-the-art performance, sample efficiency, and out-of-distribution generalization in these environments.

Under review as a conference paper at ICLR 2022
TRANSFORMERS LEARNERS
ARE
META-REINFORCEMENT
Anonymous authors Paper under double-blind review
ABSTRACT
The transformer architecture and variants presented a remarkable success across many machine learning tasks in recent years. This success is intrinsically re- lated to the capability of handling long sequences and the presence of context- dependent weights from the attention mechanism. We argue that these capabilities suit the central role of Meta-Reinforrcmeece Learning algorithm. Indeed, a meta- RL agent needs to infer the task from sequence of trajectories. Furthermore, it requires a fast adaptation strategy adapt its policy for new task which can be achieved using the self-attention mechanism. In this work, we present TrMRL (Transformers for Meta- Reinforcement Learning), a meta-RL. agent that mim- ics the memory reinstatement mechanism using the transformer architecture. It associates the recent past of working memories to build an episodic memory re- cursively through the transformer layers. This memory works a proxy the current task, and we condition a policy head on it. We conducted experiments in nigh-dimensional continuous control environments for locomotion and dexter- ous manipulation. Results show that TrMRL achieves or surpasses state-of-the-an performance, sample efficiency, and out-of-distribution generalization in these en- vironments.
INTRODUCTION
In recent years, the Transformer architecture Vaswani et al. 2017) achieved exceptional perfor- mance on many machine learning applications, especially for text (Devlin al. 2019 Raffel al, 2020) and image processing (Dosovitskiy et al.) 2021b Caron et al. 2002 Yuan al. (022)). This intrinsically relates to its few-shot learning nature Brown et al. (2020a): the attention weights work context-dependent parameters, inducing better generalization. Furthermore, this architecture par- allelizes token processing by design. This capability avoids the vanishing gradients problem, very common for recurrent models. As result, they can handle longer sequences more efficiently.
This work argues that these two capabilities are essential for a Meta-Reiinforcement Learning (meta- RL) agent. We propose TrMRL (Transformers for Meta-Reinfoo Rett-RReinfrr Learning), memory- based meta--Riifforemmett Learner which uses the transformer architecture to formulate the learning process. It works memory reinstatement mechanism (Rovec-Collies 2012) during learning, associating recent working memories to create an episodic memory which is used to contextualize the policy.
Figure illlustrates the process. We formulated each task distribution over working memories. TrMRL associates these memories using self-attention blocks create a task representation in each head. These task representations are combined in the position-wist -iiii 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) restrictions 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

---

Under review as a conference paper at ICLR 2022
Dacue Maal Head
1/.
0--1
6-1
Figure 1: Illustration of the TrMRL agent. At each timestep, it associates the recent past of working memories build an episodic memory through transformer layers recursively, We argue that the self-attention works as a fast adaptation strategy since it provides context-dependent parameters.
We conducted a series of experiments to evaluate meta-training, fast adaptation, and out-of- distribution generalization in continuous control environments for locomotion and robotic manip- ulation. Results show that TrMRL consistently achieves or surpasses the current state-of-the-an meta-RL agents performance and sample efficiency. It presents online adaptation, requiring as few as 20 timesteps to identify and achieve the desired performance test tasks. We also con- ducted 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-Lccrning is an established Machine Learning (MI.) principle to learn inductive biases from the distribution of tasks to produce a data-efficient learning system (Bengio al. 1991 Schmidhu-
ber et al.] 1996 Thrun & Pratt 1998). This principle spanned in a variety of methods in recent years,
learning different components 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 Nichol et al.
2018
Finn 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 al., 2017: Duan al. 20166] Ritter 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 this work. In this context, Wang et al. (2017); Duan et al. (2016) concurrently proposed the RL 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 to use it as replacement for RNNs. MMshra et al. (2018) proposed an architecture composed of causal convolutions (to aggregate information from past experience) and soft-attention (to pin- contrast, our work applies causal, multi-head self-attention by stabilizing the complete transformer architecture with arbitrarily large context window. Fi- nally, Ritter et al. (202I) also applied multi-head self-attention for rapid task solving for RL envi- ronments. However, in 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.
point specific pieces of information).
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 al. (2018c)
2

---

Under review as a conference paper at ICLR 2022
extended the RL² framework incorporating differentiable neural dictionary as the inductive bias for episodic memory recall. In the same line, Ritter et al. (2018b) also extended RL.² a but integrating 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, |Fortnnto 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 task proxy implemented by the transformer architecture.
RentoRReinforeementt Learning is a branch of Meta-Learning for RL agents. Some of the algo- rithms described in past paragraphs extend to the Meta-RL setting by design (Finn et al. 2017] Míshra et al. 2018: Wang et al. 2017 Duan et al., (016)). Others were explicitly designed for RL. and often aimed to create a task representation to condition the policy. PEARL (Rakelly et al. 2019) an off off-policy method that learns a latent representation of the task and explores via posterior sam- pling. 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 (aakoor et al. 2020) is also an off-policy method, but it uses a deterministic context that is not permutation invariant implemented by an RŇN. Lastly. VariBAD Zintgraf et al.] (2020) formulates the problem as a Adaative MDP and extends the RI,2 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 representa- tion: 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 pro- posed some architectural changes in the vanilla transformer, reordering layer normalization modules and replacing residual connections with expressing gating mechanisms, improving state-of-the-an performance for a set of memory environments. [oynd et al. (2020) also studied how transformer- 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 the observations.
based models can improve the performance of sequential decision-making
Finally, recent work studied how to replace RL algorithms with transformer-basee language models Janner et al. 202T] Chen al. (02T). Using a supervised prediction loss in the offline RI 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 Markov decision process (MDP) by tuple M = AS, y, H), where S
is a state space, A is an action space, P : S x A x S + [0, 20) is a transition dynamics, R. : S x A - [- Rmazz Rmaz] is bounded reward function, Po S - [0) is an initial state distribution, 7 E [0,1] is a discount factor, and H is the horizon. The standard RL objective is to maximize the cumulative reward, i.e., max E - 7
T ('((st, at)), with or ~ mo(a, 8t)) and 8, ~ P(s; 9t-1.4c-1) where Te S x A + [0, 20) is a policy parameterized by O.
3.1 PROBLEM SETUP: META-REINFORCEETE
LEARNING
In the meta-RL setting, we defin p(M) M + oo) a distribution over a set of MDPs M. During meta-training, we sample M; ~ P(M) from this distribution, where M, (S, A, Pr, R2 Po.,, 7, H... Therefore, the taske]] share similar structure in this setting, but re- ward function and transition dynamics vary. The goal is to learn a policy that, during meta-testing, can adapt to new task sampled from the same distribution pMM). In this context, adaptation means
We use the terms task and MDP interchangeahly.
3

---

Under review as a conference paper at ICLR 2022
maximizing the reward under the task in the most efficient way. To achieve this, the meta-RL. 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 few episodes.
3.2 TRANSFORMER ARCHITEECTRRE
The transformer architecture (Vaswani et al. 2017) was first proposed as an encoder-decoder ar- chitecture for neural machine translation. Since then, many variants have emerged, proposing sim- plifications or architectural changes across many ML problems (Dosovitskiy et al. 2021a [Bronn al. 2020b Parisotto et al. 2019). Here, we describe the encoder architecture it composes our memory-based meta-learner.
The transformer encoder stack of multiple equivalent layers. There are two main components in each layer: multi-head self-attention block, followed by position- position-wise 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 self- attention operation across many different heads, whose outputs are concatenated serve as input to a linear projection module, as in Equation T
MHSA K, Q. Concat(hi, b2,.... hh....,,h)Wer )
QKT
h4 = softmax(
Vd M)V,
where K, Q, V are the keys, queries, and values for the sequence input, respectively. Additionally, represents the dimension size of keys and queries representation and w the number of attention heads. M represents the attention masking operation. W, represents a linear projection operation.
The position-wise feed-forward block is 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 1 X 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 a sinusoidal function added to the sequence input /Vaswani al. 2017).
3.3 T-FIXUP INITIALIZATION
The training of transformer models is notoriously difficult, especially in the RL setting (Parisotto al.) (209).. Indeed, gradient optimization with attention layers often requires complex learning rate warmup schedules to prevent divergence (Huang et al. 2020)). Recent work suggests two main reasons for this requirement. First, the Adam optimizer (Kingma & Ba. 2017) presents high vari- ance in the inverse second moment for initial updates, proportional to divergent integral (Liu et al.) 2020). It leads to problematic updates and significantly affects optimization. Second, the backprop- agation through layer normalization can also destabilize optimization because the associated error depends on the magnitude of the input (Xiong et al. 2020).
Given these challenges, Huang al. (2020) proposed weight initialization scheme (T-Fixup) to eliminate the need for learning rate warmup and layer normalization. This is particularly important the RL setting once current RL algorithms very sensitive to the learning rate for learning and exploration.
TT-Fxxpp appropriately bounds the original Adam update make variance finite and reduce instabil- ity, regardless of model depth. We refer to Huang et al. (020)) for the mathematical derivation. We apply the T-Fixup for the transformer encoder follows:
Apply Xavier initialization GGooot & Bengio 2010) for all parameters excluding input embeddings. Use Gaussian initialization N d $), for input embeddings, where d is the embedding dimension;
Scale the linear projection matrices in each encoder attention block and position wise feed- forward block by N
4

---

Under review as a conference paper at ICLR 2022
4 TRANSFORMERS ARE CEIFFORENNO LEARNERS
In this work, we argue that two critical capabilities of transformers compose the central role of Meta-Reinforcement Learner. First, transformers can handle long sequences and reason over long- term dependencies, which is essential to the meta-RL agent to identify the MDP from a sequence of trajectories. Second, transformers present context-dependent weights from self-attention. This mechanism serves as a fast adaptation strategy and provides necessary adaptability to the meta-RL agent for new tasks.
4.1 TASK REPRESENTATION
We represent working memory the timestep 7 as a parameterized function (x(8,, ar; Tt, lt.,, where 8; the MDP state, a₁ a(a; 8t) an action, ~ R(s,, at) is the reward, and 77t is boolean flag identify whether this terminal state. Our first hypothesis is that we can define task T as a distribution over working memories, as in Equation
p(6)
T
T()) + [0, oo),
(2)
where is the working memory embedding space. In this context, one goal of a meta-RI. agent is to learn á to make distinction among the tasks the embedding space 4 Furthermore, the learned em- bedding space should also approximate the distribu- tions of similar tasks so that they share knowledge. Figure dimensional representation.
The illustration of two tasks (Ti and T2) as distributions over working memo- ries. The intersection of both densities repre- sents the ambiguity between Ti and Tz-
Figure 2:
2 illustrates this concept for
We aim to find a representation for the task given its distribution to contextualize our policy. In- tuitively, we can represent each task as a linear combination of working memories sampled by the policy interacting with it:
1T Σ ar W(oo(s,,
with
00 0
where N represents the length of segment of sampled trajectories during the policy and task inter- action. W represents an arbitrary linear transformation. Furthermore, a is a coefficient to compute how relevant particular working memory is the task representation, given the sampled working memories. Next, we show how the self-attention computes these coefficients, which we use to output an episodic memory from the transformer architecture.
4.2 AETENAITNNIION AS A FAST ADAPTATION STRATEGY
In this work, central argument that self-attention works fast adaptation strategy. The context-dependent weights dynamically compute the working memories coefficients to implement Equation We now derive how we compute these coefficients. Figure illustrates this mechanism.
Let us define of 0% and of as representatior of the working memory at timestep t in the keys, queries, and values spaces, respectively. The dimension of the queries and keys spaces is d. We aim to compute the attention operation in Equation [ for sequence of T timesteps, resulting in Equation
we slightly abuse the notation by omitting the function arguments 8t, at, and for the sake of conciseness.
5

---

Under review as a conference paper at ICLR 2022
Ch
Lnd e' 123
145
4
000
Figure 3: Illustration of causal self-attention fast adaptation strategy. In this simplified scenario (2 working memories), the attention weights æsy drives the association between the current work- ing memory and the past ones to compute task representation μt, Self-attention computes this association by relative similarity.
01,1
01,2
QKT Vd
softmax
M)V
Vd
@T,1
QT,
1
exp (0f
where
Oi.j =
En exp (of.ak)
if i
(4)
0
otherwise.
where ba) -nn ai, bjnn is the dot product between the working memories and bj. Therefore, for a particular timestep f, the self-ateention output is:
of exp (0), ot) + ... +.....6 08 exp (0),0)) Vd
14
En カー exp (1,,00)
on (or)..
(5)
Vd
n=1
Equation 5 shows that the self-attention mechanism implements the task representation in Equation by associating past working memories given that the current one Φ. It computes this association with relative similarity through the dot product normalized by the softmax operation. This induc- tive bias helps the working memory representation learning to approximate the density of the task distribution (()).
4.3 TRANSFORMERS AND MEMORY REINSTATEMENT
We now argue that the transformer model implements a memory reinstatement mechanism for episodic memory retrieval. An episodic memory system is long-lasting memory that allows agent to recall and re-experience personal events (Tulving 2002). It complements the working memory system, which is active and relevant for short periods (Baddeley 2010) and works buffer for episodic memory retrieval (Zilli & Hasselmo, 2008). Adopting this memory interaction collection of past memories. More
model, we model an episodic memory transformation concretely, we consider that a transformer layer implements this transformation:
(6) where ef represents the episodic memory retrieved from the last layer / for the timestamp and f represents the transformer layer. Equation 6 provides recursive definition, and el (the base case)
e 1-1 ......,.-11).

---

Under review as a conference paper at ICLR 2022
corresponds to the working memories 61 In this way, the transformer architecture recursively refines the episodic memory interacting memories retrieved from the past layer. We show the pseudocode for this process in Algorithm [ This refinement is guaranteed by a crucial property of the self- attention mechanism: it computes a consensus representation across the input memories associated to the sub-trajectory, representation as the memory representation that is closest on average to all likely representations (Kumar& Byrne 2004), i.e., minimizes the Bayes risk considering the of episodic memories.
stated by Theorem [
(Proof in Appendix H.. Here,
define consensus
Theorem 1. Let St = (e'...., eyy p(e,S), 01) be set of normalized episodic memory rep- resentations sampled from the posterior distribution p(elS', 0i) induced by the transformer layer I, parameterized by Or. Lel K, Q, V be the Key, Query, and Value vector spaces in the self- atlention mechanism. Then, the self-atlention in the layer I+1 computes consensus representation S e... -eep (el.9 ,
whose associated Bayes risk (in terms of negative cosine similar- ity) lower bounds the Minimum Bayes Risk (MBR) predicted from the set of candidate samples S! projected onto the V space.
J+1
Eim exp (e, "
Lastly, we condition the policy head in the episodic memory from the current timestep to sample an action. This complete process resembles a memory reinstatement operation: reminder procedure that reintroduces past elements in advance of a long-term retention test (Rovee-Collier) 20012). In our context, this "long-term retention test" identifies the task and acts accordingly to maximize rewards.
5 EXPERIMENTS AND ANALYSIS
In this section, we present an empirical validation of our method, comparing it with the current state-of-thh-ase methods. We considered digh-dimensional, continuous control tasks for locomo- tion (MuJoCaF) and dexterous manipulation (MetaWorld) We describe them in Appendix A For reproducibility (source code and hyperparameters), we refer to the released source code
5.1 EXPERIMENTAL SETUP
Meta- Training. During meta-training, we repeatedly sampled a batch of tasks to collect experience with the goal of learning to learn. For each task, we ran a sequence of E episodes. During the inter- action, the agent conducted exploration with gaussian policy. During optimization, these episodes to form single trajectory and we maximize the discounted cumulative reward of this trajectory. This is equivalent to the training setup for other on-policy meta-RL algorithms (Duan) iet al.) 20166] Zintgraf et al. 2020). For these experiments, we considered E = 2. We performed this training via Proximal Policy Optimization (PPO) (Schulman et al., 7). and the data batches mixed different tasks. Therefore, we present here an on-policy version of the TrMRL algorithm. To stabilize transformer training, we used the T-Fixup as a weight initialization scheme.
concatenate
During meta-testing, we sampled new tasks. These are different from the tasks in meta-training, but they come from the same distribution, except during Outtof-Dist (OOD) evaluation. For TrMRL, in this stage, we froze all network parameters. For each task, we ran few episodes, performing the adaptation strategy. The goal is to identify the current MDP and maximize the cumulative reward across the episodes.
Meta-Testiae Deta-nesting.
Memory Write Logic. At each timestep, we fed the network with the sequence of working memo- ries. This process works follows: the beginning of an episode (when the memory sequence is empty), we start writing the first positions of the sequence until we fill all the slots. Then, for each new memory, we removed the oldest memory in the sequence (in the "back" of this "queue") and added the most recent one (in the "front").
Comparison Methods. For comparison, we evaluated three different state-of-the-art meta-RL al- gorithms: RL Duan et al. (2016). optimized using PPO Schulman et al. (2017); PEARL (Rakelly al. y and MAML (Finn et al. (017)) whose outer-loop used TRPO (Shulman al. 2015).
We highlight that hoth MuloCo (Locomotion Tasks) and MetaWorld are built the MuloCo physies engine. We identify the set of locomotion tasks as solely for "MuloCa" to ensure simpler and concise writing during analysis of the results.
*mmiteed due to double-blind review. See supplementary material.

---

Under review as a conference paper at ICLR 2022
MeratWordd Reachvv2
Mctaword ML, Pushou
Meratoor M0.45
1.0
0.4
0.4
0.1
0.55
0:50
0.8
0.8
0.5
0.4
0.23
0.2
0.09
TiMAL
MAML-TAPG
RL2-PPO
PEARL
Figure 4: Meta- Veta-Training results for Meta World benchmarks. The plots on top represents perfor- mance training tasks, while the plots on bottom represents in the test tasks.
5.2 RESULTS AND ANALYSIS
We compared TrMRL with baseline methods in terms of meta-training, episode adaptation, and OOD performance. We also present the latent visualization for TrMRL. working memories and ablation studies. All the curves presented are averaged across three random seeds, with 95% boot- strapped confidence intervals.
Meta-Training Evaluation. Figure 4 shows the meta-training results for all the methods in the Meta World environments. All subplots show the task success rate over the training timesteps. The plots on top represent performance on training tasks, while the plots on the bottom represent on the test tasks. TrMRL outperformed all baseline methods. In the "Reach-v2", TrMRI,, RI.2. and MAMI. reached the perfect success rate, but TrMRI. was more sample efficient. For more complex scenarios, such as "Push-v2" and MLL5, TrMRL still performs consistently better. Nevertheless, we also present that all the presented methods performed poorly on ML45 for test tasks, highlighting a big improvement room. We hypothesize that this is due to the lack of meta-exploration strategy and an inductive bias to improve knowledge share among different tasks. For MuloCo locomotion tasks, refer to Appendix
Fast Adaptation Evaluation. A critical skill for meta-RI. agents is the capability of adapting new tasks given a few episodes. We evaluate this by run- ning meta-testing on 20 test tasks over sequen- tial episodes. Each agent runs its adaptation strat- egy identify the task and maximize the reward across episodes. Figure presents the results for the locomotion tasks. For AntDir and HalfCheeta Vel, TrMRL outperformed all methods. For HalfChee- tahDir, TrMRI started with better performance, PEARL outperformed after running its adaptation.
-0.5
-1.8
-.9
15
We highlight that TrMRL presented high perfor- mance since the first episode. In fact, it only requires a few timesteps to achieve high performance in test tasks. In the HalfCheetah Vel, for example, it only requires around 20 timesteps to achieve the best per- formance (Figure 5). Therefore, it presents nice
Figure 5: TrMRL's adaptation for HalfChee- tah Vel environment.
8

---

Under review as a conference paper at ICLR 2022
184
1300
MARB-7900
9,2-770
FEAN
Figure 6: Fast adaptation results MuJoCo locomotion tasks. Each curve represents the average performance over 20 test tasks. TrMRL presented high performance since the first episode due to the online adaptation nature from attention weights.
property for online adaptation. This is because the self-attention mechanism is lightweight and only quires a few working memories to achieve good performance. Hence, we can run it efficiently at each timestep. Other methods, such PEARL and MAMI., do not present such property, and they need a few episodes before executing adaptation efficiently.
OOD Evaluation. Another critical scenario is how the fast adaptation strategies perform for out-of-distribution tasks. For this case, we change the HalfCheetah Vel en- vironment to sample OOD tasks during the meta-test. In the standard setting, both training and testing target ve- locities are sampled from uniform distribution in the in- terval [0.0, 3.0). . In the OOD setting, we sampled 20 tasks in the interval [3.0, 4.0] and assessed adaptation through- out the episodes. Figure presents the results. TrMRL surpasses all the baselines methods with a good margin, suggesting that the context-dependent weights learned a robust adaptation strategy, while other methods memo- rized some aspects of the standard distribution of tasks. We especially highlight PEARL. which achieved the best performance among the methods but performed poorly in this setting. suggesting that it does generate useful la- tent representations for OOD tasks.
TWNRL
MARUT3FO
TEAS
Figure 7: OOD Evaluation in HalfChee- tah Vel environment.
6 CONCLUSION AND FUTURE WORK
In this work, we presented TrMRL, memory-based meta-RL algorithm built upon a transformer, where the multi-head self atffatttenton mechanism works as fast adaptation strategy. We designed this network to resemble memory reinstatement mechanism, associating past working memories dynamically represent task and recursively build an episode memory through layers.
TrMRL demonstrated a valuable capability learning from reward signals. On the other side, recent Language Models presented substantial improvements by designing self-supervised tasks (Devlin) ct al., 2019: Raffel et al., 2020) or even automating their generation (Shin et al.. 2020). As future work, aim to investigate how to enable these forms of self-supervision to leverage off off-policy data collected during the training and further improve sample efficiency transformers for the meta-R RL. scenario.

---

Under review as a conference paper at ICLR 2022
7 REPRODUCIBILITY STATEMENT
Code Release. To ensure the reproducibility of our research, we released all the source code as- sociated with our models and experimental pipeline. We refer to the supplementary material of this submission. It also includes the hyperparameters and the scripts to execute all the scenarios presented in this paper.
Baselines Reproducibility. We strictly reproduced the results from prior work implementations for baselines, and we provide their open-source repositories for reference.
Proof of Theoretical Results and Pseudocode. We provide detailed proof of Theorem 1 in Appendix E. containing all assumptions considered. We also provide pseudocode from TrMRL's agent to improve clarity on the proposed method.
Availability all simulation environments. All environments used in this work are freely avail- able and open-source.
REFERENCES
Marein Andrychowicz, Misha Denil, Sergio Gómez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando de Freitas. Learning to learn by gradient de- scent by gradient descent. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Asso- ciates, Inc., 2016. URL huttps://proceeeetsee htw.eeedddinnnneeesiseeeee httpp//eeeeeaaaaa...ooccoooee/aa pcpps/////ppee//////eee
pdff
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Alan Baddeley. Working memory. Current Biology, 20(4):R R14-00000. 2010. ISSN 0960-9822. doi:
URL //www. sciencedi htteer./ww....cieeed............
com/ http://ccaeceeaarteeee.s/////s22
Y. Bengio, S. Bengio, and J. Cloutier. Learning synaptic learning rule. In IJCNN-9I-Seattle International Joint Conference Neural Networks, volume ii, pp. 969 vol.2,, 1991. doi: 10.
19991UCN.11911.555660
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020a.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. CoRR, abs/2005. 14165, 2020b. URL wttp:tttt///ett wttp:/awwwwwwog..a//aaaaaaaaaaaa
www. www.g//ab///000 2000a//..0...... 14165
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers, 2021.
Lili Chen, Kevin La, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, 2021.
Yutian Chen, Matthew W. Hoffman, Sergio Gómez Colmenarejo, Misha Denil, Timothy Lilli- crap, Matt Botvinick, and Nando de Freitas. Learning leam without gradient descent by gra- dient descent. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research,
10

---

Under review as a conference paper at ICLR 2022
pp. PMLR, 06-11 Aug 2017. /tLp:ppp.//ppooedinn........sss/
Thenlee.d html
John D. Co-Reyes, Yingjie Miao, Daiyi Peng, Esteban Real, Sergey Levine, Quoe V. Le, Honglak Lee, and Aleksandra Faust. Evolving reinforcement learning algorithms, 2021.
Jacob Devlin, Ming-Waieeeee Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021a. URI. https: //openreview. wtt.:/e..nne..oooomm.e..iiiiei GNTTy
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbor,, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x scale, 2021b.
words: Transformers for image recognition at
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning, 2016.
Rasool Fakoor, Pratik Chaudhari, Stefano Soatto, and Alexander J. Smola. Meta-q-leaming, 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agmostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp- 1126-1135. PMLR, 06-11 Aug 2017. URL https: httpss///ppcceeeiigg. ml r.m.rrs/r//..
Finn17a.Ne html
Chelsea Finn, Kelvin Xu, and Sergey Levine. Probabilistic model-agnosten meta-learning In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, pp. 9537-9548, Red Hook, NY, USA, 2018. Curran Associates Inc.
Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomênech Badia, Gavin Buttimore, Charles Deck, Joel Z Leibo, and Charles Blundell. Generalization of reinforcement learners with working and episodic memory. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bus Buc, E. Fox, and R. Gareett (eds.), Advances in Neural Information Processing Sys- tems, volume 32. Curran Associates, Inc., 2019. URL https: http:::://////oeeeeeeeeem ngs nwwriis..pp... 0ttp:///w.......ee22222222222000 abDp:///222222011oaaaaaaaaaaaaaa Paw.reeaaeaaa-Ppp..a...mo pdf
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 (AISTATS' 10). Society for Artificial Intelligence and Statistics, 2010.
Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. In Proceedings of the 32nd Inter- national Conference on Neural Information Processing Systems, NIPS' 18, pp. 5307-5316, Red Hook, NY, USA, 2018. Curran Associates Inc.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition, 2015.
Abigail N. Hoskin, Aaron M. Bornstein, Kenneth A. Norman, and Jonathan D. Cohen. Refresh my memory: Episodic memory reinstatements intrude working memory maintenance. bioRxiv, 2018 doi: 10. 101/170720. URL Inttp::
201 /010//00/170720
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimiza- tion through better initialization. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 4475 4483. PMLR, 13-18 Jul 2020. URL actps: http://pcceeee.......e .ml..............mm............. mmrr preas/v119/hes 1asss//v19/hhaag2000 html
11

---

Under review as a conference paper at ICLR 2022
Frank Hutter, Lars Kotthoff, and Joaquin Vanschoren. Automated Machine Learning: Methods, Systems, Challenges. Springer Publishing Company, Incorporated, 1st edition, 2019. ISBN 3030053172.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence mod- eling problem, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Shankar Kumar and William Byme. Minimum Bayes-risk decoding for statistical machine transla- tion. In Proceedings of the Human Language Technology Conference of the North American Chapter of the Association for Computational Linguistics: HLT-NAACL 2004, pp. 169-176, Boston, Massachusetts, USA, May 2 - May 7 2004, Association for Computational Linguistics. URL BRI.https: / antho ogw. org. -gt.olg//ooo011100000000/ N0////04-11100
Ke Li and Jitendra Malik. Learning optimize, 2016.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond, 2020.
Ricky Loynd, Roland Fernandez, Asli Celikyilmaz, Adith Swaminathan, and Matthew Hausknecht. Working memory graphs, 2020.
Luckeciano C. Melo, Marcos R. O. A. Maximo, and Adilson Marques da Cunha. Bottom-up meta- policy search, 2019.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta- learner, 2018.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms, 2018.
Junhyuk Oh, Matteo Hessel, Wojciech M. Czarnecki, Zhongwen Xu, Hado P van Has- selt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. In H. Larochelle. M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1060-1070. Curran Asso- ciates, Inc., 2020. URL ptttttss:/prrceeedee F http//////aa223315555--22222 pdf
neur ips. cc.ppaper/2c22/Ei11.c.ccp
Pedro A. Ortega, Jane X. Wang, Mark Rowland, Tim Genewein, Zeb Kurth-Nelson, Razvan Pas- canu, Nicolas Heess, Joel Veness, Alex Pritzel, Pablo Sprechmann, Siddhant M. Jayakumar, Tom McGrath, Kevin Miller, Mohammad Azar, Ian Osband, Neil Rabinowitz, András György, Sil- via Chiappa, Simon Osindero, Yee Whye Teh, Hado Hasselt, Nando de Freitas, Matthew Botvinick, and Shane Legg. Meta-learning of sequential strategies, 2019.
Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learn- ing, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2020.
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off-policy meta-reiiffrreement learning via probabilistic context variables, 2019.
Sam Ritter, Ryan Faulkner, Laurent Sartran, Adam Santoro, Matt Botvinick, and David Raposo.
Rapid task-solving in novel environments, 2021.
Samuel Ritter, Jane Wang. Zeb Kurth-Nelson, Siddhant Jayakumar, Charles Blundell, Razvan Pas- canu, and Matthew Botvinick. Been there, done that: Meta-learning with episodic recall. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 4354 4363. PMLR. 10-15 Jul 2018a. URL http://ppp///poceedigge mlr. press/ /asas//800 v800riiterl1a.. htm
12

---

Under review as a conference paper at ICLR 2022
Samuel Ritter, Jane X. Wang, Zeb Kurth-Nelson, and Matthew Botvinick. Episodic control through mett-reiiforreement learning. In CogSci, 2018b. URL https: prtEps://mindmdeee /mmode ing. org/
Samuel Ritter, Jane X. Wang, Zeb Kurth-Nelson, Siddhant M. Jayakumar, Charles Blundell, Razvan Pascanu, and Matthew Botvinick. Been there, done that: Meta-learning with episodic recall, 2018c.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search, 2018.
Carolyn Govee-Collice Reinstatement of Learning, pp. 2803-2805. Springer US, Boston, MA, 2012. ISBN 978- 1 1188---4419-1428--11 doi: 10. 10.1007/978-1-44133 1-4441-14288-66 6 346. U httph://do/o.. org 100
Juergen Schmidhuber, Jieyu Zhao, and Marco Wiering, Simple principles of metalearning. Technical report, 1996.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region
policy optimization. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1889-1897, Lille, France, 07-09 Jul 2015. PMLR. URL httpps://proceet 0ttss//37/schulnane htmml
mmr..
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms, 2017.
Taylor Shin, Yasaman Razeghi, Robert L. Logan IV, Eric Wallace, and Sameer Singh. Auto- prompt: Eliciting knowledge from language models with automatically generated prompts. CoRR, abs/2010.15980) 2020. URL httpts:
org/
Sebastian Thrun and Lorien Pratt. Learning to Learn: Introduction and Overview, pp. 3-17. Kluwer Academic Publishers, USA. 1998. ISBN 0792380479.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control.
In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033,
2012. doi: 10.1109//ROS.20222.
Endel Tulving. Episodic memory: From mind to brain. Annual review of psychology, 53:1-25, 02 2002. doi: 10. 14/1///ureee..syc...... 1t0.oh..55..00001.1113114
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st Inter- national Conference on Neural Information Processing Systems, NIPS' 17, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, koray kavukcuoglu, and Daan Wierstra. Match- ing networks for one shot learning. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29. Curran As- sociates, Inc., 2016. URL URL nwwr.ips....uu.................. htt.ar/iis..cc/ppppe////////e/ee [00133783364333222500eee11333333 Paper. pdf
Jane Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn, 2017.
Jane X. Wang. Zeb Kurth-Nelson, Dharshan Kumaran, Dhruva Tirumala, Hubert Soyer, Joel Z. Leibo, Demis Hassabis, and Matthew Botvinick. Prefrontal cortex reteforrinooremmeettt leamn ing system. bioRxiv, 2018, doi: 10.1101/295964. URL http::///wwwbbbb.i..ooo
content/earlyl /0/33////////
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Yan Liu. On layer normalization in the transformer architec- ture, 2020.
13

---

Under review as a conference paper at ICLR 2022
Tianhe Yu. Deirdre Quillen, Zhanpeng He, Ryan Julian, Avnish Narayan, Hayden Shively, Adithya Bellathur, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning, 2021.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet, 2021.
Eric Zilli and Michael E Hasselmo. Modeling the role of working memory and episodic memory in behavioral tasks, Hippocampus, 18(2):133-209 2008
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep via meta- learning. 2020.
Luisa M Zintgraf, Kyriacos Shiarlis, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning, 2019.
Barret Zoph and Quoe V. Le. Neural architecture search with reinforcement learning. 2017. URL httts://arxii wwt.//gg/////ooo//o/.//ooo 61 01578
14

---

Under review as a conference paper at ICLR 2022
A METAL-RL ENVIRONMENTS DESCRIPTION
In this Appendix, we detail the environments considered in this work.
A.1 MuJoCo LOCOMOTION TASKS
This benchmark set of locomotion tasks on the MuJoCo (Todorov al. 2012) environment. It comprises different bodies, and each environment provides different tasks with different learning goals. These locomotion tasks are previously introduced by Finn et al. (2017) and Rothfuss et al. We considered different environments.
2018
AntDir:
This environment has ant body, and the goal is move forward or backward. Hence, it presents these 2 tasks.
HalfCheetahDir: This environment has a half cheetah body, and the goal is move for- ward or backward. Hence, it presents these 2 tasks.
HalfCheetal Vel: This environment also has a half cheetah body, and the goal is to achieve target velocity running forward. This target velocity comes from a continuous uniform distribution.
These locomotion task families require adaptation across reward functions.
A.2 METAWORLD
The MetaWorld (Yu et al. 2021) benchmark contains a diverse set of manipulation tasks designed for multi-task RL and meta-RL settings. MetaWorke presents variety of evaluation modes. Here, we describe the two modes used in this work. For more detailed description of the benchmark, we refer Yu al. (202T).
MLL:
This scenario considers a single robotic manipulation task but varies the goal. The meta-training "tasks" corresponds to 50 random initial object and goal positions, and meta- testing on 50 heldout positions.
MI45:
With the objective testing generalization new manipulation tasks, the bench- mark provides 45 training tasks and holds out 5 meta-testing tasks.
These robotic manipulation task families require adaptation across reward functions and dynamics.
15

---

Under review as a conference paper at ICLR 2022
B MuJoCo: META-TRAINING EVALUATION
In this Appendix, we supplement the meta-training evaluation with the results on MuloCo locomo- tion tasks. Figure g shows the average return over train tasks (on top) and test tasks (on bottom) for AntDir, VellChheetahel... and HalfCheetahDir, respectively. While PEARL failed to explore and learn in robotic manipulation tasks, presented better results on locomotion tasks, especially in sample efficiency. This is because of its off off- policc nature: it efficiently reuses trajectories sampled from previous policy versions, reducing the number of training steps needed. TrMRL achieved the same test task performance in AntDir and HalfCheetah Vel as PEARL, also keeping the metric stable over the training. When comparing with other on-policy methods, TrMRI. significantly improved performance and sample efficiency.
IaaCChenthhi
1,400
2003
2270
1.000
-200
020
An00g:
420
-350
230
2000
1500
-350
-400
107
167
167
(ours)
MAML-TRPO
RL2-PPQ
PEARL
Figure 8: Meta-Training results for MuJoCo locomotion benchmarks. The plots on top represent performance on training tasks, while the plots on bottom represent the test tasks.
16

---

Under review as a conference paper at ICLR 2022
C WORKING MEMORIES LATENT VISUALIZATION
Figure presents 3-D view of the working memories from the HalfCheetah Vel environment. We sampled some tasks (target velocities) and collected working memories during the meta-test setting. We observe that this embedding space learns a representation of each MDP as a distribution over the working memories, as suggested in Section approximately distinguish these tasks. Working memories that cross this boundary represent the ambiguity between two tasks. Furthermore, this representation also learns the similarity of tasks: for example, the cluster of working memories for target velocity 1.0 is between the clusters for " and 1.5. This property induces knowledge sharing among all the tasks, which suggests the sample efficiency behind TrMRL meta-training.
this visualization, we can draw planes that
HalfCheetahVel. Working Memories
Velocity
0.5
1.0
1.5
2.0
3.0
1.00
0.50
0.00
1.00
-0.50
0.50
-0.75
0.00
-1.00
0.50
0.25
0.50
0.75
0.75
-0.25
0.00
-0.50
1.00
-1.00
Figure 9: 3-D Latent visualization of the working memories for the HalfCheetah Vel environment. We plotted the 3 most relevant components from PCA. TrMRL learns representation of each MDP as distribution over the working memories. This representation distinguishes the tasks and approx- imates similar tasks, which helps knowledge sharing among them.
17

---

Under review as a conference paper at ICLR 2022
D ABLATION STUDY
In this section, we present an ablation study regarding the main components of TrMRL to iden- tify how they affect the performance the learned agents. For all the scenarios, we considered one environment for each benchmark to represent both locomotion (HalfCheetah Vel) and dexterous manipulation (Meta MetaWWold-MMLL 1---MLL--Rech-222 We evaluated the meta-training phase so that we could analyze both sample efficiency and asymptotic performance.
D.I T-FIXUP
In this work, employed T-Fixup address the instability from the early stages of transformer training, given the reasons deseribed in Section 3.3 In RI., the early stages training are also the moment when the learning policies are more exploratory to cover the state and action spaces better and discover rewards, preventing the convergence to sub-optimal policies. Hence, it is crucial for RL that the transformer policy learns appropriately since the beginning to drive exploration.
This section evaluated how T-Fixup becomes essential for environments where the learned behav- iors must guide exploration to prevent poor policies. For this, we present T-Fixup ablation (Figure I0 for two settings: MetaWorld-Mil MLI-Reach-y and HalfCheetah Vel. For the reach environment, we compute the reward distribution using the distance between the gripper and the target location. Hence, it is always dense and informative signal: even random policy can easily explore the environment, and T-Fixup does not interfere with the learning curve. On the other side, HalfChee- tah Vel requires functional locomotion gate to drive exploration; otherwise, it can get stuck with low rewards (e.g., cheetah is exploring while fallen). In this scenario, T-Fixup becomes crucial to prevent unstable learning updates that could collapse the learning policy to poor behaviors.
AbletionStufy Study: T-Fixup
14 Maachhcc
tuliChestarvee
0.8
-102
-200
0.0
0.2
0.4
0.8
1.5
simocrpps
Eimentage
enabled
disabled
Figure 10: Ablation results for the T-Fixup component.
D.2 WORKING MEMORY SEQUENCE LENGTH
A meta-RL agent requires sequence of interactions to identify the running task and act accordingly. The length this sequence N should large enough to address the ambiguity associated with the set of tasks, but not long to make the transformer optimization harder and less sample efficient. In this ablation, we study two environments that present different levels of ambiguity and show that they also require different lengths to achieve optimal sample efficiency.
We first analyze MetaWorld-We HILL-Recch-v2 v2. The environment defines each target location in the 3D space task. The associated reward is the distance between the gripper and this target. Hence,
18

---

Under review as a conference paper at ICLR 2022
at each timestep, the reward is ambiguous for all the tasks located on the sphere's surface with the center in the gripprr position. This suggests that the agent will benefit from long sequences. Figure I1 (left) confirms this hypothesis, the sample efficiency improves until sequences with several timesteps (N = 50).
The HalfCheetah Vel environment defines each forward velocity different task. The associated reward depends on the difference between the current cheetah velocity and this target. Hence, at each timestep. the emitted reward is ambiguous only for possible tasks. To identify the current task, the agent needs to estimate its velocity (which requires a few timesteps) and then disambiguate between these two tasks. This suggests that the agent will not benefit from very long sequences. Figure [[] (right) confirms this hypothesis: there is improvement from = to N = 5, but the performance decreases for longer sequences as the training becomes harder.
Atlation Study: Wer Working Keroory Se Sequence Length
3.6
Areraga hours
0.4
-300
3.0
-500
Cineztap
Emmssops
50
Figure 11: Ablation results for the working memory sequence length.
D.3 NUMBER OF LAYERS
Another important component is the network depth. Section 4 we hypothesized that more layers would help to recursively build more meaningful version of the episodic memory since we interact with output memories from the past layer and mitigates the bias effect from the task representations. Figure 12] shows how TrMRL behaves according to the number of layers. We observe similar pattern to the previous ablation case. For Reach-v2, more layers improved the performance by reducing the effect of ambiguity and biased task representations. For HalfCheeta Vel, we can see an improvement from a single layer to 4 or 8 layers, but for 12 layers, sample efficiency starts to decrease. On the other hand, highlight that even for deep network with 12 layers, we have stable optimization procedure, showing the effectiveness the T-Fixup initialization.
D.4 NUMBER OF ATTENTION HEADS
The last ablation case relates the number of attention heads in each MHSA block. We hypoth- esized that multiples heads would diversify working memory representation and improve network expressivity. Nevertheless, Figure shows that more heads slightly increased the performance in HalfCheetah Vel and did not interfere in Reach-v2 significantly.
19

---

Under review as a conference paper at ICLR 2022
Ablation Number
-100
3009
-302
0.0
Smeddpg
106
Umenteps
207
Figure 12: Ablation study for the number of transformer layers.
Ablatton Study: Number Attention
MotaWeadd Reachi2
KolCheschhve
Meeaage
0.4
3.0
Umaslags
1c5
Umistops
207
Figure 13: Ablation study for the number of attention heads.
20

---

Under review as a conference paper at ICLR 2022
E PROOF OF THEOREM
1
Theorem 1. Let S (Eo..... ,,y)) p(e 01) be a set of normalized episodic memory rep- resentations sampled from the posterior distribution p(elS', 01) induced by the transformer layer I, parameterized by O₁. Let K, Q V be the Key. Query, and Value vector spaces in the self- attention mechanism. Then, the self-atiention in the layer 1+1 computes consensus representation eN E e.. -exp sel ie! ; whose associated Bayes risk (in terms of negative cosine similar- Eiz exD (e, _.,))
ity) lower bounds the Minimum Bayes Risk (MBR) predicted from the set of candidate samples S projected onto the V space.
Proof. Let define Sty as the containing the projection of the elements in S' onto the V space: S = (eb ,...,e
e, ), where e Wy e' (Wv is the projection matrix). The Bayes risk of selecting ei. as representation, BR(ELV)) under a loss function C, is defined by:
BR(e = ax)
The MBR predictor selects the episodic memory @ E Sty that minimizes the Bayes Risk among I.V
the set of candidates: BM arg min BR(ê). Employing negative cosine similarity as loss function, represent MBR prediction as:
I.V. M BR
arg max peess EES)
(8)
The memory representation outputted from self-attention operation in layer 1 + 1 is given by:
I.V. E et exp (e L.K. Expexp exp (e J.K.
1+1
Σ ON. e.
(9)
10
The attention weights ON define probability distribution over the samples in S which approxi- mates the posterior distribution p(e|St, 0t). Hence, we can represent Equation 9 an expectation: 1+1 Eptelst, Finally, we compute the Bayer risk for it:
BR(c = Epiess , e
pels 91) (e. B e,)l))
[Bplesst a) (e), E meestt,o)le
2 Epleis lie, e)
= E pels o, ((e, e)) Vē e Sk.
(10)
21

---

Under review as a conference paper at ICLR 2022
F PSEUDOCODE
In this section, we present a pseudocode for TrMRL's agent during its interaction with an arbitrary MDP.
Algorithm TrMRL- Forward Pass
Require: MDP M p(M)
Require: Working Memory Sequence Length
Require: Parameterized function 6(s, r, 17)
Require: Transformer network with L layers {Si,..., Jr.)
Require: Policy Head
Initialize Buffer with N - PAD transitions: B =
e
N-..,,
t+0
Snert < 80
while episode not done do
Retrieve the
most recent transitions (8, a, r,7) from B to create the ordered subset D
D D U a AD,, (PADD:)PAD)
Compute working memories:
@ (s,, a, Tz. m)), Vis, a, T4 mh D
Set ef....,, Cx +- 01...., UN
for each / .,.,, I. do
Refine episodic memories:
el.....,ey e'x + fi(e 6 1-1)
end for
Sample a (-lef))
Collect (st+1, Te. Ne) interacting with M applying action a
8nezt t- 81+1
(st, at, Pt. Tt) t t+1
end while
22


Reviews:
[{"review_id": "P0I_aPl2yc", "reviewer": "Reviewer_irGj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "Strengths:\n* The paper addresses an important topic of Transformer architecture for meta-RL\n* The paper makes a valuable comparison of Transformer training settings and their respective impact on meta-learning\n* The preliminary results on meta-world are rather encouraging.\n\nWeaknesses:\n* Several important baselines are missing like (https://proceedings.mlr.press/v139/dance21a.html and https://arxiv.org/abs/1801.01290 )\n* The novelty of the model remains quite limited compared to the SoA Transformer current utilization in sequential decision making\n* The experimental results don't seem significant in ML45 settings which is an important transfer setting of Meta-World\n*  Maybe a connection to auto-prompting for zero-shot adaptation in controlled generation could have been interesting to develop (https://aclanthology.org/2020.emnlp-main.346.pdf).", "summary_of_the_paper": "The paper is proposing TrMRL (Transformers for Meta-Reinforcement Learning), a memory-based meta-Reinforcement Learner which uses the transformer architecture to formulate the learning process. \nIt works as a memory associating recent working information to create an episodic memory which is used to contextualize the policy.\nThe paper deals with the recently discovered limitations of Transformer training that exhibit been important to consider the particular context of sequential decision making and over-task generalization given an underlying environment dynamics.\nThe approach is evaluated on Metaworld in various task transfer settings.", "summary_of_the_review": "The paper is proposing an improvement of the transformer architecture, introducing a mechanism of memory reinstatement for meta-learning in sequential decision making.\nThe problem addressed in the paper is important and up-to-date.\nUnfortunately, several state-of-the-art approaches are missing in the evaluation part.\nMoreover, while several preliminary results are encouraging, the results depicted in ML45 remain limited compared to the proposed baseline approaches.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "O7J_hZ1wAd4", "reviewer": "Reviewer_ouZ3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "Strengths:\n1) The paper attempts to study an alternative network architecture (i.e. transformer) in RL, which is a direction that has a lot of importance and potential.\n2) The authors provide an explanation on how transformers can be a natural fit for memory-based meta RL. They argue that the attention mechanism can straightforwardly implement memory reinstatement.\n\nWeaknesses:\n1) Re the transformer use: TrMRL proposes to use transformers exactly as if they were RNNs in RL^2. The authors suggest to use the current transitions' embedding to input to the policy. This is a trivial extension of the prior work and a very obvious one. Moreover it has been done in various prior works and in several different settings. This basically makes TrMRL to be a version of RL^2 where an RNN is replaced with a transformer.\n2) Re PEARL results: the demonstrated learning curves of PEARL directly contradict the claims from the original paper and my intuition. I think that it is very unlikely that TRPO based MAML or PPO based RL^2 (both are on-policy) can outperform that dramatically SAC-based PEARL (off-policy). It is a common knowledge that off-policy methods are more sample efficient than the on-policy methods. This also makes me question the correctness of the whole experiment, as TrMRL also uses PPO as an RL backbone. \n3) Re T-Fixup initialization: the claim that the T-Fixup initialization helps is not supported in any experiment or an ablation study. It thus unclear how useful it is or essential.\n4) Re Fig 4: Why does this compare performance during meta-training rather than meta-testing? In the end we are more interested in the latter.\n5) Re Fig. 5: It would be more informative if other baselines are plotted here to server as a reference point.\n", "summary_of_the_paper": "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.\nThe paper provides an empirical study on a several high-dimensional continuous control tasks that feature locomotion and dexterous manipulation to support the aforementioned claims.", "summary_of_the_review": "While that paper studies an important direction in RL of adopting new network architectures, such as transformers, the paper, unfortunately, brings very little novelty or insight.  The proposed method largely builds off of the existing work (such as RL^2) and essentially proposes to replace RNNs with transformers. The empirical study is also questionable and,  in my opinion, misleading. Moreover, the experimental methodology is a suspect as it doesn't provide an extensive ablation or support all the claims.\n\nTaking this into account I suggest to reject the paper, as it in its current form doesn't meet the acceptance bar. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "67UrOXqY8H", "reviewer": "Reviewer_maih", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "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": "This paper investigates whether the few-shot learning capabilities of Transformers make them well-suited to the type of task adaptation required in meta RL. This is an intriguing premise, and I agree that there is a connection here that is worth studying. Unfortunately, beyond the introductory premise I found this paper somewhat difficult to parse. Many components of the algorithm are stated qualitatively, without a precise enough description to be able to reproduce, such as:\n\n> In this way, the transformer architecture recursively refines the episodic memory interacting output memories from the past layer.\nAs we compute each output memory in the perspective of a different timestep t, deeper layers reach\na consensus across the whole sub-trajectory.\n\nI generally had a hard time understanding these descriptions in the paper. For this one in particular:\n1. I am not sure what the recursive refinement refers to; could you write the recursion?\n2. \"Consensus\" is not defined or used outside of this paragraph.\n3. I think the \"working memory\" refers to feature vectors of recent transitions in the same spirit as the PEARL inference network $\\Psi( z \\mid s, a, s', r)$, and \"episodic memory\" is a function of summed working memory, as in Definition 6: $e_t = f(\\mu_t^0, \\ldots, \\mu_t^\\omega)$. However, as far as I can tell, $e_t$ is not referenced after its definition, so I do not know how it is used in the full algorithm. I also am not entirely sure what function $f$ is referring to. I assumed the memory terminology is the same as that in Fortunato et al 2019, but it would be better if these were precisely defined and motivated in this paper to make it more self-contained.\n\nI did my best to keep pedantic issues off of that list and focus on areas where I genuinely found clarity issues. (For example, \"from the perspective of\" struck me as a slightly strange anthropomorphism of Equation 5, but it is clear from context what that means, so is not as important.)\n\nWhile there is some focus given to the memory representation, there seem to be a number of missing details about the architecture and algorithm more generally. Most seriously, I could not find an actual objective aside from the general RL objective in the preliminaries. Since the experiments are in continuous control, presumably there needed to be some modification to the architecture described in Vaswani et al 2017?\n\nThe final stated contribution of this paper is a study of an important initialization method called T-Fixup. This ends up being a little underwhelming, with the ablation showing that it matters in one task but not in another. The explanation given is: \n> This result suggests that stabilizing learning since the beginning is very critical in the RL setting. This is often the\nmoment with higher exploration rates that prevent sub-optimal policies.\n\nThe first sentence seems a bit vacuously true, and I am not sure what the precise claim of the second is since the paper does not say anything about exploration.\n", "summary_of_the_paper": "This paper presents a Transformer-based approach to the meta RL problem, in which self-attention serves as a memory lookup mechanism for task adaptation. Experimental evaluation focuses on continuous control meta RL benchmarks.", "summary_of_the_review": "My main concerns about this paper are all related to clarity. It would help to have the full algorithm written precisely in one consolidated place (even if it's high-level pseudocode), rather than spread throughout the paper in informal descriptions with undefined terms.\n\nI would be happy to revisit this if the authors would like to clear things up during the discussion period.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "1aJ822mKnMu", "reviewer": "Reviewer_VA5K", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "Pros:\n(1) This paper introduces a new variant of meta-reinforcement learning by leveraging transformer, which can capture long sequence and context dependence for fast adaption in new tasks. And the intuition to leverage transformer makes sense. \n(2) It is well written and easy to understand and follow. There are few typos here and there but nothing that affects the readability of the paper.\n(3) They performed ablation study on how transformer architecture affects the performance. \n\nCons: \n(1) In the idea level, it is trivial to combine transformer model with meta-reinforcement learning to learn hidden distribution over different tasks for fast adaption. As RL^2, it builds the policy from the episodic memory generated by the transformer, instead of RNN.\n(2) The experimental results are not convincing. If I do not misunderstand it, I think N is the sequence length (correct me if it is wrong). The authors claim in the abstract that transformer can capture long sequence and context dependence for fast adaption, but why we set N=2 for all experiments? Is the transformer a good fit here? The experiment in appendix D.2 show it cannot get better result with long sequence. In Fig. 8, I also see it is not comparable to PEARL. \n\nMinor comments: \n(1)\tA typo in Eq. 5? Is it a summation from n=1 to t?\n(2)\tN=2 for all experiments? Can you try different N in Fig. 8, in order to compare with PEARL?\n\n", "summary_of_the_paper": "This paper presents TrMRL (Transformers for Meta-Reinforcement Learning), which leverages transformer to learn hidden distributions for meta-reinforcement learning. Specifically, the authors claim that the model associates the recent past of working memories to build an episodic memory recursively through the transformer layers, which can help to adapt to new tasks. The conducted experiments in high-dimensional continuous control environments for locomotion and dexterous manipulation show that TrMRL achieves comparable results on performance, sample efficiency and out-of-distribution generalization in these environments.", "summary_of_the_review": "This paper introduces a new variant of meta-reinforcement learning by leveraging transformer, which can capture long sequence and context dependence for fast adaption in new tasks. Empirically, the proposed method shows some improvements over the existing ones. However, a few major concerns are as follows.\n(1) The authors claim in the abstract that transformer can capture long sequence and context dependence for fast adaption, but why do we set N=2 for all experiments? Is the transformer a good fit here? The experiment in appendix D.2 show it cannot get better result with long sequence. \n(2) In Fig. 8, its performance is not comparable to PEARL. Is it possible to try different N and compare with PEARL?\n\nOverall, it is a good paper, but not good enough for publication. More convincing experiments needed to show the advantage of transformer here.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: At a high level, the novelty of this paper is limited: RL2 with transformers instead of RNNs. The emphasis is then placed on the experimental evaluation. Unfortunately, the reviewers felt that the experimental methodology and results were not strong enough at this stage to warrant publication. During the rebuttal, the reviewers did not engage nor discuss the author response, unfortunately, so I do not know what they think of the rebuttal. However, on evaluating the concerns of the reviewers against the updated manuscript, I think the updates do not go far enough to satisfy the concerns raised (experiments + baselines). Therefore, I recommend rejection.
Decision: Reject