Context Value Informed In-Context Reinforcement Learning

Wenhao Zhang; Shao Zhang; Xihuai Wang; Yang Li; Ying Wen

cs.LGSep 27, 2025

Context Value Informed In-Context Reinforcement Learning

Authors: Wenhao Zhang, Shao Zhang, Xihuai Wang, Yang Li, Ying Wen

Organizations: Shanghai Jiao Tong University

Abstract

In-context reinforcement learning (ICRL) has emerged as an effective paradigm for test time adaptation to unseen tasks without parameter updates. However, existing ICRL methods can exhibit brittle and unstable adaptions, and the mechanisms underlying such adaptions remain poorly understood. We provide a policy-gradient view of ICRL and argue that relying on trajectory-level return feedback can lead to high-variance updates when test time rollout is limited. Motivated by the actor--critic principle, we propose Context Value Informed ICRL (CV-ICRL), which equips the in-context policy with an explicit context value: the expected discounted return conditioned on the current state and accumulated interaction context. CV-ICRL trains a value head and writes its prediction back into the context as a value token, enabling TD-style return targets for lower-variance policy updates. Experiments on the Dark Room, Minigrid, and Procgen testbeds show that CV-ICRL substantially improves the stability of test time adaptation and achieves higher returns across tasks and environments. Code and data are available at https://github.com/Bluixe/CV-ICRL .

Figures & tables

Figure 1: From actor-only ICRL to context value informed (actor–critic-style) ICRL. ( Left ) The PG-like view models return-based credit assignment using the Monte-Carlo return

\widehat{G}_{t}=\sum_{k\geq 0}\gamma^{k}r_{t+k}

. ( Right ) CV-ICRL augments the same Transformer with a context critic that predicts a context value

v_{t+1}

and writes it back into the context, enabling a TD-style one-step target

\widehat{G}_{t}^{\mathrm{TD}}=r_{t}+\gamma v_{t+1}

for lower-variance test-time adaptation.

Figure 2: Supervised pretraining of CV-ICRL from online RL learning progress. Top: Pretraining data generation. We collect trajectories from training tasks along an online RL learning run, recording checkpoint policies

\{\pi^{(k)}\}

and annotating each transition with a value token

v_{t}=\hat{J}^{(k)}

(Monte-Carlo evaluated). Bottom: Supervised pretraining. A causal Transformer is pretrained to predict both the next action

\pi_{\theta}(\cdot\mid s_{t},\bar{C}_{t})

and the context value

v_{\theta}(s_{t},\bar{C}_{t})

from value-augmented history

\bar{C}_{t}

using a joint loss: action NLL and value MSE.

Task Type	Metric	AD	AD- $\epsilon$	IDT	CV-ICRL
SimpleCrossing S9N3	AER	$0.929\pm 0.067$	$0.152\pm 0.119$	$0.218\pm 0.181$	$\bm{0.945\pm 0.015}$
	LER	$0.941\pm 0.029$	$0.396\pm 0.252$	$0.529\pm 0.226$	$\bm{0.949\pm 0.012}$
	IF ( $\%$ )	$3.428\pm 8.615$	$81.378\pm 10.588$	$70.875\pm 15.810$	$\bm{1.561\pm 1.880}$
SimpleCrossing S11N5	AER	$0.865\pm 0.124$	$0.138\pm 0.107$	$0.323\pm 0.296$	$\bm{0.902\pm 0.076}$
	LER	$0.886\pm 0.142$	$0.329\pm 0.242$	$0.587\pm 0.331$	$\bm{0.921\pm 0.066}$
	IF ( $\%$ )	$15.754\pm 15.686$	$73.744\pm 12.312$	$67.190\pm 26.947$	$\bm{13.209\pm 13.108}$

Table 1: Minigrid results over 6 task types. We report equally weighted task means

\pm

population standard deviations over 20 unseen tasks per task type. AER and LER are higher-is-better; Instability Frequency (IF, %) is lower-is-better. IDT’s IF is omitted for BlockedUnlockPickup because it fails on the task. Best results are bolded ; second-best results are underlined .

Figure 3: Across 6 different Dark Room tasks, AD struggles to maintain performance stability, whereas CV-ICRL substantially improves stability.

Figure 4: Test time AER in Minigrid. Curves show the mean over 20 unseen tasks with

95\%

confidence intervals; the black dashed line is a single-task PPO reference trained directly on each evaluation task.

Figure 5: Test time episode reward in Dark Room and three Procgen task types. Curves show the mean over 20 unseen tasks/levels with

95\%

confidence intervals.

Figure 6: Cross-task-type generalization in Minigrid: models are trained on 6 task types and evaluated on 4 unseen task types.

Figure 7: Ablation results on Minigrid. CV-ICRL outperforms variants without value write-back, without task-dependent value information, or with RTG conditioning, showing that gains come from feeding a meaningful context-value estimate back into the policy.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: The 6 tasks for the main experiments, with different seeds corresponding to different layouts.

Figure 9: The 4 tasks for the cross-task generalization experiments, with different seeds corresponding to different layouts.

Figure 10: The three Procgen task types used in our experiments: Bigfish, Starpilot, and Miner.

Hyperparameter	Value
Learning Rate	0.0001
Policy Network Architecture	MlpPolicy
Training Steps	1000000
Save Frequency	50000
N Steps	500

Appendix

Table 2: Hyperparameters of PPO training for Dark Room

Hyperparameter	Value
Learning Rate	0.0002
Number of Epochs	20
Policy Network Architecture	CustomCNN
Training Steps	2e7
Save Frequency	64000
N Steps	1600

Appendix

Table 3: Hyperparameters of PPO training for Minigrid Pre-train

Hyperparameter	Value
Learning Rate	2e-5
Policy Network Architecture	CustomCNN
Training Steps	400000
Save Frequency	16000
N Steps	1600

Appendix

Table 4: Hyperparameters of PPO training for Minigrid Fine-tune

Environment	AD	CV-ICRL	DPT
LavaCrossingS9N2	100.00% / 0.9413	100.00% / 0.9445	85.00% / 0.8114
LavaCrossingS9N3	100.00% / 0.9319	100.00% / 0.9241	75.00% / 0.6688
SimpleCrossingS11N5	100.00% / 0.8746	100.00% / 0.9051	0.00% / 0.0000
SimpleCrossingS9N3	100.00% / 0.9308	100.00% / 0.9464	5.00% / 0.0034

Appendix

Table 5: DPT comparison on Minigrid task types used in the main experiments. Each entry reports success rate / AER.

Environment	AD	CV-ICRL	DPT
DoorKey-6x6	90.00% / 0.4473	100.00% / 0.7592	0.00% / 0.0000
DoorKey-8x8	100.00% / 0.4381	100.00% / 0.6302	0.00% / 0.0000
FourRooms	75.00% / 0.1076	65.00% / 0.1898	25.00% / 0.2292
LavaCrossingS11N5	100.00% / 0.6373	100.00% / 0.8092	0.00% / 0.0000

Appendix

Table 6: DPT comparison on unseen Minigrid task types for cross-task-type generalization. Each entry reports success rate / AER.

Method	ML1-reach-v3	ML1-push-v3
CV-ICRL	426.0	55.2
AD	253.2	18.2

Appendix

Table 7: Additional MetaWorld continuous-control results. We report Average Episode Return (AER); higher is better.

Algorithm	Mean task reward	Std across tasks	Mean within-task std	Median within-task std
AD	24.56	23.04	7.13	6.64
CV-ICRL	$\bm{31.13}$	$\bm{36.33}$	$\bm{3.95}$	$\bm{2.96}$

Appendix

Table 8: Decomposition of Dark Room variability in the late evaluation stage (steps 3k–4k). “Std across tasks” measures dispersion of task-level mean returns, while within-task std measures temporal fluctuation within each task.

Task type	Metric	CV-ICRL	IC-CQL	Reinformer
SimpleCrossingS9N3	AER	$\bm{0.945\pm 0.015}$	$0.883\pm 0.014$	$0.886\pm 0.020$
	LER	$\bm{0.949\pm 0.012}$	$0.860\pm 0.287$	$0.801\pm 0.327$
	IF (%)	$\bm{1.561\pm 1.880}$	$8.575\pm 1.381$	$10.735\pm 2.834$
SimpleCrossingS11N5	AER	$\bm{0.902\pm 0.076}$	$0.303\pm 0.192$	$0.864\pm 0.225$
	LER	$\bm{0.921\pm 0.066}$	$0.048\pm 0.210$	$0.895\pm 0.208$
	IF (%)	$13.209\pm 13.108$	$68.524\pm 17.792$	$\bm{11.711\pm 22.245}$

Appendix

Table 9: Additional baseline results over 20 unseen tasks per task type. AER and LER are higher-is-better; IF (%) is lower-is-better. Entries show mean

\pm

population standard deviation. Best mean values are bolded .

Task type	CV-ICRL	PPO-state target
SimpleCrossingS9N3	$0.945\pm 0.015$	$0.949\pm 0.001$
SimpleCrossingS11N5	$0.902\pm 0.076$	$0.916\pm 0.067$
BlockedUnlockPickup	$0.955\pm 0.010$	$0.922\pm 0.041$
Unlock	$0.966\pm 0.008$	$0.948\pm 0.015$
LavaCrossingS9N3	$0.921\pm 0.051$	$0.914\pm 0.011$
LavaCrossingS9N2	$0.943\pm 0.015$	$0.933\pm 0.052$

Appendix

Table 10: Context-indexed PPO-state supervision. AER is reported as mean

\pm

population standard deviation over 20 unseen tasks. The context selects the source PPO checkpoint, and its critic provides a state-dependent target.

Condition	SimpleCrossingS9N3	BlockedUnlockPickup
Normal	$0.942\pm 0.020$	$0.936\pm 0.027$
Zero token	$0.090\pm 0.152$	$0.097\pm 0.124$
Frozen token	$0.286\pm 0.296$	$0.241\pm 0.219$

Appendix

Table 11: Value-token interventions over 8,000 interaction steps across multiple episodes. AER is reported as mean

\pm

population standard deviation across 20 evaluator streams. These results test behavioral reliance on token content and continued write-back.

Task type	$H=400$ (main-table reference)	$H=200$
LavaCrossingS9N3	$0.921\pm 0.051$	$0.832\pm 0.282$
SimpleCrossingS9N3	$0.945\pm 0.015$	$0.037\pm 0.053$

Appendix

Table 12: Context-length comparison over 20 unseen tasks per task type. The

H=400

column uses the main-table reference;

H=200

uses a separately trained model. AER is reported as the equally weighted task mean

\pm

population standard deviation.

Environment	Stage 1 (PPO)	Stage 2 (ICRL)
Dark Room	4 h (total)	4 h (total)
Minigrid (6 task types)	24 h per task type	8 h per (method, task type)
Procgen (3 task types)	16 h per task type	10 h per (method, task type)

Appendix

Table 13: Approximate per-run compute for the original experiments on a single RTX 3090 (24 GB). Stage 1 is PPO pretraining and dataset generation; Stage 2 is supervised pretraining of an ICRL method. “per task type” multiplies over the number of task types in the environment; “per method” multiplies over AD, AD ϵ , IDT, and CV-ICRL. Added baseline, value-supervision, and deployment-intervention experiments are excluded.

In-Context RL

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

See all 42 papers

TR-ICRL: Test-Time Rethinking for In-Context Reinforcement Learning
Apr 1, 2026Wenxuan Jiang, Yuxin Zuo, Zijian Zhang +8In-Context RL RL for Language Model Reasoning
The Hong Kong Polytechnic University · Meituan-M17 · Institute of Computing Technology, Chinese Academy of Sciences +3
In-Context Reinforcement Learning under Non-Stationarity: A Survey
Jul 1, 2026A Run, Ziluo DingNon-Stationary RL In-Context RL
Independent Researcher · Beijing Innovation Center of Humanoid Robotics
Downstream-Aware Context Selection for Online In-Context Reinforcement Learning
Sep 27, 2026Ruihan A. Li, Shangtong Zhang, Rohan ChandraIn-Context RL In-Context Learning
University of Illinois Urbana-Champaign · University of Virginia

Apr 1, 2026J/K move · Enter open · S save

TR-ICRL: Test-Time Rethinking for In-Context Reinforcement Learning

Wenxuan Jiang, Yuxin Zuo, Zijian Zhang, Xuecheng Wu+7

The Hong Kong Polytechnic University · Meituan-M17 · Institute of Computing Technology, Chinese Academy of Sciences+3

In-Context RL RL for Language Model Reasoning Test-Time RL

Open paper AskPDF arXiv

In-Context Reinforcement Learning (ICRL) enables Large Language Models (LLMs) to learn online from external rewards directly within the context window. However, a central challenge in ICRL is reward estimation, as models typically lack access to ground-truths during inference. To address this limitation, we propose Test-Time Rethinking for In-Context Reinforcement Learning (TR-ICRL), a novel ICRL framework designed for both reasoning and knowledge-intensive tasks. TR-ICRL operates by first retrieving the most relevant instances from an unlabeled evaluation set for a given query. During each ICRL iteration, LLM generates a set of candidate answers for every retrieved instance. Next, a pseudo-label is derived from this set through majority voting. This label then serves as a proxy to give reward messages and generate formative feedbacks, guiding LLM through iterative refinement. In the end, this synthesized contextual information is integrated with the original query to form a comprehensive prompt, with the answer determining through a final round of majority voting. TR-ICRL is evaluated on mainstream reasoning and knowledge-intensive tasks, where it demonstrates significant performance gains. Remarkably, TR-ICRL improves Qwen2.5-7B by 21.23% on average on MedQA and even 137.59% on AIME2024. Extensive ablation studies and analyses further validate the effectiveness and robustness of our approach. Our code is available at https://github.com/pangpang-xuan/TR_ICRL.

Abstract

In-Context RL

Explore similar work

TR-ICRL: Test-Time Rethinking for In-Context Reinforcement Learning

In-Context Reinforcement Learning under Non-Stationarity: A Survey

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