Provable and Practical In-Context Policy Optimization for Self-Improvement

Tianrun Yu; Yuxiao Yang; Zhaoyang Wang; Kaixiang Zhao; Porter Jenkins; Xuchao Zhang; Chetan Bansal; Huaxiu Yao

cs.LGMar 2, 2026

Provable and Practical In-Context Policy Optimization for Self-Improvement

Authors: Tianrun Yu, Yuxiao Yang, Zhaoyang Wang, Kaixiang Zhao, Porter Jenkins, Xuchao Zhang, Chetan Bansal, Huaxiu Yao, +1 more

Organizations: Brigham Young University · University of North Carolina at Chapel Hill · Microsoft

Abstract

We study test-time scaling, where a model improves its answer through multi-round self-reflection at inference. We introduce In-Context Policy Optimization (ICPO), in which an agent optimizes its response in context using self-assessed or externally observed rewards without modifying its parameters. To explain this ICPO process, we theoretically show that with sufficient pretraining under a novel Fisher-weighted logit-matching objective, a single-layer linear self-attention model can provably imitate policy-optimization algorithm for linear bandits. Building on this theory, we propose Minimum-Entropy ICPO (ME-ICPO), a practical algorithm that iteratively uses its response and self-assessed reward to refine its response in-context at inference time. By selecting the responses and their rewards with minimum entropy, ME-ICPO ensures the robustness of the self-assessed rewards via majority voting. Across standard mathematical reasoning tasks, ME-ICPO attains competitive, top-tier performance while keeping inference costs affordable compared with other inference-time algorithms. Overall, ICPO provides a principled understanding of self-reflection in LLMs and yields practical benefits for test-time scaling for mathematical reasoning.

Figures & tables

Figure 1: The In-Context Policy Optimization (ICPO) framework. At each round $t$ , the agent leverages its history of past attempts with bandit feedback $\{(\mathbf{x}_{1},r_{1}),\dots,(\mathbf{x}_{t},r_{t})\}$ to improve its response $\mathbf{x}_{t+1}$ in order to maximize the received reward $r_{t}$ .

Figure 2: Validation of ICPO theory. (Top): Policy Matching. (Bottom): Reward-Shock Stability.

Figure 3: Performance comparison of backbone models before and after ME-ICPO.

Mean@16 (%)
Model $\backslash$ Benchmark	AIME 2024	AMC	MATH-L1	MATH-L2	MATH-L3	MATH-L4	MATH-L5
Qwen2.5-Math-7B	$11.04\pm 1.65$	$41.42\pm 0.99$	$52.62\pm 1.38$	$50.76\pm 0.95$	$49.88\pm 0.88$	$43.60\pm 0.80$	$30.58\pm 0.78$
w/ ME-ICPO	$30.42\pm 1.81$	$47.06\pm 0.84$	$62.35\pm 1.27$	$64.31\pm 1.23$	$58.87\pm 0.78$	$49.31\pm 0.97$	$38.71\pm 0.86$
$\Delta$	$+19.38$	$+5.64$	$+9.73$	$+13.55$	$+8.99$	$+5.71$	$+8.13$
Qwen2.5-Math-1.5B	$6.46\pm 0.96$	$30.42\pm 0.58$	$49.27\pm 0.81$	$48.54\pm 0.67$	$45.42\pm 0.86$	$40.28\pm 0.57$	$25.23\pm 0.45$
w/ ME-ICPO	$9.79\pm 1.11$	$33.58\pm 0.63$	$61.19\pm 0.77$	$54.93\pm 0.71$	$52.08\pm 0.69$	$46.44\pm 0.64$	$29.85\pm 0.70$

Table 1: Performance comparison on different models with Mean@16 and standard Accuracy (%).

Method	Accuracy (%)	Mean@16 (%)
w/o Reward	19.30	19.17
w/o Entropy	5.77	5.83
w/o Entropy & Reward	6.21	6.46
ME-ICPO (full)	30.05	30.42
ME-ICPO Oracle	38.19	38.12

Table 2: Ablation study of ME-ICPO. (Oracle results use groundtruth labels for reward and are shown for reference only.)

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Hyperparameter sensitivity of ME-ICPO on AIME 2024 with Qwen2.5-Math-7B.

Method	AIME-2024	AMC	MATH(Avg)	Time(s/question)
ToT (self eval)	4.38	16.19	12.51	708
ToT (Maj vote)	19.58	29.37	35.63	363
MCTR	4.60	1.20	17.20	1758
TTRL	27.20	45.18	46.83	1253
ME-ICPO (ours)	30.42	47.06	54.71	1152

Appendix

Table 3: Computational Cost and Performance Comparison (Mean@16, Qwen2.5-Math-7B) against Inference and Training Baselines.

Method	AIME-2024	AMC	MATH(Avg)
ToT (self-eval)	4.40	18.10	10.74
ToT (Maj-vote)	19.30	29.40	35.91
MCTR	23.30	2.40	33.82
TTRL	30.00	43.37	45.11
ME-ICPO (ours)	30.05	47.20	47.30

Appendix

Table 4: Computational Performance Comparison (Accuracy, Qwen2.5-Math-7B).

Model	Method	Mean@16 (%)	Acc (%)
Qwen3-4B-Instruct	Base	20.62	20.59
Qwen3-4B-Instruct	ME-ICPO (ours)	57.71	57.67
Gemini-2.5-Pro	Base	58.54	56.60
Gemini-2.5-Pro	ME-ICPO (ours)	79.17	80.00
Gemini-2.5-Flash	Base	35.21	35.42
Gemini-2.5-Flash	ME-ICPO (ours)	76.46	76.47

Appendix

Table 5: Results on frontier/long-CoT models on AIME 2024.

Method	HMMT	HMMT	APEX	APEX
	Mean@16 (%)	Acc(%)	Mean@16 (%)	Acc(%)
Base	1.04	0.67	2.55	2.61
ME-ICPO (ours)	0.42	1.33	4.59	4.57

Appendix

Table 6: Results on Harder Benchmarks (HMMT / APEX) using Qwen2.5-Math-7B.

Method	HMMT	HMMT	APEX	APEX
	Mean@16 (%)	Acc(%)	Mean@16 (%)	Acc(%)
Base	14.79	14.76	14.68	18.33
ME-ICPO (ours)	43.12	43.14	17.18	20.00

Appendix

Table 7: Results on Harder Benchmarks (HMMT / APEX) using Gemini-2.5-Flash.

$n$	$k$	Wall Time (s/question)	Tokens (k/question)	Peak VRAM (GB)
1	4	$197.52$	$99.59$	$81.50$
1	8	$228.89$	$112.25$	$81.43$
1	16	$286.74$	$379.82$	$82.24$
1	32	$468.36$	$852.47$	$82.27$
3	4	$406.65$	$247.76$	$81.43$
3	8	$469.79$	$272.56$	$81.43$

Appendix

Table 8: Empirical cost vs.

(n,k)

on AIME 2024 with Qwen2.5-Math-7B .

Explore similar work

SRPO: Self-Reflective Policy Optimization for Long-Horizon Reasoning
Aug 24, 2026Jialong Liu, Yuling Shi, Ning Yang +2RL for Language Model Reasoning Policy Optimization
School of artificial intelligence, Wuhan University, Wuhan, China · School of computer science, Shanghai Jiao Tong University, Shanghai, China · Institute of automation, Chinese Academy of Sciences, Beijing, China
$\nabla$ -Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space
Mar 5, 2026Peihao Wang, Ruisi Cai, Zhen Wang +4LLM Reasoning Test-Time Optimization
The University of Texas at Austin · UC San Diego · TTIC +1
ThinkTwice: Jointly Optimizing Large Language Models for Reasoning and Self-Refinement
Apr 2, 2026Difan Jiao, Qianfeng Wen, Blair Yang +2RL for Language Model Reasoning LLM Reasoning
Department of Computer Science, University of Toronto · Coolwei AI Lab

Aug 24, 2026J/K move · Enter open · S save

SRPO: Self-Reflective Policy Optimization for Long-Horizon Reasoning

Jialong Liu, Yuling Shi, Ning Yang, Xiaodong Gu+1

School of artificial intelligence, Wuhan University, Wuhan, China · School of computer science, Shanghai Jiao Tong University, Shanghai, China · Institute of automation, Chinese Academy of Sciences, Beijing, China

RL for Language Model Reasoning Policy Optimization RL Post-Training RL for LLM Agents LLM Reasoning LLM Self-Refinement

Open paper AskPDF arXiv

Self-reflection is a powerful mechanism for credit assignment in human learning, converting sparse outcome feedback into actionable guidance. However, its potential for post-training Large Language Models (LLMs) remains underexplored. We propose Self-Reflective Policy Optimization (SRPO), a framework that internalizes this capability. SRPO enables LLMs to analyze their own completed trajectories, synthesize errors into concise "reflection patches," and use reflection-conditioned teacher scores on student on-policy rollouts as dense token-level training signals. This process effectively transforms sparse terminal supervision into dense, token-level learning signals without requiring external critics, separate reward models, or larger teacher models. We demonstrate that SRPO achieves state-of-the-art performance across mathematical reasoning and long-horizon agentic benchmarks with exceptional data efficiency. Using a Qwen3-8B base model, SRPO attains 73.3% on AIME'24 using only 8% (0.08x) of the training FLOPs required by scaled supervised fine-tuning, while significantly improving success rates on WebShop (64.7%), ALFWorld (76.8%), and SWE-Bench-Lite (31.2%). Code is available at https://github.com/Galleons2029/SRPO

Abstract

Explore similar work

SRPO: Self-Reflective Policy Optimization for Long-Horizon Reasoning

∇\nabla∇-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space

ThinkTwice: Jointly Optimizing Large Language Models for Reasoning and Self-Refinement

$\nabla$ -Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space