In-Context Reinforcement Learning From Suboptimal Historical Data
In-context Learning; Transformer; Reinforcement Learning
Yes
No ethics review needed.
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?
- **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.
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.
**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.
Yes
No ethics review needed.
* 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).
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.
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.
* 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?
Yes
No ethics review needed.
See weakness above.
* 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.
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.
* 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.
Yes
No ethics review needed.
Please refer to the above for my questions.
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.
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.
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
We consider pretraining transformers for in-context RL from only suboptimal historical data.
@inproceedings{
anonymous2024incontext,
title={In-Context Reinforcement Learning From Suboptimal Historical Data},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=FXJm5r17Q7},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
In-context Learning; Transformer; Reinforcement Learning
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/a13be0d510531d5e2958475cb526104afc3689bc.pdf
reinforcement learning
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/3eae5ed5ce9bd4f6cc86e128849bcb690753894e.zip
In-Context Reinforcement Learning From Suboptimal Historical Data
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission