In-context Exploration-Exploitation for Reinforcement Learning

We thank the reviewer for the feedback.

How is the exploration-exploitation dilemma addressed in ICEE?

The exploration-exploitation dilemma in ICEE is primarily addressed through a method analogous to a posterior sampling style of exploration-exploitation, such as Thompson Sampling. The specific algorithm, however, depends on the return-to-go design. For instance, applying our cross-episode return-to-go to the Bayesian optimization problem, sampling actions from our learned model equates to sampling an action from the renowned probability of improvement acquisition function. This is what allows our model to deliver performance comparable to Bayesian Optimization with an expected improvement acquisition function.

Section 3 sheds more light on this dilemma by demonstrating that a generic sequence model's predictive distribution exhibits epistemic uncertainty if the model is trained with sufficient data and possesses enough model capacity. Applying this to the in-context policy learning problem, we consider a specific sequence model that predicts an action based on return-to-go, current state, and observed history: $p(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t})$. Using the notation from equation (3), $y_t$ corresponds to $a_{k,t}$, $x_t$ corresponds to $R_{k,t}$ and $o_{k,t}$, and $X_{1:t-1}$ corresponds to $H_{k,t}$, excluding past actions.

Equation (3) illustrates that the action distribution obtained from data is the marginalization of the task's posterior distribution given the observed history. Thus, sampling an action from the action distribution $p(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t})$ is equivalent to first sampling a task based on the task's posterior distribution $p(\theta | H_{k,t})$, then sampling an action conditioned on the sampled task $p(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t}, \theta)$. The latter is an explicit posterior sampling algorithm. This equivalence explains why our model can effectively balance the exploration-exploitation trade-off.

To illustrate this point further, consider the Bayesian optimization problem. Here, an action corresponds to an input location from which the unknown function's value can be taken. $\theta$ is the parameterization of the unknown function, the minimum of which we are seeking. One way to apply Thompson Sampling to BO is to use a two-step sampling process. For a scenario where there are two observed input-output pairs of the function $H_3 = (x_1, y_1, x_2, y_2)$, we can derive the posterior $p(\theta | H_3)$ of the unknown function. A specific parameterization is then sampled from the posterior $\theta’ \sim p(\theta | H_3)$. Based on $\theta'$, an action distribution is derived, proportional to the likelihood of each input location having a lower function value than both $y_1$ and $y_2$. This can be represented as $p(a_3 | R_3=true, H_3, \theta’)$. An action is then sampled from this distribution, equivalent to sampling according to the probability of improvement. This action — applied through the Thompson sampling algorithm — can enable the discovery of the function's minimum. The aforementioned relationship demonstrate that the sampling of an action from our model is equivalent to the two-step sampling process, and so, our model can also detect the minimum of the unknown function, as substantiated by our experimental results.

How does ICEE perform when trained with data collected by random policies?

ICEE can learn a good policy from data collected using random policies. In data collected using a full random policy, the chance of having good behaviors are lower. It means that ICEE will probably need more data and have a longer training time. We will include a small scale experiment to demonstrate this behavior.

Benchmark Implementation Concerns: The study's application of Algorithm Distillation (AD) diverges from the original practice, focusing on cloning a static behavior policy instead of true learning trajectories.

Our proposed method, ICEE, isn't designed to be an in-context behavior cloning algorithm and the experiments conducted do not attempt to establish ICEE complicity as a superior in-context behavior cloning algorithm to AD. Our intent was to demonstrate the advantages of employing in-context exploration-exploitation compared to in-context behavior cloning, using AD as our benchmark for the latter.

An enlightening secondary finding from our paper is that AD can learn to execute in-context policy learning from non-learning trajectory data, provided the episodes are organized according to performance. This outcome can be attributed to the epistemic uncertainty present in the action distribution, which results in naturally occurring in-context exploration-exploitation behavior.

We thank the reviewer for the feedback.

The ICEE algorithm heavily relies on importance sampling and Monte Carlo approximation, which are widely used techniques in traditional RL and not new.

While it is true that Importance Sampling and Monte Carlo approximations are common in probabilistic ML, the implementation and application in the ICEE algorithm are done in a unique and novel manner. To the best of our knowledge, the bias in action distribution learning has not previously been explored extensively in RL via supervised learning literature. Our approach, using an importance sampling-based solution, demonstrates an effective reduction in bias in the action distribution, as evidenced in the experimental section. Therefore, we firmly believe that even though these techniques are common, their use in this context does not detract from the overall contribution of the paper.

The 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?

Our cross-episode return-to-go is computed according to the sequence of episodes in the offline data. However, this does not bind it to a specific episode order. In our experimental setup, episodes were gathered independently and arranged randomly in order of creation. ICEE's ability to learn effectively from randomly ordered episodes indicates that the return-to-go design is indeed robust against episode sequence. The performance disparity observed between AD and AD-sorted provides further evidence that episode order-sensitive methods, like AD, cannot perform well in this scenario.

Our cross-episode return-to-go method generates binary values consistent across tasks. This is a clear advantage over the discounted cumulative reward design in Decision Transformer (DT). Cumulative rewards can fluctuate significantly between tasks. Multi-game Decision Transformer (MGDT) addresses this by training a specific model to predict individual sequence cumulative rewards. During inference, a bias is added to the predicted distribution to use as return-to-go. However, this technique often fails in the problems where expected cumulative rewards are hard to predict such as Bayesian Optimization (BO). In BO, a basic cumulative reward design is the distance between the function value at the current location and the function value's minimum over the sequence. This return-to-go design would not work well for BO, because predicting a function’s minimum is a notoriously hard problem. Our cross-episode return-to-go avoids this issue by generating consistent values across tasks. For instance, in the Bandit setting, the return-to-go is calculated during training by comparing each action's reward with previous ones. During inference, this action probability of the model corresponds to the probability of improving the reward when setting the return-to-go to "true". For BO, our return-to-go equates to the well-known probability of improvement acquisition function in BO literature.

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.

This paper primarily focuses on understanding the epistemic uncertainty in sequence model prediction and developing an in-context exploration exploitation algorithm. In our experiments, ICEE successfully performed well on exploration-exploitation in BO and also proved to be more efficient at solving 2D discrete RL problems than Algorithm Distillation. The paper does not claim any comparisons of ICEE's performance against traditional offline RL algorithms. Moreover, conventional offline RL methodologies would either struggle with policy learning at inference time due to the task family's POMDP setting or need online updates like RL$^2$. Such methods cannot be directly compared to ICEE for the same reason why online and offline RL methods are not comparable.

We thank the reviewer for the feedback.

As pointed below, some of the explanations of the key ideas in the paper are not very clear. 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?

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 at the time of inference by providing the projected return (return-to-go). Given that the return is the result of present and downstream actions, the distribution of action that a 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 distribution, which should be distinguished from the neural network model $p_{\psi}(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t})$. As noted in equation (5), the action's distribution is proportionate to the return's distribution, and this is subsequently weighted by the probability of action from the data collection policy. As it is the posterior distribution derived by the Bayes rule, we denote it as the “true” action posterior distribution.

To put this intuitively, if a model can accurately match equation (5), it will result in the action distribution being skewed towards the data collection policy. For instance, in a pre-recorded sequence, should an action be randomly selected with extremely low probability from the data collection policy, 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 a high probability of $p(R_{k,t}| a_{k,t}, o_{k,t}, H_{k,t})$, the resulting probability of action is weighted by the probability of action in the data collection policy, $\pi_k(a_{k,t}|o_{k,t})$, resulting in a small value. Therefore, even though (5) is the true posterior distribution of action, it is not the desirable action distribution for our model.

Ideally, the action distribution, as 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 to the data collection policy would be eliminated, thereby resolving the mentioned issue.

Between eq.(7) and eq.(8), maybe the author could provide some details on the importance sampling trick they mentioned.

As explained in the previous answer, we would like to learn an action distribution in (6) instead of the action distribution in (5). 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).

To let the model learn the action distribution in (6) instead, we first state the desirable training objective as if the data follow the distribution (6): $L_{\psi} = \sum_{k,t} \int \hat{p}(R_{k,t}, a_{k,t} |o_{k,t}, H_{k,t}) \log p_{\psi}(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t}) dR_{k,t} d a_{k,t}$

Then, we apply the importance sampling trick to introduce the true action posterior in the equation: $L_{\psi} = \sum_{k,t} \int\hat{p}(R_{k,t} |o_{k,t}, H_{k,t}) \Big( \int p(a_{k,t} |R_{k,t}, o_{k,t}, H_{k,t}) \frac{\hat{p}(a_{k,t} |R_{k,t}, o_{k,t}, H_{k,t})}{p(a_{k,t} |R_{k,t}, o_{k,t}, H_{k,t})} \log p_{\psi}(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t}) d a_{k,t} \Big) dR_{k,t}$

After some rearrangement of the equation, we get a much clearer formulation of the objective (the detailed steps can be found in appendix): $\sum_{k,t} \int p(R_{k,t}, a_{k,t} | o_{k,t}, H_{k,t}) \frac{\mathcal{U}(a_{k,t})}{\pi_k(a_{k,t}|o_{k,t})} \log p_{\psi}(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t}) d a_{k,t} dR_{k,t}$

Note the probability distribution in front is now 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 are samples from the data distribution. We get the proposed training objective. $\sum_{k,t} \frac{\mathcal{U}(a_{k,t})}{\pi_k(a_{k,t}|o_{k,t})} \log p_{\psi}(a_{k,t} | R_{k,t}, o_{k,t}, H_{k,t}), \quad R_{k,t}, a_{k,t} \sim p(R_{k,t}, a_{k,t} | o_{k,t}, H_{k,t})$

We thank the reviewer's feedback. We will fix the typos.

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.

We acknowledge that the sequence length required for our in-context policy learning indeed exceeds what needed for the Decision Transformer. However, in contrast to the preceding work, Algorithm distillation, it is already significantly shorter. In applications such as Bayesian optimization, where the sequence typically involves hundreds of evaluations, our proposed algorithm is directly applicable to real-world problems. Thanks to the continuous advancements in large language models (LLMs), the prompt size that can be processed is rapidly increasing. For example, the most recent GPT-4 Turbo can handle 128K tokens. With such improvement, we could envision that in-context policy learning can become applicable to real-world RL problems.

The last statement in page 3 seems somewhat ambiguous, ...

Indeed, the predictive distribution encapsulates both the epistemic and aleatoric uncertainty in the context of limitless data. The quoted sentence denotes that the predictive distribution can capture only the epistemic uncertainty relating to the parameters of sequences, $\theta$, and cannot include the epistemic uncertainty relating to their hyperparameters (if they exist). For example, a sequence model trained using input-output pairs taken from randomly sampled quadratic functions would fail to predict with accuracy the epistemic uncertainty about input-output pairs from a cubic function. Although this appears as model mismatch, if the model is equipped with a polynomial function hyper-prior, a Bayesian inference method can predict accurate epistemic uncertainty. This point was raised to highlight the comparative limitations of epistemic uncertainty of our method in contrast to Bayesian inference approaches. We will add clarity in the manuscript.

There could be several reasonable ways to generate suboptimal data for analysis in toy experiments, ...

We appreciate the insightful suggestion. Indeed, exploring more efficient data collection methods could be beneficial. If the data contains a sufficient proportion of optimal behaviors, the model can learn from these optimal behavior. Included suboptimal data allows the model to manage suboptimal states and learn to improve itself at inference time. A good balance of optimal and suboptimal behaviors can enhance learning efficiency. The exploration of data generation distribution parameters, such as the suggested beta distribution, could lead to more effective learning. We will explore these recommendations in the future.