P^2O: Joint Policy and Prompt Optimization

TL;DR

P^2O is a novel method that combines policy and prompt optimization to improve reinforcement learning with large language models, especially on hard samples where traditional methods struggle.

Contribution

It introduces a joint optimization framework using GEPA and context distillation to enhance LLM reasoning and out-of-distribution performance.

Findings

Restores advantage signals in RLVR on hard samples.

Outperforms standard methods with doubled rollout budgets.

Achieves up to 9.5% performance improvement.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) enhances Large Language Model (LLM) reasoning but suffers from advantage collapse on ``hard samples'' where all rollouts fail. This lack of variance eliminates crucial learning signals. For these intractable samples, simply scaling up rollout budgets offers limited gains. We introduce Joint Policy and Prompt Optimization (P$^2$O) to mitigate this collapse by alternating continuous policy updates with discrete prompt evolution. P$^2$O leverages the GEPA algorithm to discover successful reasoning prompts for intractable instances. Via context distillation, the model internalizes these prompt-induced gains directly into its parameters, removing the need for inference-time prompting. Empirically, P$^2$O restores critical advantage signals, significantly outperforming standard GRPO and surpassing baselines with doubled rollout budgets, ultimately yielding strong out-of-distribution generalization and an up to $9.5\%$ performance improvement. 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 establishes a self-reinforcing paradigm for autonomous LLM alignment.