Theoretical Foundations of $\max$@$k$ Reinforcement Learning
2026-07-20 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study how reinforcement learning works when performance is measured by the best outcome among multiple tries (called max@k), rather than just one try. They find that this setup is quite different from usual methods that focus on average returns, showing that simple memoryless strategies (Markovian policies) often don’t work well. They introduce a way to keep track of past information to make the approach optimal again and prove that learning under this criterion is more difficult. Finally, they provide a learning algorithm that uses samples efficiently to handle this harder problem.
Reinforcement Learningmax@kMarkovian PoliciesFinite-HorizonExpected ReturnPolicy OptimalitySample ComplexityState AugmentationHistory-Dependent Policies
Authors
Riccardo Poiani, Martino Bernasconi, Andrea Celli
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 $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.