Solving Harder Reinforcement Learning Tasks with Direct pass@k Optimization
Christian Walder, Deep Tejas Karkhanis
Reinforcement Learning algorithms commonly 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. Such algorithms under-utilize 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 feature a large maximum reward when considered jointly. Our primary contribution is to derive novel low variance unbiased estimators for the pass@k and its gradient, in both the binary and continuous reward settings. We show that the optimization with these estimators reduces to standard reinforcement learning with (mini-batches of) rewards that have been jointly transformed by a function that is stable and efficient to compute.While previous efforts propose transformations for $k=n$, our transformations are the first to enable robust optimization of the pass@k for any arbitrary $k \leq n$. Rather than simply 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 performance alongside significant pass@k gains.We validate our reward transformations on illustrative toy experiments, which reveal the variance reducing properties of our formulations. We also include real-world examples using the open-source model 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.