July
2026
The Tail and the Judge: Better Policy Gradients and World Models for Reinforcement Learning
Authors:
Abstract:
A quiet shift has taken place in how reinforcement learning delivers results.
Increasingly, performance comes not from a single attempt but from selection
over many: a language model is sampled repeatedly and judged by its best
verified answer, and a planning agent imagines many action sequences and
executes the one its learned world model scores highest. Selection can fail
in exactly two ways. The candidate pool may contain nothing worth choosing,
or the judge that scores the candidates may be wrong precisely where it
matters. This thesis argues that standard training objectives invite both
failures, and it redesigns them to prevent each.
The first part trains the policy for selection itself. Standard methods
maximize the average reward of a single attempt, which spreads learning
effort evenly across quality levels and neglects the rare, excellent outcomes
that selection exists to find. Extending maximum-likelihood reinforcement
learning from binary to continuous rewards, we introduce a tail-likelihood
objective that credits the policy for covering every level of quality and
concentrates its gradient on the levels it reaches only rarely. We prove that
this objective aligns training with best-of-many deployment at every sampling
budget simultaneously, and that it is optimized by an unbiased,
hyperparameter-free estimator amounting to a one-line change in standard
pipelines. Empirically, the method recovers the supervised gold-standard
gradient where popular baselines remain permanently misaligned; it lifts
maze policies from a near-zero starting success rate to solving most mazes
optimally; and it is evaluated further on object localization and
program-speed optimization.
The second part secures the judge. A planner that maximizes predicted return
will seek out exactly the regions where its world model is most optimistically
wrong. We prove that the resulting loss in policy quality is governed by the
sharpness of the model's loss landscape, and we show that sharpness-aware
training of the world model alone, leaving the planner and policy untouched,
delivers large and consistent gains that transfer across high-dimensional
continuous control and pixel-based discrete control, and across two different
model-based algorithms.
Good selection needs something worth choosing and a judge worth trusting.
Providing both, and composing them, is the program of this thesis.
Increasingly, performance comes not from a single attempt but from selection
over many: a language model is sampled repeatedly and judged by its best
verified answer, and a planning agent imagines many action sequences and
executes the one its learned world model scores highest. Selection can fail
in exactly two ways. The candidate pool may contain nothing worth choosing,
or the judge that scores the candidates may be wrong precisely where it
matters. This thesis argues that standard training objectives invite both
failures, and it redesigns them to prevent each.
The first part trains the policy for selection itself. Standard methods
maximize the average reward of a single attempt, which spreads learning
effort evenly across quality levels and neglects the rare, excellent outcomes
that selection exists to find. Extending maximum-likelihood reinforcement
learning from binary to continuous rewards, we introduce a tail-likelihood
objective that credits the policy for covering every level of quality and
concentrates its gradient on the levels it reaches only rarely. We prove that
this objective aligns training with best-of-many deployment at every sampling
budget simultaneously, and that it is optimized by an unbiased,
hyperparameter-free estimator amounting to a one-line change in standard
pipelines. Empirically, the method recovers the supervised gold-standard
gradient where popular baselines remain permanently misaligned; it lifts
maze policies from a near-zero starting success rate to solving most mazes
optimally; and it is evaluated further on object localization and
program-speed optimization.
The second part secures the judge. A planner that maximizes predicted return
will seek out exactly the regions where its world model is most optimistically
wrong. We prove that the resulting loss in policy quality is governed by the
sharpness of the model's loss landscape, and we show that sharpness-aware
training of the world model alone, leaving the planner and policy untouched,
delivers large and consistent gains that transfer across high-dimensional
continuous control and pixel-based discrete control, and across two different
model-based algorithms.
Good selection needs something worth choosing and a judge worth trusting.
Providing both, and composing them, is the program of this thesis.
Notes:
copied = false, 2000);
">
@mastersthesis{Ramasubramanian-2026-88341,
author = {Shrinivas Ramasubramanian},
title = {The Tail and the Judge: Better Policy Gradients and World Models for Reinforcement Learning},
year = {2026},
month = {July},
school = {Carnegie Mellon University},
address = {Pittsburgh, PA},
number = {CMU-RI-TR-26-90},
keywords = {RLHF, LLM training, continuous rewards, reinforcement learning},
}
author = {Shrinivas Ramasubramanian},
title = {The Tail and the Judge: Better Policy Gradients and World Models for Reinforcement Learning},
year = {2026},
month = {July},
school = {Carnegie Mellon University},
address = {Pittsburgh, PA},
number = {CMU-RI-TR-26-90},
keywords = {RLHF, LLM training, continuous rewards, reinforcement learning},
}