Reinforcement Learning for Hierarchical Reasoning Rewards: Minimax-Optimal Rates with Transformers

Naoki Nishikawa; Taiji Suzuki

cs.LGOct 6, 2026

Reinforcement Learning for Hierarchical Reasoning Rewards: Minimax-Optimal Rates with Transformers

Authors: Naoki Nishikawa, Taiji Suzuki

Organizations: The University of Tokyo, RIKEN AIP

Abstract

Reinforcement learning (RL) has become a standard tool for post-training language models on reasoning tasks, where the policy is updated by reward feedback while exploring the space of responses. Despite its empirical success, theoretical understanding of RL post-training remains limited, in particular of why on-policy exploration combined with a neural reward model is effective. In this paper, we address this question by modeling the reward as a hierarchical function on the response space: the reward consists of infinitely many local components, each of which becomes relevant only after the preceding ones have been resolved. We show that a natural Transformer-based actor--critic algorithm, which alternates between sampling from the current KL-regularized policy, fitting a Transformer critic to the observed rewards, and updating the policy, achieves the minimax optimal rates in the query budget and in the regularization strength up to logarithmic factors, and is minimax optimal for a fixed number of prompts. In contrast, we prove that sampling from the fixed reference distribution, as in offline reward modeling, can limit regret decay to a logarithmic rate. These results show that on-policy exploration progressively zooms in on the region where the reward is concentrated, and quantify its benefit for RL post-training.

Figures & tables

Figure 1: A hierarchically nested reward ( Definition 1 ), for

d=1

and zoom depths

\ell_{i}=2

. (Top) A response is a path in a tree with

2^{d}

children per node, and a prefix of depth

h

is a dyadic cube of that depth. The primitive

g_{1}

lives on

\Omega

; inside the cell

C_{1}

around its maximizer,

g_{2}

is added, rescaled to

C_{1}

, and so on. (Bottom) The resulting reward

f=\sum_{i}c_{i}h_{i}

: bumps of decreasing weight

c_{i}\asymp i^{-\gamma}

and decreasing width accumulate around the maximizer

y_{\infty}=\bigcap_{i}C_{i}

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$c_{i}=i^{-\gamma}/\zeta(\gamma)$	weight of level $i$ of the hierarchy
$T_{k}=\sum_{i>k}c_{i}$	weight of the levels deeper than $k$
$L_{i}=\sum_{a\leq i}\ell_{a}$	depth at which level $i$ is fully resolved
$C_{i}(m)$ , $E_{a}(m)$	active cell of level $i$ ; shell $C_{a-1}\setminus C_{a}$
$f^{[k]}=\sum_{i\leq k}c_{i}h_{i}$	reward truncated at level $k$
$G_{m}=\max_{y}f(y;m)$	optimal reward of prompt $m$

Appendix

Table 1: Notation used throughout the appendix.

Symbol	Meaning	Defined in
$a_{j}=c_{1+\lfloor j/\bar{\ell}\rfloor}$	accuracy scale attached to depth $j$	( 14 )
$\zeta_{M}=2+\frac{d}{2\alpha}$	exponent of $M$ in the batch sizes	( 14 )
$W_{h}$ , $\mathcal{D}_{h}(m)$	score at phase $h$ ; partition on which it is constant	Lemma 19
$\eta$	birth-accuracy level	Definition 11
$P^{\ast}_{h}(m)$ , $k(h)$	correct depth- $h$ cube; the level it has reached	Section B.3
$p_{\ast}$	Gibbs mass of the correct cube	Lemma 12

Appendix

Table 2: Notation shared by the two upper bounds.

Symbol	Meaning
$H$	number of front-phase iterations
$R$ , $S=R+1$	refinement depth of the lookahead; the lookahead itself
$\eta$ , $r_{\rm mom}$	birth-accuracy level; moment order used in the regret transfer
$\delta_{F}$ , $\delta_{h}$	total and per-phase failure budgets
$T_{\mathrm{in}}$	input length of the front-phase critic
$b$ , $r$ , $N_{\rm ref}$	batch, query depth and approximation size of the final refinement step

Appendix

Table 3: Quantities used only in the proof of the large-budget bound of Theorem 3 .

Explore similar work

Sep 29, 2026cs.LG

Unlocking the Critic: Reward-Free Policy Optimization for LLM Post-Training

Recent approaches to reinforcement learning (RL) post-training for large language models increasingly remove the critic to reduce training instability and memory overhead. Even where a critic is trained, it is discarded once training ends, although it has learned to predict outcomes. We revisit this trend and show that a pretrained critic's ability to predict future outcomes can make it a valuable asset for efficient long-horizon reasoning. First, we find that instability in critic-based RL for long chain-of-thought reasoning is largely an optimization artifact: keeping policy updates small and low in variance restores stable convergence. Second, a well-pretrained critic estimates the posterior probability of eventual success from later trajectory states and unfinished prefixes. Its predictions provide outcome-derived, dense, per-prefix learning signals that, during policy optimization, require neither completed rollouts, step-level annotations, nor external reward labels. Building on this insight, we introduce Reward-Free Policy Optimization (RFPO), which repurposes a single calibrated, frozen critic as a rollout-level reward, a value baseline for generalized advantage estimation, and a success forecaster for unfinished prefixes. We further show that binarizing the debiased score stops the policy from exploiting the critic's length bias. Binarized, RFPO matches supervised PPO without a single label in the training loop, while cutting compute and memory overhead. This makes RFPO well suited to long-horizon reasoning tasks, where outcomes arrive late and generation dominates cost: because rollouts can be rewarded before they finish, training no longer has to pay for waiting on every trajectory to complete. Our findings challenge the prevailing critic-free paradigm and establish critic-based, reward-free optimization as a scalable and computationally efficient path for LLM post-training.

Hongyang Li, Xiao Li, Caesar Wu +3

University of Luxembourg · Seafill Open-Source Community · Université Paris-Saclay

Critic-Free Reinforcement Learning Reward-Only Policy Optimization Variants Critic Long-Horizon Task Planning Position Bias

Ask PDF

May 23, 2026cs.LG

RL with Learnable Textual Feedback: A Bilevel Approach

Reinforcement learning with verifiable rewards can improve LLM reasoning, but learning remains sample-inefficient when terminal rewards are sparse. This has motivated a growing line of work on RL with textual feedback, where a critic model generates natural language feedback to guide a reasoning model (the actor), augmenting scalar rewards with richer learning signals. However, existing methods typically treat feedback as fixed or auxiliary, which misses a key property: feedback should not merely be correct, but should improve the policy (actor model) when provided in context. This motivates a paradigm of learnable textual feedback for RL. Yet the learnability and usefulness of feedback depend on the policy's ability to learn from it, making RL with learnable feedback an inherently bilevel problem. We formalize this coupling as a Stackelberg bilevel program and derive Bilevel Natural Language Actor-Critic (Bi-NAC), which jointly trains a critic to generate reward-improving feedback and an actor to exploit it. Across MATH-500, MBPP, and GPQA, Bi-NAC improves sample and parameter efficiency over RL and fixed-critic baselines: our 2B model outperforms the 3B GRPO baseline, achieving 46.6% versus 41.4% on MATH-500, while our 6B model surpasses the 7B GRPO baseline, achieving 49.3% versus 43.6% on GPQA.

Utsav Singh, Sidhaarth Sredharan, Souradip Chakraborty +1

University of Central Florida · 2Carnegie Mellon University · 3Universtiy of Maryland

LLM Reasoning Strategies Feedback Critic

Ask PDF

Jun 23, 2026cs.LG

ExTra: Exploratory Trajectory Optimization for Language Model Reinforcement Learning

Reinforcement Learning with Verifiable Rewards (RLVR) for language-model reasoning can fail at both extremes of task difficulty: easy prompts often produce all-correct, low-diversity rollout groups with little gradient signal, while hard prompts can produce all-incorrect groups with no positive reward. We introduce ExTra (Exploratory Trajectory Optimization), a GRPO-compatible framework that extracts exploration signals from the model's own rollouts. ExTra combines two mechanisms: (i) a novelty reward that adds embedding-based diversity bonuses after GRPO normalization, rewarding diverse correct solutions; and (ii) entropy-guided prefix regeneration, which scores partial trajectories using entropy signals and continues exploration from promising intermediate steps. Across six mathematical reasoning benchmarks, ExTra improves Qwen3-1.7B over GRPO by about +5 points on pass@1 and +7 points on pass@16, showing that trajectory-level exploration signals can improve both single-sample accuracy and inference-time coverage.

Wenyang Hu, Junxiang Jia, Zhen Shu +3

1National University of Singapore · 2SAP

Verifiable Rewards Trajectory-Level Credit Robust Trajectory

Ask PDF