# Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems

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

## Key facts

- Cited by 87 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 26 of its 44 references that are in the map.
- First posted on arXiv 21 May 2025.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

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

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [A General Language Assistant as a Laboratory for Alignment](https://mlsanity.com/paper/2112.00861/) (2021)
- [LeanDojo: Theorem Proving with Retrieval-Augmented Language Models](https://mlsanity.com/paper/2306.15626/) (2023)
- [Training Language Models to Self-Correct via Reinforcement Learning](https://mlsanity.com/paper/2409.12917/) (2024)
- [CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning](https://mlsanity.com/paper/2207.01780/) (2022)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [V-STaR: Training Verifiers for Self-Taught Reasoners](https://mlsanity.com/paper/2402.06457/) (2024)
- [Is Self-Repair a Silver Bullet for Code Generation?](https://mlsanity.com/paper/2306.09896/) (2023)
- [RLEF: Grounding Code LLMs in Execution Feedback with Reinforcement Learning](https://mlsanity.com/paper/2410.02089/) (2024)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Execution-based Code Generation using Deep Reinforcement Learning](https://mlsanity.com/paper/2301.13816/) (2023)
- [Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models](https://mlsanity.com/paper/2412.15287/) (2024)
- [StepCoder: Improve Code Generation with Reinforcement Learning from Compiler Feedback](https://mlsanity.com/paper/2402.01391/) (2024)
- [BOND: Aligning LLMs with Best-of-N Distillation](https://mlsanity.com/paper/2407.14622/) (2024)
- [Optimizing Language Models for Inference Time Objectives using Reinforcement Learning](https://mlsanity.com/paper/2503.19595/) (2025)
- [Variational Best-of-N Alignment](https://mlsanity.com/paper/2407.06057/) (2024)
- [InfAlign: Inference-aware language model alignment](https://mlsanity.com/paper/2412.19792/) (2024)
- [B-Coder: Value-Based Deep Reinforcement Learning for Program Synthesis](https://mlsanity.com/paper/2310.03173/) (2023)

## Built on by (most cited first, 30 of 87)

- [A Survey of Reinforcement Learning for Large Reasoning Models](https://mlsanity.com/paper/2509.08827/) (2025)
- [Pass@k Training for Adaptively Balancing Exploration and Exploitation of Large Reasoning Models](https://mlsanity.com/paper/2508.10751/) (2025)
- [POPE: Learning to Reason on Hard Problems via Privileged On-Policy Exploration](https://mlsanity.com/paper/2601.18779/) (2026)
- [The Choice of Divergence: A Neglected Key to Mitigating Diversity Collapse in Reinforcement Learning with Verifiable Reward](https://mlsanity.com/paper/2509.07430/) (2025)
- [Maximum Likelihood Reinforcement Learning](https://mlsanity.com/paper/2602.02710/) (2026)
- [Look Inward to Explore Outward: Learning Temperature Policy from LLM Internal States via Hierarchical RL](https://mlsanity.com/paper/2602.13035/) (2026)
- [Differential Smoothing Mitigates Sharpening and Improves LLM Reasoning](https://mlsanity.com/paper/2511.19942/) (2025)
- [SofT-GRPO: Surpassing Discrete-Token LLM Reinforcement Learning via Gumbel-Reparameterized Soft-Thinking Policy Optimization](https://mlsanity.com/paper/2511.06411/) (2025)
- [Trace Length is a Simple Uncertainty Signal in Reasoning Models](https://mlsanity.com/paper/2510.10409/) (2025)
- [Beyond the Sampled Token: Preserving Candidate Support in RLVR](https://mlsanity.com/paper/2510.14807/) (2025)
- [Representation-Based Exploration for Language Models: From Test-Time to Post-Training](https://mlsanity.com/paper/2510.11686/) (2025)
- [The Best of N Worlds: Aligning Reinforcement Learning with Best-of-N Sampling via max@k Optimisation](https://mlsanity.com/paper/2510.23393/) (2025)
- [Risk-Sensitive RL for Alleviating Exploration Dilemmas in Large Language Models](https://mlsanity.com/paper/2509.24261/) (2025)
- [CLPO: Curriculum Learning meets Policy Optimization for LLM Reasoning](https://mlsanity.com/paper/2509.25004/) (2025)
- [Reinforcement Learning on Pre-Training Data](https://mlsanity.com/paper/2509.19249/) (2025)
- [Pass@k Metric for RLVR: A Diagnostic Tool of Exploration, But Not an Objective](https://mlsanity.com/paper/2511.16231/) (2025)
- [Rewarding the Rare: Uniqueness-Aware RL for Creative Problem Solving in LLMs](https://mlsanity.com/paper/2601.08763/) (2026)
- [Advantage Shaping as Surrogate Reward Maximization: Unifying Pass@K Policy Gradients](https://mlsanity.com/paper/2510.23049/) (2025)
- [SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning](https://mlsanity.com/paper/2602.01062/) (2026)
- [Token Hidden Reward: Steering Exploration-Exploitation in Group Relative Deep Reinforcement Learning](https://mlsanity.com/paper/2510.03669/) (2025)
- [Understanding Diversity Collapse in RLVR via the Lens of Overtraining](https://mlsanity.com/paper/2606.15455/) (2026)
- [Thinking with Map: Reinforced Parallel Map-Augmented Agent for Geolocalization](https://mlsanity.com/paper/2601.05432/) (2026)
- [DSDR: Dual-Scale Diversity Regularization for Exploration in LLM Reasoning](https://mlsanity.com/paper/2602.19895/) (2026)
- [Why Pass@k Optimization Can Degrade Pass@1: Prompt Interference in LLM Post-training](https://mlsanity.com/paper/2602.21189/) (2026)
- [Vector Policy Optimization: Training for Diversity Improves Test-Time Search](https://mlsanity.com/paper/2605.22817/) (2026)
- [Frontis-MA1: Training an AI4AI Model towards Recursive Self-Improvement in Machine Learning Engineering](https://mlsanity.com/paper/2607.28568/) (2026)
- [LaDi-RL: Latent Diffusion Reasoning Prevents Entropy Collapse in Reinforcement Learning](https://mlsanity.com/paper/2602.01705/) (2026)
- [Unveiling Implicit Advantage Symmetry: Why GRPO Struggles with Exploration and Difficulty Adaptation](https://mlsanity.com/paper/2602.05548/) (2026)
- [Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models](https://mlsanity.com/paper/2603.24844/) (2026)
- [ResRL: Boosting LLM Reasoning via Negative Sample Projection Residual Reinforcement Learning](https://mlsanity.com/paper/2605.00380/) (2026)

## Abstract (arXiv)

Reinforcement Learning (RL) algorithms sample multiple n>1 solution attempts for each problem and reward them independently. This optimizes for pass@1 performance and prioritizes the strength of isolated samples at the expense of the diversity and collective utility of sets of samples. This under-utilizes the sampling capacity, limiting exploration and eventual improvement on harder examples. As a fix, we propose Pass-at-k Policy Optimization (PKPO), a transformation on the final rewards which leads to direct optimization of pass@k performance, thus optimizing for sets of samples that maximize reward when considered jointly. Our contribution is to derive novel low variance unbiased estimators for pass@k and its gradient, in both the binary and continuous reward settings. We show optimization with our estimators reduces to standard RL with rewards that have been jointly transformed by a stable and efficient transformation function. While previous efforts are restricted to k=n, ours is the first to enable robust optimization of pass@k for any arbitrary k <= n. Moreover, instead of trading off pass@1 performance for pass@k gains, our method allows annealing k during training, optimizing both metrics and often achieving strong pass@1 numbers alongside significant pass@k gains. We validate our reward transformations on toy experiments, which reveal the variance reducing properties of our formulations. We also include real-world examples using the open-source LLM, GEMMA-2. We find that our transformation effectively optimizes for the target k. Furthermore, higher k values enable solving more and harder problems, while annealing k boosts both the pass@1 and pass@k . Crucially, for challenging task sets where conventional pass@1 optimization stalls, our pass@k approach unblocks learning, likely due to better exploration by prioritizing joint utility over the utility of individual samples.