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

2026-07-20

AI Summary

This work addresses the theoretical foundations of $\max$@$k$ reinforcement learning in finite-horizon settings, motivated by the widespread yet theoretically unsupported use of the $\max$@$k$ evaluation criterion in tasks such as code generation and theorem proving. We uncover fundamental differences between the $\max$@$k$ objective and the classical expected return, proving that Markovian policies are generally suboptimal under this criterion. To restore optimality, we propose a compact state-augmentation scheme that effectively captures necessary historical information. Furthermore, by quantifying the performance gap between history-dependent and history-independent policies, we establish a statistical complexity lower bound for $\max$@$k$ learning and develop an algorithm that achieves optimal sample complexity.

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.

Problem

• max@k
• reinforcement learning
• theoretical foundations
• finite-horizon
• retry-aware metric

Innovation

• max@k reinforcement learning
• state augmentation
• history-dependent policies
• sample complexity
• theoretical foundations