P^2O: Joint Policy and Prompt Optimization

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as a powerful paradigm for enhancing the reasoning capabilities of Large Language Models (LLMs). However, vanilla RLVR suffers from inefficient exploration, particularly when confronting "hard samples" that yield nearzero success rates. In such scenarios, the reliance on sparse outcome rewards typically results in zero-advantage estimates, effectively starving the model of supervision signals despite the high informational value of these instances. To address this, we propose P^2O, a novel framework that synergizes Prompt Optimization with Policy Optimization. P^2O identifies hard samples during training iterations and leverages the GeneticPareto (GEPA) prompt optimization algorithm to evolve prompt templates that guide the model toward discovering successful trajectories. Crucially, unlike traditional prompt engineering methods that rely on input augmentation, P^2O distills the reasoning gains induced by these optimized prompts directly into the model parameters. This mechanism provides denser positive supervision signals for hard samples and accelerates convergence. Extensive experiments demonstrate that P^2O not only achieves superior performance on in-distribution datasets but also exhibits strong generalization, yielding substantial improvements on out-of-distribution benchmarks (+4.7% avg.).

Figures (5)

• 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 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}}$).
• Figure 3: Training Dynamics of P$^2$O. The plots illustrate the dynamics of Training Reward (Left) and Validation Accuracy (Right) throughout the optimization process using the Teacher-Ref variant on the DeepScaler-5K dataset. P$^2$O maintains a higher reward compared to the baseline by dynamically cracking hard samples. Crucially, this reward advantage translates into robust gains in validation accuracy, confirming that the context distillation mechanism effectively transfers prompt-dependent success into intrinsic model capability.
• Figure 4: Prompt Optimization Effectiveness.Left: Optimized prompts (triangles) consistently yield gains over the standard policy (circles) on both Pass@1 and Pass@6 metrics, demonstrating that GEPA effectively assists the model in bridging the performance gap. Right: With this assistance, the model continuously conquers hard samples, resulting in a steady decline in the number of intractable instances throughout the training epochs. Data points are derived from the Teacher-Ref variant on the DeepScaler-5K dataset.
• Figure 5: Qualitative Analysis: Overcoming Local Optima in Geometric Reasoning.