In-Context Reinforcement Learning From Suboptimal Historical Data

Authors: Juncheng Dong, Moyang Guo, Ethan X Fang, Zhuoran Yang, Vahid Tarokh
Primary area: reinforcement learning
Keywords: In-context Learning; Transformer; Reinforcement Learning

Abstract:
Large-scale transformer models have achieved remarkable empirical successes, largely due to their in-context learning capabilities. Inspired by this, we explore training an autoregressive transformer for in-context Reinforcement Learning (RL). In this setting, we initially train a transformer on an offline dataset consisting of trajectories collected from various RL instances, and then fix and use this transformer to create an action policy for new RL instances. Notably, we consider the setting where the offline dataset contains trajectories sampled from suboptimal behavioral policies. In this case, standard autoregressive training corresponds to imitation learning and results in suboptimal performance. To address this, we propose the Decision Importance Transformer (DIT), which emulates the actor-critic algorithm in an in-context manner. In particular, we first train a transformer-based value function that estimates the advantage functions of the behavior policies that collected the suboptimal trajectories. Then we train a transformer-based policy via a weighted maximum likelihood estimation loss, where the weights are constructed based on the trained value function to steer the suboptimal policies to the optimal ones. We conduct extensive experiments to test the performance of DIT on both bandit and Markov Decision Process problems. Our results show that DIT achieves superior performance, particularly when the offline dataset contains suboptimal historical data.

Decision: reject
Reviewer scores: 8, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1729217117263

thread_type:
official_review

<thread_object>
creator: reviewer_99qm

<content>
summary:
This paper is positioned in the in-context reinforcement learning literature. The authors proposed to mitigate the burden of using only the optimal action labels by an exponential advantage weighting. As such, trajectories of different quality can be used for update since the exponential weighting can filter out bad actions when they have low advantage values.

strengths:
This paper is well-written and clearly motivated. The authors did a great job in crafting each section so they explained the context, problem and solution clearly.

weaknesses:
However, the paper lacks novelty in my opinion. In the offline RL setting, exponential advantage weighted regression is already very popular (AWAC, AWR) and has served as basis of many more sophisticated methods, see refs [1, 2]. The KL regularized policy optimization  as in proposition 4.1 has also been extensively investigated, the result of proposition 4.2 is fairly standard in the conservative policy iteration (CPI) literature, in which [Schulman et al., 2015] is one of the variants. The literature also includes the offline case, see [3]. So this paper did not propose anything sufficiently novel imo if it simply extended to in-context case with additional conditioning on $\tau$.
I recommend the authors take into account the above-mentioned references and discuss significant theoretical differences the paper has.

I would suggest the authors go deeper in comparing the weighting scheme to enhance the paper. The exponential weighting scheme depends heavily on the estimation quality of advantage function, which itself can be unreliable. The proposed in-context estimation is not very convincing as well in this regard. And it is also due to this issue, AWAC is often found to be underperforming. Many tricks have been proposed to improve the weighting, e.g. by incorporating uncertainty [4], setting the weights of bad actions to zero to eliminate their updates [5,6]. The paper could benefit from such systematic comparison.

References:\
[1] Simple and Scalable Off-Policy Reinforcement Learning\
[2] AWAC: Accelerating Online Reinforcement Learning with Offline Datasets\
[3] Behavior Proximal Policy Optimization\
[4] Uncertainty Weighted Actor-Critic for Offline Reinforcement Learning\
[5] Offline RL with No OOD Actions: In-Sample Learning via Implicit Value Regularization\
[6] Offline Reinforcement Learning via Tsallis Regularization

questions:
Please refer to the above for my questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730255949520

thread_type:
official_review

<thread_object>
creator: reviewer_ffnq

<content>
summary:
The paper proposes the Decision Importance Transformer (DIT), a novel approach in reinforcement learning (RL) that leverages transformer models to improve decision-making using suboptimal historical data. Traditional RL models struggle with data that does not include optimal action labels, but DIT overcomes this by introducing a weighted pretraining method. This method utilizes the advantage function, which estimates the importance of each action and weights the learning process accordingly. To handle situations where advantage functions are not directly available, DIT uses a transformer-based estimator that learns to approximate these advantage values across various tasks, enhancing the model’s ability to generalize. The authors demonstrate DIT’s effectiveness on bandit and Markov Decision Process (MDP) tasks, where DIT often matches or exceeds the performance of existing models, even when trained with suboptimal trajectories.

strengths:
* By integrating a transformer-based in-context advantage estimation, the model approximates the advantage values dynamically when they are not explicitly available. This allows DIT to extend its functionality to environments where advantage functions are challenging to compute or estimate, improving its general applicability in real-world scenarios with limited labeled data.
* The authors conducted extensive testing across bandit and MDP environments, demonstrating DIT’s ability to match or exceed the performance of state-of-the-art methods. The results on tasks with noisy or suboptimal data showcase the model’s robustness and adaptability, emphasizing its potential in practical applications where perfect data is scarce.

weaknesses:
* As the optimal solution of the optimization problem in Eq.4 and the policy improvement theorem under KL constraint have been widely researched, this work seems to me the main contribution of the paper is applying it into the in-context learning RL domain. 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. 
* The weighted pretraining approach heavily relies on accurate advantage values to guide the learning process effectively. DIT uses Eq.9 to optimize the Q and V functions. However,  there is no evidence learning via this objective could avoid the overestimation problem widely exist in offline RL settings. I suggest the authors to compare the adopted training objective of Q and V with the ones with in-sample optimality guarantees such as IQL.  Or at least an illustration of Eq.9 could converge to an optimal solution should be given. 
* I suggest add some experiments between DIT and DPT using the same sub-optimal datasets to illustrate the advantages of DIT further. Current comparison such as Figure 5-7 could confuse the readers.
* The writing could be further improved regrading the notations paragraph.
* Several related works are missed which use RL or chain-of-thought to improve (in-context) LTM for decision making such as [1][2][3]. 

[1] Q-value Regularized Transformer for Offline Reinforcement Learning. ICML 2024.

[2] Rethinking Decision Transformer via Hierarchical Reinforcement Learning. ICML 2024.

[3] In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought. ICML 2024.

questions:
See weakness above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648516042

thread_type:
official_review

<thread_object>
creator: reviewer_wiqi

<content>
summary:
This paper explores the potential of large-scale transformer models for in-context learning in reinforcement learning (RL). The authors present the Decision Importance Transformer (DIT), a new framework for training transformers with suboptimal historical data, diverging from previous approaches that rely on optimal action labels. DIT leverages a weighted maximum likelihood estimation that assigns higher weights to actions with high advantage values, encouraging the model to learn near-optimal policies despite the lack of optimal trajectories in the dataset.

strengths:
Decision Importance Transformer (DIT) presents a novel approach by leveraging suboptimal historical data for in-context reinforcement learning (RL). While previous transformer-based RL approaches like Decision Pretrained Transformer (DPT) rely on datasets labeled with optimal actions, DIT uniquely addresses the challenge of training RL agents with suboptimal data by introducing a weighted maximum likelihood estimation guided by advantage functions. This framework effectively removes the constraint of needing optimal labels, broadening the applicability of transformer models in RL to real-world scenarios where optimal data is scarce or inaccessible. Additionally, the proposal of using an in-context advantage estimator to enhance policy learning from suboptimal data represents a creative blend of autoregressive modeling and advantage-weighted learning.

weaknesses:
* Baseline Comparisons: To contextualize DIT's performance on suboptimal data, the paper could be strengthened by comparing DIT to more recent offline RL methods, such as Conservative Q-Learning (CQL) or other SOTA offline RL methods.
* Add DPT experiments comparison on bandit problems: In the MDP PROBLEMS setting, many experiments are compared with DPT. Why DPT comparisons were not included for bandit problems?

questions:
* How does DIT handle scalability with larger task sets or datasets? What is the computational overhead introduced by the advantage-based weighting mechanism? such as training time or memory usage comparisons.
* 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? Can you add some discussion with these methods in the related works or more experiments comparison? 
* 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?

References：

[1] Liu, Tenglong, et al. "Adaptive Advantage-Guided Policy Regularization for Offline Reinforcement Learning." In International Conference on Machine Learning (ICML). PMLR, 2024.

[2] Hong, Zhang-Wei, et al. "Harnessing mixed offline reinforcement learning datasets via trajectory weighting." arXiv preprint arXiv:2306.13085 (2023).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699653721

thread_type:
official_review

<thread_object>
creator: reviewer_rsr6

<content>
summary:
This paper introduces the Decision Importance Transformer (DIT), a model designed to generalize to new RL tasks through in-context learning, even when pre-training data originates from sub-optimal behavior policies. DIT accomplishes this by reweighting an offline dataset using advantage functions estimated by a large-scale transformer model (LTM) that is trained autoregressively. Additionally, the authors demonstrate that the exponential reweighting technique provably ensures policy improvement. Experimental results in bandit, planning, and continuous-control tasks show that DIT outperforms baselines when generalizing to new RL instances, even when trained on data generated by sub-optimal policies.

strengths:
- **Significance**: Generalizing to new RL tasks via in-context learning is a major concern in RL, particularly for handling sub-optimal datasets in offline RL. DIT addresses both of these challenges within a unified framework.
  
- **Novelty**: The proposed DIT combines multiple methods—including LTMs, exponential reweighting, and actor-critic approaches—in a cohesive design. To the best of my knowledge, no similar approach exists in current literature.

- **Clarity**: The paper is well-organized and self-contained, presenting DIT's architecture and methodology progressively, making it accessible and providing insights to guide future work.

weaknesses:
**Suggestions for Improvement**:
1. Provide more direct empirical evidence that task information $\tau$ is being captured through in-context learning. For instance, demonstrate recovery of $\tau$ in experiments by querying with some special states.
2. The literature review is missing policy regularization methods for offline RL [1].

**Minor Issues**:
1. If $\pi(a | s ; \tau)$ and $\pi_\tau(a | s)$ refer to the same concept, choose one notation or define one in terms of the other for clarity.
2. Proposition 4.2 is labeled as "In-Context Policy Improvement." Since $\tau$ is provided explicitly, this theorem seems to justify the design of the loss function rather than "in-context learning." Renaming it to better reflect its role might avoid potential misinterpretation.

[1] Fujimoto, Scott, and Shixiang Shane Gu. "A minimalist approach to offline reinforcement learning." *Advances in Neural Information Processing Systems*, 34 (2021): 20132-20145.

questions:
1. Is DIT pre-trained in an autoregressive way or an “encoder” way, i.e., does the output of the LTMs contain only a single action token instead of a sequence of predictions?
2. In a fair comparison—such as using DIT with optimal historical data—would DIT still outperform baseline methods like DPT? If so, what do you think is the advantage of DIT over DPT when both are using optimal historical data?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732478217388

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Clarification about Contributions

comment:
Dear Reviewers,

We sincerely appreciate your help and constructive feedbacks. However, we have noticed some misunderstandings regarding our contributions. To address these concerns, we would like to provide clarification on the following points:

- Distinctions between In-context RL (ICRL) and regular RL 
- Challenges of ICRL from Suboptimal Historical Data (our main contribution)
- Technical Novelty of Our Method (our technical contribution)

---

### Distinctions between ICRL and regular RL

For standard (either online or offline) RL problems, **the goal is to learn a policy for a single RL task**, using training data collected from that task. In contrast, the goal for ICRL is much more challenging: with training data from diverse RL tasks, **ICRL aims to learn an algorithm (implemented by a transformer) for solving distinct RL tasks**. At inference time, ICRL directly applies the pretrained transformer model to new/unseen RL tasks which are unknown to the transformer, without updating the transformer’s parameters. Please see Figure 1 in https://arxiv.org/pdf/2210.14215  for a good schematic illustration. Hence, **the pretrained transformer for ICRL must generalize to new RL tasks during inference time**, conditioning on the provided context containing offline trajectories collected from these new RL tasks.  

***MAB Example***. To concretely illustrate the distinctions, consider the multi-armed bandit (MAB) problem. **For a standard RL problem, we consider a single MAB problem, and the goal is to identify the optimal bandit for this MAB problem** with a dataset of observed rewards for all bandits. For an ICRL problem, given pretraining data collected from various MAB problems with distinct reward distributions, **the goal of ICRL is to quickly identify the optimal bandit of an unseen MAB problem** whose reward distributions are different to those in the pretraining dataset. 

---
### Challenges of ICRL from Suboptimal Historical Data

Existing works for ICRL all assume access to data collected by optimal policies from diverse RL tasks, which is a strong requirement difficult to hold in practice. To this end, our work considers ICRL from **only suboptimal data**. There are two main challenges towards solving this problem:

- **Challenge 1**: How to generalize to new tasks unknown to the pretrained transformer model.
- **Challenge 2**: How to address the distribution shift so that we can use suboptimal data to infer optimal policies for distinct RL tasks.

To address Challenge 1 above, we use the in-context learning ability of transformers, which have been widely recognized and empirically proven. 

To address Challenge 2 above, we propose to use Advantage Weighted Regression (AWR). **Notably, although AWR has been studied in standard RL, it remains unclear how to generalize this approach to ICRL which, as discussed above, is a distinct and much more challenging problem than standard RL**. The main reason is that, unlike the weight function for standard RL with a single task, **the weight function for ICRL needs to be task dependent**. 

---
### Technical Novelty

Inspired by the AWR for standard RL, we also propose to use the advantage function as the weight function for ICRL. 

***Technical Challenge.*** However, because the weight function is task dependent, **we must estimate the advantage functions for all distinct RL tasks in the  pretraining dataset**. Moreover, because we don’t know the training trajectories’ source tasks and thus cannot combine the trajectories from the same RL tasks to improve estimation,  **we need to estimate the advantage functions individually for each trajectory in the pretraining dataset**. Hence, compared to AWR for standard RL where we have a significant amount of trajectories collected from the same RL task, the advantage function estimation problem for ICRL poses a formidable challenge. 

***Technical Novelty.*** Our technical novelty here is to learn the task-dependent weight function, that is the advantage functions for diverse RL tasks, with another in-context learning problem – in context value estimation. Specifically, **the proposed in-context advantage estimator can correctly estimate the task-dependent weight function, conditioned only on the pretraining trajectories without any knowledge about their source RL tasks.** This is also empirically verified in our extra experiment whose 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**.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents the Decision Importance Transformer (DIT), a model designed for in-context reinforcement learning (RL) using suboptimal historical data. DIT addresses the challenge of training RL agents with suboptimal data by using a weighted maximum likelihood estimation, which assigns higher weights to actions with high advantage values. The authors demonstrate that DIT achieves superior performance compared to other methods, particularly when the offline dataset contains suboptimal historical data. The approach is evaluated on bandit and Markov Decision Process problems. The core idea of DIT is that by emphasizing good actions during supervised pretraining, a transformer-based policy can perform well even with suboptimal data. The authors also highlight that their approach is an effective way for in-context value function estimation.

Reviewers highlighted its novelty in addressing in-context RL with suboptimal data, noting its significance for real-world applications. They praised the effective methodology that combines transformer models, exponential reweighting, and actor-critic approaches, resulting in a cohesive design. A key strength is the in-context advantage estimator, which dynamically approximates advantage values, enhancing policy learning from suboptimal data. Reviewers also noted the strong empirical results in bandit and MDP environments, showcasing DIT's robustness. Finally, they found the paper to be well-organized and clearly presented

The reviewers raised several concerns and questions about the paper. Reviewers questioned whether the advantage function estimation was reliable and if the method could avoid overestimation issues common in offline RL. Additionally, they requested more direct evidence that task information is being captured through in-context learning. 

This area chair is familiar with the ICRL, has read the paper, and offers this recommendation. The paper needs to be self-sufficient and introduce the topic to a reader that might not be fully familiar with the topic. Unfortunately, this paper fails short of doing that. The next revision of the paper should make the presentation more accessible, including some examples on ICRL in the introduction would go a long way. For example, it is not clear what authors mean by different RL instances, or what they consider actions in this context -- the part that is confusing is the mention of LLMs. Is the DIT the foundation model being trained, or is an additional model that data weighting? A system diagram depicting DIT would be very helpful. The author's should also place the work in context of the relevant literature (e.g. https://arxiv.org/pdf/2404.11018 which address ICRL with suboptimal data). Lastly, the paper claims the methodology for large transformer models. However, all evaluations are done on GPT-2. From there I am inferring that DIT is an additional model that guides the next round of in-context examples. It appears that Appendix H clarifies this point -- please make it sooner in the paper.

Overall, the paper has some interesting ideas, although its presentation should be significantly improved.