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

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement
May 7, 2026Haodong Liang, Lifeng LaiReinforcement Learning Transformer-Based RL
Department of Electrical and Computer Engineering, University of California, Davis
From Weak Data to Strong Policy: Q-Targets Enable Provable In-Context Reinforcement Learning
Sep 24, 2026Yichen Lin, Xuyuan Xiong, Xue Wang +3In-Context RL Transformer-Based RL
Shanghai Jiao Tong University, Shanghai, China · Alibaba Group · School of Management, University at Buffalo, Buffalo, NY, USA
An Optimal Control Approach To Transformer Training
Mar 10, 2026Kağan Akman, Naci Saldı, Serdar YükselTransformer Transformer Training Dynamics
Department of Mathematics, Bilkent University, Ankara, 06800, Turkey · Department of Mathematics and Statistics, Queen’s University, Kingston K7L 3N6, ON, Canada

May 7, 2026J/K move · Enter open · S save

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

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

Open paper AskPDF arXiv

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.

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