P^2O: Joint Policy and Prompt Optimization

Xinyu Lu; Kaiqi Zhang; Jinglin Yang; Boxi Cao; Yaojie Lu; Hongyu Lin; Min He; Xianpei Han

cs.LGMar 23, 2026

P^2O: Joint Policy and Prompt Optimization

Authors: Xinyu Lu, Kaiqi Zhang, Jinglin Yang, Boxi Cao, Yaojie Lu, Hongyu Lin, Min He, Xianpei Han, +1 more

Organizations: Chinese Information Processing Laboratory, Institute of Software, Chinese Academy of Sciences · University of Chinese Academy of Sciences · Institute of Information Engineering, Chinese Academy of Sciences · School of Cyber Security, University of Chinese Academy of Sciences · National Computer Network Emergency Response Technical Team/Coordination Center of China

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) enhances Large Language Model (LLM) reasoning but is suffer from advantage collapse: when all rollouts of a query receive identical rewards, the group variance vanishes, most damagingly on hard samples, where scaling rollout budgets yields little. We introduce Joint Policy and Prompt Optimization (P O) to mitigate this collapse by alternating continuous policy updates with discrete prompt evolution. P O mines hard samples with a success-rate threshold, evolves reasoning prompts for them with GEPA, and internalizes the elicited trajectories via context distillation, which optimizes each trajectory under the original query and thus removes inference-time prompting, with a Context Ratio Mask (CRM) filtering out extreme likelihood ratios. P O restores critical advantage signals and surpasses the GRPO baseline by up to 8.2 points in average accuracy on six held-out benchmarks across all training datasets and backbones, while also outperforming DAPO and other baselines. The gains are especially pronounced on hard benchmarks, reaching up to 16.3 points above GRPO on average across AIME24 and AIME25. Our findings expose the limits of standard exploration in sparse-reward environments, illuminating the potential of unifying evolutionary algorithms with reinforcement learning. This integration of discrete semantic search and continuous parameter updates provides a self-reinforcing framework that facilitates more effective LLM alignment.

Figures & tables

Figure 1: Conceptual Illustration of the P 2 O Framework. Standard policy optimization often gets trapped in local optima ( $\rho_{\text{init}}$ ) due to sparse rewards on hard samples. P 2 O bridges this exploration gap using optimized prompts (the red arrow) to reach high-reward regions that are inaccessible via standard exploration. Subsequently, the model consolidates these gains (the white arrows) by updating its parameters to master the new region ( $\rho_{\text{opt}}$ ), effectively internalizing the prompt-induced capabilities.

Figure 2: Overview of the P 2 O Framework. The training process is formulated as an alternating maximization procedure between two phases: (1) Policy Optimization with Context Distillation, where the policy $\pi_{\theta}$ is updated under a CRM to internalize reasoning patterns elicited by augmented inputs $\tilde{x}$ ; and (2) Evolutionary Prompt Optimization, where the prompt template set $\mathcal{Z}$ is evolved using GEPA to discover successful trajectories for the remaining hard samples ( $\mathcal{D}_{\text{hard}}$ ).

Base Models
Model	AIME24	AIME25	AMC	MATH500	Minerva	Olympiad	AVG
Qwen3-4B	23.8	19.4	68.0	82.8	31.6	49.9	45.9
Qwen3-8B	28.3	22.1	69.4	81.6	31.2	49.5	47.0
DeepMath-5K
Qwen3-4B-GRPO ${}_{\textnormal{Rollout-6}}$	33.8	29.0	79.5	87.6	41.5	57.5	54.8
Qwen3-4B-GRPO ${}_{\textnormal{Rollout-8}}$	35.6	28.3	79.7	89.0	40.1	56.4	54.9

Table 1: Comparative Evaluation on Mathematical Reasoning Benchmarks. Accuracy (%) of P 2 O against the baselines; gains are largest on challenging benchmarks such as AIME24 and AIME25. Best results are highlighted in bold .

Model	AIME24	AIME25	AMC	MATH500	Minerva	Olympiad	AVG
Qwen3-4B-GRPO	46.9	37.7	88.1	90.4	41.5	58.2	60.5
Qwen3-4B-P 2 O	59.8	49.4	92.2	91.4	36.4	62.2	65.2
Qwen3-4B-P 2 O ${}_{\textnormal{w/o-Context-Distillation}}$	36.5	31.5	81.1	89.6	40.1	54.8	55.6
Qwen3-4B-P 2 O ${}_{\textnormal{Same-Template-in-Group}}$	57.3	44.6	88.9	90.8	40.4	63.0	64.2

Table 2: Ablation Study of P 2 O Components. Comparison of performance on mathematical benchmarks when excluding context distillation or group prompt diversity. Data points are derived from the Teacher-Ref w/o-CRM variant on the DeepScaler-5K dataset.

Variant	P 2 O	w/o-CRM	$\Delta$
DeepMath-5K
Qwen3-4B Self-Ref	63.0	64.2	$\downarrow$ 1.2
Qwen3-4B Teacher-Ref	58.5	57.9	$\uparrow$ 0.6
Qwen3-8B Self-Ref	64.0	61.3	$\uparrow$ 2.7
Qwen3-8B Teacher-Ref	63.4	61.0	$\uparrow$ 2.4
DeepScaler-5K

Table 3: Effect of the Context Ratio Mask. Six-benchmark average accuracy (%) per training set, backbone, and reflection source;

\Delta

(P 2 O

-

w/o-CRM) marked

\uparrow

\downarrow

. See Table 9 .

Variant	P 2 O	w/o-CRM	$\Delta$
DeepMath-5K
Qwen3-4B Self-Ref	63.0	64.2	$\downarrow$ 1.2
Qwen3-4B Teacher-Ref	58.5	57.9	$\uparrow$ 0.6
Qwen3-8B Self-Ref	64.0	61.3	$\uparrow$ 2.7
Qwen3-8B Teacher-Ref	63.4	61.0	$\uparrow$ 2.4
DeepScaler-5K

Table 3: Effect of the Context Ratio Mask. Six-benchmark average accuracy (%) per training set, backbone, and reflection source;

\Delta

(P 2 O

-

w/o-CRM) marked

\uparrow

\downarrow

. See Table 9 .

$\alpha$	$\beta$	Qwen3-8B	$\Delta$ vs. w/o CRM
w/o CRM		61.3	–
0.8	1.2	58.3	$\downarrow$ 3.0
0.05	5	63.1	$\uparrow$ 1.8
0.01	10	64.0	$\uparrow$ 2.7
0.005	50	61.5	$\uparrow$ 0.2

Table 4: Hyper-parameter Ablation of the Context Ratio Mask. Six-benchmark average accuracy (%) for Qwen3-8B on the DeepMath-5K training set with Self-Ref, as the mask bounds

\alpha

and

\beta

vary; “w/o CRM” removes the mask entirely. The default interval

[0.01,10]

reaches 64.0%, 2.7 points above w/o CRM. Best result in bold .

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Policy Rollouts	GEPA Rollouts	Total Rollouts/Epoch
GRPO ${}_{\textnormal{Rollout-6}}$	$6N$	–	$6N$
GRPO ${}_{\textnormal{Rollout-8}}$	$8N$	–	$8N$
GRPO ${}_{\textnormal{Rollout-12}}$	$12N$	–	$12N$
P 2 O ${}_{\textnormal{Self-Ref}}$ ( $K=6$ )	$6N$	$2N$	$8N$
P 2 O ${}_{\textnormal{Teacher-Ref}}$ ( $K=6$ )	$6N$	$2N$	$8N$

Appendix

Table 5: Per-Epoch Rollout Budgets. All entries are rollout counts per epoch, where

N

is the number of training samples (

N=5{,}000

for DeepMath-5K and DeepScaler-5K, so that

2N=10

K). GEPA rollouts include everything charged to the budget of Algorithm 3 : mini-batch rollouts, reflection calls (including all Kimi-K2 calls), and all dev evaluations; GRPO uses none. Because the GEPA budget is a hard cap, P 2 O spends

8N

rollouts per epoch at

K=6

and therefore matches GRPO

{}_{\text{Rollout-8}}

Bucket	# Samples (%)	# Improved (%)	w/o Template	w/ Template	w/o Template	w/ Template
			Reward Variance		Reward Mean
All hard samples	1980 (100%)	740 (37.4%)	0.2494	0.2491	0.4765	0.5295
All-incorrect ( $\sum_{k=1}^{K}r_{ik}=0$ )	559 (28.2%)	202 (36.1%)	0	0.1205	0	0.1401
Near-all-incorrect ( $\sum_{k=1}^{K}r_{ik}/K\leq 1/6$ )	745 (37.6%)	321 (43.1%)	0.0399	0.1617	0.0416	0.2029

Appendix

Table 6: Advantage collapse on the hard samples of DeepMath-5K at the end of epoch 0. Self-Ref w/o-CRM variant with

K=6

rollouts per query. A sample is counted as Improved if its average reward over the

K

rollouts is higher when the template is applied than under the bare query. Variance and mean are averaged over the samples of each bucket; the two extreme buckets are nested inside the set of all hard samples.

Metric	Self-Ref	Teacher-Ref	$\Delta$ (Teacher $-$ Self)
Mean dev-hard score	16.2	24.4	$\uparrow$ 8.2
Best dev-hard score	26.7	31.7	$\uparrow$ 5.0
Final downstream 6-benchmark AVG	64.2	57.9	$\downarrow$ 6.3

Appendix

Table 7: Template quality versus downstream performance DeepMath-5K, Qwen3-4B, w/o-CRM. Dev-hard scores are accuracies (%) of the template set produced by GEPA at the end of epoch 0, averaged over all candidate templates ( mean ) and taken at the best template ( best ), where both variants evaluate on the same dev-hard split with the same policy weights. The last row reports the final six-benchmark average of the fully trained models (Table 9 ).

\Delta=\textnormal{Teacher-Ref}-\textnormal{Self-Ref}

Model	AIME24	AIME25	AMC	MATH500	Minerva	Olympiad	AVG
Qwen3-4B-P 2 O τ=2/6	56.7	45.4	86.9	92.4	40.4	63.1	64.2
Qwen3-4B-P 2 O τ=1/6	50.0	42.5	87.3	90.2	40.8	61.6	62.1
Qwen3-4B-P 2 O τ=4/6	57.3	44.6	88.0	91.6	42.6	60.4	64.1

Appendix

Table 8: Hard Sample Success-rate Threshold

\tau

Ablation. Results on DeepMath-5K using the Self-Ref w/o-CRM variant with

K=6

. Thresholds are reported as rates. Our default

\tau=2/6

treats samples with at most one correct rollout as hard.

Model	AIME24	AIME25	AMC	MATH500	Minerva	Olympiad	AVG
DeepMath-5K
Qwen3-4B-P 2 O ${}_{\textnormal{Self-Ref}}$	55.3	40.0	86.2	91.2	43.0	62.4	63.0
Qwen3-4B-P 2 O ${}_{\textnormal{Self-Ref, w/o-CRM}}$	56.7	45.4	86.9	92.4	40.4	63.1	64.2
Qwen3-4B-P 2 O ${}_{\textnormal{Teacher-Ref}}$	46.5	32.5	82.3	88.6	41.5	59.3	58.5
Qwen3-4B-P 2 O ${}_{\textnormal{Teacher-Ref, w/o-CRM}}$	44.6	34.4	81.9	90.2	36.0	60.0	57.9
Qwen3-8B-P 2 O ${}_{\textnormal{Self-Ref}}$	57.1	41.7	86.7	92.2	43.0	63.4	64.0

Appendix

Table 9: Full Results of the Context Ratio Mask Ablation. Benchmark-wise accuracy (%) of P 2 O and its w/o-CRM counterpart with the Qwen3-4B and Qwen3-8B backbones, on the DeepMath-5K and DeepScaler-5K training sets and the two reflection sources. AVG averages the six benchmarks.

Explore similar work

The Easy, the Hard, and the Learnable: Confidence and Difficulty-Adaptive Policy Optimization for LLM Reasoning
Jun 6, 2026Zhanke Zhou, Xiangyu Lu, Chentao Cao +4Frictive Policy Optimization LLM Reasoning Strategies
1TMLR Group, Department of Computer Science, Hong Kong Baptist University · 2Stanford University · 3Sydney AI Centre, The University of Sydney
MCPO: Mastery-Consolidated Policy Optimization for Large Reasoning Models
Apr 18, 2026Zhaokang Liao, Yingguo Gao, Yi Yang +2Reinforcement Learning With Verifiable Reward Frictive Policy Optimization
Ant Group · Zhejiang University
Nonsense Helps: Prompt Space Perturbation Broadens Reasoning Exploration
May 7, 2026Langlin Huang, Chengsong Huang, Jinyuan Li +3Large Language Model Reinforcement Learning LLM Reasoning Strategies
Washington University in St. Louis

Jun 6, 2026, cs.LGJ/K move · Enter open · S save

The Easy, the Hard, and the Learnable: Confidence and Difficulty-Adaptive Policy Optimization for LLM Reasoning

Zhanke Zhou, Xiangyu Lu, Chentao Cao, Brando Miranda+3

1TMLR Group, Department of Computer Science, Hong Kong Baptist University · 2Stanford University · 3Sydney AI Centre, The University of Sydney

Frictive Policy Optimization LLM Reasoning Strategies Hard

Open paper AskPDF arXiv

RL with verifiable rewards can substantially improve LLM reasoning, yet standard GRPO-style training often treats easy, hard, and learnable questions alike through uniform sampling and weighting, leading to inefficient compute allocation. We study GRPO by tracking token log-probabilities, group-normalized advantages, and the induced token-level update weights. This reveals three recurring dynamics as training proceeds: (1) confidence inflation, (2) advantage contraction, and (3) hierarchical convergence. These findings suggest that the utility of each update depends strongly on both question difficulty and the model's current competence. Motivated by this, we propose Confidence and Difficulty-adaptive Policy Optimization (CoDaPO), which assigns each question a bounded value from rollout confidence and empirical difficulty. CoDaPO then uses this value to reweight policy updates and resample high-value learnable questions within mini-batches, thereby increasing discovery within the learnable band under a fixed compute budget. Across twelve benchmarks, CoDaPO consistently improves accuracy over existing RL methods. Our code is publicly available at https://github.com/tmlr-group/CoDaPO.

Abstract

Explore similar work

The Easy, the Hard, and the Learnable: Confidence and Difficulty-Adaptive Policy Optimization for LLM Reasoning

MCPO: Mastery-Consolidated Policy Optimization for Large Reasoning Models

Nonsense Helps: Prompt Space Perturbation Broadens Reasoning Exploration