In-context reinforcement learning (ICRL) exploits the impressive in-context learning ability of large language models to zero-shot generalize to new, unseen tasks at inference time without any parameter updates. Among many approaches, the Decision-Pretrained Transformer (DPT) is particularly appealing for its theoretical grounding: under idealized assumptions such as a sufficiently expressive model and infinite pretraining data, DPT can perform Bayesian posterior sampling. Whether this characterization survives in practice remains open. Therefore, we ask: is ICRL, specifically DPT and our proposed DPT variants, a Bayesian posterior sampler in practice? To approach this question, we establish two necessary conditions that a Bayesian posterior sampler must satisfy: sufficient-statistic invariance and the martingale property. We then construct statistical hypothesis tests to statistically assess if practically-trained ICRL models satisfy these conditions. Applying these tests in Gaussian, Bernoulli, and linear bandit settings, we find that both conditions are consistently violated, suggesting that ICRL may not behave as Bayesian posterior samplers in practice. Beyond this negative result, our tests provide a principled toolkit for testing the Bayesian behavior of ICRL.
Figures & tables
Figure 1: DPT versus Thompson Sampling in Gaussian bandits: cumulative regret and total variation (TV) distance between action distributions.
Linear bandits
Offline
Online
n
Sort
Replace
Sort
Replace
25
1.56×10−18
7.60×10−7
4.08×10−22
2.31×10−4
50
1.76×10−49
1.65×10−12
2.95×10−59
1.31×10−8
100
4.88×10−76
1.21×10−21
1.05×10−87
4.79×10−26
150
1.81×10−67
1.74×10−42
8.36×10−69
7.27×10−33
Table 1: p -values for sufficient-statistic invariance tests on DPT in Gaussian, Bernoulli, and linear bandits ( N=5000 , α=0.05 ). Here, n denotes the context length.
Gaussian
Bernoulli
n
m
Offline
Online
Offline
Online
50
100
<10−300
<10−300
<10−300
<10−300
200
<10−300
<10−300
<10−300
<10−300
300
<10−300
<10−300
<10−300
<10−300
400
<10−300
<10−300
<10−300
<10−300
450
<10−300
<10−300
<10−300
<10−300
Table 2: p -values for martingale tests on DPT ′ in Gaussian, Bernoulli, and linear bandits ( N=5000 , α=0.05 ). Here, n and m denote the starting-context and continuation lengths, respectively.
Gaussian
Bernoulli
n
m
Offline
Online
Offline
Online
50
100
6.85×10−110
2.89×10−248
2.02×10−5
0.099
200
1.33×10−192
<10−300
0.004
0.014
300
1.20×10−61
2.04×10−157
3.09×10−4
0.028
400
1.90×10−53
9.77×10−121
0.002
0.042
450
5.02×10−57
8.08×10−246
0.001
0.013
Table 3: p -values for martingale tests on GDPT in Gaussian, Bernoulli, and linear bandits ( N=5000 , α=0.05 ). Here, n and m denote the starting-context and continuation lengths, respectively.
Appendix figures & tables31 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Gaussian
Bernoulli
Linear
Environment and data
Number of arms K
5
5
10
Feature dimension d
—
—
2
Environment prior
μa∼U[0,1]
pa∼Beta(1,1)
θ∼N(0,I2/2)
Reward distribution
N(μa,0.09)
Bernoulli(pa)
N(xa⊤θ,0.09)
Context horizon H
500
500
200
Appendix
Table 4: Environment, model, and statistical test hyperparameters for DPT, DPT ′ , and GDPT.
Figure 2: Performance comparison of DPT and Thompson Sampling in Gaussian bandits under online and offline deployment.
Figure 3: Performance comparison of DPT and Thompson Sampling in Bernoulli bandits under online and offline deployment.
Figure 4: Performance comparison of DPT and Thompson Sampling in linear bandits under online and offline deployment.
Linear
Offline
Online
n
Sort
Replace
Sort
Replace
25
1.96×10−102
5.71×10−28
1.15×10−86
4.19×10−29
50
6.30×10−96
2.33×10−27
7.27×10−93
2.22×10−28
100
9.44×10−108
9.99×10−18
3.07×10−99
1.10×10−21
150
1.78×10−93
3.22×10−16
8.76×10−84
3.62×10−14
Appendix
Table 5: p -values for sufficient-statistic invariance tests on DPT ′ in Gaussian, Bernoulli, and linear bandits ( N=5000 , α=0.05 ). Here, n denotes the context length.
Figure 5: Performance comparison of DPT ′ and Thompson Sampling in Gaussian bandits under online and offline deployment.
Figure 6: Per-arm martingale increments for DPT ′ in Gaussian bandits with offline starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 7: Per-arm martingale increments for DPT ′ in Gaussian bandits with online starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 8: Performance comparison of DPT ′ and Thompson Sampling in Bernoulli bandits under online and offline deployment.
Figure 9: Per-arm martingale increments for DPT ′ in Bernoulli bandits with offline starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 10: Per-arm martingale increments for DPT ′ in Bernoulli bandits with online starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 11: Performance comparison of DPT ′ and Thompson Sampling in linear bandits under online and offline deployment.
Figure 12: Per-arm martingale increments for DPT ′ in linear bandits with offline starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 13: Per-arm martingale increments for DPT ′ in linear bandits with online starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Linear
Offline
Online
n
Sort
Replace
Sort
Replace
25
1.37×10−107
3.64×10−36
6.02×10−91
3.81×10−36
50
5.32×10−91
1.83×10−41
1.22×10−81
1.34×10−32
100
1.99×10−49
2.37×10−51
2.73×10−42
1.91×10−39
150
2.60×10−31
2.88×10−49
3.15×10−29
9.43×10−45
Appendix
Table 6: p -values for sufficient-statistic invariance tests on GDPT in Gaussian, Bernoulli, and linear bandits ( N=5000 , α=0.05 ). Here, n denotes the context length.
Figure 14: Performance comparison of GDPT and Thompson Sampling in Gaussian bandits under online and offline deployment.
Figure 15: Per-arm martingale increments for GDPT in Gaussian bandits with offline starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 16: Per-arm martingale increments for GDPT in Gaussian bandits with online starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 17: Performance comparison of GDPT and Thompson Sampling in Bernoulli bandits under online and offline deployment.
Figure 18: Per-arm martingale increments for GDPT in Bernoulli bandits with offline starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 19: Per-arm martingale increments for GDPT in Bernoulli bandits with online starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 20: Performance comparison of GDPT and Thompson Sampling in linear bandits under online and offline deployment.
Figure 21: Per-arm martingale increments for GDPT in linear bandits with offline starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 22: Per-arm martingale increments for GDPT in linear bandits with online starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Gaussian
Offline
Online
n
Sort
Replace
Sort
Replace
50
Equal
6.65×10−36
Equal
6.60×10−90
100
Equal
2.79×10−94
Equal
2.54×10−62
200
Equal
5.38×10−217
Equal
1.46×10−46
300
Equal
8.01×10−305
Equal
2.11×10−44
Appendix
Table 7: p -values for sufficient-statistic invariance tests on permutation-invariant DPT in Gaussian bandits ( N=5000 , α=0.05 ). “Equal” denotes identical action probabilities and no p -value is reported.
Figure 23: Performance comparison of permutation-invariant DPT and Thompson Sampling in Gaussian bandits under online and offline deployment.
Gaussian
Offline
Online
n
Sort
Replace
Sort
Replace
50
Equal
2.46×10−233
Equal
<10−300
100
Equal
1.09×10−293
Equal
<10−300
200
Equal
<10−300
Equal
<10−300
300
Equal
<10−300
Equal
<10−300
Appendix
Table 8: p -values for sufficient-statistic invariance tests on permutation-invariant DPT ′ in Gaussian bandits ( N=5000 , α=0.05 ). “Equal” denotes identical action probabilities and no p -value is reported.
Gaussian
n
m
Offline
Online
50
100
<10−300
<10−300
200
<10−300
<10−300
300
<10−300
<10−300
400
<10−300
<10−300
450
<10−300
<10−300
Appendix
Table 9: p -values for martingale tests on permutation-invariant DPT ′ in Gaussian bandits ( N=5000 , α=0.05 ). Here, n and m denote the starting-context and continuation lengths, respectively.
Figure 24: Performance comparison of permutation-invariant DPT ′ and Thompson Sampling in Gaussian bandits under online and offline deployment.
Figure 25: Per-arm martingale increments for permutation-invariant DPT ′ in Gaussian bandits with offline starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
Figure 26: Per-arm martingale increments for permutation-invariant DPT ′ in Gaussian bandits with online starting contexts, where m and t denote the lengths of the martingale continuation and the starting contexts, respectively.
In-context reinforcement learning (ICRL) refers to the ability of RL agents to adapt to new tasks at inference time without parameter updates by conditioning on additional context. Recent empirical studies further demonstrate that Chain-of-Thought (CoT) generation can amplify this ICRL capability. This paper is the first to provide a theoretical understanding on how CoT interacts with ICRL. We conduct our analysis in a policy evaluation setup with linear Transformer. We prove that with specific Transformer parameters, the CoT generation process is equivalent to repeatedly executing temporal difference learning updates. Additionally, we provide finite sample convergence analysis showing that the policy evaluation error decreases geometrically with CoT length and eventually saturates at a statistical floor determined by the context length. We also prove that the desired Transformer parameters are a global minimizer of the pretraining loss, providing a theoretical understanding on the empirical emergence of those parameters.
In-Context Reinforcement Learning (ICRL) enables Large Language Models (LLMs) to learn online from external rewards directly within the context window. However, a central challenge in ICRL is reward estimation, as models typically lack access to ground-truths during inference. To address this limitation, we propose Test-Time Rethinking for In-Context Reinforcement Learning (TR-ICRL), a novel ICRL framework designed for both reasoning and knowledge-intensive tasks. TR-ICRL operates by first retrieving the most relevant instances from an unlabeled evaluation set for a given query. During each ICRL iteration, LLM generates a set of candidate answers for every retrieved instance. Next, a pseudo-label is derived from this set through majority voting. This label then serves as a proxy to give reward messages and generate formative feedbacks, guiding LLM through iterative refinement. In the end, this synthesized contextual information is integrated with the original query to form a comprehensive prompt, with the answer determining through a final round of majority voting. TR-ICRL is evaluated on mainstream reasoning and knowledge-intensive tasks, where it demonstrates significant performance gains. Remarkably, TR-ICRL improves Qwen2.5-7B by 21.23% on average on MedQA and even 137.59% on AIME2024. Extensive ablation studies and analyses further validate the effectiveness and robustness of our approach. Our code is available at https://github.com/pangpang-xuan/TR_ICRL.
Wenxuan Jiang, Yuxin Zuo, Zijian Zhang +8
The Hong Kong Polytechnic University · Meituan-M17 · Institute of Computing Technology, Chinese Academy of Sciences +3
In-context reinforcement learning (ICRL) enables large language model agents to adapt to new environments using their interaction history without updating model parameters. However, repeatedly conditioning on growing histories can lead to substantial token cost. We propose a bounded-history context-management framework that predicts the task-dependent downstream effect of removing historical interactions to guide history selection and determine a decision-dependent context budget. Formally, our framework uses the full rolling history as a reference. The predictor evaluates removal effects, defines a deletion ordering, and applies a shared selection criterion to determine how much history to retain at each decision. We evaluate the method in closed-loop SUMO driving under held-out in-distribution, unseen-domain, and unseen-route settings, and in ScienceWorld under a continual ICRL protocol. Relative to a baseline using the full context, our method reduces total token usage by 25.7%, 25.8%, and 23.2% across the three driving settings while maintaining comparable closed-loop driving performance. In ScienceWorld, it reduces total token usage by 52.1% compared to full context and uses 30.2% and 37.8% fewer tokens than the Recent and Similarity baselines, respectively, while maintaining performance.
Ruihan A. Li, Shangtong Zhang, Rohan Chandra
University of Illinois Urbana-Champaign · University of Virginia