PaperlayerView exact revision on arXiv

Source-linked AI summary

Tail-Likelihood Reinforcement Learning

Shrinivas Ramasubramanian, Daman Arora, Fahim Tajwar, Guanning Zeng, Qingyang Wu, Zhongzhu Zhou, Chenfeng Xu, Haiwen Feng, Yuda Song, Aarti Singh, Ruslan Salakhutdinov, J. Andrew Bagnell, Jeff Schneider, Andrea Zanette

arXiv:2609.02987v1cs.LGstat.ML

TL;DR

Average reward can miss differences in rare high-reward rollout coverage, which matters when additional samples are used during training or inference. TailRL optimizes the log-probability of exceeding randomly chosen reward thresholds, emphasizing rare upper-tail outcomes. Across four settings, it improves learning and inference-time scaling, including strong results when high-reward rollouts are rare or suboptimal behaviors dominate.

  • Problem

    Mean reward does not fully characterize generative policies because similarly averaged policies can have different probabilities of producing rare, exceptionally good rollouts.

  • Method

    TailRL views continuous rewards as threshold-indexed binary events and maximizes the expected log-likelihood of exceeding a uniformly chosen threshold.

  • Results

    Across object localization, maze navigation, GUI grounding, and code optimization, TailRL matches or exceeds relevant baselines, including 128–256× fewer inference rollouts for GUI grounding and a 7.7× Best-of-1024 speedup for code optimization.

  • Takeaways & Limitations

    TailRL provides a practical critic-free modification to existing RL pipelines that leverages rare high-reward samples and improves benefits from additional inference samples.

  • Takeaways & Limitations

    The main formulation assumes a deterministic scalar reward r(x, z) in [0, 1], with more general bounded ranges treated separately.

Abstract

from arXiv · show

Reinforcement learning typically optimizes average reward. For generative policies, the average can hide an important distinction: two policies can achieve the same mean reward while having very different chances of producing a rare but high-reward rollout. This matters as sampling increases during training and inference, since its benefit depends on retaining probability mass on high-reward outcomes. We propose to optimize this coverage directly. Rather than considering only expected reward, we consider all of its upper tails: for each reward threshold, how likely is the policy to exceed it? This turns a continuous reward into a family of binary success events. We introduce Tail-Likelihood Reinforcement Learning (TailRL), which maximizes the log-probability of exceeding a randomly chosen reward threshold. Its gradient gives more weight to rare, high-reward rollouts and can be interpreted as a mixture of Best-of-(k) gradients. TailRL requires only a simple modification to the advantage function, making it compatible with existing reinforcement learning pipelines. Across object localization, maze navigation, GUI grounding, and code optimization, TailRL leverages rare high-reward training samples to avoid suboptimal solutions and yields models that benefit more from additional samples at inference time.

1 Introduction

TailRL addresses the loss of rare high-reward rollout coverage that average-reward training can overlook. It optimizes upper-tail likelihoods and shows gains across four settings, especially when high-reward outcomes are rare or training favors suboptimal behavior.

  • Motivation: Generative policies with similar mean rewards can differ substantially in their probability of producing rare, exceptionally good rollouts, affecting training and inference-time sampling.Standard RL can lose coverage over rare high-reward rollouts, limiting Best-of-k performance and gains from additional samples.
  • Method: TailRL treats each continuous reward threshold as a binary success event and maximizes the log-likelihood of exceeding a uniformly chosen threshold.For rewards in [0, 1], this preserves information across reward levels rather than selecting one success threshold.
  • Method: TailRL emphasizes rarer, higher-reward outcomes by optimizing upper-tail probabilities across reward levels instead of only the reward mean.Its gradient assigns greater weight to thresholds that are difficult for the policy to reach.
  • Method: The TailRL gradient decomposes harmonically over Best-of-k gradients, linking training to inference-time scaling without choosing one inference budget in advance.This provides a direct connection between the training objective and coverage of high-reward rollouts.
  • Method: A finite group of N rollouts defines an order-N objective interpolating from expected reward at N = 1 to the population objective as N →∞, with an unbiased gradient estimator.Unlike REINFORCE, increasing the rollout budget changes the objective being optimized rather than only reducing variance.
  • Results: Across object localization, maze navigation, GUI grounding, and code optimization, TailRL matches or outperforms baselines and avoids suboptimal solutions using rare high-reward samples.It learns from initial policies with 0.01% success where expected-reward baselines fail, matches RLOO’s Pass@1024 with 128–256× fewer inference rollouts, and reaches a 7.7× Best-of-1024 speedup on code optimization.

2 Preliminaries

The paper models a policy as a generator of rewards and contrasts expected reward with sampling-sensitive metrics that emphasize high-reward outcomes. Binary rewards motivate MaxRL, which increases the emphasis on rare successes.

  • Inference-Time Scaling: Best-of-k is the expected maximum reward from k independent continuous-reward rollouts, while Pass@k is the probability of at least one success for binary rewards.Both metrics reduce to mean reward at k = 1 and approach the maximum supported reward as k increases.
  • Policy as a Generative Model of Rewards: A policy and deterministic reward function induce a reward distribution whose randomness comes from policy sampling.Expected-reward RL shapes this distribution using only its mean.
  • Inference-Time Scaling: Policies with the same mean reward can differ substantially in Best-of-k and Pass@k because they assign different probability to high-reward outcomes.This motivates objectives that act on the full reward distribution rather than only its mean.
  • Threshold Events: Continuous rewards can be decomposed into threshold events, with each rollout clearing every threshold below its reward.Groups of rollouts estimate tail probabilities across reward thresholds.
  • Binary Rewards and MaxRL: For binary rewards, expected reward equals success probability, whereas MaxRL maximizes its log-probability and therefore emphasizes rare successes.Its gradient is a harmonic mixture of Pass@k gradients.

3 Tail-Likelihood Reinforcement Learning

TailRL extends likelihood maximization from binary success to all continuous-reward thresholds. It emphasizes difficult upper-tail events and has an exact harmonic connection to Best-of-k sampling, while recovering MaxRL for binary rewards.

  • TailRL Objective: TailRL maximizes the expected log-probability that a rollout exceeds a uniformly sampled reward threshold.It applies the MaxRL likelihood principle at every threshold rather than selecting one cutoff.
  • TailRL Objective: Expected reward aggregates tail probabilities arithmetically, whereas TailRL aggregates their logarithms, making small tail probabilities more influential.Higher thresholds receive greater gradient weight because they are harder to reach.
  • Harmonic Decomposition over Best-of-k: TailRL’s gradient decomposes into a harmonic mixture of Best-of-k gradients, combining learning signals across all inference budgets.The logarithm supplies harmonic weights 1/k, so no target inference budget is selected in advance.
  • Recovery of MaxRL for Binary Rewards: For binary rewards, every nontrivial threshold represents the same success event, so TailRL reduces exactly to MaxRL.Best-of-k also coincides with Pass@k in this setting.
  • TailRL Objective: Using only a chosen success threshold can make lower and higher rewards equivalent after crossing the threshold, so TailRL preserves the continuous reward signal.TailRL uses the full reward support and introduces no weighting hyperparameter under uniform threshold sampling.
  • Probabilistic Interpretation of TailRL: The probabilistic interpretation combines threshold events using independent rollouts, whose audit probability factorizes across thresholds.Independence is necessary because reusing one rollout collapses the nested events to the hardest threshold.

4 Estimating the TailRL Gradient

Finite rollout groups define tractable truncated TailRL objectives and unbiased gradient estimators. Increasing the rollout budget changes the optimized objective, progressively incorporating stronger Best-of-k signals rather than merely reducing variance.

  • From Finite Rollouts to a Finite-Order Objective: A finite rollout budget N defines an order-N truncation of the population TailRL objective and admits an unbiased policy-gradient estimator.The population objective cannot be estimated directly from finite rollouts.
  • From Finite Rollouts to a Finite-Order Objective: The finite-order threshold weight equals 1 at T = 1 and approaches 1/pθ(x, τ) as T →∞.Increasing T smoothly shifts optimization from expected-reward RL toward population TailRL.
  • Finite-Rollout Estimator: Sampled rollout weights are computed exactly after sorting rewards, with thresholds cleared by fewer rollouts receiving more weight.This concentrates credit on harder-to-reach reward levels.
  • Finite-Rollout Estimator: TailRL differs from a standard critic-free policy-gradient method only in how sampled rewards become rollout weights.The resulting estimator retains the standard score-function form.
  • Finite-Rollout Estimator: One rollout recovers the expected-reward gradient, while larger groups incorporate progressively higher Best-of-k learning signals.Unlike REINFORCE, increasing rollout count changes the optimized objective rather than only reducing estimation variance.
  • Finite-Rollout Estimator: Centering rollout weights reduces variance and makes TailRL a drop-in replacement for standard policy-gradient advantage calculation.The centered estimator remains unbiased for the corresponding truncated objective.

5 Unifying Gradient Weight View

The paper unifies expected-reward and TailRL objectives by the gradient weights assigned to tail probabilities. Expected reward weights all tails equally, whereas TailRL increasingly emphasizes rare, high-threshold outcomes.

  • Objective Comparison: Figure 3 compares Order-T TailRL’s approach to the population weight with PKPO’s capped weight that vanishes as tail probability grows.The comparison highlights distinct weighting behavior across the objectives.
  • Gradient Weighting: TailRL and expected-reward maximization differ only in how they weight gradients of tail probabilities.The scalar function ϕ determines the gradient weight ϕ′ assigned to each tail probability.
  • Expected Reward: Expected reward uses ϕ(p)=p and ϕ′(p)=1, assigning equal weight regardless of how often each tail occurs.This provides the baseline weighting in the unified objective family.
  • TailRL: TailRL uses ϕ(p)=log p and ϕ′(p)=1/p, emphasizing higher reward thresholds reached rarely.Because tail probability decreases with the threshold, inverse-probability weighting favors rarer upper tails.
  • Truncated TailRL: Order-T TailRL interpolates between expected reward at T=1 and population TailRL as T approaches infinity.Its multiplier is 1 at T=1 and approaches 1/p as T→∞.

6 Experiments

The experiments evaluate TailRL’s population objective, finite-rollout approximation, rare-success learning, inference-time scaling, and resistance to moderate-reward shortcuts across four tasks. Results show advantages over expected-reward baselines and competitiveness with supervised objectives in the reported settings.

  • Experimental Questions: The experiments test population-level learning, finite-rollout convergence, rare high-reward learning, inference-time scaling, and resistance to suboptimal shortcuts.The four settings are ImageNet localization, Text-Maze navigation, GUI grounding, and Code Optimization.
  • ImageNet Object Localization: On ImageNet localization, TailRL matches or exceeds objectives that directly supervise ground-truth coordinates using only scalar IoU rewards.It achieves higher CorLoc@0.5 and CorLoc@0.75 than L1+GIoU while obtaining comparable mean IoU.
  • ImageNet Object Localization: Increasing training rollouts moves finite-rollout TailRL toward the exact population objective across reported localization metrics.The same convergence appears at the gradient level as the sampled gradient approaches the population-level gradient.
  • ImageNet Object Localization: At matched N=1024, TailRL outperforms GRPO and RLOO across CorLoc@0.5, mean IoU, and Best-of-1024 IoU.With N=16, TailRL also exceeds both baselines trained with N=1024 and outperforms PKPO on CorLoc@0.5 and mean IoU.
  • Text-Maze Navigation: TailRL outperforms expected-reward methods when the initial policy has poor coverage of rare excellent maze rollouts.RLOO and GRPO fail to learn reliably in the low-initial-success regime, while TailRL benefits from rare exceptional rollouts.
  • GUI Grounding: TailRL’s Pass@k continues rising with additional GUI-grounding samples, unlike RLOO’s earlier plateau and GRPO’s inferior scaling.On ScreenSpot-Pro, TailRL reaches RLOO’s mean Pass@1024 with 8 rollouts at 3B and 4 rollouts at 7B.
  • Code Optimization: In Code Optimization, TailRL maintains higher entropy and raises mean reward to 2.92, nearly three times the copying value.Expected-reward baselines collapse toward a copying shortcut with correctness above 98% and mean reward just below 1.0.

7 Related Work

The paper situates TailRL among objectives that optimize inference-time success or tail-sensitive behavior rather than expected reward. It distinguishes TailRL by extending likelihood-based tail optimization from binary to continuous rewards without adding a separate diversity regularizer.

  • Expected-Reward RL: GRPO and RLOO are critic-free score-function policy-gradient variants that alter baselines or normalization without changing the objective.RLOO changes the baseline, while GRPO applies common per-input normalization.
  • Binary-Reward Objectives: MaxRL maximizes log success probability for binary rewards and decomposes its gradient into a harmonic sum of Pass@k gradients.TailRL extends this likelihood-based perspective to continuous rewards.
  • Inference-Time Objectives: PKPO directly targets a selected Pass@k or Best-of-k objective using unbiased finite-rollout estimators.Other related methods likewise optimize selected inference-time metrics or monotone transformations of success probability.
  • Coverage-Sensitive RL: TailRL addresses coverage loss by weighting reward levels inversely by their tail probabilities through the objective itself rather than adding a separate diversity regularizer.This targets the narrowing of policy distributions that can remove rare, high-reward behavior.

8 Conclusion

TailRL reframes continuous rewards as thresholded success events and optimizes their likelihood, connecting training to high-reward coverage and Best-of-k performance. It provides finite-rollout estimators and theoretical links to MaxRL, while experiments show benefits when rare high-reward outcomes or suboptimal behaviors dominate.

  • 8 Conclusion: TailRL maximizes the expected log-probability of exceeding a uniformly sampled reward threshold for continuous rewards.It recovers MaxRL exactly for binary rewards.
  • 8 Conclusion: Its gradient decomposes harmonically over Best-of-k gradients, linking the objective to coverage of high-reward rollouts and inference-time scaling.The finite-rollout estimator is critic-free and changes existing reinforcement-learning pipelines through the advantage calculation.
  • 8 Conclusion: Across four settings, TailRL gains most when high-reward rollouts are rare or training favors a common but suboptimal behavior.In these regimes, it learned more reliably from rare outcomes and benefited more from additional training and inference samples.
  • 8 Conclusion: Threshold weighting is equivalent to choosing a monotone coordinate system for reward quality, while nonuniform weighting remains outside the work’s scope.Uniform weighting uses the original normalized reward axis and adds no function or hyperparameter.

C.4 Harmonic Best-of-k Expansion of the Tail-Likelihood

The TailRL objective is the limit of harmonic partial sums of Best-of-k objectives, with the first truncation matching expected-reward reinforcement learning up to a constant. Its finite-rollout estimator is unbiased and reduces to MaxRL for binary rewards.

  • C.4 Harmonic Best-of-k Expansion of the Tail-Likelihood: Each order-T TailRL truncation is a partial sum of Best-of-k objectives with harmonic coefficients, converging to the full tail-likelihood.The truncations are upper bounds that tighten as T grows.
  • C.4 Harmonic Best-of-k Expansion of the Tail-Likelihood: TailRL at order 1 equals Best-of-1 and expected-reward reinforcement learning up to an additive constant, sharing its gradient.Thus, increasing rollout order changes the optimized objective rather than merely reducing estimator variance.
  • C.4 Harmonic Best-of-k Expansion of the Tail-Likelihood: The truncated gradient is a positively weighted combination of Best-of-k gradients for k = 1, . . . , T, with coefficient 1/k.This gives the gradient its harmonic Best-of-k interpretation.
  • C.4 Harmonic Best-of-k Expansion of the Tail-Likelihood: The finite-rollout estimator is unbiased for the order-N truncated gradient under the stated assumptions.Its proof interprets the estimator as MaxRL applied at every reward threshold on one shared rollout group.
  • C.4 Harmonic Best-of-k Expansion of the Tail-Likelihood: For binary rewards, successful rollouts receive weight 1/M, unsuccessful rollouts receive zero, and the uncentered estimator reduces to MaxRL.After mean-centering, successful weights are 1/M −1/N and unsuccessful weights are −1/N.

D TailRL on a General Reward Range

TailRL extends from rewards in [0, 1] to any bounded reward range by averaging log-tail-probabilities over uniformly sampled thresholds. Shifting or rescaling rewards preserves the objective’s results and, after normalization, leaves the estimator algorithm unchanged.

  • D TailRL on a General Reward Range: For rewards in [rmin, rmax], TailRL averages log-tail-probability over thresholds uniformly drawn from that range.The normalization factor prevents the objective scale from depending on reward units.
  • D TailRL on a General Reward Range: Affine reward rescaling is equivalent to applying the unit-interval objective to the rescaled reward ˜r := (r −rmin)/(rmax −rmin).Shifting or rescaling rewards only relabels thresholds, so the paper’s results carry over directly.
  • D TailRL on a General Reward Range: The normalized finite-rollout estimator operates on raw rewards from any bounded range without changing the algorithm.Reward rescaling multiplies every weight by rmax −rmin, and the normalization cancels this factor.

E Connection to Ordinal Cross-Entropy

TailRL is exactly ordinal cross-entropy against the ideal target r = 1, viewing each reward threshold as a binary event. Its finite-rollout procedure estimates this objective from sampled rewards, while the broader evaluation compares TailRL across reward-driven tasks and sampling metrics.

  • E Connection to Ordinal Cross-Entropy: Threshold decomposition represents a continuous target as binary events, one for each threshold, and fits those events with binary classification.The paper identifies TailRL as the reinforcement-learning counterpart of this decomposition.
  • E Connection to Ordinal Cross-Entropy: Maximizing TailRL exactly minimizes ordinal cross-entropy against the ideal target r = 1.For binary rewards, the integral collapses to the MaxRL objective log qθ(x).
  • E Connection to Ordinal Cross-Entropy: With N rollouts, TailRL truncates the threshold integral at the resolution supported by the group and estimates its gradient unbiasedly.As N →∞, the finite procedure recovers the ordinal cross-entropy objective at target r = 1.
  • E Connection to Ordinal Cross-Entropy: RLOO, GRPO, and TailRL share a critic-free update template and differ in how rewards are mapped to advantages.TailRL weights rare threshold clearances more heavily and reduces to MaxRL advantages for binary rewards.
  • E Connection to Ordinal Cross-Entropy: Pass@k measures binary success probability, whereas Best-of-k measures expected maximum reward for continuous rewards; they coincide for binary rewards.Both evaluate sampling k rollouts and retaining the best one.

H.3 Additional Results

Additional object-localization results show that TailRL improves with larger training rollout budgets, approaches the population objective, and preserves useful continuous-reward information. In Text-Maze, its performance depends on encountering rare successful rollouts during training.

  • ImageNet object localization: TailRL improves steadily as the training rollout budget N increases and approaches the exact population objective, unlike the expected-reward baselines.At every evaluated budget, both baselines perform worse than TailRL at its smallest reported budget.
  • ImageNet object localization: Difficulty bands partition localization validation images by ground-truth box area into easy, medium, and hard subsets.Easy covers areas from 0.30 to 0.70; medium covers 0.10 to 0.30 or 0.70 to 0.95; hard is below 0.10 or above 0.95.
  • ImageNet object localization: At CorLoc@0.9, TailRL reaches approximately 0.39, compared with 0.24 for the 0.75-threshold MaxRL variant and 0.02 for the 0.5-threshold variant.Binarized methods perform well near their chosen thresholds but degrade at stricter quality levels.
  • Text-Maze: Text-Maze rewards combine formatting, progress toward the goal, and path efficiency, with reward 1 reserved for shortest-path success.The policy is trained from supervised checkpoints spanning very low to higher initial success rates.

I.5 Additional Results

Additional Text-Maze results show that TailRL gains more from inference and training rollouts than expected-reward baselines, especially when initial shortest-path success is rare. Its higher entropy and longer generated paths indicate continued exploration during training.

  • Inference-time scaling: As inference rollouts increase, Pass@k and Best-of-k reward improve more rapidly for TailRL than for GRPO or RLOO.Single-sample evaluation understates the differences among the learned rollout distributions.
  • Training rollout budget: For an initial shortest-path success rate of approximately 0.02%, TailRL’s Pass@k and Best-of-k reward increase sharply between N = 4 and N = 16.GRPO remains near the floor, while RLOO remains at zero shortest-path success across the evaluated budgets.
  • Training dynamics: The expected-reward baselines lose entropy early and stop changing their rollout distributions, whereas TailRL retains substantially more entropy throughout training.TailRL’s generated paths grow steadily longer during training.
  • Evaluation caveat: Published ScreenSpot-Pro results are not directly comparable because they use a different tool-call format.The study instead compares methods under the same pipeline.

J.7 Additional Results

Additional GUI-grounding and code-optimization results show that TailRL’s advantage emerges with larger inference budgets and persists throughout training. In code optimization, it preserves entropy and produces high-speedup rewrites rather than converging on copying.

  • GUI-grounding inference scaling: At k = 128, TailRL leads the better baseline in eleven of twelve scale-category cells, by up to 12.0 points at 3B and 11.5 points at 7B.The advantage appears across task categories rather than through a single category.
  • GUI-grounding training dynamics: At every inference budget larger than one, TailRL separates from both baselines within the first epoch and maintains that separation throughout training.At k = 1, the three post-trained methods remain close together.
  • GUI-grounding reward decomposition: The format reward saturates within the first epoch for every method, so method differences are carried by the point term.This decomposition separates parsing compliance from localization quality.
  • GUI-grounding training dynamics: TailRL maintains higher policy entropy than expected-reward baselines at both model scales, matching its continued Pass@k growth with larger inference budgets.The same training-time pattern appears in Text-Maze.
  • Code runtime optimization: The code-optimization reward has an unbounded upper tail, while copying is a safe but suboptimal reward-1 shortcut.This creates a setting where preserving probability mass on exceptional rewrites matters.
  • Code runtime optimization: TailRL’s Best-of-k density is centered near 5× at k = 1 and shifts rightward as k grows, while GRPO and RLOO remain pinned at the copy line.The pretrained model improves toward 2× but does not reach TailRL’s region.

K.7 Sample Generations

The sample generations contrast TailRL’s verified high-speedup rewrite with baseline and pretrained outputs on a held-out edit-distance problem. The example illustrates that selection exposes substantial differences among the learned rollout distributions.

  • Program example: The displayed program computes edit distance with a two-dimensional dynamic program over the input strings.The rewritten sample preserves the stated input/output behavior and algorithmic structure.
  • Baseline samples: GRPO and RLOO reproduce the input program byte for byte inside their sampled code blocks.Their outputs therefore exemplify the copying shortcut described in the code-optimization results.
  • Held-out sample: At step 300, TailRL’s best of 64 rollouts achieves a verified 27.0× speedup on held-out problem p02318.The pretrained model’s best of 64 rollouts reaches 5.05×.
  • Program example: The sample rewrite changes headers, input routines, and loop presentation while claiming identical output for valid inputs.The accompanying explanation characterizes these as performance-oriented implementation changes.
Loading 2609.02987v1…