In-context Exploration-Exploitation for Reinforcement Learning
Authors: Zhenwen Dai
Federico Tomasi
Sina Ghiassian
Venue: ICLR 2024 poster (2024)
Keywords: Offline reinforcement learning
in-context learning
decision transformer
Primary area: reinforcement learning

Abstract: In-context learning is a promising approach for online policy learning of offline reinforcement learning (RL) methods, which can be achieved at inference time without gradient optimization. However, this method is hindered by significant computational costs resulting from the gathering of large training trajectory sets and the need to train large Transformer models. We address this challenge by introducing an In-context Exploration-Exploitation (ICEE) algorithm, designed to optimize the efficiency of in-context policy learning. Unlike existing models, ICEE performs an exploration-exploitation trade-off at inference time within a Transformer model, without the need for explicit Bayesian inference. Consequently, ICEE can solve Bayesian optimization problems as efficiently as Gaussian process biased methods do, but in significantly less time. Through experiments in grid world environments, we demonstrate that ICEE can learn to solve new RL tasks using only tens of episodes, marking a substantial improvement over the hundreds of episodes needed by the previous in-context learning method.

Published as a conference paper at ICLR 2024
IN-CONTEXT EXPLORATION-EXPPLOTTTTNN FOR RE- INFORTEMENT LEARNING
Zhenwen Dai, Federico Tomasi, Sina Ghiassian
Spotify Research
{zhenwend, federicot, sinag] it..
ABSTRACT
In-context learning is a promising approach for online policy learning of offline reinforcement learning (RL) methods, which can be achieved at inference time without gradient optimization. However, this method hindered by significant computational costs resulting from the gathering of large training trajectory sets and the need to train large Transformer models. We address this challenge by introducing In-context Auplooaation-Exloitatiiii (ICEE) algorithm, designed to optimize the efficiency of in-context policy learning. Unlike existing mod- ceplooaatio-explottttiii trade-off at inference time within
els, ICEE performs a Transformer model, without the need for explicit Bayesian inference. Conse- quently, ICEE can solve Bayesian optimization problems as efficiently Gaus- sian process biased methods do, but in significantly less time. Through experi- ments in grid world environments, we demonstrate that ICEE can learn to solve new RL tasks using only tens of episodes, marking a substantial improvement over the hundreds of episodes needed by the previous in-context learning method.
INTRODUCTION
Transformers represent highly effective approach sequence modelling, with applications extend- ing across multiple domains like text, image, and audio. In the field of reinforcement learning (RL), Chen al.] (2021) and Janner et al. (2021) have suggested the concept of addressing offline RL as a sequential prediction problem using Transformer. This method has proven successful in tackling a range of tasks using solely offline training when combined with large-scale sequence modelling techniques (Lee et al. 2022 Reed et al. 2022a).. A notable shortcoming lies with the policy's inabil- ity to self-improve when employed in online environments. To overcome this, fine-tuning methods have been introduced, like (Zheng et al. 2022), which enable continued policy improvement. How- ever, these methods often rely on slow, computationally expensive gradient- based optimization.
Conversely, in-context learning, characteristic observed in large language models (LLMs), can handle new tasks by providing the task details through language prompts, effectively eliminating the need for fine-tuning. [Lskin et al. (2023) suggest in-context learning algorithm for RI,, which uses a sequence model to distill a policy learning algorithm from RL training trajectories. The resulting model is capable of conducting policy learning inference-time through an iterative process of action sampling and prompt augmentation. This method incurs significant computational costs in collecting extensive training trajectory sets and training large Transformer models that are needed to model substantial part of a training trajectory. The key reason for this high computational cost the lengthy RL training trajectories resulting from the slow trial-and-erroo process of RL policy learning algorithms.
This paper aims to improve the efficiency of in-context policy learning by eliminating the need to learn from policy learning trajectories. In an ideal scenario, efficient policy learning could be achieved through an efficient trial-and-error process. For simplified RL problems such as multi- armed bandits (MAB), efficient trial-and-error process such Thompson sampling and upper con- fidence bounds are proven to exist. This process, often referred to as the the epistemic uncertainty derived from Bayesian belief. However, it
(EE) trade-off, relies heavily intractable to infer the exact epistemic uncertainty of sequential RL problems using conventional Bayesian methods. In light of recent studies uncertainty estimation of L.LMs (Yin et al.) 2023)

---

Published as a conference paper at ICLR 2024
examine predictive distributions of sequence models, demonstrating that, by training with purely supervised learning on offline data, a sequence model can capture epistemic uncertainty in sequence prediction. This suggests the potential for implementing EE in offline RL.
Based on this observation, we develop an in-context policy learning. ICHE takes inputs a sequence of multiple episodes of the same task and predicts the corresponding action at every step conditioned on some hindsight information. This offline RL. design resembles the Decision Transformer (DT. Chen et al., 202T). but ICEE tackles in-context policy learning by modeling multiple episodes of a task while DT only models a single episode. Furthermore, these episodes do not need to originate from a training trajectory, thereby averting the high computational costs associated with generating and consuming learning trajectories. The action distribution learned in DT biased towards the data collection policy, which may not be ideal when it is sub-optimal. To address this bias, we introduce an unbiased objective and develop a specific form of hindsight information for efficient EE across episodes.
ceporaation-pploiaattinn (ICEE) algorithm for
With experiments, we demonstrate that the EE behavior emerges in ICEE during inference thanks to epistemic uncertainty in action prediction. This is particularly evident when applying ICEE to Bayesian optimization (BO), as the performance of ICEE is on par with a Gaussian process based method on discrete BO tasks. We also illustrate that ICEE can successfully improve the policy for a new task with muls--ndd-rroo from scratch for sequential RL problems. To the best of our knowl- edge, ICHE is the first method that successfully incorporates in-context exploration-explotaatio into RL through offline sequential modelling.
2 RELATED WORK
Meta Learning. Interest in meta learning learning to learn algorithms has recently increased. While learner is an agent that learns to solve a task using observed data, a learning to learn algorithm includes meta-learner that continually improves the learning process of the learner (Schmidhuber
et al. 1996 Thrun& & Pratt 2012 Hospedales et al. [202]] Sutton 2022]. J. Lots of work has been done within the space of meta learning. For example. Finn et al. (2017) proposed general model-agnostic meta-learning algorithm that trains the model's initial parameters such that the model has maximal performance on new task after the model parameters have been updated through few gradient steps computed with a small amount of data from the new task. Other works on meta learning include improving optimizers (Andrychowicz et al., 2016 Li & Malik] 2016 Ravi & Larochelle 2016 Wichrowska et al., 2017). improving few-shot learning (Mishra et al. 2017 Duan et al., 2017
learning to explore (Stadie et al., 2018), and unsupervised learning (Hsu et al. 2018).
In the space of deep meta RL (Wang et al. 2016), a few works focused on a special form of meta- learning called meta- gradienss In meta- geettgggasietss the meta-learner is trained by gradients by mea- suring the effect of meta-parameters on learner that itself is also trained using a gradient algorithm (Xu et al.] 2018). In other work, Zheng et al. (2018) used meta-gradients to learn rewards. Gupta et al. (2018) 8) focused on automating the process of task design in RL, to free the expert from the bur- den of manual design of meta- weea-lee tasks. Similarly, Veeriah et al. (2019) presented a method for reinforcement learning agent to discover questions formulated as general value functions through the use of non-myopic meta-gradients. More recently, meta gradient reinforcement learning has seen substantial breakthroughs from performance gains on popular benchmarks hybrid offline-online algorithms for meta RL (Xu et al., |020 | Zahavy et al. 22220] Flennerbag et al. [202]] Mitchell et al.,
2021 Yin et al. 2023 Pong et al. (022). The role of uncertainty in meta RL has been studied
by Zintgraf et al. (2021). which result in an efficient online meta RL method. This work is then
extended by Doriman al. (202I) to the off-policy setting.
Offline Reinforcement Learning, In general. reinforcement learning was proposed as a fundamen- tally online paradigm (Sutton) 1988) [Sutton et al.. 1999) [Sutton& & Barto, 2018). This online learning nature comes with some limitations such as making it difficult to adopt it to many applications for which it is impossible to gather online data and learn at the same time, such as autonomous driving and is also sometimes not as data efficient as it could bem since it might choose to learn from sam- ple and then dispose the sample and move on to the next (Levine et al. 2020). One idea for getting more out of the collected experience is to replay buffers. When using buffers, part of the samples are kept in a memory and then are re-used multiple times so that the agent can learn more from them Lin 992 Mnih et al. 2015). A variant of reinforcement learning, referred as offline RI. studies
2

---

Published as a conference paper at ICLR 2024
RL algorithms that can learn fully offline, from fixed dataset of previously gathered data without gathering new data at the time of learning (Ernst et al.) 2005
2005; Lange et al. 2012 Fujimoto et al. 2019; Siegel et al. 2020 Gulcehre et al. 2020 Nair et al. 2020). Recent literature
Riedmiller
on decision transformers also focuses on offine RI. (Chen et al. 2021 because it needs to calculate
the return to go training time, which in turn necessitates previously gathered data.
In-context learning. In-context RL algorithms are the ones that improve their policy entirely in- context without updating the network parameters or without any fine-tuning of the model(Lu al.) 2021). There has been some work studying the phenomenon of in-context learning trying to explain how learning in-context might be possible Abernethy et al. (2023); Min et al. (2022). The "Gato" agent developed by Reed et al.] (2022b) works as model, multi-task, multi mulli-embodiment gen- eralist agent, meaning that the same trained agent can play Atari, caption images, chat, stack blocks with real robot arm only based on its context. By training an RL agent at large scale, [Team et al. (223)) showed that an in-context agent can adapt to new and open ended 3D environments. Of spe- cial interest to us is Algorithm Distillation (AD), which is a meta-RL method (Laskin et al. 2023 More specifically, AD is an in-context offline meta-RL method. Basically, AD is gradient- adapts to downstream tasks without updating its network parameters.
3 EPISTEMIC UNCERTAINTY IN SEQUENCE MODEL PREDICTION
DT, also known as Upside-down RL., treats the offline policy learning problem as a sequence mod- elling problem. In this section, consider a generic sequencing model and analyze its predictive uncertainty.
Let X1:7 (0) ( TT) be sequence of inputs of length T and YI.T (y): ..... . yr) be the cor- responding sequence of outputs, Assume that the output sequence is generated following a step-wise probabilistic distribution parameterized by 0, y/ p(yr [x, X1:--1, Yi:t-1, 0). Each sequence is generated with a different parameter sampled from its prior distribution, O ~ p(0). These define generative distribution of sequence:
T
P(ET:: 0 x p(0)p(y) m1) II p(y art X1--11, Y1=-1,
t=2
A sequence modeling task is often defined training autoregressive model parameterized by Xice, [ Yia-1), given dataset of sequences D {x(1, yO), generated from the above unknown generative distribution. At the limit of infinite data, the objective of maximum likelihood learning of the above sequence model can be formulated as マ arg max Cxp.
sb, Py
Cxp - - p( p(Yi:t- EEEE(Yii
111-1) 11
DEL (p(yolz,, X Y1st-1) (yrl2,, X1a-11 Yia-1)) dY +C, where (-I(- denotes the Kullback Leibler divergence and C is constant with respect to sb.
The left hand side distribution the cross entropy term p(y ær,, Xia-1, Yaa-)) is the true pre- dictive distribution of y: æ. conditioned on the observed history Yu--1 and X Xia-1, which can be written as
p XI--1, V1u-1) / plyelor, XI--1, Y 01=-1,,)p(00 Y Y
(3)
where
p(0 p(0)p( Y 1,
p(0 X X YI.-
(4)
p((0') (p( Y 1:--1 O')=1,'''d'"
As shown above, the tree predictive distribution of y: 24 contains both aleatoric uncertainty and epistemic uncertainty, in which the epistemic uncertainty is contributed by p(0) Yi+-)). With sufficient data and model capacity, the generative distribution in the sequence model Pe (ys Lr, XI=-1, II= will be trained to match the true predictive distribution. As result, we can expect epistemic uncertainty to be contained in the predictive distribution a sequence model. Note that the predictive distribution
only capture the epistemic uncertainty with spect to the parameters of sequences 6, but does not include the epistemic uncertainty regarding the hyper-parameters (if exist).
3

---

Published as a conference paper at ICLR 2024
4 IN-CONTEXT POLICY LEARNING
Epistemic uncertainty is the essential ingredient of EE. With the observation that the predictive distribution sequence model contains epistemic uncertainty, design an in-context policy learning algorithm with EE.
Consider the problem of solving family of RL games based on offline data. From each game, a set of trajectories are collected from a number of policies, where (i) (6)
(6) T k.k a k,1 ", k,1"...,, Ok. k.T. - is the trajectory of the k-th episode of the i-th game and o, a, denote the observed state, action and reward respectively. The policy used to collect 7/6) is denoted as (1) (i) (i) We concatenate all the episodes of the i-th game into a single se- (a k.
(6) (4) (()
quence 7 = (T] ..... T For convenience, the superscript (9) will be omitted in the following text unless the i-th game is explicitly referred to.
We propose a sequence model that is trained to step-wisely predict Py: (ax,,) Ok.t, Hg.e). where Rg, is the return-to-go at the k-th episode and the time step and H = 17 -i:k-1) the history until the time step including the past episode The above model formulation is similar to DT but a sequence in DT only contains one episode. Note that, different from AD, the concatenated trajectories do not need to be from a RL learning algorithm.
As shown in the previous section, by doing maximum likelihood learning on the collected trajecto- ries, the predictive distribution will be trained to match the true posterior distribution of action of the data collection policy,
pl RR. ann, Ok.t H HR, (a. plas, Rk, Out H&.e) I ((RR,,la Ok,t, H T (ay (k..loo,))
(5)
where p( p(R ax, Ok, HK.)) is the distribution of return after the time step following Tk-
As shown in ()). the posterior distribution of action biased towards the data collection policy. Following such action distribution allows collection policy but will lead to recreation of suboptimal trajectories if the data collection policy is not optimal. A more desirable action distribution is the action distribution that corresponds to the specified return without the influence of data collection policy, i.e.
reproduce the trajectories generated by the data
pl R.. Ok.t, HA, ((.e))(ann))
p(ax,e RE, Ok. HA.1) = J p( Re a Ok,t, HA, (aa()
(6)
where ((xx,)) is the uniform random policy, which gives all the actions equal probabilities. To let the sequence model learn the unbiased action distribution, the maximum likelihood objective needs to be defined as
Exp E/ p(Rr., ax ok,,, HR.)) log Pup (ax, Ru,,, Ok,,, Hk.
After applying the importance sampling trick, the Monte Carlo approximation of the above objective can be derived as
Σ U(ax.)) log Pub Pus (ax, Re,e. Ok,t, HR.)),
E ~
(8)
Te (ak,,,(oo).))
k.t
where are Te (ap.2l (k,r) and Rest ~ p(Ru, O Hx,)), i.c., ax, and RR are sampled from the data collection policy nTk
DESIGN OF RTOGR
Return-to-go is crucial component DT for solving RL tasks at inference using trained sequence model. Its return-to-go scheme is designed to calculate the expected return signal from a single episode. In order to achieve in-context policy learning, design the return-to-ge across episodes.
Ithough the learned action distribution does not need condition the return-to-go of the past actions, for the convenience of causal Transformer implementation, Hy contains the return-to-go of the past actions.
4

---

the cumulative rewards starting Published as a conference paper at ICLR 2024
Algorithm 1: In-context Exploration-EExllottttoo (ICEE) Action Inference
Input: trained ICEE model Pv
H1,1 =
for each episode 1,., K do
Reset the environment;
Select the episode return-to-oo ex;
for each step t ., T' do
Get an observation oz.i
Sample the step return-to-go CKtt q(c,,))
Sample an action asst ~ Po Ra.,, Ok.t: Hk,)):
Take the sampled action ar and collect the reward rhet Expand the history H4.++1 HK : U {Re, OR,t, ag.t Tk,
end
Compute the true return-to-go based on the whole episode (Rv, Update the history HA,T, with the true return-to-go (R
end
The return-to to-go of ICEE consists of two components: one for the individual steps within episode and the other for the cross-episode behavior, Rx,: = (ck,2, ūx). The in-episode return-to-go ck follows the design used in (Chen al. 2021), which is defined from the current step C,.t = "E>> rk,er. This design borrows the concept of the cumulative reward of RL and has the benefit of encapsulating the information of the future rewards following the policy. This is very useful when the outcomes of future steps strongly depend on the state and action of the current step. It allows the action that leads action that leads to a bad future outcome in the sequence model. The downside is that with a non- expert data collection policy, the optimal return-to- toorn-ttoooo-ao each state is often unobserved. This will limits the ability of the sequence model to achieve better performance than the data collection policy inference time.
good future outcome be differentiated from the
In the cross-episode return-to-go design, the situation is different. The initial states of individual episodes are independent each other. What determines the cumulative rewards of individual episodes are the sequence of actions. If we consider the whole policy space as the action space for each episode, the cross-episode decision making is closer to MAB, where the policy is the action and the return of an episode is the MAB reward. Driven by this observation, define the return- ttturn-to-go based on the improvement the return of current episode compared to all the previous episodes. Specifically, we define the cross-episode retur-to-go as
Fk > maxisjsk-1 otherwise.
where Fx = Ee Pke is the cumulative reward of the k-th episode. Intuitively, at inference time, by conditioning on 0g 1, take actions from a policy that is "sampled" according the probability of being better than all the previous episodes. This encourages the sequence model to deliver better performance after collecting more and more episodes. This design avoids the limitation of the need of observing the optimal policy learning trajectories.
Action Inference. After training the sequence model, the model can be used perform policy learning from scratch. At each step, we sample an action from the sequence model conditioned on the trajectory so far and a return-to-go go for the step. The return-to-ge -o for action sampling is defined as follows. The cross-episode return-to-go Cg is always set to one to encourage policy improvements. For the in-episode return-to-go, we follow the action inference proposed by by al. (2022). During training of ICEE, separate sequence model is trained to predict the discretized return from the trajectories, Pw ( Ck.e [x, Ok,t, Hk,)). At inference time, an in-episode return-to-go for each step is sampled from augmented distribution
g(ck,) x (o(Ch,el Ok,t HR.)) Ck.t Caiin y*. Cmax Cmin
(10)
This augmentation biases the return-to-go distribution towards higher values, which encourages the agent take actions that leads better returns. The return prediction is not combined into the main
5

---

Published as a conference paper at ICLR 2024
EI
KDFF
Clpp aa Miaaum
KCEB CEEE-biened
Capptt Purn: Mrirmm
ICEE ICEB-cmeed
Function Entuadoon
Euppod (Beconds)
(a)
(b)
Figure 1: Discrete BO results on 2D functions with 1024 candidate locations. 16 benchmark func- tions are used with five trials each. The y-axis shows the average distance between the current best estimate and the true minimum. The x-axis of (a) is the number of function evaluations and the x-axis of (b) is the elapsed time.
sequence model as in (Lee et al. 2022), because in the main sequence model the returns are also fed into the inputs. In this way, the model quickly figures out Ct. can be predicted from C This is problematic because at inference time the true return cannot be observed until the end of an episode.
After CR is sampled, an action sampled conditioned on the combined return-to-ge RR,-- The resulting state and reward are concatenated to the trajectory for the next step action prediction. At end of an episode, we will recalculate the true ca,, and č. based on the rewards from the whole episode and update the return-to-go in the trajectory with the recalculated ones. This makes the trajectory at inference time be as close to the training trajectories as possible. The description of the action inference algorithm can be found in Algorithm
BAYESIAN OPTIMIZATION EXPERIMENTS
Bayesian optimization (BO) is a very successful application of apploratio-exxllotattiin (EE). is able to search for the optimum of a function with the minimum number of function evaluations. The Gaussian process (GP) based BO methods have been widely used in various domains like hyper- parameter tuning, drug discovery, aerodynamic optimization. To evaluate the EE performance of ICEE, we apply it to BO and compare it with a GP-based approach using one of most widely used acquisition functions, expected improvement (EI).
We consider a discrete BO problem. The task is to find the location from fixed set points that has the lowest function value as few number of function evaluations as possible. BO can be viewed as a special type of multi-armed bandit (MAB), where each action is associated with a location in bounded space. To solve BO with ICEE, encode the iterative search trajectory of function as a single sequence, where 0, is the location of which a function value is collected at the step 4, r. is the corresponding function value and R2 the return-to-go. The observations (o.) can be used to encode any side information that is known about the function and the minimum. As no such information is available for generic BO, do not use }or} in our experiment. As the function value can seen as the available immediate reward for, we treat each action as a different episode and use only the episode return-to-go as the R, here. As each action is associated with a location, we embed actions by learning a linear projection between the location space and the embedding space. When decoding a Transformer output for an action, the logit of an action generated using a MLP that takes as input the Transformer output together with the embedding associated with the action. The design to tackle the challenge that the set of locations of each function may be different.
To train ICEE to solve the discrete BO problem, need generate training data that consists of onput-output pairs randomly sampled functions. During training, the input-output pairs at random locations are generated on-the-fly. We GP with the Matérn 5/2 kernel to sample 1024 points for each function. The locations of these points are sampled from a uniform distribution on [0, 1]. The length scales of the kernel sampled from uniform distribution on [0.05, 0.3].
Ideally, for each sampled function, we can use a BO algorithm to solve the search problem and take the sequence of actions the training data. This is very computationally expensive due to

---

Published as a conference paper at ICLR 2024
the high computational cost of BO algorithms. Thanks to the EE property of ICEE, the training data do not need to come from an expert inference algorithm. This is very different from behavior cloning/ / algorithm distillation. We use cheap data collection algorithm that randomly picks actions according to a probability distribution that is computed based on the function values of actions. The probability distribution is defined follows: The actions are sorted in ascending order according to their function values and the action at the i-th position gets the probability p(a = a;) x exp( 77) y is sampled from a uniform distribution between and 10 for each function. Intuitively, the locations with lower function values are more likely to be picked.
After training ICEE, we evaluate its performance on a set of 2D benchmark functions. We use 16 21) functions implemented in (Kim & Choi (017). The input space of every function is normalized to be between O and 1. We randomly sample different set of 1024 points from each function and normalize the resulting function values be zero-mean and unit-variance. Each function was given five trials with different initial designs. At each step of search, we compute the difference of the function value between the current best estimate and the true minimum. The average performance of all the evaluated functions is shown in Fig. ! The performance of ICEE is compared with the GP-based BO method using the El acquisition function and the random baseline that picks locations according to a uniform random probability. "ICEE-biased" denotes a variant of ICEE that does not use the action bias correction objective as shown in (8). The search efficiency of ICEE is on par with the GP-based BO method with EL. Both of them significantly better than random and are noticeably better than ICEE-biased. The performance gap between ICEE and ICEE-biased shows the loss of efficiency due to the biased learned action distribution. Being able to perform par with a state-of-the thet-errtttttr art BO method demonstrates that ICEE is able to perform state-of-the-ar EE with in-context inference.
A clear advantage of ICEE is that the whole search is done through model inference without need of any gradient optimization. In contrast, GP-based BO methods need fit GP surrogate function at each step, which results into significant speed difference. Fig. Ib shows the same search results with the x-axis being the elapsed time. All the methods run on a single A100 GPU. Thanks the in-context inference, ICEE is magnitude faster than conventional BO methods. More details of the BO experiments can be found in Appendix
7 RL EXPERIMENTS
We investigate the in-context policy learning capability of ICEE on sequential RL problems. To demonstrate the capability of in-context learning, we focus on the families of environments that cannot be solved through zero-shot generalization of a pre-trained model, so in-context policy learn- ing is necessary for solving the tasks. This implies that some information that are important to the success of a task is missing from the state representation and will need to be discovered by the agent. We use the two grid world environments in (Lee et al. 2022 : dark room and dark key-to-door. The details of them are follows.
Dark Room. The experiment takes place in a 2D discrete POMDP, where an agent is placed inside a room to locate a goal spot. The agent has access to its own location coordinates (at, y). but is unaware of the goal's placement requiring it to deduce it from the rewards received. The room's 's dimensions are 9x9 with the possible actions by the agent including moving one step either left, right, up or down, or staying idle, all within an episode length of 20. Upon completion, the agent gets placed back the mid mid-point of the map. Two environment variants are considered for this experiment: The Dark Room case where the agent obtains a reward (r=1) each time the goal is achieved, and the Dark Room Hard case where the rewards are sparse (r=1 only once for attaining the goal). Whenever the reward value is not 1, it will be considered as O. Different from (Lee et al. 2022). we keep the room size of the hard case to be 9 x 9.
Dark Key-to-Door. This setting is similar to the Dark Room, but with added challenging features. The task of the agent is to locate an invisible key to receive a one-time reward of r=1, and subse- quently, identify invisible door to gain another one-time reward of r=1. Otherwise, the reward remains at r=0. The agent's 's initial location in each episode resets randomly. The room size is still 9 x 9 but the episode length increases to 50 steps.
7

---

Published as a conference paper at ICLR 2024
To collect data for offline training, we sample set of new games for each mini-batch. We collect K episodes from each game. Thanks to the EE capability of ICEE, the training data do not need to be from a real RL learning algorithm like Deep Q-Network (DQN), which is expensive to run. Instead, we let a cheap data collection policy act for K' episodes independently and concatenate the resulting episodes into single sequence. We
an e-greedy version of the "cheating" optimal policy. The policy knows the goal location that is unknown to the agent and will move straightly towards the goal with 1 probability and with probability it will take an action that does not make the agent closer to the goal. For each episode, is sampled from a uniform distribution between O and 1. Intuitively, this policy has some chance to solve a game efficiently when c is small but on average it does not deliver good performance. For Dark Room experiments, each sequence consists of 50 episodes and for Dark Key-To-Door, it consists of 20 episodes.
Dark Room (Biased). To demonstrate the benefits of EE for sequential RL problem when the data collection policy cannot be optimal, we create a variant of the dark room environment. At each step, the data collection policy takes the "left" action with the 2/3 probability and with the 1/3 probability acts as described above. At training time, the goal can be at any where in the room and, at evaluation time, the goal will only appear on the right hand side where x > 5.
For sequential RL problems, ICEE consists of two sequence models: one for action prediction and the other for the in-episode return-toggo prediction. The return-to-go sequence model takes as inputs the sequence of state, action, reward triplets and predicts the in-episode retur-to-go. The action prediction model takes as inputs the sequence of the triplets and the two return-to-go Rg, and predicts the action sequence. The two models are trained together with the same gradient optimizer. To encourage ICEE to solve the games quickly, when calculating the in-episode return-to-go, negative reward, -1/T, is given to each step that does not receive a reward, where T is the episode length. Both o and C, are discrete and tokenized.
7.1 BASELINE METHODS
Source. We use the data collection policy baseline for comparison. As the data collection policy solves each episode independently, we calculate the average return across multiple episodes.
Algorithm Distillation (AD, Laskin al.) 2023). The in-context learning algorithm that distills a RL algorithms from RL training trajectories. AD predicts action based only on the current states and the history of state, action and reward triplets. We replicate the implementation of AD using the Transformer architecture as ICEE. We apply AD to the same training data as ICEE uses (ignoring the return-to-o signals), despite they are generated from RL learning trajectories.
AD-sorted. AD is designed to be trained on RL learning trajectories. An important property of RL learning trajectories is that the performance of the agent gradually increases throughout training. To mimic such trajectories using our data, we sort the episodes in a sequence according to the sampled e of the data collection policy in the descent order. deterriiee how close the data collection policy is to the optimal policy. In this order, the episodes in latter position of sequence tend to have higher return. We train AD using this sorted sequences instead of the original ones.
Multi-game Decision Transformer (MGDT, et al. 2022). MGDT is not an in-context learning algorithm. We train MGDT using only one episode from each sampled game. The performance of MGDT shows that what the performance of the agent is when there is no in-context policy learning.
7.2 EVALUATION & RESULTS
After training, ICEE will be evaluated on solving a set of sampled games. The inference algorithm described in Alg. [ No online model update is performed by ICEE and all the baseline methods at evaluation time. For cach sampled game, ICEE and two variants of AD will act for K episodes consecutively. In each episode, the trajectories from the past episodes are used as in the history representation. Ideally, good performing agent identifies the missing information with as few number of episodes as possible and then maximizes the return in the following episodes. For each problem, we sample 100 games and K is 50 for Dark Room and 20 for Key-To- Key-To-Door.
The experiment results are shown in Fig. 2 ICEE is able to solve the sampled games efficiently compared to the baseline methods. The EE capability allows ICEE search for the missing infor-
8

---

Published as a conference paper at ICLR 2024
Dark Room
Dark Room (Hardd)
Dark Key-To-Dons To-Door
15.0
125-
10.0
825m
75-
Poau
ICEE
0.41
AD
AD sorted
0.2
sorted
AD sorted
source
Sourco
MGOT
MEDT
MGDT
90
40
00
10
50
10
10
15
Env Ep sodes
Ern Expaxdes
Ern Epbosches
(a)
(b)
Dark Room (Biased)
Average Action Entropy
1.6
Dack Raom
Dark Room (Hand)
Key-Tbb-Too
Rekin
1.0
ICEE
ICEE blased
Episodes
(d)
(e)
Figure 2: Experimental results of in-context policy learning on grid world RL problems. An agent is expected to solve a game by interacting with the environment for K episodes without online model updates. The y-axis of (a-d) shows the average returns over 100 sampled games after each episode. The y-axis of (e) shows the entropy of the action distribution. The x-axis shows the number of episodes that an agent experiences.
mation efficiently and then acts with confidence once the missing information is found. An indicator of such behavior is the continuous decrease of the action entropies as the agent experiences more episodes (see Fig. 2ep.
As expected, the original AD learns mimic the data collection policy, which results an average performance slightly below the data collection policy. MGDT fails to solve most of the games due to the missing information. Interestingly, despite that the training data not generated from a RL learning algorithm, AD-sorted is able to clone the behavior of the data collection policy with different e at different stages, which allows
solve the games at the end of the sequence.
ICEE-biased is not shown in Fig. 2a Fig. 2b and Fig. 2c as it achieves similar performance as ICEE does. The reason is that there is no clear bias the action distribution of the data collection policy. However, as shown in Fig. 2d for the Dark Room (Biased) environment, ICEE clearly outperforms (iae-biased, as it can overcome the bias in the data collection policy and keep sufficient uncertainty in action distribution to explore the right hand side of the room. AD-sorted fails the task because clones the data collection policy, which is unlikely to solve the tasks due to the action bias.
8 CONCLUSION
In this paper, we analyze the predictive distribution of sequence models and show that the predic- tive distribution can contain epistemic uncertainty, which inspires the creation of an EE algorithm. We present an in-context EE algorithm by extending the DT formulation to in-context policy learn- ing and deriving unbiased training objective. Through the experiments on BO and discrete RL problems, we demonstrate that: (i) ICEE can perform EE in in-context learning without the need of explicit Bayesian inference; (ii) The performance of ICEE is par with state-of-the-an art BO meth- ods without the need of gradient optimization, which leads to significant speed-up; (ii) New RL tasks can be solved within tens of episodes.
9

---

Published as a conference paper at ICLR 2024
REFERENCES
Jacob Abemeethy, Alekh Agarwal, Teodor V Marinov, and Manfred K Warmuth. A mechanism for sample-efficiet in-context learning for sparse retrieval tasks. arXiv preprint (tt1:2305.170440 2023.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
Maximilian Balandat, Brian Karrer, Daniel R. Jiang, Samuel Daulton, Benjamin Letham, An- drew Gordon Wilson, and Eytan Bakshy. BoTorch: A Framework for Efficient Monte-Carlo Bayesian Optimization. Advances in Neural Information Processing Systems 33, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lec, Aditya Grover, Misha Laskin, Picter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, pp. 15084-15097, 2021.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning identifiability challenges and effective data collection strategies. In Advances in Neural Information Processing Systems, pp.
2021.
Yan Duan, Marcin Andrychowicz. Bradly Stadie, OpenAl Jonathan Ho. Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. Advances in neural information processing systems, 30, 2017.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6, 2005.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learming, pp. 1 1126-1 135. PMLR, 2017.
Sebastian Flennerhag, Yannick Schroecker, Tom Zahavy, Hado van Hasselt, David Silver, and Satin- der Singh. Bootstrapped meta-learning arXiv preprint 1tw.::200090050000 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. International conference on machine learning, pp. 2052-2062. PMLR, 2019.
Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel Mankowitz, Cosmin Paduraru, et al. RI unplugged: A suite of benchmarks for offline reinforcement learning. Advances Neural Information Process- ing Systems, 33:7248-7255 2020.
Abhishek Gupta, Benjamin Eysenbach, Chelsea Finn, and Sergey Levine. Unsupervised meta- learning for reinforcement learning. arXiv preprint arXiv: 180,000000 2018
Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. Meta-learning in neural networks: survey. IEEE transactions on pattern analysis and machine intelligence, 44(9): 5149-5169, 2021.
Kyle Hsu, Sergey Levine, and Chelsea Finn. Unsupervised learning via meta-learning arXiv preprint 4rX1v:1810.02334. 2018.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances Neural Information Processing Systems. pp. 2273-1286, 2021.
Jungtack Kim and Seungjin Choi. BayesO; A Bayesian optimization framework in Python. https:
7/bayeso. org, 2017.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforce- ment learning: State-of-the-art, pp. 45-73. Springer, 2012.
10

---

Published as a conference paper at ICLR 2024
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, and Volodymyr Mnih. In-context reinforcement learning with algorithm distillation. In The Eleventh International Conference on Learning Representations, 2023.
Kuang- Huen Lee, Ofir Nachum, Mengjiao (Sherry) Yang, Lisa Lee, Daniel Freeman, Sergio Guadar- rama, lan Fischer, Winnie Xu, Erie Jang, Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers. In Advances in Neural Information Processing Systems, pp. 27921-27936, 2022.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01665 2020.
Ke Li and Jitendra Malik. Learning optimize. arXiv preprint arXiv: arXiv: (600::6800.. 2016.
Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8:293-321, 1992.
Kevin Lu, Aditya Grover, Pieter Abheel, and Igor Mordatch. Pretrained transformers as universal computation engines. arXiv preprint arXiv:2103.05247, 1, 2021.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202. 122.1i::2202.183.7, 12837, 2022.
Nikhil Mishra, Mostafa Rohaninejad. Xi Chen, and Pieter Abbeel. A simple neural attentive meta- learner. arXiv preprint arXiv:170.00314, 2017.
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. Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, nature, 518(7540):529-5333 2015.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online rein- forcement learning with offline datasets. arXiv preprint (0XX::2000609355,, 2020.
Vitchyr H Pong, Ashvin V Nair, Laura Smith, Catherine Huang, and Sergey Levine. Offline meta- reinforcement learning with online self-supervision. In International Conference on Machine Learning, pp. 17811-17829. PMLR, 2022.
Sachin Ravi and Hugo Larochelle. Optimization as model for few-shot learning. In International conference on learning representations, 2016.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom EEccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. May 2022a.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXhv:2205.06173 2022b.
Martin Riedmiller. Neural fitted q iteration first experiences with a data efficient neural reinforce- ment learning method. In Machine Learning: ECML 2005: 16th European Conference Ma- chine Learning, Porto, Portugal, October 3-7, 2005. Proceedings 16, pp. 317-328. Springer, 2005.
Juergen Schmidhuber, Jieyu Zhao, and MA Wiering. Simple principles of metalearning. Technical report IDSIA, 69:1-23, 1996.
11

---

Published as a conference paper at ICLR 2024
Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Ne- unert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.0839% 2020.
Bradly C Stadie, Ge Yang, Rein Houthooft, Xi Chen, Yan Duan, Yuhuai Wu, Pieter Abbeel, and llya Sutskever. Some considerations learning to explore via meta-reinforrement learning. arXiv preprint arXiv:1803.01118,, 2018.
Richard Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3:9-44, 1988.
Richard S Sutton. A history of meta-gradient: Gradient methods for meta-leariigg arXiv preprint arXiv:2202.09701, 2022.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopehand, Nathalie timesclee adaptation an open-ended task space. arXiv preprint arXiv:2301.07008 2023.
Michael Chang, Natalie Clay, Adrian Collister, et al.
Sebastian Thrun and Lorien Pratt. Learning to learn. Springer Science & Business Media, 2012.
Vivek Veeriah, Matteo Hessel, Zhongwen Xu, Janarthanan Rajendran, Richard L Lewis, Junhyuk Oh, Hado P van Hasselt, David Silver, and Satinder Singh. Discovery of useful questions as auxiliary tasks. Advances in Neural Information Processing Systems, 32, 2019.
Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint
1055633, 2016.
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and gen- cralize. International conference on machine learning, pp. 3751-3760. PMLR, 2017.
Zhongwen Xu, Hado P van Hasselt, and David Silver. Meta- deta-gradient reinforcement learning. Ad- vances in neural information processing systems, 31, 2018.
Zhongwen Xu, Hado P van Hasselt, Matteo Hessel, Junhyuk Oh, Satinder Singh, and David Silver. Meta gradient reinforcement learning with an objective discovered online. Advances in Neural Information Processing Systems, 33:152554-11 12215555526444,, 2020.
Zhangyue Yin, Qiushi Sun, Qipeng Guo, Jiawen Wu. Xipeng Qiu, and Xuanjing Huang. Do large language models know what they don't ton't know? In Findings of the Association for Computational Linguistics: ACI. 2023, pp. 8653-8665, 8665, July 2023.
Tom Zahavy, Zhongwen Xu, Vivek Veeriah, Matteo Hessel, Junhyuk Oh, Hado P van Hasselt, David Silver, and Satinder Singh. A self self-tuning actor-critic algorithm. Advances in neural information processing systems, 33:20913 12-2099329924... 2020.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, pp. 27042-27059, 2022.
Zeyu Zheng, Junhyuk Oh, and Satinder Singh. On learning intrinsic rewards for policy gradient methods. Advances in Neural Information Processing Systems, 31, 2018.
Luisa Zintgraf, Sebastian Schulze, Cong Lu, Leo Feng, Maximilian Igl, Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: Variational bayes- adaptive deep rl via meta-learning. Journal of Machine Learning Research, 22(289):1-39, 2021.
12

---

Published as a conference paper at ICLR 2024
Appendix
A DERIVATION OF THE OBJECTIVE OF A SEQUENCE MODEL
The maximum likelihood objective sequence model (2 can derived in the following steps.
c = / p(Yrrr,0/ 1)x1 Po Y1:1 X 1 d0
p(Y1:r, 10fir,010 0|x111) 1:1) log II (v or X1:-1, 1
=
Σ
11 X X1a-1, Y1s-- (id=-))dYi:rdl
Σ
p (o:pollx X1.-1, Y1st-1ddY1.
Σ
/ P(YY.c-1 (YY X X Y1 log Pe (36 X14-11
X XI.-1, Y1.-1) log Per (11 X1--11 Y1--1)
11 P X X Y dY
+ P( X X Y1u-1) log pl X14-1, -1
T
Σ
p p(Y X1n-1) DKL (p(u) xc, X14-1. Yu 1 X14-1, Y(a-1)) dY
P(Y::111 1 plu X1--11 Y1--- dY
E 1x (plyn le, X X1r-1, Y Pol X1a-1, (1) 1Y + C,
=
B DERIVATION OF THE UNBIASED OBJECTIVE
The Decision Transformer's approach to offline RL includes the training of an action distribution conditioned the return, thereby allowing for the sampling of an action the time of inference by providing the projected return (return-to-go). Given that the return is the result of present and down- stream actions, the distribution of action that model tries to learn can be reframed as an action's posterior distribution, as is presented in equation (5). Note that equation (5) outlines the data dis- tribution, which should be distinguished from the neural network model p& Pa.eeee/
Ok. HR.)). As noted in equation (5), the action's 's distribution is proportionate to the return's 's distribution, and this is subsequently weighted by the probability of action from the data collection policy. As is the posterior distribution derived by the Bayes rule, we denote it as the "true" action posterior distribution.
To put this intuitively, if model can accurately match equation (5). it will result in the action dis- tribution being skewed towards the data collection policy. For instance, in a pre-recorded sequence, should action be randomly selected with extremely low probability from the data collection pol- icy, yet yield a high return, the subsequent action distribution in equation 5 would attribute a minute probability to the relevant action, given the high return. Despite the observation of a high return following the action, which would suggest high probability of RR (a Ok,ts He.)). the resulting probability of action is weighted by the probability of action in the data collection policy,
13

---

Published as a conference paper at ICLR 2024
Th (dg.elok..) resulting in small value. Therefore, even though ( is the true posterior distribution of action, it is not the desirable action distribution for our model.
Ideally, the action distribution, demonstrated in equation (6) should be proportional only to the return's distribution and would be unaffected by the data collection policy. With such a distribution, the undesired devaluation due the data collection policy would be eliminated, thereby resolving the mentioned issue.
As explained above, would like to learn action distribution in (6 instead of the action dis- tribution in ( However, since (5) is the true action distribution of data, the common maximum
likelihood training objective will let the model match the action distribution in (5).
The unbiased objective of learning the action distribution in (8) can be derived the following steps.
To let the model learn the action distribution in 6 instead, first state the desirable training objective as the data follow the distribution
Lup
Σ
p Ra,,, ao HI.e) log pro (ax,e Ra,e Ou,3, Hag.d
k,t
p( (Ra,, Hr.e) (f Siax..Re.. Ou,t, HR,)) log Pa Pw (ax, Rk,t, Ox,t Ha,e)dan.d d
Σ
=
Then, we apply the importance sampling trick to introduce the true action posterior in the equation:
Bi ok HEa) ( planaal Ru.e ou,, HK.)) Max Rig, Ok,t, HK.)) plan, PE Ok,t, Hi,))
Lu
Σ
R,t
log Pep (ak REg, Ok,t HE,))ax, dR
After some rearrangement the equation, we get a much clearer formulation of the objective:
pl Rx. lag. H., RRe..lo.......
Las = Σ k.t
p( p(Rk,loN Had)(
plan.el Ok, Ha.)) pl lar., o,,,, He. (e.)). ox.e)
P(R..loe...H... He.))
log Pap (ax,1 Pac,es Ok,t, , HA HR,dd dR
p( i(Re,elox, HEd) (f planse Raass Ok HK.)) Ty Tx(a lok, ) Rk, Ok.t. H Hk.))
= =
log Dop (ax,t RK,es Ok,t H HKS)dauR) dRa.
/ p P(Reolon, Har) (f Spaast Onee Ha.)) u(an,)
= =E Rrd
Tx =
log (ak. Re,e, Ok,t, H Ha,d)daz.) dR
= =E k,t
U(ar,)) log Pvs Pw (R Ok,t, Hu,)daa,
p( an, ok,, Ha.) m=(ax.ooox.))
Note the probability distribution in front is the joint distribution of return and action in the data distribution. We apply the Monte Carlo approximation to the integral by considering the recorded data samples from the data distribution. We get the proposed training objective.
U(ax.)
Lep Σ k,t
log Pep (ar, RK,,, Ok1t HEd)s Rk,, axet pl PRR,,, ar. H
Th (ok,)))
14

---

Published as a conference paper at ICLR 2024
C IMPLEMENTATION DETAILS
ICEE is implemented based on nanoGPT For the RL experiments, ICEE contains 12 layers with 128 dimensional embeddings. There are heads in the multi-head attention. We use the Adam optimizer with the learning rate 10-5
D BAYESIAN OPTIMIZATION EXPERIMENTS
We consider a discrete BO problem. The task is to find the location from fixed set of points that has the lowest function value few number of function evaluations as possible. At the beginning of search, the function values of some randomly picked locations are given. The BO algorithm will be asked suggest a location where the function minimum lies and then, the function value will be collected from the suggested location. The suggestion-evaluation loop will be repeated with fixed number of times. The performance of a BO algorithm evaluated by how quickly it can find the function minimum.
The list of 2D functions used for evaluations are: Branin, Beale, Bohachevsky, Bukin6, DeJong5, DropWave, Eggholder, GoldsteinPrice, HolderTable, Kiml, Kim2, Kim3, Michalewicz, Shubert, SixHumpCamel, Cameell
El
ICEE
Gap to Func Minimum
ICEE-biased
109
random
10-
0
20
40
60
80
100
Function Evaluations
Figure 3: Discrete BO results on 2D functions with 1024 candidate locations. shows the 95% confidence interval of the mean.
The shading area
The expected improvement baseline is implemented using BOTorch (Balandat al. 2020 We used the "SingleTTaste TaskGP" class for the GP surrogate model, which uses Matern 5/2 kernel with Gamma prior the length scales.
E QUANTITATIVE RESULTS OF THE DISCRETE RL EXPERIMENTS
Please find the quantitative comparison of the RL experiments shown in Fig. below. The shown values are the average returns over 100 sampled games and the values in the parentheses are the confidence intervals of the mean estimates, which correspond to the shaded area in the figure. We take three time points along trajectories of in-context policy learning. MGDT not able to update policy at inference time, we only estimate single mean return for each game.
15

---

Published as a conference paper at ICLR 2024
Dark Room (10th Episode)
Dark Room (30th Episode)
Dark Room (50th Episode)
ICEE
8.15 (1.29)
12.37 (1.14)
13.61 (0.86)
AD
3.74 (1.15)
4.51 (1.17)
4.03 (1.15)
AD-sorted
0.05 (0.05)
3.83 (0.87)
12,48 (1.37)
MGDT
1.86 (0.93)
1.86 (0.93)
1.86 (0.93)
source
5.13 (1.19)
5.13 (1.19)
5.13 (1.19)
Dark Room (Hard) (10th)
Dark Room (Hard) (30th)
Dark Room (Hard) (50th)
ICEE
0.48 (0.10)
0.74 (0.09)
0.79 (0.08)
AD
0.33 (0.09)
0.43 (0.10)
0.43 (0.10)
AD-sorted
0.08 (0.05)
0.55 (0.10)
0.75 (0.08)
MGDT
0.09 (0.06)
0.09 (0.06)
0.09 (0.06)
source
0.51 (0.10)
0.51 (0.10)
0.51 (0.10)
Dark Key-To-Doom (5th)
Dark Key-To-Doo To-Door (10th)
Dark Key- Key-To-Door (20th)
ICEE
1.04 (0.15)
1.50 (0.12)
1.84 (0.08)
AD
0.67 (0.15)
1.02 (0.17)
0.94 (0.17)
AD-sorted
0.17 (0.08)
0.84 (0.14)
1.77 (0.09)
MGDT
0.34 (0.11)
0.34 (0.11)
0.34 (0.11)
source
1.10 (0.19)
1.10 (0.19)
1.10 (0.19)
16


Reviews:
[{"review_id": "YvREXTTLxr", "reviewer": "Reviewer_5QaL", "summary": "This paper proposes a novel way to leverage the powerful autoregressive sequence learning capabilities in transformer architectures to decision making problems by finding a way to bring in exploration-exploitation for in context policy learning. This goes beyond recent work along with a similar theme for decision making problems (Decision transformers, Algorithm distillation) in two important ways. First, it learns to explore and exploit across multiple episodes concatenated as one sequence, and secondly it is designed to exploit data generated from arbitrary policies that can, in principle, be highly inefficient. \n\nExperiments are conducted on two well motivated settings -- (1) bayesian optimization and (2) a family of episodic toy RL tasks in a 2d discrete maze. The authors also investigate controlled variations in the data collection policy to highlight how the proposed technique compares favorably to Algorithm Distillation.", "questions": "- The last statement in page 3 seems somewhat ambiguous, could the authors clarify what they mean by this? e.g. _\"...Note that the epistemic uncertainty associated with cross sequence latent variables if exist, e.g., the posterior of parameters in $p(\\theta)$, would be not learned into the sequence model.\"_ But shouldn't the true predictive distribution of $y_t$ implicitly encode both the epistemic and aleatoric in the limit of infinite data?\n\n- There could be several reasonable ways to generate suboptimal data for analysis in toy experiments, and the authors have investigated an $\\epsilon-$greedy version of the (oracle) optimal where $\\epsilon$ is sampled independently for each _episode_ with full $[0,1]$ support. As long as there are enough samples covering a large enough $\\epsilon$, a non-trivial part of the episode collection would effectively be oracle policy demos? If that's the case, a more effective stress might be more meaningful e.g. as one possible minimal change, one could sample $\\epsilon$ more generally as a beta distribution with two params and rerun the experiments across a range of those params.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Simple and well articulated but technically precise characterization of the issues surrounding in context policy learning, especially related to epistemic uncertainty.\n- The experimental results, while in simple domains at small scale, are nevertheless well motivated and designed to illustrate the promise of the proposed ideas.\n- For the BO benchmark, the proposed approach is competitive with EI, but at a small fraction of the cost. For the sequential RL tasks, it achieves better performance compared to other comparable meta learning approaches which were not designed to do exploration in context.", "weaknesses": "The context length requires a sequence of episodes for in-context learning, which can make it fundamentally quite challenging in terms of scale to go beyond small dimensional problems. \n\nNits/minor typos:\n- GP biased -> GP based\n- pg 3: indefinite -> infinite\n- Sec 6: wildly -> widely", "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": "", "extra_scores": {}}, {"review_id": "7srm4SIQf1", "reviewer": "Reviewer_1NXH", "summary": "This paper studies in-context RL and proposes a method that addresses the exploration-exploitation tradeoff. Specifically, during the pre-training, the method assigns a reward to each trajectory so that the model can learn to improve the policy. Furthermore, the paper applies the method on discrete Bayesian Optimization and several RL environments, which show better performance than those of some baselines such Algorithm Distillation.", "questions": "1.\tThe return-to-go for cross-episode seems simple and straightforward. Is it possible to have a more sophisticated design such that the performance can be further improved? For example, current return-to-go seems depends on the order of the offline trajectories. Is there any possible way to avoid this and make design to be order-insensitive?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ In-context RL is a timely topic in the study of RL and this work provides some interesting idea about the design of in-context RL algorithms.\n+ The design of cross-episode reward successfully removes the requirement that the offline dataset is generated from some RL learning algorithms. This design improves the applicability of in-context RL.\n+ The experimental results show promising performance of ICEE in the early episodes.", "weaknesses": "- The experimental evaluation is not sufficient. There is no comparison between in-context RL algorithms and traditional offline RL algorithms, multi-task RL algorithms or posterior sampling based algorithms. \n- The proposed ICEE algorithm lacks theoretical performance guarantees.\n- The ICEE algorithm heavily relies on importance sampling and Monte Carlo approximation, which are widely used techniques in traditional RL and not new.", "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": "", "extra_scores": {}}, {"review_id": "6owOnQeR0k", "reviewer": "Reviewer_5MSL", "summary": "The authors develop a in-context exploration-exploitation algorithm for reinforcement learning. This algorithm can be viewed as an extension and refinement of Decision Transformer by generalizing the model from a single episode to multiple episode, and modified the objective to eliminate the bias. The algorithm is extensively tested against Bayesian optimization problems and multiple reinforcement\nlearning tasks, the proposed algorithm is able to solve sampled games much faster than baseline methods.", "questions": "On page 4, below eq (5), the statement \"true posterior distribution of action is biased towards the data collection policy\" is not that straightforward from eq (5), could you add more explanations?\n\nBetween eq.(7) and eq.(8), maybe the author could be provide some details on the importance sampling trick they mentioned.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The experimental results are very promising. While the algorithm is an extension of Decision Transformer(CT), it certainly outperforms DT significantly.\n\nThe application of in-context learning in RL is innovative and has potentials to bring advancement to learning.", "weaknesses": "The explanation provided in the paper on why the proposed algorithm is not very convincing. There is no quantitative analysis provided for identifying the difference between ICEE and DT. As pointed below, some of the explanations of the key ideas in the paper are not very clear.", "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": "", "extra_scores": {}}, {"review_id": "SgRY0DZzsj", "reviewer": "Reviewer_cBRg", "summary": "The paper presents a novel algorithm called In-context Exploration-Exploitation (ICEE). Learning to learn in context is promising but demands extensive learning trajectory data and expansive training of large Transformer models. The ICEE algorithm is designed to address these challenges by augmenting decision transformers with an extra cross-episode return-to-go and training on concatenated episodic trajectories. It learns to balance exploration versus exploitation directly within the Transformer model during inference, eliminating the need for separate Bayesian inference processes. This enables ICEE to tackle Bayesian optimization tasks as effectively as methods based on Gaussian processes but at a notably faster rate. Experimental results show that ICEE can adeptly learn and solve new RL tasks with just tens of episodes, corroborating its efficiency.", "questions": "1. How is the exploration-exploitation dilemma addressed in ICEE?\n2. How does ICEE perform when trained with data collected by random policies?\n\nI will be pleased to raise my score if these questions, especially the first one, can be properly answered.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "**Learning from Episodic Trajectories**: ICEE learns in-context reinforcement learning using merely concatenated episodic trajectories, which makes it a versatile choice for scenarios where abundant learning trajectories are unavailable and constitutes a novel and significant contribution to field of reinforcement learning.\n\n**Connection to Thompson Sampling**: ICEE resolves the exploration-exploitation dilemma in a way analogous to posterior sampling methods, like Thompson Sampling. Given the history, it samples a move that could possibly lead to an improved outcome, gaining either a high return or new information. This method could be the most scalable and efficient implementation among all Thompson Sampling variants.", "weaknesses": "**Ambiguous Presentation**: Section 3 of the paper feels disjointed from its succeeding sections, leading to some confusion regarding the role and definition of the parameter $\\theta$ within ICEE's framework.\n\n**Questionable Experimental Setup**: The experiments predominantly utilize noisy optimal policy variants for data collection. This leaves an open question regarding ICEE's performance when relying on trajectories from random policies.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)