cs.LGJul 20, 2026

Theoretical Foundations of max\max@kk Reinforcement Learning

Authors: Riccardo PoianiMartino BernasconiAndrea Celli

Organizations: Bocconi University

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.

Explore similar work

CardsList
  1. On Advantage Estimates for Max@K Policy Gradients

    Jun 4, 2026Shota Takashiro, Soichiro Nishimori, Paavo Parmas +6Policy GradientInference-Time