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

TL;DR

PrAg-PO is a novel reinforcement learning method that enhances mathematical reasoning in language models by promoting diverse prompts and formats during training, leading to improved accuracy and robustness.

Contribution

It introduces Prompt Augmented Policy Optimization (PrAg-PO), a simple approach that mixes prompt templates with format rewards to increase diversity and robustness in reasoning models.

Findings

PrAg-PO outperforms existing methods like GRPO and DAPO in reasoning accuracy.

PrAg-PO mitigates premature training collapse.

PrAg-PO achieves competitive results on mathematics benchmarks.

Abstract

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. The code and model checkpoints are available at https://github.com/wenquanlu/PrAg-PO.