Agentic In-Context Learning: What Algorithms Can Transformers Learn in Closed-Loop Interaction?

Shiguang Wu; Yaqing Wang; Quanming Yao

cs.AIOct 16, 2026

Agentic In-Context Learning: What Algorithms Can Transformers Learn in Closed-Loop Interaction?

Authors: Shiguang Wu, Yaqing Wang, Quanming Yao

Organizations: Tsinghua University · Beijing Institute of Mathematical Sciences and Applications

Abstract

The learning algorithm view of in-context learning (ICL) studies how fixed-parameter Transformers learn from supplied examples. Language model agents extend this setting through interaction: their actions determine what evidence enters context next. We call this regime Agentic In-Context Learning (Agentic ICL). Existing studies largely depend on externally supplied contexts or reward-bearing interaction histories, while work on general tool-using agents lacks a learning algorithm view. We represent an interactive learner by its state, update, and decision rule, covering active learning, bandits, and online reinforcement learning (RL). For expressivity, we construct causal Transformers that approximate these algorithms over a finite horizon, and bound the resulting loss in return. For training, we characterize when optimizing process rewards preserves behavior that maximizes the final outcome. Distillation experiments show that Transformers can learn to execute three representative recursions. We then train models from outcome or process rewards without oracle action or state targets, as in RL post-training. Different rewards yield distinct decision policies. Improving the alignment of process supervision with outcome-optimal action preferences shifts learned decisions accordingly. Such behavioral differences persist in a pretrained language model across textual paraphrases.

Figures & tables

Figure 1: Illustrative comparison of Agentic ICL, ICL and ICRL. Agentic ICL studies how agents solve general agentic tasks.

Setting	State $s_{t}$	Update $U_{t}$	Decision $G_{t}$ / examples
Active learning	Labeled set, query pool, posterior	Add queried label, update posterior	BALD, value-of-information query selection
Bandits	Posterior or reward statistics	Posterior or statistical update	Thompson sampling, UCB
Online RL	Values, policy parameters, or learned model	TD, policy-gradient, or model update	Q-learning, SARSA, actor-critic, learned-model planning

Table 1: Illustrative settings and algorithms covered by Theorem 4.1 .

Figure 2: Learned execution from supervised distillation. Each panel compares a distilled model with its corresponding random Transformer (null baseline).

x-

ticks are different metrics. Full values are provided in Appendix Tables 12 and 13 .

Figure 3: Distances to candidate policies under four training objectives. Each panel reports policy TV to the four candidates on a shared diagnostic set.

Figure 4: Varying the credit discount

\gamma

\alpha=\tau_{v}=1

. The gray band marks the common ordering/decision crossover; the green band marks zero oracle violations, after which learned agreement is nearly constant.

Figure 5: Natural-language tool use under three training objectives. Top: paired bars compare distances to Bayes DP and KG references; the closer candidate is unchanged under unseen paraphrases. Bottom: task returns decrease modestly under paraphrases and sharply under shuffled cards, and move all policies far from both references. Appendix Table 24 reports all values.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Algorithmic instance	Returned record $x_{t}\in\mathcal{X}_{t}$	Action $a_{t}\in\mathcal{A}_{t}$	Environment kernel $K_{M}$	Information state $s_{t}\in\mathcal{S}_{t}$
Setting: sequential information acquisition
Pool-based active learning (BALD)	$(z_{t},y_{t})\in\mathcal{Z}\times\mathcal{Y}$	$\operatorname{query}(z_{t})$ , $z_{t}\in\mathcal{P}_{t}$	$p_{M}(y_{t}\mid z_{t})$ , with the queried $z_{t}$ copied into the record	$(\mathcal{D}_{t},\mathcal{P}_{t},b_{t})$ : labeled set, unlabeled pool, posterior
Bayesian sequential design (VOI/KG)	$x_{t}\in\mathcal{X}_{t}(a_{t})$	query or experiment	$K_{M}(x\mid h_{t-1},a_{t})$	$(h_{t-1},b_{t})$ , $b_{t}(m)=\mathbb{P}(M=m\mid h_{t-1})$
Setting: bandits
Thompson sampling	reward $r_{t}\in[0,1]$	arm $i_{t}\in[K]$	$r_{t}\sim P_{M,i_{t}}$	posterior $b_{t}(M)$ over arm-reward models
UCB	reward $r_{t}\in[0,1]$	arm $i_{t}\in[K]$	$r_{t}\sim P_{M,i_{t}}$	$(t,N_{t},S_{t})$ , with counts and reward sums in $\mathbb{R}^{K}$

Appendix

Table 2: Observations, actions, environment kernels, and learner states for each algorithm, grouped by interaction setting.

N_{t},S_{t}

are bandit statistics;

N_{t},C_{t},Z_{t}

in the model-based rows are transition counts, reward sums, and visit counts. Terminal flags

d_{t}

make bootstrapping conventions explicit.

Split	$H$	$q$	Cost law	Purpose
Training	$\operatorname{Unif}\{2,3,4\}$	$\operatorname{Unif}\{2,\ldots,8\}$	$c_{e}\sim\operatorname{Unif}\{1,2\},\,c_{\min}=1$	short-horizon training
In-distribution test	$\{2,4\}$	same	same	short-horizon reference
Nominal-horizon control	$\{6,8\}$	$q\leq 4$	same	no calls beyond four
Feasible-long test	$6$ / $8$	$q\geq 6$ / $q=8$	same	longer calls feasible

Appendix

Table 3: Joint budget-horizon protocol for symbolic BTS. The two longer-horizon rows are prespecified strata of a test distribution whose budget values and cost support are already present during training. Feasibility does not imply execution; actual policy-dependent call lengths are reported separately .

Policy	Mean	Median	90th pct.	$\Pr(N>4)$	$\Pr(N=6)$	Gap to Bayes $\downarrow$
Hard Bayes oracle	2.88	3	4	8.0%	2.0%	0
Outcome reward, 6L/256W	2.73	3	4	5.0%	1.0%	0.019 $\pm\,0.004$

Appendix

Table 4: Autonomously executed interaction lengths on the prespecified feasible-long stratum. The distribution is

H_{0}=6,q\in\{6,7,8\}

, with unchanged likelihood, prior, cost, and discount laws. Value gaps use the same unfiltered stratum; a nonzero long-call tail supports limited autonomous length extrapolation, not a claim that most episodes require six calls.

Policy	Mean	Median	90th pct.	$\Pr(N>4)$	$\Pr(N=8)$	Gap to Bayes $\downarrow$
Hard Bayes oracle	3.011	3	5	10.3%	0.2%	0
Outcome reward, 6L/256W	2.799	3	4	6.5%	0.0%	0.028 $\pm\,0.005$

Appendix

Table 5: Autonomous call counts on the separate

H_{0}=8,q=8

feasible-long stratum. The tail beyond four calls supports limited autonomous length extrapolation; most episodes still terminate earlier. Return gaps are mean

\pm

SD across five seeds.

Training condition	Bayes agreement $\uparrow$	KG agreement $\uparrow$	Neither
Outcome reward	88.7% $\pm\,1.9\%$	8.1% $\pm\,1.4\%$	3.2% $\pm\,0.7\%$
One-step process reward	11.4% $\pm\,2.2\%$	86.5% $\pm\,2.0\%$	2.1% $\pm\,0.9\%$
Outcome-preserving	87.9% $\pm\,2.0\%$	8.8% $\pm\,1.5\%$	3.3% $\pm\,0.8\%$
Bayes DP-score distillation	93.4% $\pm\,1.1\%$	4.3% $\pm\,0.8\%$	2.3% $\pm\,0.5\%$

Appendix

Table 6: Agreement with Bayes DP and KG on the shared diagnostic set. Every entry uses

D_{BK}

, with unique, different Bayes and KG maximizing actions and a margin of at least

0.02

under both reference score functions. The set contains one history from each of 1,024 qualifying tasks. Values are mean

\pm

standard deviation across five seeds, and the three columns sum to 100% within each seed.

Condition	Objective audited	Opt. mass $\uparrow$	Regret $\downarrow$	Nearest / runner-up	Gap $\uparrow$
Outcome reward	outcome reward	0.803 $\pm\,0.018$	0.010 $\pm\,0.002$	Bayes / two-step	0.038 $\pm\,0.008$
One-step process reward	one-step verifier score	0.812 $\pm\,0.019$	0.007 $\pm\,0.002$	KG / entropy	0.091 $\pm\,0.010$
Bayes DP-score distillation	outcome, audit only	0.771 $\pm\,0.014$	0.013 $\pm\,0.002$	Bayes / two-step	0.057 $\pm\,0.007$

Appendix

Table 7: Objective alignment and candidate separation on frozen diagnostic set

D

. Optimal-action mass and regret use unregularized maximizing sets and action values; candidate distances use the specified behavioral kernels and include receding two-query lookahead. Differences between nearest and runner-up distances are computed within seed before summarization. Distillation is audited against outcome optimality but was trained against a score-softened teacher.

Condition	Bayes	Two-step	KG	Entropy	Open-loop
Outcome reward	.043	.081 †	.182	.249	.286
One-step process reward	.219	.174	.047	.138	.275
Outcome-preserving	.046	.084	.184	.254	.289
Bayes DP-score distillation	.018	.075 †	.214	.267	.301

Appendix

Table 8: Mean distances to the expanded candidate family on

D

, using the reference policies defined in Appendix G.2 .

\dagger

: means obtained by adding the paired Bayes/two-step gap in Table 7 to the Bayes distance, with the same histories and weights; no SD is inferred from this addition. The verifier’s two-step distance is

0.174\pm 0.013

; the preserving condition’s is

0.084\pm 0.010

Condition / audit	Top-action agreement $\uparrow$	OptMass $\uparrow$	Regret $\downarrow$
Outcome / outcome	96.9% $\pm\,1.0\%$	.923 $\pm\,.014$	.006 $\pm\,.001$
One-step process reward / immediate score	95.7% $\pm\,1.4\%$	.916 $\pm\,.017$	.005 $\pm\,.001$
Distillation / outcome	98.1% $\pm\,0.8\%$	.891 $\pm\,.012$	.008 $\pm\,.001$
Preserving / outcome	96.6% $\pm\,1.2\%$	.918 $\pm\,.015$	.006 $\pm\,.001$

Appendix

Table 9: Objective alignment after excluding near-ties, using objective-specific margin

\geq 0.02

. Values are mean

\pm

SD across five seeds. Distillation has the highest top-action agreement but lower optimal-action mass, consistent with learning a soft teacher.

Paired intervention	Bayes ref.	Exact verifier	Outcome model	One-step process reward model
Remaining horizon $2\to 4$	.204	0	.187 $\pm\,.020$	.021 $\pm\,.005$
Remaining budget $2\to 4$	.173	0	.161 $\pm\,.018$	.018 $\pm\,.004$

Appendix

Table 10: Action TV between paired resource prompts, with current legal actions unchanged. Reference entries are deterministic; model entries are mean

\pm

SD across five seeds. The outcome model responds to continuation opportunities, whereas the model trained with one-step process reward responds much less.

Diagnostic	Outcome reward	One-step process reward
posterior log-odds NRMSE $\downarrow$	0.021 $\pm\,0.003$	0.024 $\pm\,0.004$
Bayes-update identity NRMSE $\downarrow$	0.028 $\pm\,0.004$	0.031 $\pm\,0.004$
action TV on state-equivalent order swaps $\downarrow$	0.014 $\pm\,0.002$	0.016 $\pm\,0.003$
unseen $H=6$ value gap $\downarrow$	0.019 $\pm\,0.004$	0.052 $\pm\,0.007$
likelihood-intervention TV to score reference $\downarrow$	0.058 $\pm\,0.006$	0.073 $\pm\,0.009$

Appendix

Table 11: State, update, and generalization diagnostics. Values are mean

\pm

standard deviation across five seeds; hard Bayes/KG agreement is reported in Table 6 .

Reference	action TV $\downarrow$	state / update NRMSE $\downarrow$	long-horizon TV $\downarrow$	value gap $\downarrow$
Bayes DP	0.018 $\pm\,0.002$	0.013 $\pm\,0.002$ / 0.019 $\pm\,0.003$	0.031 $\pm\,0.004$	0.006 $\pm\,0.001$
UCB bandit	0.009 $\pm\,0.001$	0.006 $\pm\,0.001$ / 0.008 $\pm\,0.001$	0.019 $\pm\,0.003$	0.004 $\pm\,0.001$
Q-learning	0.016 $\pm\,0.003$	0.011 $\pm\,0.002$ / 0.017 $\pm\,0.002$	0.067 $\pm\,0.007$	0.008 $\pm\,0.002$
Null range	0.408-0.760	0.895-1.000 / 0.216-0.691	0.481-0.732	0.031-0.533

Appendix

Table 12: Algorithm-compilation tests. Task rows report distilled models. Null ranges are the minimum and maximum task means for untrained Transformers.

Oracle	action TV $\downarrow$	state / update NRMSE $\downarrow$	long-horizon TV $\downarrow$	value gap $\downarrow$
Bayes DP	0.760 $\pm\,0.001$	1.000 $\pm\,0.001$ / 0.691 $\pm\,0.002$	0.732 $\pm\,0.002$	0.533 $\pm\,0.002$
UCB bandit	0.490 $\pm\,0.001$	1.000 $\pm\,0.001$ / 0.308 $\pm\,0.001$	0.481 $\pm\,0.003$	0.031 $\pm\,0.002$
Q-learning	0.408 $\pm\,0.002$	0.895 $\pm\,0.001$ / 0.216 $\pm\,0.002$	0.490 $\pm\,0.004$	0.035 $\pm\,0.001$

Appendix

Table 13: Random Transformer baselines for Table 12 . Values are mean

\pm

standard deviation; lower is better for every metric. Long-horizon evaluation uses Bayes DP horizons five through eight and doubles transcript length for UCB and Q-learning, as specified above.

Probe evaluation	Metric	Bayes DP	UCB	Q-learning
unmodified	state	0.013 $\pm\,0.002$	0.006 $\pm\,0.001$	0.011 $\pm\,0.002$
unmodified	update	0.019 $\pm\,0.003$	0.008 $\pm\,0.001$	0.017 $\pm\,0.002$
feedback shuffled	state	1.153 $\pm\,0.061$	0.741 $\pm\,0.048$	1.062 $\pm\,0.057$
feedback shuffled	update	1.089 $\pm\,0.055$	0.697 $\pm\,0.043$	0.986 $\pm\,0.051$
targets permuted	state	1.003 $\pm\,0.024$	0.994 $\pm\,0.021$	1.005 $\pm\,0.026$
targets permuted	update	0.989 $\pm\,0.022$	1.010 $\pm\,0.025$	0.997 $\pm\,0.023$

Appendix

Table 14: Negative controls for the post hoc probes. Metrics are NRMSE. The feedback shuffle breaks the correspondence between a history and its returned records; the target permutation breaks the correspondence between frozen representations and trainer-side states. Target permutation leaves the trained policy unchanged.

$H_{0}$	$k=0$	1	2	3	4	5	6	7	8	Depth mean
5	.019	.020	.021	.024	.028	.032	—	—	—	.024 $\pm\,.003$
6	.020	.021	.024	.027	.032	.035	.037	—	—	.028 $\pm\,.004$
7	.021	.024	.027	.029	.035	.039	.043	.046	—	.033 $\pm\,.005$
8	.024	.026	.030	.034	.038	.042	.047	.052	.058	.039 $\pm\,.006$

Appendix

Table 15: Bayes distillation: action TV to the score teacher by initial horizon and prefix depth. Individual depth entries are means; SDs are available for the horizon-level aggregates. Prefixes with only one query remaining test execution rather than Bayes/KG planning separation.

Paired with $H_{0}=4$	5	6	7	8
Model action TV	.010 $\pm\,.002$	.011 $\pm\,.002$	.013 $\pm\,.003$	.014 $\pm\,.003$

Appendix

Table 16: Changes in the Bayes-distilled model when only the inactive horizon field is increased, mean

\pm

SD across five seeds.

Condition	Trainer objective	Entropy protocol	Aligned reference
Outcome reward	terminal BTS return	$0.01\to 0$ in first $30\%$ ; then $0$	hard Bayes DP
Outcome-preserving process reward	outcome plus perturbed telescoping potential	same	hard Bayes DP
Primary one-step process reward	$\widetilde{r}_{t}^{\rm 1V}$ , $\beta=0$	$0.01\to 0$ in first $30\%$ ; then $0$	exact one-step process reward optimum
Bayes DP-score distillation	cross-entropy to score teacher	none	$\pi_{{\rm B},0.02}^{\rm score}$

Appendix

Table 17: Alignment between the primary experimental objectives and policy references. The three reward conditions are optimized with PPO, have an entropy-free final phase, and are evaluated against hard unregularized references. The credit-assignment sweep is specified separately below.

$\alpha$	$\tau_{v}$	Violation fraction on $D$	$\min_{h\in D}g_{h}$
.50	.50	52.4%	-.124
.50	.75	48.5%	-.101
.50	1.00	45.1%	-.086
.50	1.25	47.9%	-.112
.75	.50	45.7%	-.117
.75	.75 ⋆	41.3%	-.096

Appendix

Table 18: Immediate-credit verifier-specification audit on the fixed 4,096-history set

D

. Entries are deterministic oracle quantities.

\star

marks the primary setting;

\square

marks the exact KG specification. Its violation fraction is 100% on

D_{BK}

, where Bayes and KG optima differ by construction.

$(\alpha,\tau_{v})$	Nearest candidate	Candidate TV	Hard-objective TV	Verifier regret	Return
(.50,.50)	Entropy	.091 $\pm\,.010$	.238 $\pm\,.024$	.012 $\pm\,.003$	.668 $\pm\,.014$
(.50,.75)	Entropy	.079 $\pm\,.009$	.211 $\pm\,.021$	.009 $\pm\,.002$	.687 $\pm\,.012$
(.50,1.00)	Hybrid †	.087 $\pm\,.010$	.198 $\pm\,.020$	.008 $\pm\,.002$	.696 $\pm\,.011$
(.50,1.25)	Entropy	.098 $\pm\,.012$	.226 $\pm\,.025$	.011 $\pm\,.003$	.679 $\pm\,.014$
(.75,.50)	KG	.076 $\pm\,.009$	.219 $\pm\,.022$	.010 $\pm\,.002$	.696 $\pm\,.013$
(.75,.75) ⋆	KG	.047	.188 $\pm\,.019$	.007 $\pm\,.002$	.710 $\pm\,.009$

Appendix

Table 19: Learned policies in the immediate-credit verifier-specification sweep. Candidate TV uses behavioral references; hard-objective TV uses the uniform kernel on exact verifier optima. Entries with uncertainty are mean

\pm

SD across five seeds; the primary KG distance is available as a mean only.

\dagger

: the displayed distance is to KG, not a distance to a hybrid reference.

Condition	Bayes DP-score TV $\downarrow$	KG-score TV $\downarrow$	Two-step-score TV $\downarrow$	Return $\uparrow$	Post. NRMSE $\downarrow$	Update NRMSE $\downarrow$
Outcome-preserving process reward	0.046 $\pm\,0.007$	0.184 $\pm\,0.015$	0.084 $\pm\,0.010$	0.752 $\pm\,0.009$	0.023 $\pm\,0.004$	0.030 $\pm\,0.005$

Appendix

Table 20: Outcome-preserving process-reward control with the fixed perturbed potential

\Phi_{\epsilon}

\epsilon=0.05

. The perturbation changes intermediate reward values while leaving the unregularized outcome-optimal policy set unchanged. Architecture, training budget, seed identities, and evaluation tasks are matched to the outcome-reward condition. Values are mean

\pm

standard deviation across five seeds.

$\gamma$	Violations	$G_{D}$	Violations	$G_{W}$
	$D$ , 4,096 histories		$W$ , 1,024 histories
.00	1540	-.078000	1024	-.103180
.25	1428	-.069700	1001	-.101249
.50	1076	-.043800	836	-.087732
.70	628	-.020600	482	-.060792
.85	231	-.006280	164	-.027286

Appendix

Table 21: Ordering violations and minimum action-value gaps at

\alpha=\tau_{v}=1

. Counts and cost-normalized gaps are computed on the indicated diagnostic sets. Zero counts denote no violations, rather than a rounded small percentage. Values bracket crossings on

W

(.90,.95)

and on

D

(.98,.99)

, subject to the numerical precision scope stated below.

$\gamma$	Bayes TV	Two-step TV	KG TV	Entropy TV	Open-loop TV
.00	.218 $\pm\,.012$	.172 $\pm\,.013$	.038 $\pm\,.006$	.142 $\pm\,.012$	.277 $\pm\,.014$
.25	.211 $\pm\,.012$	.166 $\pm\,.012$	.048 $\pm\,.007$	.151 $\pm\,.012$	.278 $\pm\,.013$
.50	.179 $\pm\,.011$	.129 $\pm\,.011$	.073 $\pm\,.009$	.172 $\pm\,.013$	.281 $\pm\,.014$
.70	.128 $\pm\,.010$	.077 $\pm\,.008$	.118 $\pm\,.011$	.205 $\pm\,.013$	.283 $\pm\,.013$
.85	.081 $\pm\,.010$	.063 $\pm\,.009$	.157 $\pm\,.013$	.226 $\pm\,.014$	.285 $\pm\,.015$
.90	.061 $\pm\,.008$	.069 $\pm\,.010$	.168 $\pm\,.014$	.239 $\pm\,.014$	.286 $\pm\,.014$

Appendix

Table 22: Complete five-candidate profiles for the credit-assignment sweep at

\alpha=\tau_{v}=1

D

, mean

\pm

SD across five seeds. These are distances to behavioral references, not hard cumulative-objective optima.

$\gamma$	Nearest	Return	OptMass	Regret	Hard agreement on $W$
					Bayes	KG	Neither
.00	KG	.711 $\pm\,.009$	.830 $\pm\,.018$	.005 $\pm\,.001$	6.9%	90.5%	2.6%
.25	KG	.716 $\pm\,.010$	.824 $\pm\,.019$	.006 $\pm\,.002$	9.1%	87.3%	3.6%
.50	KG	.725 $\pm\,.010$	.809 $\pm\,.020$	.008 $\pm\,.002$	21.0%	72.1%	6.9%
.70	Two-step	.737 $\pm\,.011$	.786 $\pm\,.023$	.009 $\pm\,.002$	49.0%	42.4%	8.6%
.85	Two-step	.744 $\pm\,.011$	.792 $\pm\,.022$	.010 $\pm\,.002$	76.7%	17.5%	5.8%

Appendix

Table 23: Performance across credit discounts, agreement with the cumulative objective on

D

, and agreement with Bayes DP and KG on

W

. Return, OptMass and raw-unit regret are mean

\pm

SD across five seeds. Agreement columns are means only and sum to 100%. Except at

\gamma=.90

, the same candidate is nearest in all five seeds. At

\gamma=.90

, three seeds are closest to Bayes DP and two to two-step, so the condition is labeled hybrid/unidentified.

Training	Evaluation	Bayes-ref. TV $\downarrow$	KG-ref. TV $\downarrow$	return $\uparrow$
Outcome reward	training templates	0.131 $\pm\,0.012$	0.191 $\pm\,0.015$	0.676 $\pm\,0.018$
	unseen paraphrases	0.178 $\pm\,0.014$	0.245 $\pm\,0.019$	0.641 $\pm\,0.021$
	shuffled cards	0.674 $\pm\,0.023$	0.698 $\pm\,0.020$	0.296 $\pm\,0.015$
One-step process reward	training templates	0.268 $\pm\,0.018$	0.114 $\pm\,0.011$	0.629 $\pm\,0.020$
	unseen paraphrases	0.277 $\pm\,0.021$	0.153 $\pm\,0.013$	0.602 $\pm\,0.023$
	shuffled cards	0.671 $\pm\,0.019$	0.685 $\pm\,0.022$	0.290 $\pm\,0.014$

Appendix

Table 24: Natural-language tool use under three training objectives. The shuffled-card control breaks the mapping between tool descriptions and environment kernels.

Training	Evaluation	Two-step TV	Entropy TV	Open-loop TV
Outcome reward	training templates	.155	.232	.288
	unseen paraphrases	.206	.282	.334
	shuffled cards	.683	.708	.721
One-step process reward	training templates	.229	.199	.300
	unseen paraphrases	.244	.220	.313
	shuffled cards	.678	.704	.720

Appendix

Table 25: Additional candidate-distance means for the text conditions in Table 24 , evaluated with the same text histories and weights as its Bayes/KG columns. These three distances are available as means only; they do not certify a seed-stable nearest-candidate label.

In-Context Learning

In-context learning studies how models adapt to new tasks from examples or instructions supplied at inference time, without changing their weights.

See all 402 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

Mar 11, 2026

Implicit Statistical Inference in Transformers: Approximating Likelihood-Ratio Tests In-Context

In-context learning (ICL) allows Transformers to adapt to novel tasks without weight updates, yet the underlying algorithms remain poorly understood. We adopt a statistical decision-theoretic perspective by investigating simple binary hypothesis testing, where the optimal policy is determined by the likelihood-ratio test. Notably, this setup provides a mathematically rigorous setting for mechanistic interpretability where the target algorithmic ground truth is known. By training Transformers on tasks requiring distinct geometries (linear shifted means vs. nonlinear variance estimation), we demonstrate that the models approximate the Bayes-optimal sufficient statistics from context up to some monotonic transformation, matching the performance of an ideal oracle estimator in nonlinear regimes. Leveraging this analytical ground truth, mechanistic analysis via logit lens and circuit alignment suggests that the model does not rely on a fixed kernel smoothing heuristic. Instead, it appears to adapt the point at which decisions become linearly decodable: exhibiting patterns consistent with a voting-style ensemble for linear tasks while utilizing a deeper sequential computation for nonlinear tasks. These findings suggest that ICL emerges from the construction of task-adaptive statistical estimators rather than simple similarity matching.

Faris Chaudhry, Siddhant Gadkari

Department of Computer Science Imperial College London

Likelihood Ratio Testing In-Context Learning Transformer Interpretability Decision-Making under Uncertainty

Ask PDF

May 8, 2026

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

In-context reinforcement learning (ICRL) studies agents that, after pretraining, adapt to new tasks by conditioning on additional context without parameter updates. Existing theoretical analyses of ICRL largely rely on linear attention, which replaces the softmax function in the standard attention with an identity mapping. This paper provides the first theoretical understanding of ICRL without making the unrealistic linear attention simplification. In particular, we consider the standard softmax attention used in practice. We show that, with certain parameters, the layerwise forward pass of a Transformer with such softmax attention is equivalent to iterative updates of a weighted softmax temporal difference (TD) learning algorithm. Here, weighted softmax TD is a new RL algorithm that performs policy evaluation in kernel space and adopts both linear TD and tabular TD as special cases. We also prove that under a certain contraction condition, the policy evaluation error decays as the number of layers grows, with the identified parameters above. Finally, we prove that those parameters are a global minimizer of a pretraining loss, explaining their emergence in our numerical experiments.

Zixuan Xie, Xinyu Liu, Claire Chen +3

University of Virginia · California Institute of Technology · Purdue University

Reinforcement Learning Transformer-Based RL Temporal-Difference Learning Transformer Attention In-Context RL Contextual RL

Ask PDF

Abstract

In-Context Learning

Explore similar work

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

Implicit Statistical Inference in Transformers: Approximating Likelihood-Ratio Tests In-Context

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