PrAg-PO: Prompt Augmented Policy Optimization for Robust and Diverse Mathematical Reasoning

Introduction

Reinforcement learning algorithms such as group-relative policy optimization (GRPO) have shown strong potential for improving the mathematical reasoning capabilities of large language models. While a growing body of work seeks to improve training entropy, rollout diversity, and exploration, most existing methods still train models with a single fixed reasoning prompt or template, which can encourage prompt-specific overfitting and unstable training dynamics. In this work, we introduce Prompt Augmented Policy Optimization (PrAg-PO), a simple policy optimization method that mixes prompt templates with template-specific format rewards during training. By encouraging models to generate reasoning traces under diverse instructions and output formats, PrAg-PO increases rollout diversity and improves robustness. Compared with GRPO and DAPO, PrAg-PO achieves significantly higher reasoning accuracy while mitigating premature training collapse. Empirically, experiments on DeepSeek-R1-Distill-Qwen-1.5B, Qwen2.5-Math-1.5B, and Qwen3-1.7B show that PrAg-PO consistently outperforms strong baselines and achieves competitive performance against recent methods on mathematics benchmarks, using only a fixed MATH Level 3-5 training set of 8.5K problems.

Results

Weights 🏋🏻‍♂️

• PrAg-PO Qwen2.5-Math-1.5B Step 2720

• PrAg-PO Qwen2.5-Math-1.5B Step 2480 (80+ on MATH500!)

• PrAg-PO Qwen2.5-Math-1.5B Step 2820 (outdated)

• PrAg-PO DeepSeek-R1-Distill-Qwen-1.5B Step 1100

• PrAg-PO DeepSeek-R1-Distill-Qwen-1.5B Step 1160

• PrAg-PO Qwen3-1.7B Step 1520

• PrAg-PO Qwen3-1.7B Step 720

Training

Wandb training logs for experiments: here.

First install required packages for verl following their installation instructions. Our project is developed on top of this version of verl, using conda environment. We have provided our environment.yml.

To train PrAg-PO on Qwen2.5-Math-1.5B:

sbatch examples/grpo_trainer/run_qwen2_5_math_1_5b.sh

To train PrAg-PO on DeepSeek-R1-Distill-Qwen-1.5B:

sbatch examples/grpo_trainer/run_deepseek_r1_qwen_1_5b.sh

To train PrAg-PO on Qwen3-1.7B:

sbatch examples/grpo_trainer/run_qwen3_1_7b.sh

To convert checkpoints to huggingface format

sbatch convert_model.sh

Evaluation

We use greedy decoding with average of 5 inference results as our final reported result due to the stochasticity of vLLM inference engine.

To evaluate trained Qwen2.5-Math-1.5B:

sbatch eval_qwen2.5.sh

To evaluate trained DeepSeek-R1-Distill-Qwen-1.5B:

sbatch eval_deepseek.sh

To evaluate trained Qwen3-1.7B:

sbatch eval_qwen3.sh

Comments

Our codebase builds heavily on verl, understand-r1-zero and TreePO. Thanks to their great works!

Citation

If you find this repository helpful, please consider giving this repo a star ⭐ and citing:

@misc{lu2026pragpopromptaugmentedpolicy,
      title={PrAg-PO: Prompt Augmented Policy Optimization for Robust and Diverse Mathematical Reasoning}, 
      author={Wenquan Lu and Hai Huang and Enqi Liu and Randall Balestriero},
      journal={arXiv preprint arXiv:2602.03190},
      url={https://arxiv.org/abs/2602.03190},
      year={2026},
}