In-Context Reinforcement Learning From Suboptimal Historical Data

Dear Reviewer 99QM,

We appreciate your comments and hope that our response below can help address your concerns.

Novelty.

We observe some misunderstanding regarding our contribution and in-context RL. Please see our general response for a detailed discussion about our contributions, particularly why extending AWR to in-context RL is technically novel.

Extra Experiments

To further corroborate the effectiveness of our approach (other than the significantly improved performance with reweighted pretraining), we have followed a suggestion from Reviewer fFnq to conduct an extra experiment to empirically prove that our method can accurately learn the value functions for distinct behavioral policies and RL tasks so that advantage-based reweighting can be instantiated. It can be observed in Figure 8 in Appendix I that the prediction error of the trained in-context advantage estimator is very small.

Comparisons with Advanced Reweighting Techniques

We would like to note that our main contribution is an effective approach for in-context value function estimation and to prove that even the most straightforward reweighting technique enables in-context RL with only suboptimal data whereas all existing works for in-context RL requires access to data generated by the optimal policies. While we believe that all the referred more advanced weighing mechanisms for standard RL can be applied to further improve our method’s performance, they are tangential to our key contribution and thus should not be considered as evidence to criticize the novelty of our work.

Insights and Value of DIT for the Community

Our key insight is that as long as the “relatively good” actions can be identified and emphasized during supervised pretraining, a transformer-based policy can perform as good as those LTM policies pretrained with much more expensive pretraining data such as DPT.

In particular, although the weighted action labels for one trajectory can be misleading, the weighted MLE objective averaged over trajectories across diverse environments does lead to high-quality LTM policies with strong in-context RL capability to generalize to unseen RL tasks. Moreover, with our in-context advantage estimator, we also observe that LTMs can generate reliable in-context estimations for value functions of distinct behavioral policies and RL tasks with only supervised learning and a single trajectory as context. We believe these insights and observations are important to share with the community so that the potential of in-context RL can be fully investigated.

On the practical aspect, our method is easy to implement and significantly increases the feasibility of in-context RL as sub-optimal trajectories are much easier to collect (large companies often have tremendous databases consisting of historical trajectories collected by non-expert users), creating opportunities for in-context RL to be applied in much broader applications.

Hi, thanks for the clarification and extra experiments, they are appreciated.

I understand the authors have contributed to in-context RL by taking inspirations from standard RL literature. But the current paper lacks many important pieces to make the methodology convincing. Let me rephrase my questions:

1. as the authors said, key to the proposed method is the advantage function because they control the weight of update. But for this critical problem, as the authors put it, was only given the space around Eq. 8 stating that "estimation by $D^i$ is unreliable and we propose to estimate it by two models". Why is this the case? I didn't see any qualitatively new methods have been put forward to resolve this unreliability issue nor theoretical insights of estimating this critical quantity. It is appreciated the authors have added Figure 8 to show it indeed worked, but why, when, and how it worked, whether the standard advantage value could have worked, these questions remain unanswered.

2. because the authors are tackling a more challenging problem than solving a fixed task, it is reasonable to expect that the advantage functions are more difficult to estimate. A known issue of exp weighting is that even for actions with very bad advantage value, it still gives positive weights because the exp function is positive everywhere. Since the proposed agent learns in an offline manner, the agent can only see a limited range of samples. So even if a state-action pair is bad, its likelihood is still increased and others unseen pairs implicitly decreased. How did the authors deal with this issue?

3. regarding theoretical novelty. I do not agree with the authors' claims that this extension brings new insights. By looking at the proof of Theorem 4.1 and 4.2, it is quite clear that nothing significantly new was used. Take theorem 4.2 for example: the authors started by using a lemma from [Achiam et al. 2017] followed by Pinksker's inequality and Jensen's inequality. Up to this point the techniques were very common and general-purpose in the CPI, TRPO types of algorithms; only the last step was really relevant to this in-context setting where an expectation was taken. Let me ask the authors: how is it qualitatively different from the CPI-based algorithms like the ones I mentioned?

Dear Reviewer fFnq,

Thanks for your constructive comments and we will include the referred works in our literature review section. Please see our response to your other comments below.

Thus I suggest the authors to conclude the contribution more from the aspect that what's the new insight of combing in-context RL and advantage weighted regression.

This is a valuable suggestion. Please see the general response for a detailed clarification of our contributions.

In terms of insights, our key insight is that as long as the “relatively good” actions can be identified and emphasized during supervised pretraining, a transformer-based policy can perform as good as those LTM policies pretrained with much more expensive pretraining data such as DPT.

In particular, although the weighted action labels for one trajectory can be misleading, the weighted MLE objective averaged over trajectories across diverse environments does lead to high-quality LTM policies with strong in-context RL capability to generalize to unseen RL tasks. Moreover, with our in-context advantage estimator, we also observe that LTMs can generate reliable in-context estimations for value functions of distinct behavioral policies and RL tasks with only supervised learning and a single trajectory as context. We believe these insights and observations are important to share with the community so that the potential of in-context RL can be fully investigated.

Notation Paragraph

We have followed your suggestion to improve our notation paragraph in Section 4.

I suggest add some experiments of DPT using the same sub-optimal datasets

We appreciate this insightful suggestion. Indeed, in all experiments, we compare DIT with a baseline BC, which is an unweighted version of our method DIT. Given that DIT and DPT use the same transformer model architecture, BC is equivalent to DPT trained with sub-optimal data. Hence, DIT’s effectiveness over DPT under sub-optimal data is clearly demonstrated as DIT significantly outperforms BC.

Overestimation Problem of the Value Functions: However, there is no evidence learning via this objective could avoid the overestimation problem widely exist in offline RL settings.

We would like to note that DIT relies on the behavioral policies’ value functions, instead of the optimal value functions. To this end, the overestimation problem commonly observed in offline RL is much less a concern for DIT because

1. we don’t need to evaluate a unseen policy
2. we don’t need to estimate a greedy policy

The reason for point 1 is that we use the learned value functions of the behavioral policies to generate weights for actions taken by the behavioral policies. In other words, we are always in-distribution with respect to the behavioral policies. The reason for point 2 is that we are only estimating the value functions for the behavioral policies, not for a greedy policy or optimal policy. Thus, we don’t need to bootstrap with a max operator which is the main cause of overestimation problems.

Extra Experiments for Validation. To present further evidence that the value functions are estimated well, we conduct an extra experiment to test the error of learned transformer-based value functions. The results are presented in Figure 8 in Appendix I. It can be observed that the prediction error of the trained in-context advantage estimator is very small.

Lastly, we would like to note that the ultimate criterion to evaluate the effectiveness of the learned value functions should be the in-context RL performance on unseen RL tasks. To this end, DIT significantly outperforms its unweighted variation, corroborating the effectiveness of the learned in-context advantage function estimator.

Dear Reviewer Wiqi,

We appreciate these valuable comments and hope that the response below can help clarify your concerns.

This DIT makes actions with high advantage values receive more weight, leading to guaranteed policy improvements over the behavior policies. It is very similar to some offline RL methods, such as A2PR[1], AW/RW[2]. How DIT specifically differs from or improves upon these methods?

Please see our general response about a detailed clarification of our contributions.

Baselines of In-context RL and Offline RL: To contextualize DIT's performance on suboptimal data, the paper could be strengthened by comparing DIT to more recent offline RL methods

At a high level, in-context RL is similar to a meta-RL problem using transformers. In the testing/inference stage, a new task (unknown to the transformer) is sampled from the family of desired tasks, and the pretrained transformer model is deployed to assess the performance. Thus, direct comparison with offline RL methods may not be appropriate for contextualizing the performance of in-context RL algorithms.

To comprehensively evaluate the performance of our method, we compare with several SOTA in-context RL methods. Moreover, we also compare with DPT pretrained with the optimal action labels, which is not accessible to our method DIT. Thus, DPT can be considered as an oracle upper bound for DIT. Notably, although without any optimal action labels, the performance of DIT matches that of DPT on most problem instances. This corroborates the effectiveness of our method.

Extra Experiments: Inclusion of DPT in the bandit experiments

Thank you for this very constructive comment. For the bandit problems, our primary goal was to compare with the theoretically optimal bandit algorithms such as the Thompson Sampling.

Following your suggestion, we have conducted experiments to include the performance of DPT. Please see Figure 3 for the updated results. As expected, DIT is slightly outperformed by DPT since DPT has access to optimal bandits during pretraining whereas DIT do not. However, please note that DIT still outperforms the theoretically optimal bandit algorithms, and the regret curves of DIT and DPT demonstrate similar trends. These results prove that the effectiveness of DIT.

Computational Overhead: How does DIT handle scalability with larger task sets or datasets? What is the computational overhead introduced by the advantage-based weighting mechanism?

The DIT framework can be decomposed into 3 steps:

• Step (1): training the in-context advantage function estimator;
• Step (2): labeling all the state-action pairs with the trained advantage estimator;
• Step (3): weighted supervised pretraining.

The extra computational overhead of DIT comes from Step (1) and Step (2). In particular, the overhead of Step (1) is about the same as Step (3). Step (2) is computationally simple and its cost is negligible compared to Step (1) and (3). Thus, the overall pretraining computational cost is about 2x the unweighted pretraining method.

For specific computation time, our model uses exactly the same architecture as DPT and we include this part in Appendix F. Inference time for our model is also minimal.

We would like to also note that for methods based on historical data, the computational cost is often not the main concern. Thus, given the significantly improved performance, the doubled pretraining time is very often acceptable.

For different tasks, how many does DIT train the in-context task identification transformer, the Q and Value models? Is there one for all tasks?

To facilitate in-context advantage function estimation, we train only one pair of Q and V transformers for all tasks. When combined, they estimate the advantage functions for all tasks.

Dear Reviewer Rsr6,

Thank you for these valuable comments! Following your advice, we have updated the name of Proposition 4.2 and included the referred work in our literature review. Please see below our response to your other questions.

If $\pi(a|s;\tau)$ and $\pi_\tau(a|s)$ refer to the same concept$

We use $\pi(a|s;\tau)$ to refer to some meta policy that takes the task identity $\tau$ as input to generate distinct policies for varying tasks. On the other hand, we use $\pi_\tau(a|s)$, e.g., $\pi^b_\tau(a|s)$, to represent some policy that is fixed for task $\tau$.

We appreciate this question and have highlighted this in our manuscript to avoid confusion.

Provide more direct evidence that task information is being captured

Conceptually, the reason why DPT and DIT can construct task information in the inference time is that transformer models have the in-context learning ability — that is, these models can solve an unseen task during inference time, by only looking at a few demonstration examples. There are a few prior works showing that the key to the success of in-context learning is that transformers construct representations of the task during inference time, based on the demonstration examples. Such representations are named “task vectors” or “function vectors” in these works Function Vectors in Large Language Models [1] and In-Context Learning Creates Task Vectors [2].

Following the same rationale, we expect that DPT/DIT uses the in-context learning ability to infer the desired task during inference time. And inside these models, representations of the task are constructed during inference time, based on the trajectories sampled from the new task during inference time.

An Extra Experiment as Evidence. To validate this hypothesis, we conduct an extra experiment where, during inference time, we do not condition on the offline trajectories sampled from the task of interest. Instead, we draw offline trajectories from a different task, and use these trajectories as ``demonstration examples” in the prompt and test the performance of DPT/DIT on the task of interest. We compare this experiment setting with our experiment where the offline trajectories are indeed sampled from the desired task.

If our hypothesis that DPT/DIT extracts task-relevant information, we anticipate that, conditioning on trajectories from a different task, we will see a huge performance drop. The reason is that the offline trajectories mislead DPT/DIT in terms of which task the model is applied to. This is shown in Figure 9 in Appendix J.

From Figure 9, we can see that the extracted task information is crucial for the success of DPT/DIT: their performance degrades significantly if the task information is misleading. This shows that DPT/DIT’s transformer models heavily rely on the extracted task information and, given the significant performance boost when offline trajectory is sampled from the desired task, they indeed learn useful task information.

Pretraining of DIT

We follow the pretraining procedure of DPT to pretrain DIT in an autoregressive way where the output actions only depend on the prior inputs. This is achieved by using a standard GPT2 transformer architecture with casual attention masks.

Both DIT/DPT pretrained with optimal data.

We have the results for both DIT and DPT training on optimal historical data. In this way, both methods perform very well and there is no discernible difference. This observation is expected because the goal of DIT is to recover the missing information about the optimal action labels and to close the performance gap with DPT pretrained with optimal action labels. With optimal historical data, there is no need for this goal, and given that DPT and DIT use the same transformer architecture, their performance should also be very similar.

We choose to not include these results given that the contribution of this work is in-context RL with only sub-optimal data, which is an important research problem since in most real-world environments, optimal data is not commonly accessible.

---

[1] Todd, E., Li, M., Sharma, A. S., Mueller, A., Wallace, B. C., & Bau, D. (2024). Function Vectors in Large Language Models. The Twelfth International Conference on Learning Representations. https://openreview.net/forum?id=AwyxtyMwaG

[2] Hendel, R., Geva, M., & Globerson, A. (2023). In-Context Learning Creates Task Vectors. The 2023 Conference on Empirical Methods in Natural Language Processing. https://openreview.net/forum?id=QYvFUlF19n

Dear Reviewer Rsr6,

These are great points for discussion, and we appreciate them.

How do you think DIT could balance the trade-off between exploration and exploitation during online deployment?

From a theoretical perspective, the supervised pretraining employed by DPT/DIT is shown to be equivalent to implicit posterior sampling [1]. That is to say, conditioned on the context consisting of trajectories collected from a desired RL task, transformer implicitly builds a posterior distribution over the tasks (with the distribution of pretraining RL tasks as the prior distribution), sample a task from the posterior distribution, and use the optimal policy for the sampled task to take action. See, for example, Appendix C.1 of [1] for more details. This posterior sampling is a generalization of the Thompson Sampling algorithm, a theoretical optimal solution for the multi-armed bandit problems. In particular, posterior sampling is proved to be sample-efficient with online Bayesian regret guarantees [2]. Thus, the supervised pretraining framework employed by DPT/DIT allows for strong Exploration vs. Exploitation capabilities.

Is the same distribution used during evaluation?

Yes, we follow the conventional setting of in-context RL to evaluate on in-distribution yet unseen RL tasks.

how does DIT maintain its performance when evaluated under different distributions?

This is a great question. From the perspective of posterior sampling, testing on out-of-distribution RL tasks results in a wrong prior for the postering distribution. Thus, similar to almost all Bayesian methods, it would require more samples / trajectories to identify the true task and act optimally. While algorithms with a strong out-of-distribution performance is an important research question, this work primarily considers the standard in-distribution setting as in-context RL is still in its early stages. With that being said, we will include this important discussion into our discussion of weakness. Thank you again for this insightful point.

---

[1] Lee, Jonathan, et al. "Supervised pretraining can learn in-context reinforcement learning." Advances in Neural Information Processing Systems 36 (2024).

[2] Osband, Ian, Daniel Russo, and Benjamin Van Roy. "(More) efficient reinforcement learning via posterior sampling." Advances in Neural Information Processing Systems 26 (2013).