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

arXiv: https://arxiv.org/abs/2602.03190 · HTML: https://mlsanity.com/paper/2602.03190/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 25 of its 60 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)

## Builds on (most cited first, 25 of 25)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvement](https://mlsanity.com/paper/2409.12122/) (2024)
- [Group Sequence Policy Optimization](https://mlsanity.com/paper/2507.18071/) (2025)
- [The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models](https://mlsanity.com/paper/2505.22617/) (2025)
- [Reasoning with Exploration: An Entropy Perspective](https://mlsanity.com/paper/2506.14758/) (2025)
- [GDPO: Group reward-Decoupled Normalization Policy Optimization for Multi-reward RL Optimization](https://mlsanity.com/paper/2601.05242/) (2026)
- [Geometric-Mean Policy Optimization](https://mlsanity.com/paper/2507.20673/) (2025)
- [Soft Adaptive Policy Optimization](https://mlsanity.com/paper/2511.20347/) (2025)
- [SEED-GRPO: Semantic Entropy Enhanced GRPO for Uncertainty-Aware Policy Optimization](https://mlsanity.com/paper/2505.12346/) (2025)
- [Defeating the Training-Inference Mismatch via FP16](https://mlsanity.com/paper/2510.26788/) (2025)
- [TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling](https://mlsanity.com/paper/2508.17445/) (2025)
- [Beyond Pass@1: Self-Play with Variational Problem Synthesis Sustains RLVR](https://mlsanity.com/paper/2508.14029/) (2025)
- [QuestA: Expanding Reasoning Capacity in LLMs via Question Augmentation](https://mlsanity.com/paper/2507.13266/) (2025)
- [On Entropy Control in LLM-RL Algorithms](https://mlsanity.com/paper/2509.03493/) (2025)
- [RePO: Replay-Enhanced Policy Optimization](https://mlsanity.com/paper/2506.09340/) (2025)
- [Latent Chain-of-Thought? Decoding the Depth-Recurrent Transformer](https://mlsanity.com/paper/2507.02199/) (2025)
- [MO-GRPO: Mitigating Reward Hacking of Group Relative Policy Optimization on Multi-Objective Problems](https://mlsanity.com/paper/2509.22047/) (2025)
- [Interleaved Reasoning for Large Language Models via Reinforcement Learning](https://mlsanity.com/paper/2505.19640/) (2025)
- [Harder Is Better: Boosting Mathematical Reasoning via Difficulty-Aware GRPO and Multi-Aspect Question Reformulation](https://mlsanity.com/paper/2601.20614/) (2026)
- [DRA-GRPO: Your GRPO Needs to Know Diverse Reasoning Paths for Mathematical Reasoning](https://mlsanity.com/paper/2505.09655/) (2025)
- [Transformation-Augmented GRPO for Enhancing Exploration in Reasoning of Large Language Models](https://mlsanity.com/paper/2601.22478/) (2026)

## Built on by (most cited first, 1 of 1)

- [Nonsense Helps: Prompt Space Perturbation Broadens Reasoning Exploration](https://mlsanity.com/paper/2605.05566/) (2026)

## Abstract (arXiv)

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.