2026/07/20 by Riccardo Poiani, Martino Bernasconi, Andrea Celli
#cs.LG
Reinforcement Learning is a cornerstone technique for modern large reasoning models. Usually, for difficult tasks such as code generation and theorem proving, the agent is evaluated by generating K responses rather than sampling a single response, and performance is then measured using a retry-aware metric such as max@k. Despite their practical importance, the theoretical foundations of learning under such criteria remain limited. In this work, we provide a theoretical study of the max@k learning problem in finite-horizon reinforcement learning. We show that optimizing the max@k objectives is fundamentally different from standard expected-return maximization. In particular, we prove that Markovian policies are in general insufficient, identify a compact state augmentation that restores optimality, and explicitly characterize the performance gap that can arise between history-dependent and non-history-dependent policies. Moreover, we show that learning max@k-optimal policies is statistically harder than standard reinforcement learning and provide an efficient algorithm that achieves the optimal sample complexity rate.