Paper: In-context Exploration-Exploitation for Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The experiments predominantly utilize noisy optimal policy variants for data collection, leaving an open question regarding ICEE's performance when relying on trajectories from random policies.

Author response:
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.

Reviewer comment (Experiments):
How does ICEE perform when trained with data collected by random policies?

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
> 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.

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

Author response:
> 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.

Reviewer comment (Experiments):
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?

Author response:
> 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.

Reviewer comment (Experiments):
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.

Author response:
> 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.

Reviewer comment (Writing):
Nit/typo: "GP biased" should be "GP based".

Author response:
We thank the reviewer's feedback. We will fix the typos.

Reviewer comment (Writing):
Nit/typo: On page 3, "indefinite" should be "infinite".

Author response:
We thank the reviewer's feedback. We will fix the typos.

Reviewer comment (Writing):
Nit/typo: In Section 6, "wildly" should be "widely".

Author response:
We thank the reviewer's feedback. We will fix the typos.

Reviewer comment (Theory):
The last statement on page 3 seems somewhat ambiguous and needs clarification. Specifically, _"...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 uncertainty in the limit of infinite data?

Author response:
> 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.

Reviewer comment (Experiments):
There are several reasonable ways to generate suboptimal data for analysis in toy experiments. The authors examined 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 demonstrations. If that is the case, a more effective stress test might be more meaningful; for example, as a minimal change, one could sample $\epsilon$ more generally as a beta distribution with two parameters and rerun the experiments across a range of those parameters.

Author response:
> 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.