In-context Reinforcement Learning with Algorithm Distillation
Authors: Michael Laskin
Luyu Wang
Junhyuk Oh
Emilio Parisotto
Stephen Spencer
Richie Steigerwald
DJ Strouse
Steven Stenberg Hansen
Angelos Filos
Ethan Brooks
maxime gazeau
Himanshu Sahni
Satinder Singh
Volodymyr Mnih
Venue: ICLR 2023 notable top 5% (2023)
Keywords: Reinforcement Learning
Transformers
Learning to Learn
Large Language Models

Abstract: We propose Algorithm Distillation (AD), a method for distilling reinforcement learning (RL) algorithms into neural networks by modeling their training histories with a causal sequence model. Algorithm Distillation treats learning to reinforcement learn as an across-episode sequential prediction problem. A dataset of learning histories is generated by a source RL algorithm, and then a causal transformer is trained by autoregressively predicting actions given their preceding learning histories as context. Unlike sequential policy prediction architectures that distill post-learning or expert sequences, AD is able to improve its policy entirely in-context without updating its network parameters. We demonstrate that AD can reinforcement learn in-context in a variety of environments with sparse rewards, combinatorial task structure, and pixel-based observations, and find that AD learns a more data-efficient RL algorithm than the one that generated the source data.

Published as a conference paper at ICLR 2023
IN-CONTEXT REINFORCEMENT LEARNING WITH ALGORITHM DISTILLATION
Michael Laskin"! Luyu Junhyukk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, Maxime Gazeau, Himanshu Sahni, Satinder Singh, Volodymyr Mnih DeepMind
ABSTRACT
We propose Algorithm Distillation (AD), method for distilling reinforcement learning (RL) algorithms into neural networks by modeling their training his- tories with causal sequence model. Algorithm Distillation treats learning reinforcement learn as an across-episode sequential prediction problem. A dataset of learning histories is generated by a source RL algorithm, and then causal transformer is trained by autoregressively predicting actions given their preceding learning histories as context. Unlike sequential policy prediction architectures that distill post-learning expert sequences, AD is able to improve its policy entirely in- in-context without updating its network parameters. We demonstrate that AD can reinforcement learn in-context in variety of environments with sparse rewards, combinatorial task structure, and pixel-based observations, and find that AD learns more data-efficient RL algorithm than the one that generated the source data.
Data Generation
Task
RL algorithm learning histories
Task Yl
(n)
(Og, do, ro, 01, d1, Pl...... Or, ar, PT)
learning progress
Model Training
Predict actions using across-episode contexts
o
dg
To
O--1
Og-1
Fe-1
Ot
Causal Transformer
Po 0t)
Figure 1: Algorithm Distillation (AD) has two steps
datasct of learning histories is collected from individual single-task RL algorithms solving different tasks; (ii) causal transformer predicts actions from these histories using across-episodis contexts. Since the RL policy improves throughout the learning histories, by predicting actions accurately AD learns to output an improved policy relative to the one seen in its context. AD models state-action-reward tokens, and does not condition returns.
INTRODUCTION
Transformers have emerged as powerful neural network architectures for sequence modeling (Vaswani et al., 2017). A striking property of pre-trained transformers is their ability to adapt to downstream tasks through prompt conditioning or in-context learning. After pre-training on large offline datasets, large transformers have been shown to generalize to downstream tasks in text completion (Brown al., 2020), language understanding (Devlin al., 2018), and image generation (Yu al., 2022).
'Equal contribution. Corresponding authors: {ml askin, vmni h 2deeppi ndd.iiddcom com.

---

Published as a conference paper at ICLR 2023
Recent work demonstrated that transformers can also learn policies from offline data by treating offline Reinforcement Learning (RL) as a sequential prediction problem. While Chen et al. (2021) showed that transformers can learn single-task policies from offline RL data via imitation learning, subsequent work showed that transformers can also extract multi-task policies in both same-domain (Lee et al., 2022) and cross-domail settings (Reed et al., 2022). These works suggest a promising paradigm for extracting generalist multi-task policies first collect a large and diverse dataset of environment interactions, then extract a policy from the data via sequential modeling. We refer the family of approaches that learns policies from offline RL data via imitation learning as Offline Policy Distillation, or simply Policy Distillation' (PD).
Despite its simplicity and scalability, a substantial drawback of PD is that the resulting policy does not improve incrementally from additional interaction with the environment. For instance, the Multi- Game Decision Transformer (MGDT, Lee et al., 2022) learns a return-onditioned policy that plays many Atari games while Gato (Reed et al., 2022) learns policy that solves tasks across diverse environments by inferring tasks through context, but neither method can improve its policy in-context through trial and error. MGDT adapts the transformer to new tasks by finetuning the model weights while Gato requires prompting with an expert demonstration to adapt to Distillation methods learn policies but not Reinforcement Learning algorithms.
task. In short, Policy
We hypothesize that the reason Policy Distillation does not improve through trial and error is that it trains on data that does not show learning progress. Current methods either learn policies from data that contains no learning (e.g. by distilling fixed expert policies) or data with learning (e.g. the replay buffer of an RL agent) but with a context size that is too small to capture policy improvement.
Our key observation that the sequential nature of learning within RL algorithm training could, in principle, make it possible to model the process of reinforcement learning itself as a causal sequence prediction problem. Specifically, if transformer's 's context is long enough include policy improvement due to learning updates it should be able to represent not only fixed policy but policy improvement operator by attending to states, actions and rewards from previous episodes. This opens the possibility that any RL. algorithm can be distilled into sufficiently powerful sequence model such as a transformer via imitation learning. converting it into an in-context RL. algorithm. By in-context RL we mean that the transformer should improve its policy through trial and error within the environment by attending to its context, without updating its parameters.
We present Algorithm Distillation (AD), a method that learns an in-context policy improvement operator by optimizing a causal sequence prediction loss on the learning histories of an RL algorithm. AD has two components. First, large multi-task dataset is generated by saving the training histories of an RL. algorithm on many individual tasks. Next, a transformer models actions causally using the preceding learning history as its context. Since the policy improves throughout the course of training of the source RL algorithm, AD is forced to learn the improvement operator in order to accurately model the actions any given point in the training history. Crucially, the transformer context size must be sufficiently large (i.e. across-episodic) to capture improvement in the training data. The full method is shown in Fig.
We show that by imitating gradient- RL algorithms using a causal transformer with sufficiently large contexts, AD can reinforcement learn new tasks entirely in-context. We evaluate AD across a number of partially observed environments that require exploration, including the pixel-based Watermaze (Morris, 1981) from DMLab (Beattie al., 2016). We show that AD is capable of in-context exploration, temporal credit assignment, and generalization. We also show that AD learns a more data-efficient algorithm than the one that generated the source data for transformer training. To the best of our knowledge, AD is the first method to demonstrate in-context reinforcement learning via sequential modeling of offline data with an imitation loss.
2 BACKGROUND
Partially Observable Markov Decision Processes: A Markov Decision Process (MDP) consists of states 8 E S, actions E A, rewards R. discount factor 7, and a transition probability function p( p($t11) 1 at), where t is an integer denoting the timestep and (S, A) are state and action spaces. In
What we refer to as Policy Distillation is similar to Rusu et al. (2016) but the policy is distilled from offline data, not teacher network.
2

---

Published as a conference paper at ICLR 2023
environments described by an MDP, each timestep t the agent observes the state st, selects action a ~ (<([s)) from its policy, and then observes the next state $141 ~ ((-]st, Qe) sampled from the transition dynamies of the environment. In this work, we operate in the Partially Observable Markov Decision Process (POMDP) setting where instead of states E S the agent receives observations E O that only have partial information about the true state of the environment. Full state information may be incomplete due to missing information about the goal in the environment, which the agent must instead infer through rewards with memory, or because the observations are pixel-based, or both.
Online and Offline Reinforcement Learning: Reinforcement Learning algorithms aim to maximize the return, defined as the cumulative sum of rewards Ee 2 throughout an agent's lifetime or episode of training. RL algorithms broadly fall into two categories: on-policy algorithms (Williams, 1992) where the agent directly maximizes a Monte-Carlo estimate of the total returns or off-policy (Mnih et al., 2013; 2015) where an agent learns and maximizes a value function that approximates the total future return. Most RL algorithms maximize returns through trial-and-error by directly interacting with the environment. However, offline RL (Levine et al., 2020) has recently emerged as an alternate and often complementary paradigm for RL where an agent aims to extract return maximizing policies from offiine data gathered by another agent. The offline dataset consists of (s, a, r) tuples which are often used to train an off- off-policy agent, though other algorithms for extracting return maximizing policies from offline data are also possible.
Atte-ionnnionn and Transformers The self-attention (Vaswani et al., 2017) operation begins by projecting input data X with three separate matrices onto D-dimensional vectors called queries Q. keys K, and values V. These vectors are then passed through the attention function:
(1) The QK" term computes an inner product between two projections of the input data X. The inner product is then normalized and projected back to a D D-dimensional vector with the scaling term V. Transformers (Vaswani al., 2017; Devlin et al., 2018: 8; Brown et al., 2020) utilize self-attention as core part of the architecture to process sequential data such as text sequences. Transformers are usually pre-trained with a self-supervised objective that predicts tokens within the sequential data. Common prediction tasks include predicting randomly masked out tokens (Devlin et al., 2018) or applying a causal mask and predicting the next token (Radford et al., 2018).
Attention Q, K, V) = softmax(QK' /VDV.
We refer to the family of methods that treat offline Reinforcement Learning as a sequential prediction problem as Offline Policy Distillation, or Policy Distillation (PD) for brevity. Rather than learning value function from offline data, PD extracts policies by predicting actions in the offline data (Lee. behavior cloning) with a sequence model and either return conditioning (Chen et al., 2021; Lee et al., 2022) or filtering out suboptimal data (Reed et al., 2022). Initially proposed to learn single-task policies (Chen et al., 2021; Janner et al., 2021), PD was recently extended to learn multi-task policies from diverse offline data (Lee et al., 2022: Reed et al., 2022).
Offline Policy Distillation:
In In-Context Learning: In-context learning refers to the ability to infer tasks from context. For example, large language models like GPT-3 (Brown et al., 2020) or Gopher (Rae al., 2021) can be directed at solving tasks such as text completion, code generation, and text summarization by specifying the task through language as prompt. This ability to infer the task from prompt is often called in-context learning. We use the terms 'in-weights learning' and 'in-context learning' from prior work on sequence models (Brown et al., 2020; Chan et al., 2022) to distinguish between pradient-based learning with parameter updates and gradient-free learning from context, respectively.
3 METHOD
Over the course of its lifetime a capable reinforcement learning (RL) agent will exhibit complex behaviours, such as exploration, temporal credit assignment, and planning. Our key insight is that an agent's 's actions, regardless of the environment it inhabits, its internal structure, and implementation, can be viewed function of its past experience, which we refer to as its history. Formally, we write:
H. 3 b4 == (Oo: ap, Tu..... ot- 41-1, Tt-1, oz, Qe re) = (0st, ag:) and we refer to a long² history-conditiooned policy as an algorithm:
(2)
P P:HUO- (AA),
(3)
2I Long enough to span learning updates, e.g. across episades.
3

---

Published as a conference paper at ICLR 2023
where ((A) denotes the space of probability distributions the action space A. Eqn. (3) suggests that, similar to policy, an algorithm can be unrolled in an environment to generate sequences ob- servations, rewards, and actions. For brevity, denote the algorithm P and environment (i.e. task) as M, such that the history of learning for any given task M is generated by the algorithm PM-
(Oo, Ao. Ro...., Ar A,, ~ PM.
(4)
Here, we're denoting random variables with uppercase Latin letters, egg. O, A, R. and their values with lowercase Latin letters, e.g. O, a, т. By viewing algorithms long history-conditioned policies, we hypothesize that any algorithm that generated a set of learning histories can be distilled into neural network by performing behavioral cloning actions. Next, present method that, provided agents' lifetimes, learns a sequence model with behavioral cloning to map long histories to distributions over actions.
3.1 ALGORITHM DISTILLATION
Suppose the agents' lifetimes, which we also call learning histories, are generated by the source algorithm psource for many individual tasks {M, producing the dataset D:
D { O0 .o
D=
(n)
(m) (n)
(x) (n)) of aT TT
(n)
prource n=1
(5)
Then we distill the source algorithm's 's behaviour into a sequence model that maps long histories to probabilities over actions with negative log likelihood (NLL) loss and refer to this process as algorithm distillation (AD). In this work, we consider neural network models Po with parameters 0 which we train by minimizing the following loss function:
( - Σ
log Po(A at (n) (n) 1 (n)).
Intuitively, sequence model with fixed parameters that is trained with AD should amortise the source RL algorithm psource and by doing so exhibit similarly complex behaviours, such as exploration and temporal credit assignment. Since the RL policy improves throughout the learning history of the source algorithm, accurate action prediction requires the sequence model to not only infer the current policy from the preceding context but also infer the improved policy, therefore distilling the policy improvement operator.
Algorithm AAlgorithm Distillation
Require: Train {MMiinn and test (MM tasks, observations o E O, actions a E A, and rewards E R.
Require: Network parameters di for i 1,... N source RI. algorithms.
Require: Network parameters 0 for causal transformer Po that predicts actions.
Require: An empty buffer to store data D.
i... do
>> Part 1: Dataset Generation
2: Sample task MMtaainn randomly from the train task distribution.
3: Train the source RL algorithm di until it converges to the optimal policy.
4: Save the learning history h (oo, ao, To, or, ar, rT), to the dataset D < DUU (8)
5: end for
6: while Pa converged do
Part Algorithm Distillation 7: Randomly sample multi-episodic subsequence n = (o,, a, ....... Oytes asscr Pzte), of length
8: Autoregressively predict the actions with Po compute the NLL loss in Eq.
9: Backpropagate update transformer parameters.
10: end while
11:
for ds == 1 Meest do
Part 3: Autoregressive Evaluation Sample task Ma randomly from the test task distribution. Initialize empty context queue C. Unroll the transformer (=)) ) the environment storing sequential transitions (i.e. histories) C.
12:
14:
Measure the return accumulated by the agent for each episode of evaluation.
15: end for
3.2 PRACTICAL IMPLEMENTATION
In practice, we implement AD as a two-step procedure. First, a dataset of learning histories is collected by running an individual gradient-based RI. algorithm many different tasks. Next,
4

---

Published as a conference paper at ICLR 2023
sequence model with multi-episodic context is trained to predict actions from the histories. We describe these two steps below and detail the full practical implementation in Algorithm 1.
Dataset Generation: A dataset of learning histories is collected by training N individual single-task gradient-based RL algorithms. To prevent overfitting to any specific task during sequence model training, task M is sampled randomly from a task distribution for each RL run. The data generation step is RL algorithm agnostic any RL algorithm can be distilled. We show results distilling UCB exploration (Lai & Robbins, 1985), an on-policy actor-critic (Mnih et al., 2016), and an off-policy DQN (Mnih et al., 2013), in both distributed and single-stream settings. We denote the dataset of learning histories as D Eq. 5.
Training the Sequence Model: Once dataset of learning histories D is collected, a sequen- tial prediction model is trained predict actions given the preceding histories. We utilize the GPT (Radford et al., 2018) causal transformer model for sequential action prediction, but AD is compatible with any sequence model including RNNs (Williams & Zipser, 1989). For instance, we show in Appendix L that AD can also be achieved with LSTM (Hochreiter & Schmidhuber, 1997), though less effectively than AD with causal transformers. Since causal transformer train- ing and inference are quadratic in the sequence length, we sample across-episodic subsequences h = (o ag Of+er +ee aj +e) of length c < T from rather than training full histories.
EXPERIMENTAL SETUP
4.1 ENVIRONMENTS
To investigate the in-context RL capabilities of AD and the base- lines (see next section), we focus on environments that cannot be solved through zero-shot generalization after pre-training. Specif- ically, we require that each environment supports many tasks, that the tasks cannot be inferred easily from the observation, and that episodes are short enough to feasibly train across-episodic causal transformers for more details regarding environments see Appendix B. We list the evaluation environments below:
Adversarial Bandit:
a multi-armes bandit with 10 arms and 100 trials similar to the environment considered in RL. (Duan et al., 2016). However, during evaluation the reward is out of distribu- tion. Reward is more likely distributed under odd arms 95% of the time during training. At evaluation, the opposite happens reward appears more often under even arms 95% of the time.
Figure 2: Agent view from the DM- Lab Watermaze environment. The task find hidden platform that elevates once found.
Dark Room: a 2D discrete POMDP where an agent spawns in a room and must find a goal location. The agent only knows its own (z. coordinates but does not know the goal location and must infer it from the reward. The room size is 9 x 9, the possible actions are one step left, right, up, down, and no-op, the episode length is 20, and the agent resets at the center of the map. We test environment variants - Dark Room where the agent receives r = 1 every time the goal is reached and Dark Room Hard. hard exploration variant with a 17 x 17 size and sparse reward (r 1 exactly once for reaching the goal). When not r = 1, then r = 0.
Dark Key-to-Door: similar to Dark Room but this environment requires an agent to first find an invisible key upon which it receives a reward of r = 1 once and then open an invisible door upon once again. Otherwise, the reward making the task space combinatorial with 812 = 6561 possible tasks. This environment is similar to the one considered in Chen et al. (2021) except the key and door are invisible and the reward is semi- sparse (r = 1 for both key and the door). The agent is randomly reset. The episode length is 50 steps.
which it receives a reward of
The room size is 9
DMLab Watermaze: partially observable 3D visual DMLab environment based on the classic Morris Watermaze (Morris, 1981). The task is navigate the water maze to find a randomly spawned trap-door. The maze walls have color patterns that can be used to remember the goal location. Observations pixel images of size 72 x 96 3. There are 8 possible actions in total, including going forward, backward, left, or right, rotating left or right, and rotating left or right while going forward. The episode length is 50, and the agent resets the center of the map. Similar to Dark
5

---

Published as a conference paper at ICLR 2023
Bandit Scores
Normalized Score
Cdd (Seen)
Even (Unseen)
0.5
Uniform (Unseen)
-0.5
AD
RL?
RL' nL--Transfermer
Figure 3: Adversarial Bandit (Section 5): AD, RL?, and ED evaluated on a 10-arm bandit with 100 trials. The source data for AD comes from learning histories from UCB (l.i & Robbins, 1985). During training. the reward distributed under odd 95% of the time and under even RL can in-context learn in-distribution lasks, AD generalizes better out distribution. Running RL with a transformer generally doesn't I offer an advantage over the original LSTM variant. ED performs poorly both in and out of distribution relative to AD and RL Scores are normalized relative to UCB.
95% of the time during evaluation. Both AD
Room, the agent cannot see the location of the goal from the observations and must infer through the reward of r if reached and r 0 otherwise; however, the goal space is continuous and therefore there are an infinite number of goals.
4.2 BASELINES
The main aim of this work investigate to what extent AD reinforcement learns in-context relative to prior related work. AD is mostly closely related to Policy Distillation, where policy is learned with a sequential model from offline interaction data. In-context online meta-RL is also related though not directly comparable to AD, since AD is in-context offline meta-RL method. Still, we consider both types of baselines to better contextualize our work. For more detailed discussion of these bascline choices refer the reader to Appendix C. Our basclines include:
Expert Distillation (ED): this baseline is exactly the same as AD but the source data consists expert trajectories only, rather than learning histories. ED is most similar Gato (Reed al., 2022) except ED models state-action-rewarr sequences like AD, while Gato models state-action sequences.
Source Algorithm: compare AD to the gradient-basee source RL algorithm that generates the training data for distillation. We include running the source algorithm from scratch as a baseline to compare the data-efficiency of in-context RL the in-weights source algorithm.
RI.2 (Duan et al., 2016): an online meta-RL. algorithm where exploration and fast in-context adapta- tion are learned jointly by maximizing multi-episodic value function. RL is not directly comparable to AD for similar reasons to why online and offline RL algorithms are not directly comparable RL gets to interact with the environment during training while AD does not. We use RL asymptotic performance as approximate upper bound for AD.
4.3 EVALUATION
After pre-training, the AD transformer P5 can reinforcement learn in-context. Evaluation is exactly the same as with in-weishts RL algorithm except the learning happens entirely in-context without updating the transformer network parameters. Given an MDP (or POMDP), the transformer interacts with the environment and populates its own context (i.e. without demonstrations), where the context is a queue containing the last e transitions. The transformer's 's performance is then evaluated in terms of ability to maximize return. For all evaluation runs, we average results across training seeds with 20 evaluation seeds each for total of 100 seeds. A task M is sampled uniformly from the test task distribution and fixed for each evaluation seed. The aggregate statistics reported reflect multi-tast performance. We evaluate for 1000 and 160 episodes for the Dark and Watermaze environments respectively and plot performance as a function of total environment steps at test-time.
EXPERIMENTS
The main research question of this work is whether an in-weights RL algorithm can be amortized into an in-context via Algorithm Distillation. The in-context RL. algorithm should behave in similar

---

Published as a conference paper at ICLR 2023
Dark Room
Dark Room (Hard)
Dark Key to-
Watermaee
AD
FD Source
02
RI2 (3 billion)
12 15
20
02
12
20
30 26 30 Stepe (housard))
10
50
Steas (Thcosard))
Steps (hcusand))
Figure 4: Main results: we evaluate AD, RL3, ED, and the source algorithm on environments that require memory and exploration. In these environments, an agent must reach goal location that can only be inferred through binary reward. AD is consistently able to in-context reinforcement learn across all environments and is more data-efficient than the A3C ("Dark" environments) (Mnih al., 2016) or DQN (Watermazee (Mnih al., 2013) source algorithm it distilled. We report the mean return +: standard deviation over training seeds with 20 test seeds each.
way as the in-weights one and exhibit exploration, credit assignment, and generalization capabilities. We begin our analysis
clean and simple experimental setting where all three properties are required to solve the task- the Adversarial Bandit described in Sec. 4.
To generate the source data, we sample a set of training tasks {M run the Upper Confidence Bound algorithm (Lai & Robbins, 1985), and save its learning histories. We then train a transformer to predict actions as described in Alg. We evaluate AD, ED, and RL and normalize their scores relative to UCB and random policy (r Trand.) //
Trand.). The results are shown in Fig. 3. We find that both AD and RL2 can reliably in-context learn tasks sampled from the training distribution while ED cannot, though ED does do better than random guessing when evaluated in-distribution. However, AD can also in-context learn to solve out of distribution tasks whereas the other methods cannot. This experiment shows that AD can explore the bandit arms, can assign credit by exploiting an arm once reached, and can generalize well out of distribution nearly as well as UCB. We now move beyond the bandit setting and investigate similar research questions more challenging RL environments and present our results as answers to a series of research questions.
Does Algorithm Distillation exhibit in-context reinforcement learning? To answer this question, we first generate source data for Algorithm Distillation. In the Dark Room and Dark Key-to-Door environments we collect 2000 learning histories with an Asynchronous Advantage Actor-Critic (A3C) (Mnih et al., 2016) with 100 actors, while in DMLab Watermaze we collect 4000 learning histories with a distributed DQN with 16 parallel actors (see Appendix F for asymptotic learning curves of the source algorithm and Appendix O for hyperparameters) Shown in Fig. 4, AD in-context reinforcement learns in all of the environments. In contrast, ED fails to explore and learn in-context in most settings. We use RI.² trained for billion environment steps as a proxy for the upper bound of performance for meta-RI. method, RI.2 achieves a near-optimal asymptotic score in all the environments except for Dark Room (Hard). Despite learning from offline data, AD matches asymptotic RL on the Dark environments and approaches it (within 13%) on Watermaze.
Credit-assinmmnt:: In Dark Room, the agent receives r = 1 each time it visits the goal location. Even though AD is trained to condition only on single timestep reward and not episodic return tokens, it is still able maximize the reward, which suggests that AD has learned to do credit assignment.
Exploration: Dark Room (Hard) tests the agents exploration capability. Since the reward is sparse (r exactly once), most of the learning history has reward values of r 0. Nevertheless, AD infers the goal from previous episodes in its context which means it has learned to explore and exploit.
Generalization: Dark Key-to-Door tests in-distribution generalization with combinatorial task space. While the environment has a total of ~ 6.5k tasks, less than 2k were seen during training. During evaluation, AD both generalizes and achieves near-optimal performance on mostly unseen tasks.
Can Algorithm Distillation learn from pixel-based observations? DMLab Watermaze a pixel- based environment that is larger than the Dark environments with tasks sampled from continuous uniform distribution. The environment is partially observable in two ways the goal is invisible until the agent has reached it and the first-person view limits the agent's 's field of vision. Shown in Fig. 4, AD maximizes the episodic return with in-context RL while ED does not learn.
Can AD learn more data-efficient RL algorithm than the one that produced the source data? In Fig. 4, AD is significantly more data-efficient than the source algorithm. This gain is a byproduct of distilling multi-stream algorithm into single-stream one. The source algorithms
7

---

Published as a conference paper at ICLR 2023
Algorithm Distillation
Expert Distillation
15
15
5% optimal
5% optimal
Return
20
25% optimal
Return
10
25% optimal
optimal
optimal
75% optimal
75% optimal
optimal
95% optimal
%
20
20
Staps (Thousands)
Staps
Figure 5: AD and ED conditioned on partial demonstrutionn: We compare the performance of AD and ED when prompted with demonstration from the source algorithm's 's training history on Dark Room (semi-dense). While ED slightly improves and then maintains performance from the input policy, AD is able improve it in-context until the policy is optimal nearly optimal.
(A3C and DQN) are distributed, which means they run many actors in parallel to achieve good performance. 3 A distributed RL algorithm may not be very data-efficient in aggregate but each individual actor can be data-efficient. Since the learning history for each actor is saved separately, AD achieves similar performance to the multi-stream dis- tributed RL algorithm, but is more data-efficient as a single-stream method.
Single-stream AD
2.0
These data-efficiency gains are also evident for distilling single-stream algorithms. In Fig 6 show that by sub- sampling every k-th episode (where k = 10) from a single stream A3C learning history, AD can still learn more data-efficient in-context RL algorithm (for more detail, see Appendix J). Therefore, AD can be more data-efficient than both a multi and single-stream source RL algorithm.
Return
1.0
Source
AD
0.0
10
20
30
40
50
Env Steps (Thousands))
While AD is more data-efficient, the source algorithm achieves slightly higher asymptotic performance (see Ap- pendix F). However, the source algorithm produces many single-task agents with unique set of weights on per task Mn. while AD produces a single generalist agent with weights 0 that are fixed across all tasks.
Figure 6: Single-stream Algorithum Distil-
lation: AD trained on the learning history from an A3C agent with only one actor (i.e. single-stream). By training on subsampled learning histories (see Sec. 5), AD learns more data-efficient in-context RL algorithm.
Is it possible to accelerate AD by prompting it with demonstrations? Although AD can reinforcement learn without relying on demonstrations, it has the added benefit that, unlike the source algorithm, it can be conditioned on or prompted with external data. To answer the research question, we sample policies from the hold-out test-set data along different points of the source algorithm history from a near-random policy to near-expert policy. We then pre-fill the context for both AD and ED with this policy data, and run both methods in the environment in Dark Room (Fig. 5). While ED maintains the performance of the input policy, AD improves every policy in-context until it is near-optimal. Importantly, the more optimal the input policy the faster AD improves until it is optimal.
Dark Key-to-Door
20 epl
Return
1.5
30 epl
40 epl
1.0
50 (1 op)
100 (2 eps)
200 (4 eps)
10
20
30
Env Steps (Thousands)
Figure 7: Context size: AD Dark Key-to- Door with different context sizes. In-context RL only emerges once the context size is large enough and across-episodic.
What context size required for in-context RL to emerge? We've hypothesized that AD requires sufficiently long (i.e. across-episodic) contexts to in in-context reinforcement learn. We test this hypothesis by training several AD variants with different context sizes on the Dark Room environment. We plot the learning curves of these different variants in Fig. 7 and find that multi-episodic contexts of 2-4 episodes are necessary to learn a near-optimal in-context RL algorithm. Initial signs of in-context RL begin to emerge when the context size is roughly the length of episode. The reason for this is likely that the context is large enough to retrain across-episodic information e.g. at the start of a new episode, the context will be filled with transitions from most of the previous episode.
Indeed, current state-of-the-art RI. algorithms such as MuZero (Schrittwieser et al., 2019) and Muesli (Hessel al., 2021) rely distributed actors.
8

---

Published as a conference paper at ICLR 2023
6 RELATED WORK
Offline Policy Distillation: Most closely related to our work the recent advances in learning policies from offline environment interaction data with transformers, which we have been referring to as Policy Distillation (PD). Initial PD architectures such Decision Transformer (DT) (Chen et al., 2021) and Trajectory Transformer (Janner et al., 2021) showed that transformers can learn single-task policies from offline data. Subsequently the Multi-Game Decision Transformer (MGDT) (Lee et al., 2022) and Gato (Reed et al., 2022) showed that PD architectures can also learn multi-task same domain and cross-domain policies, respectively. Importantly, these prior methods use contexts substantially smaller than an episode length, which is likely the reason in-context RL was not observed in these works. Instead, they rely on alternate ways to adapt to new tasks MGDT finetunes the model parameters while Gato gets prompted with expert demonstrations to adapt to downstream tasks. AD adapts in-context without finetuning and does not rely on demonstrations. A number of recent works have explored more generalized PD architectures (Furuta et al., 2021), I). prompt conditioning (Xu et al., 2022), and online srudient-basee RL (Zheng et al., 2022). Some PD architectures such as DT and MGDT are instantiations of Upside Down RL (UDRL) Schmidhuber (2019); Srivastava et al. (2019) where rather than learning a value function, a policy is conditioned directly the desired return. However, AD is not explicitly doing UDRL since it is not conditioned on returns. In fact, return maximization is not specified anywhere in the AD objective but rather emerges implicitly by modeling the learning histories of an RL algorithm.
Meta Reinforcement Learning: AD falls into the category of methods that learn to reinforcement learn, also known as meta-RL. Specifically, AD an in-context offline meta-RL method. This general idea of learning the policy improvement process has a long history in reinforcement learning, but has been limited to meta-learning hyper- seper-parameters' until recently (Ishii et al., 2002). In-context deep meta-RL. methods introduced by Wang et al. (2016) and Duan et al. (2016) are usually trained in the online setting by maximizing multi-episodic value functions with memory- based architectures through environment interactions (Ni et al., 2022). Meta-RL through multi-episodic value functions has been done in both on-policy (Duan et al., 2016) and off-policy (Rakelly et al., 2019; Fakoor et al., 2020) settings. Another common approach online meta-RL RL includes optimization-bart methods that find good network parameter initializations for meta-RL (Hochreiter et al., 2001; Finn et al., 2017: Nichol et al., 2018) and adapt by taking additional gradient steps. Like other in-context meta-RL approaches, AD is gradient-fres its network parameters. Recent works have proposed learning to reinforcement learn from offline datasets, or offline meta-RL, using Bayesian RL (Dorfman et al., 2021) and optimization-bases meta-RL (Mitchell et al., 2021). Given the difficulty of offline meta- -L, Pong et al. (2022) proposed a hybrid offline-online strategy for meta-RL.
it adapts to downstream tasks without updating
i--ontext Learning with Transformers: In this work, we make the distinction between in-context learning and incremental or in-context learning. In-context learning involves learning from a provided prompt or demonstration while incremental in-context learning involves learning from one's own behavior through trial and error. While many recent works have demonstrated the former, it is much less common to see methods that exhibit the latter. Arguably, the most impressive demonstrations of in-context learning to date have been shown in the text completion setting (Radford et al., 2018; Chen et al., 2020; Brown et al., 2020) through prompt conditioning. Similar methodology was recently extended to show powerful composability properties in text-conditioned image generation (Yu et al., 2022). Recent work showed that transformers can also learn simple algorithm classes, such as linear regression, in-context in small-scale setting (Garg al., 2022). Like prior in-context learning methods, Garg et al. (2022) required initializing the transformer prompt with expert examples. While the aforementioned approaches were examples of in-context learning, recent work (Chen al., 2022) demonstrated incremental in-context learning for hyperparameter optimization by treating hyperparameter optimization as a sequential prediction problem with a score function.
7 CONCLUSION
We have demonstrated that can distill RL algorithms by modeling their learning histories causally with imitation learning and that AD can learn more data-efficient algorithms than those that generated the source data. We hope that AD inspires further investigation into in-context reinforcement learning from the research community.
FFor a recent example of meta-optimization in RL. see Flennerhag al. (2022).
9

---

Published as a conference paper at ICLR 2023
REFERENCES
Charles Beattie, Joel Z. Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, Julian Schrittwieser, Keith Anderson, Sarah York, Max Cant, Adam Cain, Adrian Bolton, Stephen Gaffney, Helen King, Demis Hassabis, Shane Legg, and Stig Petersen. DeepMind Lab. CoRR, abs/1612.03801, 2016.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
Stephanie CY Chan, Adam Santoro, Andrew K Lampinen, Jane X Wang, Aaditya Singh, Pierre H Richemond, Jay McClelland, and Felix Hill. Data Distributional Properties Drive Emergent In-Context Learning in Transformers. arXiv preprint 2rXiv:2205.05055, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084- 5084-5 2021.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and llya Sutskever. Generative pretraining from pixels. In International Conference on Machine Learning, pp. 1691 1703. PMLR. 2020.
Yutian Chen, Xingyou Song. Chansoo Lee, Zi Wang. Qiuyi Zhang, David Dohan, Kazuya Kawakami, Greg Kochanski, Arnaud Doucet, Marc' aurelio Ranzato, Sagi Perel, and Nando de Freitas. Towards Learning Universal Hyperparameter Optimizers with Transformers, 2022.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint wrXiv:181000480. 2018.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline Meta Reinforcement Learning-Leentifinnn Identifiability Challenges and Effective Data Collection Strategics. Advances Neural Information Processing Systems, 34:4607-4613 2021.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Picter Abbeel. RL3: Fast Reinforcement Learning via Slow Reinforcement Learning, 2016.
Rasool Fakoor, Pratik Chaudhari, Stefano Soatto, and Alexander Smola. Meta-q-learning. In International Conference Learning Representations, 2020. URL
net/forum: S
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-1135. PMLR, 2017.
Sebastian Flennerhag, Yannick Schroecker, Tom Zahavy, Hado van Hasselt, David Silver, and Satinder Singh. Bootstrapped meta- learnig. In International Conference on Learning Representations, 2022. URL https: /ttp::/:////peeeeeeeee net wttp:/oooomoododo/o/oooe--e-- 1E5.
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:21 2021.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes, 2022.
Matteo Hessel, Ivo Danihelka, Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre, Theophane Weber, David Silver, and Hado van Hasselt. Muesli: Combining Improvements in Policy Optimiza- tion. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 Proceedings of Machine Learning Research, pp. 4214 4226. PMLR, 2021.
Sepp Hochreiter and Jürgen Schmidhuber. Long Short- Term Memory. Neural Computation, 9())
1735-1780, 1997.
10

---

Published as a conference paper at ICLR 2023
Sepp Hochreiter, A Steven Younger, and Peter R Conwell. Learning to learn using gradient descent. In International conference on artificial neural networks, pp. 87-94 94. Springer, 2001.
Shin Ishii, Wako Yoshida, and Junichiro Yoshimoto. Control of pnploiiaaton-eeploraatoo meta- parameter in reinforcement learning. Neural networks, 15(4-6)x665663) 2002.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement Learning as One Big Sequence Modeling Problem. arXiv preprint urXiv:2106.02039 2021.
T.L. Lai and Herbert Robbins. Asymptotically Efficient Adaptive Allocation Rules, Adv. Appl. Math., 6(1):4-22, mar 1985. ISSN 0196-8858. doi: 00.1066019--88885599000000
Kuang-Huee Lee, Ofr Nachum, Mengjiao Yang. Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, lan Fischer, Eric Jang, Henryk Michalewski, and Igor Mordatch. Multi-Game Decision Transformers, 2022.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement leamnigg Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01644 2020.
Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pp. 7780-7791. PMLR, 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 5187754):55295555 533, 2015.
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous Methods for Deep Reinforcement Learning. In Maria-Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY. USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 1928-1937 JMLR.org. 2016.
Richard G.M. Morris. Spatial localization does not require the presence of local cues. Learning and Motivation, (22)):399260, 1981, doi: 10. 10.10160002-969006 1900020-5
Rafael Müller, Simon Korbllihh, and Geoffrey E Hinton. When does label smoothing help? In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL https: ptttpa://pproceedingge neurips wwt.::e....aaaaappppepeeee wttpppppcccpppp//../////// 9/file/ .df..
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free RL can be a strong baseline for many pomdps. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 16691-16723. PMLR, 2022. URL https: ppwpo/eeeediieeeee ngs mmrr. press/vv62/ wwt.i///........ html.
Alex Nichol, Joshua Achiam, and John Schulman. On First-Order Meta-Learning Algorithms, 2018.
Vitchyr H Pong. Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline meta-reiiffocemment learning with online self-supervision. In International Conference on Machine Learning, pp. 1781 1-17899. PMLR, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training, 2018.
11

---

Published as a conference paper at ICLR 2023
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann. Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang. Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre. Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean-Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake Hechtman, Laura Weidinger, Iason Gabriel, William Isaac, Ed Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling Language Models: Methods, Analysis & Insights from Training Gopher, 2021.
Kate Rakelly, Aurick Zhou. Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-pollcy mett-reinforreement learning via probabilistic context variables. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 5331-5340. PMLR, 2019. URL http: http:::/ppoceeeingg wtr.ppres/vve v97 rakellyllle...ee 9a. html.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A Generalist Agent, 2022.
Andrei A. Rusu, Sergio Gomez Colmenarejo, Çaglar Gülgehre, Guillaume Desjardins, James Kirk- patrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukeuoglu, and Raia Hadsell. Policy distillation. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL httt:://// //arxiv...///////aa/ org/ abs / 11111.11111 06295.
Jürgen Schmidhuber. Reinforcement learning upside down: Don't predict rewards just map them to actions. CoRR, abs/1912.00275, 2019. URL http: / / a rx wxw.ioooooo wwwpxivoorggaaaoo httgg///////// / 912.. 02875.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model. CoRR, abs/1911 1.19911.000 1.08665, 2019.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaskowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv: 191200877, 2019.
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 28 2818-2826, 2016. doi: 10. 10.1109/CVPR.22166.33
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Jane X Wang. Zeb Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn, 2016.
Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos,
Ronald Williams. Simple Statistical Fralient-Flllowing Algorithms for Connectionist Reinforce- ment Learning. Mach. Learn., 8:229-256, 1992. doi: 10. 100//F9999999.
Ronald J Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
12

---

Published as a conference paper at ICLR 2023
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. International Conference on Machine Learning, pp. 24631-24645. PMLR, 2022.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang. Burcu Karagol Ayan, Ben Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge, and Yonghui Wu. Scaling Autoregressive Models for Content-Riee Text-to-lmage Generation, 2022.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. arXiv preprint arXiv:2202.05607 2022.
13

---

Published as a conference paper at ICLR 2023
A DISCUSSION AND LIMITATIONS
Discussion: While AD is more data-efficient than the source algorithm, we found that the source algorithm achieves a slightly higher asymptotic score in harder environments. This presents a trade- off between data-efficiency and asymptotic optimality when using AD that should be taken into consideration when applying AD to a specific problem. In terms of long-term consequences, AD presents path for converting narrow single-task RI agents into multi-task generalist ones. To date, deep RL research has mostly focused on powerful single-task agents Mnih et al. (2015); Schrittwieser al. (2019); Hessel al. (2021). These algorithms have produced powerful but fata-ineefficiet agents, which has limited their applicability beyond narrow domains. AD offers path toward training substantially more data-efficient, though perhaps less optimal, generalist agents by distilling narrow RL algorithms into sequence models like transformers.
Limitations: While AD is a general method, in this work we'v shown that AD generalizes to new tasks within the same domain. Showing the ability to distill cross-domain RL algorithms and generalize to new domains would be an interesting line if inquiry for future work. AD also requires storing many learning histories which could take up significant memory, though this may not be too much of an issue since the learning histories can be stored on disk rather than RAM. Perhaps the main limitation of AD is that most RL environments of interest have long episodes and modeling multi- episodic context requires more powerful long-horizon sequential models than the ones considered in this work. We believe this is promising direction for future research.
B ENVIRONMENT CONSIDERATIONS
In this work, consider environments where zero-shot generalization is difficult, the agent must learn through trial and error. We also want environments where overfitting to any particular task is difficult to ensure our method is general. A final practical consideration is that consider environments wher across-episodic histories can be feasibly modeled with a causal transformer. Given these considerations, our evaluation environments need to satisfy three criteria:
1. Supports many tasks: The environments must be multi-task to ensure that our agent and the baselines do not overfit to any single task and instead is able to in-context reinforcement learn across many tasks within a given domain.
2. Task must be hard infer: To ensure that the downstream tasks are hard generalize to in zero-shot, we use environments that require exploration. Namely, we require environments where either the task can only be inferred from the reward and not the observation, or tasks that are partially observable.
3. Supports malti-episodic contexts: Lastly, we impose practical constraint the environment episodes must be short enough such that a normal GPT-like transformer can fit multiple episodes in its context. Since this work introduces AD as method, we wish investigate in the cleanest possible setting using canonical architecture. We leave investigating AD with more complex architectures that scale to longer sequences for future work.
Prior related works (Chen et al., 2021; Lee et al., 2022; Janner et al., 2021; Reed et al., 2022) evaluated on Atari, OpenAl gym, and as well as other environments. However, Atari and OpenAl gym don't satisfy at least one of the above criteria. Atari and OpenAl gym episodes are often long and can contain thousands or more transitions per episode, so it's technically challenging to populate causal transformer's context with across-episode histories. Indeed, the prior related works only considered within eeisode context lengths. Additionally. it is often easy to infer the task from either the observation or the dense reward alone in both Atari and OpenAl gym, which reduces the need for exploration. For these reasons, we evaluate in environments that satisfy all three criteria instead.
C CLOSELY RELATED PRIOR METHODS
In our main results we use Expert Distillation (ED) as a baseline. Here, we discuss how the most closely related methods differ from AD and why ED is sufficient to support the paper's claims.
14

---

Published as a conference paper at ICLR 2023
Expert Distillation (ED): ED is most similar to Gato (Reed et al., 2022), which models expert sequences from a converged RL policy using a causal transformer. ED also trains a causal transformer to predict actions using expert policy data. There are two key differences between ED and Gato. First, unlike Gato which utilizes small (relative to an episode length) within-episode contexts, ED is trained on the same across-episode contexts as AD, so the architectures used by ED and AD are the same. contexts alone but also learning
The benefits of AD cannot therefore be attributed to across progress in the offline data used train AD. Second, ED models state-action-rewanl sequences while Gato models only state state-action sequences. The main difference between ED and AD is that AD is trained on full multi-task learning histories rather than expert policy data.
Decision Transformer (DT) (Chen et al., 2021) and Multi-Gama Decision Transformer (Lee et al., 2022): DTs learn return-conditioned policies from single-task offline data collected by an RI. agent. While the training data itself (an RI agent's replay buffer) contains learning, the context sizes used DT are too small to capture any learning progress or identify the task using across-episode information. For instance, the Atari experiments use context of length 30 50 tokens, or 10 - 17 transitions. Atari games can have hundreds or thousands of transitions in a single episode, which means these contexts capture mostly within-episode information. Additionally, very little learning progress happens in the underlying replay buffer data within that many transitions.
Another difference between DT and AD ED is that DT learns retuun-conditioned model whereas AD / ED are both condiiionodd In our setting return-conditioning alone cannot yield an optimal policy since the agent does not know the task until after it explores the environment and can identify it using across-episode contexts. Since (i) DT uses small within-episodic contexts and (ii) return-conditioning would not help in the environments considered, this baseline similar to ED with a small within-episode context which is strictly weaker than the long across-episode context variant of ED we consider.
Trajectory Transformer (TTO) (Janner et al., 2021): Like AD, TTO also models selee-actioo-rewaaie tokens but in addition to predicting actions also learns world model by predicting states and rewards. To maximize return, TTO then uses beam search to select high-reward actions. However, in our setting, TTO will run into the same problem as DT. To model rewards accurately it will need longer across-episodic contexts since one environment supports many tasks. Similar to DT, MGDT, and Gato, TTO uses smaller within-episode contexts. For this reason, TTO will fare no better than DT, MGDT, ED in the settings consider. We also note that in contrast to TTO, AD is model-free. In AD, actions are sampled from the transformer cistor--conditioned predictions and return maximization emerges from modeling the learning histories of an RL algorithm.
To summarize, AD differs from prior methods mainly because its context is across-episodic and hence large enough capture learning progress and task information. AD could further be augmented by learning world models like TTO or conditioning on returns like DT, but these investigations would be well suited for future work since they are tangential the main research question addressed in this work whether in-context RL can emerge by imitating the learning histories of an RL algorithm with long across-episodic contexts.
AD is also closely related to prior work in in-context meta-RL, While both AD and in-context meta-RL model across-episodie histories with memory- architectures, prior in-context me algorithms, such as RL2 (Duan et al., 2016) are trained online and rely on learning multi-episodic value functions with TD learning while AD trained offline and uses a supervised imitation learning objective.
15

---

Published as a conference paper at ICLR 2023
D EXPERT DISTILLATION MAIN RESULTS
We elaborate further on the main results in Fig. 4 and provide intuition regarding the behaviors of the ED baseline. In Dark Room, Dark Room (Hard), and Watermaze, ED performance is either flat or it degrades. The reason for this is that ED only saw expert trajectories during training, but during evaluation needs to first explore (i.e. perform non-expert behavior) identify the task. This required behavior is out-of out-of-distriburi distribbutiio for ED and for this reason it does not reinforcement learning in-context. In Dark Key-to-Door the agent reset randomly at the beginning of each episode, whereas in all of the environments the agent's s starting position is fixed. Due to random resets, the ED agent sometimes reset by the first goal in Dark Key-to-Door which allows it to occasionally identify the first goal of the task, which is why it shows slight improvement.
E MODEL SIZE
We investigate how transformer capacity affects performance in Fig. 8. While in-context RL emerges across all model sizes investigated, we find that increasing the model depth, the model width in terms of embedding dimension, and (to lesser extent) the number of attention heads improves performance on Dark Key-to-Doos
3.0
8.5
Figure 8: Model size investigations: We investigate how increasing model capacity affects AD. While in-context RL with AD emerges regardless the model capacity, increasing the model depth and width helps improve AD until achieves optrmal performance.
16

---

Published as a conference paper at ICLR 2023
F SOURCE ALGORITHM TRAINING RUNS
Dark Room (Hard)
Key-lo-Dorr
Watermaze
0.75
0.75
0.50
0:5
0.5
0.00
0.23
0.50
0.75
1.03
0.00
0.25
0.50
0.75
1.00
0.25
0.53
0.75
1.50
Env Staps
1w5
Ere Steps
Em Steps
165
Figure 9: Asymptotic performance of the A3C (Maih al., 2016) and Q-Å variant of the DQN (Mnih al., 2013) algorithms used to produce learning histories Dark Watermaze environments. These curves show the learning histories AD is trained on. The source algorithms plotted Fig. 4 are the same as in these plots.
G LABEL SMOOTHING ABLATION
For the harder exploration task of Dark Room (Hard), we found that adding label smoothing regular- ization (Szegedy et al., 2016; Müller et al., 2019) improved the in-context learning ability of AD In Figure 10 we ablate the benefit of using label smoothing for different a values well as with it turned off. Each curve in the figure denotes average performance over 5 training seeds. We can see that adding label smoothing up point improves the in-context learning ability of Algorithm Distillation, with performance continually increasing with the number of evaluation episodes.
Label Smoothing Ablation Dark Room (Hard)
0.6
alpha=0.0
Return
0.4
alpha=0.1
alpha=0.2
alpha=0.3
0.2
0.0
200
400
600
800
1000
Episodes
Figure 10: AD trained with different amounts label smoothing on Dark Room (Hard).
H RL NETWORK ARCHITECTURE: TRANSFORMER VS LSTM
We compared using a transformer as the architecture for RL instead of an LSTM. Figure 11, we both transformer and LSTM RL. agents over the Dark Room environment. The curves shown are the best from a sweep over learning rate and unroll length byperparameters. The transformer architecture is 4-yeesss with model size of 256 and pre-norm layer normalization placement. While both agents reached similar level of final performance, all RL transformer models trained tended be more unstable with the average return not as consistent as with LSTM architecture. Given the poor performance of the transformer-bast RL 2 on the simpler Dark Room setting, our other experimental settings used the LSTM.
17

---

Published as a conference paper at ICLR 2023
LSTM v.s. Transformer
Final Episodic Reward
LSTM
0.8
Transformer
0.6
0.4
0.2
250
500
750
1000
Training Step (1e6)
Figure 11: Comparison of LSTM and Transformer architecture for RL agent on Dark Room. Each curve averaged over 5 training seeds with the shaded area representing the standard error.
NUMBER OF TRAINING TASKS IN SOURCE DATA
# of tasks during training
2.0
tasks
tasks
18 tasks
1.5
37 tasks
Return
75 tasks
151 tasks
303 tasks
606 tasks
0.5
1212 tasks
2424 tasks
Env Steps (Thousands)
Algorithm distillation trained different numbers of training tasks Dark Key-to-Door evaluated on a fixed set of test tasks for 300 episodes of evaluation.
Figure 12:
One interesting question is how many tasks AD needs to be trained on to learn algorithm that generalizes to held out tasks. We trained AD on different numbers of Dark Key-to-Door training tasks and evaluated the resulting models the same set of test tasks. Figure 12 shows the in- context learning plots for the resulting AD models on the set test tasks. As a reminder, there are 812 6561 unique Dark Key-to- ooor tasks. Models trained 1, any in-context learning on test tasks. While models trained on 37, 75 and 151 tasks did not achieve good performance overall, they did exhibit some in-context learning over the course of 300 episodes. The best models were trained on 1212 and 2424 tasks which corresponds to roughly 18% and 37% of the total number of tasks in the Dark Key-to-Door domain.
18 training tasks did not show
SINGLE-STREAM ALGORITHM DISTILLATION
We provide more details around the experimental setup for the single-stream result shown in Fig. 6. We showed in Fig. 4 that when AD is trained on data from a subset of the actors of a distributed source RL algorithm, the resulting model is more data efficient than the source algorithm. Here we confirm that AD can produce a faster algorithm than the one it was trained on in the single-stream setting. For this experiment we trained A3C on 2048 Dark Key-to-Door tasks for 2000 episodes each. We then trained AD on the resulting data while subsampling the learning histories by factor of 10. More concretely, we took every 10th episode from each of the learning history, which resulted in 200 episode compressed learning trajectory for each task. Figure 6 compares the resulting AD model evaluated on a set of test tasks to the performance of the source algorithm on these tasks. The model learned by AD learss much faster than the source algorithm confirming that Algorithm Distillation can turn a slow gradient- gradient-based algorithm into a much more data efficient in-context learning algorithm.
18

---

Published as a conference paper at ICLR 2023
K
RANDOM MASK
Dark Key-to-Door
0.1
1.5
Return
1.0
0.5
10
15
Env Steps (Thousands)
Figure 13: Downstream performance of Algorithm Distillation with different values of random masking during training in 9x9 1 goal gridworld.
During training, input tokens were randomly masked to avoid overfitting training data. This plot shows the downstreams results on 9x9 Dark Key-to-Door domain with different values of this random masking. Values of 0.3 0.5 perform the best with the value of 0.3 chosen for all experiments.
L AD NETWORK ARCHITECTURE: TRANSFORMER VS LSTM
Here we consider the importance of the Transformer architecture to the success of algorithm dis- tillation (AD) comparing to AD based off of LSTM (Hochreiter & Schmidhuber, 1997). Specifically, the LSTM receives the concatenated embeddings of (o;, 0g. r:) triplets to the most recent time step 1. The output of the LSTM is then concatenated with the current observation or embedding and both are then fed through multi-layer perceptron (MLP) policy torso to produce distribution over the present action az The LSTM hidden size (512), MLP depth (2), and MLP width (256) were swept and tuned by grid search based on downstream reward attainment.
Comparing Transformer AD and LSTM AD on the Dark Key-to-Dov Door task (Figure 14), we find that both agents are capable of in-context learning, demonstrating that the success of AD is not tied to the underlying network architecture. However, we also find that the Transformer variant consistently outperforms the L.TM variant, which is why all other experiments in this paper employ the Transformer variant. This finding is consistent with the recent wider success Transformer based architectures recurrent neural network (RNN)-based architectures in sequence prediction tasks.
M DISTILLING ACTIONS VS PROBABILITIES
In the rest of the paper, we use (one-hot) actions taken by the source policy as the prediction target for AD. Here, we compare that choice to predicting the source policy probabilities from which that action was sampled. In other words, is it better to distill actions or probabilities?
Figure 15 compares all combinations of: 1) distilling actions vs probabilities, and conditioning those predictions on past actions, probabilities, or both. The plot labels represent various input/output combinations, e.g. sa sar->a indicates observing states, actions, and rewards while predicting actions (i.e. the main variant of AD in the rest of the paper), sar->p represents using the same observations but instead predicting probabilities, and sapr->a represents observing states, actions, probabilities, and rewards while predicting actions, etc.
Two conclusions stand out from this plot. First, the original action distillation variant of AD (sar->a) performs best, followed closely by probability distillation (sar->p). In the paragraphs and experiments below, we explore why action distillation outperforms probability distillation. Second, all variants that condition past probabilities catastrophically fail. We speculate
19

---

Published as a conference paper at ICLR 2023
Dark Key-to-Door
2,00
1.75
1.50
Return
1.25
1.00
0.50
Transformer
0.25
LSTM
0.00
50
100
150
200
250
300
Episodes
Figure 14: Comparison between algorithm distillation with a Transformer and LSTM architecture on Dark Key-to-Daw Mean + 1 standard deviation over 5 training seeds and 20 evaluation seeds. 300 episodes corresponds to 15k environment steps.
One-Hot Actions vs Probabilities
2.00
1.75
1.50
sar->-a
Return
1.25
sar-2-p
spr->a
1.00
sapr->p
0.75
sapr->a
0.50
spr->p
0.25
0.00.0
0.2
0.4
0.6
0.8
1.0
Episodes
163
Figure 15: Comparison between algorithm distillation with various input/output combinations on Dark Key-to-Door, Mean + standard deviation over 5 training seeds and 20 evaluation seeds. 1000 episodes corresponds to 50k environment steps.
two somewhat contradictory reasons why this might oceur: 1) The AD prediction task involves a combination of inferring the current policy, well predicting when policy updates will occur. Since observing policy probabilities provides more information about the current policy than actions do, it is possible that this leads AD fully focusing inferring the current policy and ignoring the prediction of policy updates. In other words, observing only actions taken may act as a useful information bottleneck. 2) On the other hand, observing probabilities may inadvertently leak information about when policy updates occur. If the transformer context includes multiple visits to the same state, AD could learn compare the policy probabilities to infer whether policy update occurred in between the two visits. If AD learned during training to rely on this information, then during autoregressive evaluation, it may be "waiting" for policy update than never occurs. If this is the souree of the issue, then potential solution would be retraining with explicit policy update tokens and including them between evaluation episodes, however we leave this investigation future work.
Now we return to the question of why action distillation (sar->a) outperforms probability distillation (sar->p). Intuitively, distilling one-hot actions will result in more deterministic policy than distilling probabilities when training on finite data. Indeed, the red curves in the top two plots of Figure 16 show that probability distillation (top left) converges to a higher entropy (~1122 bits) policy than does action distillation (top right, ~1 bit). Noting this discrepancy, we speculated that encouraging probability distillation towards a more deterministic policy might lead to increased performance. To do so, experimented with entropy penalty regularizer added to the AD loss
20

---

Published as a conference paper at ICLR 2023
from equation 6, leading to the modified objective:
E(0,a)== Σ log Po a!") (1 of")) all (o(AA
where H[] is the Shannon entropy in bits, and a regularization weight. The top two panels decreasing policy entropy for action distillation
of Figure 16 show the effect of increasing (top right) and probability distillation (top left), throughout training. The bottom two panels show the corresponding changes in evaluation returns. Notably, entropy penalization indeed improves performance for probability distillation (bottom left), with the best performing regularization value (a 0.1) achieving similar return to the original unregularized action distillation (bottom right, red). Interestingly, this is the amount regularization that leads probability distillation to have the most similar entropy to the unregularized action distillation variant (~1 bit) as well. Further entropy penalization for action distillation, on the other hand, does not lead to increased performance (bottom right). Together, these results suggest that action distillation naturally leads the "optimal" amount of entropy regularization on its own, at least in the environments we study.
sar->p p policy entropy
sar->a policy entropy
2.0
entropy penalty
Entropy (bits)
1.6
Entropy (bits)
0.05
1.4
0.2
0.3
1.2
1.0
0.8
0.8
0.6
0.25
0.50
0.75
1.00 Episodes
1.25
1.50
1.75
0.6
0.25
0.50
0.75
1.00
1.25
1.50 1.75
Episodes
165
sar->p reward
sar->a reward
Return
1.6
Return
1.4
1.2
1.0.0
18.00
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Episodes
163
Episodes
163
Figure 16: Policy entropy during training and reward during evaluation for the action distillation (sar->a) and probability distillation (sar->p) variants of on Dark Key-to-Door. Colors indicate the strength of the weight on the entropy penalty regularizer from equation 7. Mean + standard deviation over 5 training seeds and 20 evaluation seeds. 1000 episodes corresponds 50k environment steps.
In addition to closing gap between probability and action distillation with an entropy penalty, we also explored retuning various hyperparameters for sa sar->p rather than reusing those tuned for sar->a. We found that many hyperparameters, such as the dropout rate, attention dropout rate, and sequence mask probability, had similar optimal values in the two cases, and so retuning them did not help. However, one hyperparameter that did have effect was the size of the transformer context window (Figure 17). While action distillation performance increased only up to a context window size of 200 steps (4 episodes) and then plateaued (Figure 17, right), probability distillation performance continued to increase up to a context window size of 300 steps (6 episodes) before plateauing (Figure 17, left). Thus, while comparison our default context window size of 200
21

---

Published as a conference paper at ICLR 2023
steps favored action distillation, for larger context window sizes, action and probability distillation performed similarly. We leave any further explanation of why probability distillation requires larger context windows than action distillation to future work.
sar->p
sar->a
2.00
2.00
1.75
1.75
1.50
Return
1.25
Return
1.25
1.00
1.00
context window size
0.75
0.75
100 steps (2
0.50
0.50
200 steps (4
300 steps (6
0.25
0.25
400 (8
0.08.0
0.08.0
Episodes
Episodes
Figure 17: Context window size dependence for the action distillation (sar->a) and probability distillation (sar->p) variants of AD on Dark Key-to-Door. Mean + standard deviation over training seeds and 20 evaluation seeds. 1000 episodes corresponds to 50k environment steps.
22

---

Published as a conference paper at ICLR 2023
N ALGORITHM DISTILLATION HYPERPARAMETERS
Hyperparameter
Dark Room
Dark Room (Hard)
Dark Key-to-Door
Watermaze
Embedding Dim.
64
Number of Layers
4
Number Hcads
Feedforward Dim.
2048
Position Encodings
Absolute Post Norm
Layer Norm Placement
Dropout Rate
0.1
Context Window
600 tokens (200 timesteps)
Attention Dropout Rate
0.5
0
0.5
0.5
Sequence Mask Prob
0.3
0.5
0.3
0.3
Label Smoothing
0
0.2
0
0
Table 1: Algorithm Distillation Architecture Hyperparameters.
Hyperparameter
Value
Batch Size
128
Optimizer
Adam
32
0.99
Gradient Clip Norm Threshold
Learning Rate Schedule
Cosine Decay 2c-6
Initial Value
Peak Value
3c-4
Table 2: Algorithm Distillation Optimization Hyperparameters.
Layer
Hyperparameter
Value
Conv Block
Channel
128
Conv
Kernel
Stride
Decay Rate
0.999
BatchNorm
cps
le-5
Activation
ReLU
Kernel Stride
Max Pooling
Dropout
Rate
Network
Conv Blocks
Final Linear Layer
Units
256
Table 3: Watermaze Image Encoder Hyperparameters.
23

---

Published as a conference paper at ICLR 2023
O SOURCE RL ALGORITHM HYPERPARAMETERS
0.1 DARK ENVIRONMENTS
Hyperparameter
Value
Batch Size (Num. Actors)
100
À
0.95
Agent Discount
0.99
Entropy Bonus Weight
0.01
MLP Layers
MLP Hidden Dim
128
Optimizer
Adam
B1
0.9
82
0.999
c
le-6
Learning Rate
le-4
Table 4: Source A3C Algorithm Hyperparameters for Dark Environments.
O.2 DMLAB WATERMAZE
Hyperparameter
Value
Batch Size
8
Rollout Length
40
Rollout Overlap
31
Number of Actors
Reply Buffer Capacity
le5
Offline Data Fraction
0,7
À
0.75
0.01
Agent Discount
Target Update Period
ResNet Channels
[32, 64, 64]
ResNet Kemness
[3. 3]
ResNet Strides
[1, 1,
Pool Kernels
[3, 3, 3]
Pool Strides
[2, 2, 2]
Optimizer
Adam
81
0,9
f22
0.999
le-6
Gradient Clip Norm Threshold
Learning Rate
lc-4
Table 5: Source DQN(Q-A) Algorithm Hyperparameters for Watermaze.
24

---

Published as a conference paper at ICLR 2023
P RL HYPERPARAMETERS
Hyperparameter
Value
RL Algorithm
A3C
Learning Rate
3e-4
Batch Size
256
Unroll Length
20
LSTM Hidden Dim.
256
LSTM Number of Layers
Episodes Per Trial
10
Table 6: RL Hyperparameters used in "Dark" Environments.
Hyperparameter
Value
RL Algorithm
DQN(Q-A)
Learning Rate
le-4
Batch Size
96
Unroll Length
40
LSTM Hidden Dim.
256
LSTM Number of Layers 1
Episodes Per Trial
30
Table RL Hyperparameters used in the Watermaze Environment.
25

---

Published as a conference paper at ICLR 2023
Q
ATTENTION MAPS
1700
100
1800
1900
200
1700
1800
1900
1700
100
1800
200 0
100
200
1900
1700
1800
1900
Figure 18: Attention maps for AD from five separate seeds. White and gray colors correspond to low and high attention. Red and blue colors indicate that those transitions correspond to an episode restart and positive reward token, respectively. The left column plots attention for an AD transformer after 200 time-steps of evaluation (when the context is initially filled). The right column plots attention after 1900 steps (38 episodes) of evaluation. Each episode has length of 50 steps. From these patterns, it is evident that AD attends to tokens across several episodes to predict its next action.
26


Reviews:
[{"review_id": "jJ_WT6w3U0", "reviewer": "Reviewer_E88f", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "", "summary_of_the_paper": "The paper proposes Algorithm Distillation, a transformer-based method for imitating the learning process implied by a reinforcement learning algorithm. The learning process is emulated by simply imitating the sequence of actions produced by an algorithm given a history of observations with the environment.", "summary_of_the_review": "In short, I believe this is a well-executed investigation of a really intriguing idea. I recommend acceptance, but I encourage the authors to elaborate on the possible long-term consequences or more general implications of the development of such a class of methods.", "strength_and_weaknesses": "Strengths:\n- The approach presented in this paper is at the same time very simple and intellectually stimulating;\n- The empirical investigation is detailed and answers interesting questions one might have about this type of approaches;\n- The presentation of the ideas is clear and concise.\n\nWeaknesses:\n- While I believe the careful investigation of such an idea is definitely worth even without an immediate implication on how to leverage it in a more standard context (i.e., when should I both to imitate my RL algorithm instead of simply running it?). It would be better to have a direct discussion on which ones would be the implication of improving this type of methods in the long run;\n- Is the context length the only blocker in scaling this method to more complex tasks? From the current paper, it is not clear whether more complex transformer models (with longer context sizes) could actually be enough to scale to the usual reinforcement learning benchmarks.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "V6pFYI3Sr_", "reviewer": "Reviewer_GFbm", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "This paper proposes an offline meta-RL method that uses a causal transformer as a backbone model to predict actions . In particular, this method learns policies conditioned on a long context using a casual transformer where the context is built based on trajectories from multiple offline tasks. To show the effectiveness of their method, this paper runs various experiments on Adversarial Bandit, Dark Room, and DMLab Watermaze benchmarks. ", "summary_of_the_review": "This is an interesting paper and the results are promising. However, I'd make my final recommendation after rebuttal as I require some clarifications from others. In general, the proposed method is not original and based on previous meta-RL works. Using the decision transformer with context is the main contribution of this paper. ", "strength_and_weaknesses": "While I find the result of this work interesting as it shows that it is possible to do more efficient offline meta-RL when the problem is formulated in the right way (e.g. using long context), I have a couple of concerns about this work:\n\n- Meta-RL:\nUsing context/history in meta-RL has been studied in previous works [1,2,3] and all those works show that encoding past histories leads to significant improvement in performance (see results in [2]). Although the setting of those papers are online-RL, the similarity is uncanny and this work should have discussed those extensively in the paper (I'd suggest adding a section to discuss this). The only difference that I see with this paper and those works is use of transformer vs recurrent models to encode the context aside from online vs offline setting. Also, is it fair to say that the main contribution of this paper is to adopt the use of context/history in the decision transformer? \n\n- Choice of baseline:\nThis paper uses RL^2 as a baseline. However, previous works [1,2] showed that RL^2 has poor performance in comparison to other meta-RL methods (see figure 2  of [1] ). Hence, choosing a better baseline method could provide a better conclusion about this paper.\n  \n- Improvement over behavioral policy:\nAuthors claim in multiple parts of the paper that \"AD learns to output an improved policy relative to the one seen in its context.\" I am confused about this as I don't see the results support this claim. Take figure 4 as an example, AD doesn't outperform RL^2. Am I missing something here? \n\n- Ablation studies:\n\"Can AD learn a more data-efficient RL algorithm than the one that produced the source data?\" experiment shows that AD outperforms source algorithm. The setting of this experiment can be inconclusive as the source method and AD have different model sizes (?) and it appears to me AD uses larger network sizes. If that is the case, it is not a fair comparison. Moreover, \"Training the Sequence Model\" in page 5 can be a good experiment as long as it is a fair comparison in terms of model sizes. Finally, it would be useful if authors add a table in the appendix illustrating AD and baselines network/model sizes for each experiment in the paper to provide a better picture about the setting of the experiments. \n\n- Context/History:\nThe paper says dataset D contains a set of \"learning\" histories. However lines 3 and 4 of Algorithm 1 show otherwise. In particular, Algorithm 1 shows that D contains ONLY expert data as optimal policies are used to collect data. This weakens the entire discussion of policy improvement in the paper as data contains only expert data.\n\n\n[1] Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables, ICML 2019 http://proceedings.mlr.press/v97/rakelly19a/rakelly19a.pdf\n\n[2] Meta-Q-Learning, ICLR 2020 https://openreview.net/pdf?id=SJeD3CEFPH\n\n[3] Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs, ICML 2022 https://proceedings.mlr.press/v162/ni22a/ni22a.pdf\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "ScFolvEIo2", "reviewer": "Reviewer_8EkG", "summary": "", "questions": "", "limitations": "", "rating": 10, "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": "", "summary_of_the_paper": "The authors propose AD.  AD works by training a transformer (or other causal sequence model) on RL histories of agents that are learning (exhibiting improving performance during the history).  Then after this training, when allowed to run on the environment, the transformer continues to improve performance, resulting in improving learning curves without any parameter updates.\n", "summary_of_the_review": "This paper explores a topic that I believe is extremely important; the result (the fact that this approach worked) was surprising to me.  It is well-written, I have no significant concerns, and I have very little to say other than to recommend an accept.  I would consider recommending a “strong accept” if other reviewers felt similarly.\n", "strength_and_weaknesses": "Edits:\n- intro: “does not improve through trial an error”, “and” missing a “d”\n\nStrengths:\n- Well written; this paper was a pleasure to read.\n- To the best of my knowledge, this work is novel.\n- The fact that this approach worked was quite surprising to me.  This “surprise factor” makes this an important and interesting topic in my opinion.\n- I also think that this kind of “in-context” learning (in other words, “learning” only within what-can-be-viewed-as-analogous-to the model’s short-term memory, without updating parameters) is an important thing to study. I hope this will inspire future work that better studies the intersection of this kind of learning with more traditional learning.\n\nWeaknesses:\n- The paper would greatly benefit by a more clear early definition of “in-context”, instead of making the reader infer it through context.  (There is sort of a definition in the abstract, but even that is more of an implied definition than a definition.  I suggest adding a more clear definition to the early introduction.)\n- One might argue that the contribution is slightly limited: the authors simply propose a new idea and show that it works.  However, given the novelty and importance of this topic, I am not very concerned about the extent of the contribution.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "IhP1hpH05Z", "reviewer": "Reviewer_HwSZ", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "This paper proposes to train a transformer to imitate offline data coming from a RL algorithm. When trained on a distribution of tasks, it can then generalise the behaviour of the RL algorithm into the new task. This is in comparison to recent papers using transformers to imitate already-trained RL policies, instead of algorithms. In this regard, it is closer to meta-RL algorithms like learning to reinforcement learn or RL^2, with the difference of being offline rather than online. The paper then evaluates on a wide variety of environments and analyses many different factors involved in the proposed approach.", "summary_of_the_review": "The community will benefit from this paper as it is a very thorough exploration of a relevant idea, particularly in the current LLM-for-everything climate. The writing is also clear and it reads very well.", "strength_and_weaknesses": "- Strengths\n    - I think the evaluation of this paper is very very good. The diversity of environments is great, and most of all, I was impressed by the many different questions the author asked themselves about the proposed approach, from out-of-distribution, to training on partial demonstrations or the importance of the length of the context. Great job!\n    - The paper reads very clearly. In part this is because it sets out to do one job and do it well, and I appreciate that.\n    - The idea is insightful and potentially useful in the long run for meta-RL.\n- Weaknesses\n    - I think the originality of the proposed approach is good enough to pass the bar for acceptance, but I wouldn't consider it great. It is relatively close to GATO and RL2. It is also close to multiple works that are not cited but very relevant: Bootstrapped Meta-learning [Flennerhag et al., Outstanding Paper award last ICLR] as well as Upside-Down RL [Schmidhuber] and, less closely, the learning-to-optimize community [Li&Malik '16, Andrychowicz et al. '16].\n    - In multiple places, the comparison to the standard (non-learned) RL algorithm needs to be described more in favor of the latter rather than the proposed approach. For instance, it is argued that the \"algorithm\" coming from AD generalizes out-of-distribution, but this wouldn't apply to widely different tasks: an RL algorithm applies to image-input problems and torque-based inputs just as fine, but the algorithm from AD wouldn't have nearly that level of generality (and which other methods in meta-RL have started to achieve). Furthermore, the sentence \"the source algorithm produces many single-task agents [...] while AD produces a single generalist agent\" while technically true is quite incomplete, as AD needs to store each individual history to be able to represent each specialist agent.\n    - A point that concerns me from the evaluation was that in the main figure, RL2 is shown only in its asymptotic value, as if it was an Oracle whose performance couldn't be put as just another curve (which may surpass AD). One may argue that it could be unfair for AD, as it is online rather than offline, but the same can be said about the Source being on a non-meta setting and not exploiting meta-training data, being at a disadvantage w.r.t. AD. I would really appreciate a clarification on this.", "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:

Decision: Accept: notable-top-5%