Downstream-Aware Context Selection for Online In-Context Reinforcement Learning

Ruihan A. Li; Shangtong Zhang; Rohan Chandra

cs.LGSep 27, 2026

Downstream-Aware Context Selection for Online In-Context Reinforcement Learning

Authors: Ruihan A. Li, Shangtong Zhang, Rohan Chandra

Organizations: University of Illinois Urbana-Champaign · University of Virginia

Abstract

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.

Figures & tables

Figure 1: Mean selected context size

K

across normalized decision progress under held-out ID, Severe-3x, and unseen-route evaluation. Dashed and dotted lines denote the matched fixed-

K

baselines.

Evaluation Type	Method	Mean $K$	Post-selection Tokens $\downarrow$	Avg Total $\downarrow$	Safety $\uparrow$	Comfort $\uparrow$	Efficiency $\uparrow$
Held-out ID	Full-H20	$20.00$	$130$ K	$218$ K	$0.88$	$0.76$	$0.79$
	Recent-13	$13.00$	$92$ K	$181$ K	$\mathbf{0.88}$	$0.76$	$0.80$
	Similarity-13	$13.00$	$91$ K	$180$ K	$0.88$	$\mathbf{0.77}$	$0.80$
	Ours	$13.63$	$\mathbf{73}$ K	$\mathbf{162}$ K	$0.88$	$0.76$	$\mathbf{0.81}$
Severe- 3x	Full-H20	$20.00$	$151$ K	$247$ K	$0.82$	$0.78$	$0.83$
	Recent-13	$13.00$	$103$ K	$201$ K	$0.87$	$0.77$	$0.84$

Table 1: Evaluation results across the held-out in-distribution, unseen OOD, and unseen-route settings. Token counts are rounded to the nearest thousand (K), and performance scores to two decimals. Lower token usage is better, while higher Safety, Comfort, and Efficiency scores are better. Bold indicates the best value based on the unrounded results.

Method	Mean $K$	Late-Round Total/Trial $\downarrow$	Avg Total /Trial $\downarrow$	Return $\uparrow$
Full-H20	$20.00$	$352{,}809$	$282{,}879$	$95.50$
Recent-12	$12.00$	$202{,}710$	$194{,}188$	$\mathbf{96.06}$
Similarity-12	$12.00$	$227{,}433$	$217{,}706$	$94.78$
Ours	$12.70$	$\mathbf{87{,}100}$	$\mathbf{135{,}470}$	$\mathbf{96.06}$

Table 2: ScienceWorld evaluation. Late-round statistics are averaged over rounds 19–40, when context selection is active for our method, whereas all-round statistics are averaged over rounds 1–40. Because histories grow over the episodes, late-round token usage can exceed the all-round average. “Total” includes both input and output tokens. ”Return” are average over all-round.

Metric	Ours	LongLLMLingua
Avg Total / Trial $\downarrow$	$135{,}470$	$\mathbf{131{,}467}$
Return $\uparrow$	$\mathbf{96.060}$	$94.780$
History Proc. Latency (s) $\downarrow$	$\mathbf{1.244}$	$27.715$

Table 3: ScienceWorld comparison with LongLLMLingua.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Scenario	Method	Mean $K$	Post-selection Tokens $\downarrow$	Performance $\uparrow$
Held-out ID	Recent-13	$13.00$	$91{,}730$	$\mathbf{0.8838}/0.7617/0.7981$
	Recent-13 + LORC	$13.00$	$74{,}458$	$0.8808/0.7602/0.7825$
	Ours	$13.63$	$\mathbf{73{,}233}$	$0.8836/\mathbf{0.7645}/\mathbf{0.8071}$
Severe-3x	Recent-13	$13.00$	$102{,}971$	$\mathbf{0.8673}/0.7715/\mathbf{0.8351}$
	Recent-13 + LORC	$13.00$	$98{,}165$	$0.8561/\mathbf{0.7802}/0.7996$
	Ours	$13.02$	$\mathbf{85{,}747}$	$0.8625/0.7759/0.8256$

Appendix

Table 4: Comparison of Recent, Recent+LORC, and Ours across evaluation scenarios. For driving, total tokens denote post-selection tokens and performance is Safety/Comfort/Efficiency (S/C/E). For ScienceWorld, total tokens denote post-selection tokens per trial and performance is return. Lower token usage is better; higher performance is better.

Setting	$\bm{\tau}$	Mean $K$	Estimated Post-selection Input Tokens $\downarrow$	Change vs. $\tau=0.05$
Held-out ID	$0.02$	$17.58$	$81{,}734$	$+17.79\%$
	$0.03$	$16.50$	$78{,}402$	$+12.99\%$
	$\mathbf{0.05}$	$\mathbf{13.63}$	$\mathbf{69{,}387}$	$\mathbf{0.00\%}$
	$0.07$	$11.64$	$63{,}183$	$-8.94\%$
	$0.10$	$10.64$	$59{,}685$	$-13.98\%$
	$0.15$	$10.14$	$57{,}931$	$-16.51\%$

Appendix

Table 5: Sensitivity to the feasibility threshold

\tau

. Post-selection input tokens are estimated from the 21st decision onward on fixed trajectories and averaged per episode. Changes are reported relative to

\tau=0.05

Method	Mean Return $\uparrow$	Running-Max Return $\uparrow$
Full-H20	$74.34$	$95.50$
Recent-12	$\mathbf{76.02}$	$95.78$
Similarity-12	$69.37$	$94.78$
Recent-12 + LORC	$70.27$	$94.78$
LongLLMLingua (matched-budget)	$73.74$	$94.67$
Ours	$74.22$	$\mathbf{95.94}$

Appendix

Table 6: ScienceWorld Post-selection return comparison. Higher is better.

Figure 2: ScienceWorld performance over 40 interaction trials. Left: mean trial return. Right: mean running-max return. Shaded regions denote SEM across tasks.

In-Context RL

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

See all 42 papers

Do LLMs Learn from Rewards in Context? : Rethinking the role of reward in In-Context RL
Oct 8, 2026Minchan Kwon, Seunghee Koh, Sunghyun Baek +2In-Context RL Contextual RL
Korea Advanced Institute of Science and Technology (KAIST)
Remember Less, Adapt Longer: Bounded-State In-Context Reinforcement Learning with Recurrent Summary Memory
Oct 19, 2026Jiawen Wang, Cheng Xue, Peng Zhang +1Reinforcement Learning In-Context RL
The Australian National University
Context Value Informed In-Context Reinforcement Learning
Sep 27, 2025Wenhao Zhang, Shao Zhang, Xihuai Wang +2In-Context RL Test-Time RL
Shanghai Jiao Tong University

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

Do LLMs Learn from Rewards in Context? : Rethinking the role of reward in In-Context RL

Minchan Kwon, Seunghee Koh, Sunghyun Baek, Minsung Bae+1

Korea Advanced Institute of Science and Technology (KAIST)

In-Context RL Contextual RL RL for LLM Agents In-Context Learning

Open paper AskPDF arXiv

LLM agents increasingly improve at inference time by accumulating experience in context rather than by updating parameters. This process is often described as in-context reinforcement learning (ICRL). Whether in-context learning (ICL) can actually play the role of RL, however, has not been tested. We study this question in its simplest form, direct ICRL, where the model conditions directly on raw trajectory-reward pairs, and ask whether the reward acts as a learning signal. Through controlled experiments on four benchmarks across six models, we find that the reward is read, but its effect is small: flipping, randomizing, or removing the reward leaves the improvement curve almost unchanged, and this holds even under meta-prompts that explicitly instruct the model to explore, exploit, or reason over rewards. Trajectories drive improvement, but not through their semantic content: shuffled or corrupted trajectories work as well as real ones. These patterns closely mirror those known in ICL, suggesting that direct ICRL is better understood as a special case of ICL than as inference-time RL. This reframing has implications for agent memory design: ICL factors such as input distribution and demonstrations may matter more than RL elements such as reward shaping and exploration.

Abstract

In-Context RL

Explore similar work

Do LLMs Learn from Rewards in Context? : Rethinking the role of reward in In-Context RL

Remember Less, Adapt Longer: Bounded-State In-Context Reinforcement Learning with Recurrent Summary Memory

Context Value Informed In-Context Reinforcement Learning