From Weak Data to Strong Policy: Q-Targets Enable Provable In-Context Reinforcement Learning

Yichen Lin; Xuyuan Xiong; Xue Wang; Xiangfu Meng; Mike Mingcheng Wei; Tao Yao

cs.LGSep 24, 2026

From Weak Data to Strong Policy: Q-Targets Enable Provable In-Context Reinforcement Learning

Authors: Yichen Lin, Xuyuan Xiong, Xue Wang, Xiangfu Meng, Mike Mingcheng Wei, Tao Yao

Organizations: Shanghai Jiao Tong University, Shanghai, China · Alibaba Group · School of Management, University at Buffalo, Buffalo, NY, USA

Abstract

Existing in-context reinforcement learning methods mainly pretrain Transformers with supervised behavior-prediction objectives. This enables task inference from context, but makes the learned policy strongly depend on the quality of offline actions: when trajectories are weak or suboptimal, imitation itself becomes a biased learning signal. We propose Q-Target Pretrained Transformers (QTPT), which keeps the context-conditioned Transformer architecture but replaces behavior cloning with a Bellman-style Q-target objective. QTPT therefore learns to use rewards and transitions in the context to estimate action values, rather than simply imitating the behavior policy. We theoretically analyze QTPT in stochastic linear bandits and finite-horizon MDPs, showing stronger robustness to data quality than supervised pretraining. Empirically, QTPT improves over supervised behavior prediction on controlled RL benchmarks with random or suboptimal data, and we examine extensions to D4RL Kitchen and AntMaze. Supplementary experiments evaluate backbone robustness, meta-RL comparisons, task-coherent context, and unsupported-action value overestimation. These comparisons distinguish the benefits of Q-target pretraining from the remaining limitations of offline coverage.

Figures & tables

Figure 1: Comparison of Q-target Pretrained Transformers (QTPT) and Supervised Pretrained Transformers (SPT) for in-context RL. While SPT requires expert demonstrations for good performance, QTPT maintains robust performance even when trained on suboptimal data.

Panel of Figure 2 — Figure 2: (a) Suboptimality gap over time for QTPT and SPT under random and LinUCB pretraining data, averaged over 1000 runs. (b) Supplementary mixture evaluation: mean gap and one standard deviation across five independent training seeds, evaluated on 1,000 task instances. The horizontal axis is the fraction of expert trajectories. These are finite-horizon results at $T=200$ , not an empirical verification of asymptotic rates or a basis for extrapolating curve crossings.

Variant	Random data	Expert data
QTPT (full)	14.46 $\pm$ 3.27	18.29 $\pm$ 4.06
QTPT-MC	9.33 $\pm$ 3.01 (-35%)	13.77 $\pm$ 3.68 (-25%)
QTPT-NoCtx	0.38 $\pm$ 0.61 (-97%)	0.34 $\pm$ 0.65 (-98%)

Table 1: Controlled ablations on the Darkroom benchmark under matched architectures and offline datasets. Historical submitted evaluation: mean reward with 95% confidence-interval half-widths over 300 evaluation episodes. These intervals measure episode-level variation, not variability across independent training seeds; higher is better.

Panel of Figure 3 — Figure 3: (a) Average cumulative reward in Darkroom. (b) Average cumulative reward in Dark Key-to-door. QTPT is usually stronger under weak data, while extremely sparse expert-only settings remain difficult for TD learning.

Method	Kitchen-partial-v0	Kitchen-mixed-v0	AntMaze-umaze-v2
QTPT (Ours)	13.23 $\pm$ 4.03	14.62 $\pm$ 3.36	35.94 $\pm$ 13.52
SPT	10.84 $\pm$ 4.18	13.25 $\pm$ 2.68	26.96 $\pm$ 11.39
Q-SFT	11.37 $\pm$ 1.14	9.44 $\pm$ 4.86	19.05 $\pm$ 9.21
QDT	12.58 $\pm$ 4.28	7.28 $\pm$ 5.23	20.00 $\pm$ 8.66

Table 2: Main-text D4RL extension. QTPT is compared against matched Transformer baselines under the same revised experimental protocol. Kitchen entries report D4RL normalized score (%), and AntMaze reports success rate (%). Values report active-stage mean

\pm

standard deviation across 3–4 seeds; active-stage performance averages non-zero periodic evaluation checkpoints. This conditional checkpoint statistic excludes failed checkpoints and is descriptive, not the conventional final-policy D4RL evaluation. Higher is better; boldface is omitted because no seed-level significance claim is made.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Random Data Results	Default	Double DQN	Softmax + Hardmax	ReLU-Attn + Softmax-Attn
Results	30.9	34.01	32.48	197.26
Expert Data Results	Default	Double DQN	Softmax + Hardmax	ReLU-Attn + Softmax-Attn
Results	39.4	45.76	53.73	203.74

Appendix

Table 3: Comparison of results for different configurations of QTPT .

Embedding Dimension	16	32	64	128	256
Results	135.7	54.0	40.5	33.3	30.9

Appendix

Table 4: Results for varying embedding dimensions.

Number of Heads	1	2	4	8	16
Results	34.2	30.2	30.9	29.3	31.4

Appendix

Table 5: Results for varying number of attention heads.

Number of Layers	1	2	4	8
Results	122.3	33.6	31.0	30.9

Appendix

Table 6: Results for varying number of layers.

Figure 4: Comparison of Pretrained Transformer Performance Using Q-learning and Supervised Learning in Stationary and Non-stationary Environments. The shaded regions represent the standard deviation of the suboptimality estimates based on 1000 simulation runs.

Figure 5: Accuracy (%) on the GSM8K test set across different mixtures of suboptimal and expert data.

Figure 6: Performance comparison against strong baselines on the Stochastic Linear Bandit task

Method	$G_{\mathrm{OOD}}\downarrow$	Harmful rate $\downarrow$	OOD regret $\downarrow$	Rollout regret $\downarrow$
QTPT	$0.2569\pm 0.1021$	$0.2064\pm 0.0362$	$0.1828\pm 0.0754$	$153.15\pm 58.99$
Naive hard-max QTPT	$0.3351\pm 0.0387$	$0.2588\pm 0.0090$	$0.2856\pm 0.0157$	$211.67\pm 7.81$
QTPT + CQL penalty	$0.2582\pm 0.1021$	$0.2144\pm 0.0390$	$0.1903\pm 0.0675$	$139.86\pm 55.74$
Naive hard-max + CQL penalty	$0.3330\pm 0.0386$	$0.2514\pm 0.0084$	$0.2840\pm 0.0176$	$205.40\pm 10.36$

Appendix

Table 7: Support-controlled diagnostic. Mean

\pm

standard deviation across five paired training/data seeds. Lower is better; harmful rate is a fraction.

Expert mix	QTPT reward $\uparrow$	QTPT + CQL penalty reward $\uparrow$
0.0	$8.3053\pm 8.7520$	$9.3333\pm 8.7971$
0.5	$15.7753\pm 10.7997$	$14.0393\pm 7.8813$
1.0	$6.2247\pm 1.9315$	$8.9653\pm 4.2767$

Appendix

Table 8: Darkroom diagnostic: five-seed held-out reward, mean

\pm

standard deviation, validation-best checkpoints. This protocol is separate from Table 1 .

Method / acquisition	Transitions	Pseudo-regret $\downarrow$	Bayes progress $\uparrow$
QTPT / fixed random corpus	1M	$1.915\pm 0.085$	$85.4\%\pm 5.1\%$
RL 2 / on-policy trials	1M	$2.356\pm 0.112$	$58.9\%\pm 6.7\%$
QTPT / fixed random corpus	5M	$1.784\pm 0.041$	$93.3\%\pm 2.5\%$
RL 2 / on-policy trials	5M	$1.780\pm 0.056$	$93.6\%\pm 3.4\%$

Appendix

Table 9: Bernoulli bandit, mean

\pm

standard deviation over three independent training seeds.

Method	Pseudo-regret $\downarrow$	Mean training time (s)
QTPT	$0.969\pm 0.140$	842.7
PEARL (fixed-offline adaptation)	$2.472\pm 0.074$	1062.7

Appendix

Table 10: Fixed-offline continuous meta-bandit. Mean

\pm

standard deviation over seven training seeds; pseudo-regret is

\sum_{t=1}^{10}(a_{t}-M)^{2}

Backbone	Expert mix	QTPT gap $\downarrow$	SPT gap $\downarrow$
GPT-2	0.0	$33.52\pm 3.85$	$222.20\pm 2.35$
GPT-2	0.5	$34.42\pm 2.10$	$133.41\pm 13.97$
GPT-2	1.0	$35.53\pm 3.19$	$46.92\pm 1.36$
Llama	0.0	$33.18\pm 1.44$	$221.93\pm 2.19$
Llama	0.5	$40.59\pm 2.49$	$141.29\pm 11.15$
Llama	1.0	$38.11\pm 2.07$	$39.88\pm 0.93$

Appendix

Table 11: Cumulative gap at

T=200

, mean

\pm

standard deviation over five training seeds. Mix is expert-trajectory fraction. Lower is better.

Expert mix	QTPT gap $\downarrow$	SPT gap $\downarrow$
0.00	$31.97\pm 1.45$	$220.85\pm 1.24$
0.10	$31.58\pm 0.41$	$210.40\pm 1.24$
0.20	$32.57\pm 0.66$	$196.89\pm 11.14$
0.25	$33.92\pm 1.25$	$173.79\pm 9.33$
0.50	$36.24\pm 1.57$	$139.98\pm 11.59$
0.75	$36.43\pm 1.75$	$78.30\pm 2.56$

Appendix

Table 12: Values used for Figure 2(b) ; mean

\pm

standard deviation over five training seeds per mixture.

Input protocol	Regret $\downarrow$	Success $\uparrow$	Action flip rate
Current-episode history	$0.0033\pm 0.0003$	$0.8808\pm 0.0097$	$0.2036\pm 0.0120$
History + unrelated-task support	$0.0037\pm 0.0002$	$0.8694\pm 0.0096$	$0.2220\pm 0.0074$
History + same-task support	$0.0002\pm 0.0000$	$0.9638\pm 0.0069$	$0.0475\pm 0.0099$

Appendix

Table 13: Diagnostic A: extra-context source, five-seed mean

\pm

standard deviation. Regret is true reward loss against the best grid action; success is the fraction choosing it.

Input protocol	Regret $\downarrow$	Success $\uparrow$	Action flip rate
No history	$0.2386\pm 0.0000$	$0.1172\pm 0.0000$	$0.0000\pm 0.0000$
Current-episode history	$0.0033\pm 0.0003$	$0.8847\pm 0.0100$	$0.1978\pm 0.0203$
Unrelated history (same length)	$0.2386\pm 0.0000$	$0.1172\pm 0.0000$	$0.0000\pm 0.0000$
History + same-task support	$0.0002\pm 0.0001$	$0.9536\pm 0.0112$	$0.0502\pm 0.0052$
Oracle goal token	$0.0001\pm 0.0001$	$0.9547\pm 0.0110$	$0.0000\pm 0.0000$

Appendix

Table 14: Diagnostic B: task relevance versus token count, five-seed mean

\pm

standard deviation. Displayed

0.0000

values are rounded.

In-Context RL

In-context reinforcement learning studies how agents adapt their policies from interaction histories without updating model parameters.

See all 42 papers

From Policy Evaluation to Off-Policy Control: Provable Emergence of In-Context Q-Learning
Oct 13, 2026Yingjun Mo, Yajie Zhao, Fang KongIn-Context RL Transformer-Based RL
Southern University of Science and Technology (SUSTech) · University of Electronic Science and Technology of China (UESTC)
Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning
May 8, 2026Zixuan Xie, Xinyu Liu, Claire Chen +3Reinforcement Learning Transformer-Based RL
University of Virginia · California Institute of Technology · Purdue University
Q-learning Penalized Transformer for Safe Offline Reinforcement Learning
Sep 28, 2026Shengchao Hu, Peng Wang, Jifeng Hu +5Reinforcement Learning Transformer-Based RL
Shanghai Jiao Tong University, China · Shenzhen Campus of Sun Yat-sen University, China · Jilin University, China +1

Oct 13, 2026J/K move · Enter open · S save

From Policy Evaluation to Off-Policy Control: Provable Emergence of In-Context Q-Learning

Yingjun Mo, Yajie Zhao, Fang Kong

Southern University of Science and Technology (SUSTech) · University of Electronic Science and Technology of China (UESTC)

In-Context RL Transformer-Based RL Q-Learning Off-Policy RL

Open paper AskPDF arXiv

Transformers can implement reinforcement-learning algorithms in context, yet how off-policy control emerges within Transformers remains poorly understood. We study Q-learning, a foundational model-free control algorithm that combines off-policy bootstrapping with greedy policy improvement. We prove that a two-layer attention-only Transformer can implement a finite-temperature form of Q-learning, where the temperature

τ>0

controls the approximation to greedy action selection, recovering window Q-learning as

τ\to0

. We train the structured Transformer on randomly generated MDP tasks using output-only supervision. Experiments show that the theoretically predicted parameter structure emerges during training, including the internal greedy action-selection mechanism, and that the trained Transformer achieves performance comparable to Q-learning on unseen tasks. Our theoretical analysis further shows that the equivalent realizations form a solution manifold and that, under suitable task-diversity conditions, the training iterates converge locally to this manifold at an exponential rate. For the random-task model analyzed here, the local worst-case number of gradient steps needed to reduce the distance to the manifold by a factor

ε

Θ\left(τ^{-5}\log(1/\varepsilon)\right)

, quantifying the increasing optimization cost as the learned update approaches window Q-learning.

Abstract

In-Context RL

Explore similar work

From Policy Evaluation to Off-Policy Control: Provable Emergence of In-Context Q-Learning

Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning

Q-learning Penalized Transformer for Safe Offline Reinforcement Learning