# HeaPA: Difficulty-Aware Heap Sampling and On-Policy Query Augmentation for LLM Reinforcement Learning

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 42 of its 55 references that are in the map.
- First posted on arXiv 30 Jan 2026.

## Placed in

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

## Builds on (most cited first, 30 of 42)

- [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)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [Kimi k1.5: Scaling Reinforcement Learning with LLMs](https://mlsanity.com/paper/2501.12599/) (2025)
- [OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problems](https://mlsanity.com/paper/2402.14008/) (2024)
- [Group Sequence Policy Optimization](https://mlsanity.com/paper/2507.18071/) (2025)
- [Policy Distillation](https://mlsanity.com/paper/1511.06295/) (2015)
- [Seed1.5-Thinking: Advancing Superb Reasoning Models with Reinforcement Learning](https://mlsanity.com/paper/2504.13914/) (2025)
- [SimpleTIR: End-to-End Reinforcement Learning for Multi-Turn Tool-Integrated Reasoning](https://mlsanity.com/paper/2509.02479/) (2025)
- [A Minimalist Approach to LLM Reasoning: from Rejection Sampling to Reinforce](https://mlsanity.com/paper/2504.11343/) (2025)
- [Act Only When It Pays: Efficient Reinforcement Learning for LLM Reasoning via Selective Rollouts](https://mlsanity.com/paper/2506.02177/) (2025)
- [Reinforcement Learning with Verifiable Rewards: GRPO's Effective Loss, Dynamics, and Success Amplification](https://mlsanity.com/paper/2503.06639/) (2025)
- [Efficient Reinforcement Finetuning via Adaptive Curriculum Learning](https://mlsanity.com/paper/2504.05520/) (2025)
- [From Automation to Autonomy: A Survey on Large Language Models in Scientific Discovery](https://mlsanity.com/paper/2505.13259/) (2025)
- [Not All Rollouts are Useful: Down-Sampling Rollouts in LLM Reinforcement Learning](https://mlsanity.com/paper/2504.13818/) (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)
- [Black-Box On-Policy Distillation of Large Language Models](https://mlsanity.com/paper/2511.10643/) (2025)
- [Can Prompt Difficulty be Online Predicted for Accelerating RL Finetuning of Reasoning Models?](https://mlsanity.com/paper/2507.04632/) (2025)
- [No Prompt Left Behind: Exploiting Zero-Variance Prompts in LLM Reinforcement Learning via Entropy-Guided Advantage Shaping](https://mlsanity.com/paper/2509.21880/) (2025)
- [KDRL: Post-Training Reasoning LLMs via Unified Knowledge Distillation and Reinforcement Learning](https://mlsanity.com/paper/2506.02208/) (2025)
- [NGRPO: Negative-enhanced Group Relative Policy Optimization](https://mlsanity.com/paper/2509.18851/) (2025)
- [SPEED-RL: Faster Training of Reasoning Models via Online Curriculum Learning](https://mlsanity.com/paper/2506.09016/) (2025)
- [Beyond Correctness: Harmonizing Process and Outcome Rewards through RL Training](https://mlsanity.com/paper/2509.03403/) (2025)
- [Improving Sampling Efficiency in RLVR through Adaptive Rollout and Response Reuse](https://mlsanity.com/paper/2509.25808/) (2025)
- [CAR: Conceptualization-Augmented Reasoner for Zero-Shot Commonsense Question Answering](https://mlsanity.com/paper/2305.14869/) (2023)

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

- [Better, Faster: Harnessing Self-Improvement in Large Reasoning Models](https://mlsanity.com/paper/2605.24998/) (2026)
- [CuSearch: Curriculum Rollout Sampling via Search Depth for Agentic RAG](https://mlsanity.com/paper/2605.11611/) (2026)
- [UI-Mate: Advancing Open-Weight Foundation GUI Agents with In-Context Demonstrations](https://mlsanity.com/paper/2608.15930/) (2026)
- [Rollout Efficiency in Reinforcement Learning for Reasoning Large Language Models: A Taxonomy and Future Directions](https://mlsanity.com/paper/2609.25463/) (2026)

## Abstract (arXiv)

RLVR has become a standard recipe for training LLMs on reasoning tasks with verifiable outcomes, but when rollout generation dominates the cost, efficiency hinges on which prompts are sampled and when. In practice, prompt pools are often static or only weakly coupled to policy progress, so uniform sampling fails to track the moving capability frontier and wastes rollouts on regions that are already solved or still unreachable. Prior methods improve efficiency via filtering, curricula, adaptive rollout allocation, or teacher guidance, but they often assume a fixed pool, which does not support stable on-policy pool growth, or they introduce additional teacher cost and latency. In this work, we propose HeaPA (Heap Sampling and On-Policy Query Augmentation), which maintains a bounded, evolving pool, tracks the frontier with heap-based boundary sampling, grows the pool via on-policy augmentation under lightweight asynchronous validation, and stabilizes correlated queries via topology-aware pool statistics re-estimation and controlled reinsertion. Across two training corpora, two training recipes, and seven benchmarks, HeaPA consistently improves accuracy and reaches target performance with fewer computations at comparable wall-clock time. Analyses attribute the gains to frontier-focused sampling and on-policy pool growth, with more pronounced improvements at mid-to-large model scales. Our training code is publicly available at https://github.com/horizon-llm/HeaPA.