Using Reward Uncertainty to Induce Diverse Behaviour in Reinforcement Learning

Abstract

-cross Abstract: Classical reinforcement learning (RL) typically seeks a deterministic policy that maximizes the expected sum of a scalar reward. Yet, modern applications such as language model fine-tuning or scientific discovery demand diversity. Existing remedies such as entropy regularization or diversity bonuses often require fragile trade-offs that sacrifice performance for stochasticity or rely on heuristic metrics that can misalign policy rankings. We argue that diversity is more naturally understood as the rational response to uncertainty in the reward. When the reward function is not perfectly known--as is the case with ambiguous preferences or imperfect reward models--committing to a single action can be sub-optimal. Building on this, we propose a fundamental reformulation of the RL objective by replacing the scalar reward with a distribution over reward functions, and applying a non-linear objective over sets of actions. The result is a framework in which calibrated behavioural diversity emerges naturally, remains controllable through the reward function distribution, and is obtained without sacrificing expected reward. Focusing on the contextual bandit setting as commonly used in large language model (LLM) post-training, we derive a principled gradient estimator for this objective and prove that our formulation naturally generalizes both vanilla policy gradient and more recently developed action-set approaches. We provide didactic experiments which complement our theoretical results, and our large-scale empirical results in LLM reasoning further demonstrate that this framework offers a robust and theoretically grounded alternative for complex RL tasks where the traditional formulation of the problem fails to induce the desired breadth of agent behaviour.

Authors 10

Anthony GX-Chen, Ankit Anand, Gheorghe Comanici, Zaheer Abbas, Eser Ayg\"un, David Smalling, Shibl Mourad, Doina Precup, Andr\'e Barreto, Mark Rowland

Specification

Official page

arxiv.org/abs/2606.03962

Source:arXiv (Atom API + RSS)·T1·observed 2 h ago·high

Arxiv announce type

replace

Source:arXiv (Atom API + RSS)·T1·observed 2 h ago·high

arXiv id

2606.03962

Source:arXiv (Atom API + RSS)·T1·observed 2 h ago·high

Categories

cs.LG, cs.AI

Source:arXiv (Atom API + RSS)·T1·observed 2 h ago·high

PDF

arxiv.org/pdf/2606.03962

Source:arXiv (Atom API + RSS)·T1·observed 2 h ago·high

Primary category

cs.LG

Source:arXiv (Atom API + RSS)·T1·observed 2 h ago·high

Published

11 Sept 2026

Source:arXiv (Atom API + RSS)·T1·observed 2 h ago·high

Each value shows its source, tier and observation time. Conflicting claims are kept side by side and flagged — never averaged. How AI Atlas records facts →

Provenance