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

Yingjun Mo; Yajie Zhao; Fang Kong

cs.LGOct 13, 2026

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

Authors: Yingjun Mo, Yajie Zhao, Fang Kong

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

Abstract

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 $ε$ is $Θ\left(τ^{-5}\log(1/\varepsilon)\right)$ , quantifying the increasing optimization cost as the learned update approaches window Q-learning.

Figures & tables

Panel of Figure 1 — Figure 1: Final learned $\mathbf{P}_{1}$ , $\mathbf{V}_{1}$ , $\mathbf{P}_{2}$ and $\mathbf{V}_{2}$ matrices after training at $\tau=0.2$ , $B=32$ , and $n=20$ . the non-zero structure matches the theoretical constructions $(\mathbf{P}_{1},\mathbf{V}_{1},\mathbf{P}_{2},\mathbf{V}_{2})$ in ( 9 ) and ( 11 ).

Figure 2: Left: Performance of Transformers on unseen tabular MDPs (sixteen valid parameter blocks). Right: The learned scalar projection

\beta_{\mathrm{att}}

and its target

1/\tau

(six support parameter blocks).

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning and default
$S,m,d$	State count $S=9$ , action count $m=4$ , Q-feature dimension $d=36$ .
$\mathbf{w},{\boldsymbol{\theta}}$	Q parameter $\mathbf{w}\in\mathbb{R}^{d}$ ; distinct Transformer parameter vector ${\boldsymbol{\theta}}$ .
$\tau$	Teacher temperature.
$\alpha,\gamma,\epsilon,\eta$	Q-update step $0.2$ , discount $0.5$ , behavior exploration $0.1$ , Adam learning rate.
$B,U$	Independent contexts per optimizer update; number of optimizer updates.
$n$	Transitions within one context, default $20$ ; not minibatch size $B$ .

Appendix

Table 1: Notation. A task is one independently sampled MDP/feature realization; a context is a valid trajectory window generated in that task. In the training, one task supplies one context.

Tasks ( $10^{6}$ )	Soft error (%)	Hard error (%)	Greedy tr (%)	Greedy ctx (%)
0.32	$9.0157\pm 0.6508$	$9.0465\pm 0.6716$	$95.314\pm 0.102$	$56.587\pm 0.776$
5.12	$7.1287\pm 0.2225$	$7.1792\pm 0.2118$	$96.708\pm 0.157$	$67.205\pm 0.895$
10.24	$4.4694\pm 0.6952$	$4.5328\pm 0.6750$	$98.246\pm 0.087$	$80.769\pm 1.045$
20.48	$0.0657\pm 0.0002$	$0.6655\pm 0.0019$	$99.920\pm 0.005$	$98.840\pm 0.082$
25.60	$0.0625\pm 0.0002$	$0.6633\pm 0.0020$	$99.927\pm 0.003$	$98.936\pm 0.035$
37.76	$0.7393\pm 0.0557$	$0.9834\pm 0.0444$	$99.654\pm 0.012$	$96.003\pm 0.197$

Appendix

Table 2: Dense held-out contexts: update errors and greedy agreement

Tasks ( $10^{6}$ )	Soft error (%)	Hard error (%)	Greedy tr (%)	Greedy ctx (%)
0.32	$55.7678\pm 4.5493$	$57.4641\pm 4.8244$	$95.752\pm 0.151$	$66.636\pm 0.899$
5.12	$44.0258\pm 1.9607$	$45.6216\pm 1.9971$	$97.076\pm 0.085$	$75.928\pm 0.833$
10.24	$27.8956\pm 4.6804$	$29.2407\pm 4.7820$	$98.406\pm 0.067$	$86.035\pm 0.483$
20.48	$0.4427\pm 0.0066$	$5.1372\pm 0.0071$	$99.985\pm 0.002$	$99.834\pm 0.023$
25.60	$0.4181\pm 0.0073$	$5.1137\pm 0.0097$	$99.987\pm 0.003$	$99.875\pm 0.022$
37.76	$4.5953\pm 0.3353$	$6.7774\pm 0.2200$	$99.681\pm 0.012$	$97.070\pm 0.077$

Appendix

Table 3: One-hot held-out contexts: update errors and greedy agreement

Tasks ( $10^{6}$ )	$E_{A}$	$\beta_{\rm att}$	$E_{A,\rm shape}$
0.32	$0.729464\pm 0.000924$	$1.353730\pm 0.004596$	$0.017515\pm 0.000474$
5.12	$0.435038\pm 0.001273$	$2.828350\pm 0.006350$	$0.024805\pm 0.000613$
10.24	$0.303830\pm 0.000436$	$3.483203\pm 0.002233$	$0.016895\pm 0.000379$
20.48	$0.081884\pm 0.000192$	$4.590583\pm 0.000960$	$0.000274\pm 0.000009$
25.60	$0.077982\pm 0.000189$	$4.610089\pm 0.000944$	$0.000148\pm 0.000007$
37.76	$0.023377\pm 0.000188$	$4.885535\pm 0.001009$	$0.004731\pm 0.000092$

Appendix

Table 4: Attention residual and scale across training.

Figure 3: Terminal detail of all 16 initialized blocks for seed 0. Each block has its own color range, shared across seeds for that block. This view complements the full matrices; it does not replace structural zeros by cropped or normalized values in the full-matrix figures.

Method	Mean $J(\pi)$	95% interval
Transformer	0.479225	$[0.461365,0.497690]$
Soft window-Q	0.480136	$[0.462090,0.498943]$
Hard window-Q	0.465381	$[0.447874,0.483684]$
Online Q, constant	0.449844	$[0.426708,0.473394]$
Online Q, visit count	0.522897	$[0.502016,0.544255]$
Oracle	0.529490	$[0.508685,0.550915]$

Appendix

Table 5: Exact greedy-policy return after 20,000 interactions on 512 unseen MDPs.

Figure 4: Frozen dense-feature-trained Transformers on unseen tabular MDPs and measures exact policy returns under a common interaction budget.

$\tau$	Soft error (%)	Hard error (%)	Teacher bias (%)	Greedy (%)
1	$0.00283\pm 0.00107$	$10.03718\pm 0.00020$	10.0370	99.9985
0.8	$0.00261\pm 0.00071$	$7.114048\pm 0.000049$	7.1140	99.9974
0.5	$0.00276\pm 0.00078$	$3.21972\pm 0.00013$	3.2198	99.9962
0.2	$0.00393\pm 0.00019$	$0.63950\pm 0.00014$	0.6385	99.9990
0.1	$0.07575\pm 0.00025$	$0.25509\pm 0.00040$	0.1960	99.9707

Appendix

Table 6: Update accuracy on 8,192 common dense-feature contexts.

$B$	Soft error (%)	Hard error (%)	Greedy ( $\%$ )	$\widehat{\tau}_{\rm eff}$
1	$0.0202\pm 0.0008$	$0.6348\pm 0.0007$	99.994	0.200113
2	$0.0135\pm 0.0013$	$0.6347\pm 0.0004$	99.995	0.200128
4	$0.0110\pm 0.0014$	$0.6349\pm 0.0004$	99.995	0.200161
8	$0.0077\pm 0.0009$	$0.6351\pm 0.0002$	99.999	0.200223
16	$0.0052\pm 0.0005$	$0.6352\pm 0.0003$	99.999	0.200299
32	$0.0041\pm 0.0003$	$0.6355\pm 0.0001$	100.000	0.200373

Appendix

Table 7: Final checkpoints on 8,192 new shared contexts.

Condition	$N_{task}$	Soft error (%)	Hard error (%)	Greedy (%)
Constant Adam	10,240,000	$3.242\pm 0.732$	$3.332\pm 0.717$	$98.100\pm 0.142$
Uniform cosine	10,240,000	$0.04028\pm 0.00419$	$0.64309\pm 0.00076$	$99.980\pm 0.006$
Low-rate extension	20,480,000	$0.00415\pm 0.00023$	$0.63428\pm 0.00012$	$100.000\pm 0.000$

Appendix

Table 8: Results of Soft error, Hard error, and Greedy agreement with three learning rate schedules.

Panel of Figure 5 — Figure 5: Final learned $\mathbf{P}_{1}$ , $\mathbf{V}_{1}$ , $\mathbf{P}_{2}$ and $\mathbf{V}_{2}$ matrices after training at $\tau=0.2$ , $B=32$ , and $n=20$ . the non-zero structure matches the theoretical constructions $(\mathbf{P}_{1},\mathbf{V}_{1},\mathbf{P}_{2},\mathbf{V}_{2})$ in ( 9 ) and ( 11 ).

Panel of Figure 6 — Figure 6: All six actual trained blocks, seed 0. $\mathbf{A}_{\boldsymbol{\theta}},\mathbf{L},\mathbf{C}$ are $36\times 36$ ; $\mathbf{B}_{-},\mathbf{B}_{+}$ are $36\times 37$ ; $\mathbf{B}_{r}$ is $1\times 37$ . The fitted scalar is $\widehat{c}=0.967893$ . Dense $\mathbf{L}$ and $\mathbf{B}_{+}$ are compatible with a nontrivial invertible gauge and are not replaced by identity-shaped pictures.

Panel of Figure 7 — Figure 7: Normalized relations for the same model. The first four targets are $I$ , the reward-row target is $[1,0,\ldots,0]$ , and the two bias targets are zero. These are display-only diagnostics, not raw weights. Common limits are used for each named relation across seeds and methods, including the failed control.

Figure 8: Parameter recovery versus cumulative fresh tasks at

B=32

n=20

, and

\tau=0.2

. Left: attention residual

E_{A}

. Right: the aggregate of the six remaining TD and bias residuals,

E_{\rm TD}

$\tau$	$1/\tau$	$\beta_{att}$	$E_{A}$	$E_{\rm TD}$	$E_{\rm all}$
1	1	1.000000	$3.808\times 10^{-5}$	$5.410\times 10^{-5}$	$(6.624\pm 2.102)\times 10^{-5}$
0.8	1.25	1.250000	$3.150\times 10^{-5}$	$4.741\times 10^{-5}$	$(5.715\pm 1.049)\times 10^{-5}$
0.5	2	2.000002	$2.424\times 10^{-5}$	$4.882\times 10^{-5}$	$(5.466\pm 1.323)\times 10^{-5}$
0.2	5	4.992236	$1.553\times 10^{-3}$	$6.384\times 10^{-5}$	$(1.554\pm 0.017)\times 10^{-3}$
0.1	10	8.174396	$0.1826$	$6.590\times 10^{-4}$	$0.18256\pm 0.00020$

Appendix

Table 9: Final parameter recovery at 20.48 million tasks,

B=32

, and

n=20

. Entries are five-seed means;

E_{\rm all}

includes a pointwise 95% interval half-width. Residuals are dimensionless. The scalar projection

\beta_{att}

is compared with

1/\tau

without removing nonscalar distortion.

Figure 9: Temperature-dependent recovery at

B=32

and

n=20

. Left: attention residual versus cumulative tasks, with five paired seeds per temperature. Right: the learned scalar projection

\beta_{att}

and its target

1/\tau

. The target is shown only for evaluation; no matrix target is supplied during training.

Figure 10: Batch-size effects under a fixed task budget. Left: attention scale and traceless-shape residuals, whose squares sum to

E_{A}^{2}

for each model. Right: aggregate residuals at four checkpoints of the same training runs. Intermediate curves are not independently retrained short-budget experiments.

$n_{\rm train}$	Transitions (M)	$E_{A}$	$E_{\rm TD}$	$E_{\rm all}$	95% interval half-width
1	20.48	$0.0128$	$2.033\times 10^{-4}$	$0.0128$	$1.969\times 10^{-4}$
5	102.4	$2.155\times 10^{-3}$	$8.059\times 10^{-5}$	$2.156\times 10^{-3}$	$8.661\times 10^{-5}$
10	204.8	$1.699\times 10^{-3}$	$6.440\times 10^{-5}$	$1.700\times 10^{-3}$	$5.381\times 10^{-5}$
20	409.6	$1.553\times 10^{-3}$	$6.384\times 10^{-5}$	$1.554\times 10^{-3}$	$1.728\times 10^{-5}$
40	819.2	$1.739\times 10^{-3}$	$6.393\times 10^{-5}$	$1.740\times 10^{-3}$	$5.290\times 10^{-5}$
80	1638.4	$2.019\times 10^{-3}$	$7.050\times 10^{-5}$	$2.020\times 10^{-3}$	$3.846\times 10^{-5}$

Appendix

Table 10: Training-length comparison at

\tau=0.2

and

B=32

. Every model receives 20.48 million tasks and 640,000 updates; supplied transitions differ. Residuals are five-seed means, with the 95% interval half-width for

E_{\rm all}

in the final column. The default is shared across sweeps.

$n_{\rm train}$	Dense matched (%)	Dense common 20 (%)	Tabular matched (%)
1	$0.01893\pm 0.00032$	$0.01573\pm 0.00025$	$0.08467\pm 0.00095$
5	$0.00561\pm 0.00039$	$0.00525\pm 0.00038$	$0.01558\pm 0.00077$
10	$0.00431\pm 0.00043$	$0.00417\pm 0.00042$	$0.01284\pm 0.00049$
20	$0.00393\pm 0.00019$	$0.00393\pm 0.00019$	$0.01222\pm 0.00036$
40	$0.00367\pm 0.00020$	$0.00384\pm 0.00021$	$0.01413\pm 0.00060$
80	$0.00380\pm 0.00012$	$0.00411\pm 0.00012$	$0.01679\pm 0.00044$

Appendix

Table 11: Length-specific models on matched-length and common-

n_{\rm test}=20

dense tests, and matched-length one-hot tests. Entries are update NRMSE percentages with pointwise 95% seed-interval half-widths. The tabular column measures feature transfer without one-hot training.

Figure 11: Training context length at

\tau=0.2

and

B=32

, with 20.48 million tasks per model. Left: final parameter residuals. Right: update errors on matched-length and common-length dense-feature tests.

Figure 12: Online Q-learning is included in the comparisons; the variant with decaying visit-count learning rate yields the best performance.

$\tau$	Transformer [95% CI]	Hard Q (window)	Soft Q (window)	Oracle	Random
1	0.5036 [0.4835, 0.5237]	0.4740	0.5036	0.5374	0.0033
0.8	0.5033 [0.4829, 0.5227]	0.4740	0.5025	0.5374	0.0033
0.5	0.5003 [0.4804, 0.5198]	0.4740	0.4996	0.5374	0.0033
0.2	0.4880 [0.4686, 0.5074]	0.4740	0.4861	0.5374	0.0033
0.1	0.4833 [0.4644, 0.5022]	0.4740	0.4795	0.5374	0.0033

Appendix

Table 12: Final returns on 512 unseen tabular tasks:

n_{\rm test}=20

, zero Q initialization.

Panel of Figure 13 — Figure 13: Tabular transfer at $\tau=1$ , $\tau=0.8$ , $\tau=0.5$ , and $\tau=0.1$ , with $n_{\rm test}=20$ , zero Q initialization, and 20,000 interactions. Error bars denote pointwise 95% intervals. Transformer weights are frozen throughout. All experiments use the same 512 tasks and 20,000-interaction budget.

Figure 14: Separating training-length and deployment-length effects at

\tau=0.2

. Left: models trained at

n_{\rm train}=20

and deployed with different windows. Right: models trained at different lengths and all deployed at

n_{\rm test}=20

. Both evaluations use the same task bank and 20,000-interaction budget.

In-Context RL

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

See all 42 papers

May 7, 2026

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

We investigate the ability of transformers to perform in-context reinforcement learning (ICRL), where a model must infer and execute learning algorithms from trajectory data without parameter updates. We show that a linear self-attention transformer block can provably implement policy-improvement methods, including semi-gradient SARSA and actor-critic, via explicit parameter constructions. Beyond existence, we design a teacher-mimicking training procedure, analyze its gradient-flow dynamics, and establish the first convergence guarantee in the ICRL literature: under suitable richness conditions on the training MDP distribution, gradient flow converges locally and exponentially to an optimal parameter manifold corresponding to the desired RL update. Empirically, training transformers on randomly generated tabular MDPs confirms these predictions: the learned models recover the parameter structure of our explicit constructions and, when deployed on unseen MDPs, deliver strong in-context control performance. Together, these results illuminate how transformer architectures internalize and execute classical reinforcement learning algorithms in context, bridging mechanistic understanding and training dynamics in ICRL.

Haodong Liang, Lifeng Lai

Department of Electrical and Computer Engineering, University of California, Davis

Reinforcement Learning Transformer-Based RL Policy Optimization In-Context RL Policy Learning

Ask PDF

Sep 24, 2026

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

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.

Yichen Lin, Xuyuan Xiong, Xue Wang +3

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

In-Context RL Transformer-Based RL Offline RL Temporal-Difference Learning

Ask PDF

Mar 10, 2026

An Optimal Control Approach To Transformer Training

In this paper, we develop a rigorous optimal control-theoretic approach to Transformer training that respects key structural constraints such as (i) realized-input-independence during execution, (ii) the ensemble control nature of the problem, and (iii) positional dependence. We model the Transformer architecture as a discrete-time controlled particle system with shared actions, exhibiting noise-free McKean-Vlasov dynamics. While the resulting dynamics is not Markovian, we show that lifting it to probability measures produces a fully-observed Markov decision process (MDP). Positional encodings are incorporated into the state space to preserve the sequence order under lifting. Using the dynamic programming principle, we establish the existence of globally optimal policies under mild assumptions of compactness. We further prove that closed-loop policies in the lifted is equivalent to an initial-distribution dependent open-loop policy, which are realized-input-independent and compatible with standard Transformer training. To train a Transformer, we propose a triply quantized training procedure for the lifted MDP by quantizing the state space, the space of probability measures, and the action space, and show that any optimal policy for the triply quantized model is near-optimal for the original training problem. Finally, we establish stability and empirical consistency properties of the lifted model by showing that the value function is continuous with respect to the perturbations of the initial empirical measures and convergence of policies as the data size increases. This approach provides a globally optimal and robust alternative to gradient-based training without requiring smoothness or convexity.

Kağan Akman, Naci Saldı, Serdar Yüksel

Department of Mathematics, Bilkent University, Ankara, 06800, Turkey · Department of Mathematics and Statistics, Queen’s University, Kingston K7L 3N6, ON, Canada

Transformer Transformer Training Dynamics Optimal Control

Ask PDF

Abstract

In-Context RL

Explore similar work

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

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

An Optimal Control Approach To Transformer Training