In-Context Reinforcement Learning From Suboptimal Historical Data

Thank you for your thoughtful and constructive feedback. We hope our responses below address your concerns.

Experiments

Bandit (Figure 2). Our bandit experiments are designed to be comprehensive and aligned with established evaluation standards. As bandit problems are well-understood, we follow the evaluation protocol used in DPT [1], comparing DIT against strong baselines such as the known optimal algorithms Thompson Sampling and UCB/LCB, as well as DPT itself — which serves as an oracle for DIT due to its access to additional optimal bandit information during pretraining. These comparisons provide a rigorous evaluation of DIT in this setting.

MDPs (Figures 3(a) & 4(a)). The goal of these experiments is to evaluate adaptation speed — how quickly each method improves reward on new tasks. While DIT has not fully converged to its optimal performance in Figure 3, the performances of all baselines have plateaued, and DIT already surpasses them in terms of performance. We believe this clearly demonstrates the superior adaptability of DIT.

To provide additional empirical insights, we follow your advice to continue evaluation until the performance of DIT fully converges. As DIT’s performance in Figure 4(a) (Miniworld) has already converged, we focus on Darkroom and we run 10 additional episodes to see the final converged reward, it turns out that DIT will continue advance the score a bit then keep stable (around 50), compared to 40 at episode 40, expanding its performance gain over baseline methods.

Method

Normalizing Constant and Policy Reinforce Baselines. We prove in Appendix D.3 that, when using the exponentiated advantage function for reweighting, the normalizing constant (partition function) can be safely ignored and set to 1.

To further illustrate the effectiveness of DIT, we follow the reviewer’s suggestion to conduct an extra ablation study to compare DIT to several policy-reinforced baselines with different reweighting schemes. We evaluated three weight alternatives:

1. Cumulative reward
2. Exponentiated cumulative reward
3. Batch-normalized DIT, where advantage values are first estimated and then normalized across all trajectories in the pretraining dataset.

We consider offline testing from expert trajectory on Meta-world and summarize the key results in the following table.

We observe that DIT significantly outperforms all baselines. In particular, batch-normalized DIT performs worse than standard DIT. This aligns with our theoretical analysis, which highlights the importance of using trajectory-specific weights. Normalizing across the entire dataset violates this principle, which likely degrades performance.

We will include these results and a discussion of the policy reinforcement methods into our experiment section to further improve the quality of our work. We thank the reviewer for this constructive suggestion.

Training Objective (on-policy value function)

Thank you for this insightful observation. We agree that the value function equals the expected Q-function only under on-policy rollouts.

Indeed, in our case, DIT performs pretraining using trajectories generated by behavioral policies, and the learned value functions are used to weight the actions taken by those same behavioral policies. Thus, all training rollouts are on-policy with respect to the value functions being used, and there are no off-policy discrepancies in the objective.

We believe this is one key strength of DIT and will clarify this point in the revised version. Additionally, we appreciate the note about the time indices in the Bellman equations and will correct them accordingly.

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

Thank you for your constructive comments. Please see our responses below to address your concerns.

References

We appreciate these shared references. We will include a discussion of AD-\epsilon, AWR, and AMAGO-2 into our manuscript. In particular, we highlight that AD-\epsilon still requires sampling actions from the good policies to perform noise distillation. Moreover, to create noise distilled dataset, they need to actively sample trajectories following a noise schedule. In comparison, our work assumes a setting of only historical data, which is more realistic and easier to satisfy.

Theoretical Results

The authors deeply appreciate the reviewer’s detection of this typo, and we will correct it in our updated manuscript. Yes, the expectation in Equation (3) should be with respect to the meta-policy $\pi(a|s\tau)$ and the expectation in Equation (4) is with respect to the behavioral policy $\pi^b_{\tau}(a|s)$: we use data collected by the behavioral $\pi^b_{\tau}(a|s)$ to learn a meta-policy $\pi(a|s;\tau)$.

Experiment Setup

Comprehensive Benchmarks. Our goal is to provide a more comprehensive empirical evaluation than prior works, such as DPT (which focuses on bandits and navigation tasks like Dark Room and Mini-world) and prompt-DT (which focuses solely on continuous control). To that end, we include representative benchmarks from all three domains: bandits, navigation, and continuous control. While we agree that Key-to-Door is a valuable benchmark, it shares significant structural similarities with tasks like Dark Room and MiniWorld. As such, we believe it would offer limited additional insight beyond the settings already covered. Due to the time constraints of the rebuttal phase, we are unable to include new results at this stage. However, we are more than happy to run additional experiments on Key-to-Door and include them in the final manuscript to further strengthen the empirical evaluation.

Meta-World Benchmark. We follow the same environment setting of Prompt-DT[1] to use 20 tasks from the Meta-World reach-v2 suit (specifically ML1-pick-place-v2), while we choose the task horizon to be 100. We choose the converged SAC policy as the optimal policy. The cumulative reward of a random policy is around 1.5 within a horizon of 100 and around 20 for the optimal policy. We will update the manuscript to make this clear.

Extra Experiments

To provide additional insights, we conduct extra experiments to compare DIT with several policy reinforce [2] baselines where the reweighting is based on cumulative reward rather than the advantage function. Specifically, we consider reweighting with cumulative rewards and exponentiated cumulative rewards for offline testing from expert trajectory on Meta-world and summarize the key results in the following table.

The significantly improved performance of DIT over these two baselines further proves the effectiveness of our proposed method.

Performance of DPT.

The difference in DPT's performance between our work and the original paper stems from the nature of the pretraining datasets. In the original DPT paper, pretraining trajectories were collected using uniformly random policies, which ensured broad coverage of the state-action space. In contrast, as noted in our manuscript, we use suboptimal policies (achieving only ~30% of optimal performance) to collect pretraining data. This choice naturally reduces coverage and impacts DPT’s performance.

However, we believe it better reflects realistic scenarios—especially in practical applications where historical data are typically collected by suboptimal or heuristic policies, not uniformly random ones.

To ensure a fair comparison, both DPT and our method (DIT) are trained on the same set of suboptimal trajectories. However, we provide DPT with additional privileged information: a set of randomly sampled out-of-trajectory query states with corresponding optimal action labels sampled from optimal policies. This positions DPT as a strong oracle-style upper bound relative to DIT, further underscoring the strength of DIT’s performance under more constrained assumptions.

[1] Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., & Gan, C. (2022, June). Prompting decision transformer for few-shot policy generalization. In international conference on machine learning (pp. 24631-24645). PMLR.

[2] Williams, R. J. (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8, 229-256.

Thank you for your helpful and constructive comments. We will incorporate the referenced work into the final manuscript. In addition, we will include citations supporting the availability of abundant historical data in real-world settings (e.g., using [1] as a standard reference). Please see below for our responses to your other concerns.

Why DPT and ICRL Work

As discussed in the original DPT paper, on a high level, DPT internally conducts posterior sampling (PS) for MDP, where the transformer model is pretrained to infer the target MDP from the given context dataset and to take action according to the optimal policy for the inferred target MDP. As the context size increases, the inferred MDP becomes closer to the true target MDP, thus leading to improved performance. Meanwhile, there are also theoretical works establishing the efficacy of the supervised pretraining approach taken by DPT, e.g., [2].

Prediction without a context dataset. The DPT model can still predict an action using only the query state as input (without a context dataset), as the prediction is only based on the token corresponding to the query state.

Method

Choice of hyperparameter. As shown in our theoretical result, $\eta$ represents the penalty for the KL divergence. Thus, if the historical data is less suboptimal (high quality), it should be high as we would like to stay close to the behavior policies; on the contrary, if the data is low quality, it should be chosen as a small value so that we can improve more over the behavior policies.

Value transformer model structure. Our design follows the in-context learning setup, where transformers are trained for regression tasks $y = f(x)$ using sequences like $x_1, y_1, ..., x_n, y_n$. In our case, value estimations (for V and Q) can be seen as regression problems, e.g., set $x = s$ and $y = G$ (cumulative reward) when estimating V. Based on this, we design our value transformer as shown in Figure 6. Figure 7 provides empirical evidence that it can accurately estimate value functions in context.

Cost of training in-context advantage estimator. As detailed in the manuscript, we train them with standard supervised training objectives for transformers, which are stable and straightforward to optimize. While this moderately increases the total pretraining time, we emphasize that the primary bottleneck in ICRL lies in collecting high-quality pretraining data, not in training time. To this end, DIT reduces this data requirement considerably, trading off for a modest increase in training complexity. We believe this tradeoff is worthwhile and offers strong practical benefits, making ICRL more accessible and scalable in real-world applications.

Experiments

Historical Data Suboptimality. In our experiments, we use behavioral policies with performance 30%-50% of the optimal policies to collect historical data, a commonly used setting for offline RL. We would like to clarify that we use the intermediate training checkpoints of SAC to collect trajectories so that they are indeed suboptimal. We will update the manuscript to make this clear in the introduction to avoid confusion.

Experiment Setup. For Meta-world, we follow the setting of Prompt-DT to use Meta-World reach v2. We use Mini-world mainly for ablation study to understand the importance of the proposed reweighting mechanism and the effect of lack of optimal action labels. Given this purpose, we compare DIT with DPT and a variation of DIT without reweighting.

Performance of DPT. The observed performance difference of DPT between this work and the original paper is due to the difference of the pretraining datasets. Due to space constraint, please refer to our response to Reviewer Gczk for more details.

DIT sometimes outperforms DPT

We appreciate the reviewer’s insightful observation. This is indeed a compelling research direction.

At a high level, DPT uses a single randomly sampled query state paired with an optimal action label for each trajectory during pretraining, whereas DIT leverages multiple reweighted suboptimal action labels. Interestingly, these many reweighted suboptimal labels can collectively provide a stronger learning signal than a single optimal label, resulting in better pretraining objectives in some settings. Of course, if DPT were given access to many optimal action labels per trajectory, it would likely surpass DIT, as DIT does not use any optimal supervision.

[1] Levine, S., Kumar, A., Tucker, G., & Fu, J. (2020). Offline reinforcement learning: Tutorial, review. and Perspectives on Open Problems, 5.

[2] Lin, Licong, Yu Bai, and Song Mei. "Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining." arXiv preprint arXiv:2310.08566 (2023).

Thank you for your insightful comments. We hope our response can address your concerns.

Presentation

Style Improvement. We will follow the reviewer’s suggestion to remove most of the colored text, keeping only minimal phrases with key importance.

Clarification on the term “In-Context.” We appreciate the reviewer highlighting this potential ambiguity. While “in-context learning” originated in the LLM literature, the term has recently been adopted by several works in reinforcement learning to describe similar ideas—namely, models that adapt to new tasks by conditioning on sequences of past transitions, without gradient updates. Our use of “in-context RL” follows this growing convention. Nonetheless, we agree that the dual usage can be confusing. In the revised manuscript, we will explicitly distinguish in-context RL from ICL in language models to ensure clarity and avoid potential misunderstandings.

Methods.

Zero-shot Adaptation. DIT builds on the supervised pretraining framework of DPT and can act as an online meta-policy in new environments. Specifically, it can predict actions without conditioning on a context dataset, as the prediction is only based on the token corresponding to the query state. In the zero-shot setting—i.e., with no context tokens—DIT and DPT default to the behavior learned during pretraining, effectively leveraging a strong prior. Cold-start performance is demonstrated in the online testing experiments (Figures 2, 3, and 4). These results show that DIT quickly improves its performance in-context with more information (trajectories), adapting efficiently even in the cold-start setting.

Learned Advantage Function and Distribution Shift

We appreciate the reviewer’s insightful comment. We would like to clarify that DIT does not require robustness to distribution shift during the pretraining.

This is because the learned advantage functions are only applied to actions taken by the same behavioral policies that generated the trajectories—meaning all computations remain in-distribution with respect to those behavioral policies.

This design ensures stability and avoids the challenges typically associated with off-policy corrections. We consider this to be a key strength of DIT, and we will make this point more explicit in the revised manuscript.

Clarification of Contributions

While similar theoretical results have been established in the standard RL setting, our contribution lies in extending these insights to the in-context reinforcement learning (ICRL) framework. Although theory is not the primary focus of this work, we include the analysis to better motivate and ground the design of DIT, helping the reader understand why the method is effective.

Our core insight is that if "relatively good" actions can be identified and emphasized during supervised pretraining, a transformer-based policy can match the performance of meta-policies trained with significantly more expensive data, such as DPT. Although individual weighted action labels may be noisy, our analysis and experiments show that the weighted MLE objective—when averaged over diverse environments—can yield high-quality meta-policies capable of generalizing to unseen tasks.

Furthermore, with the proposed in-context advantage estimator, we observe that transformer models can learn to generate reliable value function estimates across tasks and behavioral policies using only supervised learning. We believe these empirical insights are important to share, as they highlight the practicality and potential of ICRL.

On the practical side, DIT is simple to implement and significantly improves the feasibility of in-context RL, as suboptimal trajectories are much easier to collect in real-world systems, especially in industry settings where large amounts of historical data are already available. This opens the door for broader adoption of ICRL methods in practical applications.