English

Theoretical Foundations of $\max$@$k$ Reinforcement Learning

Machine Learning 2026-07-20 v1

Abstract

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 KK responses rather than sampling a single response, and performance is then measured using a retry-aware metric such as max\max@kk. 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\max@kk learning problem in finite-horizon reinforcement learning. We show that optimizing the max\max@kk 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\max@kk-optimal policies is statistically harder than standard reinforcement learning and provide an efficient algorithm that achieves the optimal sample complexity rate.

Cite

@article{arxiv.2607.17823,
  title  = {Theoretical Foundations of $\max$@$k$ Reinforcement Learning},
  author = {Riccardo Poiani and Martino Bernasconi and Andrea Celli},
  journal= {arXiv preprint arXiv:2607.17823},
  year   = {2026}
}