TR-ICRL: Test-Time Rethinking for In-Context Reinforcement Learning
Organizations: The Hong Kong Polytechnic University · Meituan-M17 · Institute of Computing Technology, Chinese Academy of Sciences · Xi’an Jiaotong University ·
Abstract
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.
Figures & tables
| Setting | Acc.(MATH) | Acc.(AMC) |
|---|---|---|
| TR-ICRL (Step 1) | 71.00 | 53.01 |
| BoN (N=16) | 69.60 | 43.37 |
| Self-Refine (Iteration=8) | 61.20 | 42.17 |
| Reflexion | 66.60 | 55.42 |
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| Model | Params | Avg. Tokens ( ) | Relative Cost (units) | Accuracy (%) | |||
|---|---|---|---|---|---|---|---|
| (B) | MATH | AMC | MATH | AMC | MATH | AMC | |
| Qwen2.5-72B | 72 | 0.62 | 0.92 | 89.28 | 132.48 | 62.10 | 41.11 |
| TR-ICRL (step 1) | 7 | 0.48 | 0.81 | 80.64 | 136.08 | 71.00 | 53.01 |
| Dataset | Train Num | Test Num | Options Num |
|---|---|---|---|
| MATH500 | 0 | 500 | N/A |
| AMC | 0 | 83 | N/A |
| GSM8K | 7473 | 1319 | N/A |
| AIME2024 | 0 | 30 | N/A |
| AIME2025 | 0 | 30 | N/A |
| MedQA | 10178 | 1273 | 4 |
In-Context RL
In-context reinforcement learning studies how agents adapt their policies from interaction histories without updating model parameters.
See all 42 papers