Using Reward Uncertainty to Induce Diverse Behaviour in Reinforcement Learning

2026-06-02

AI Summary

Traditional reinforcement learning relies on deterministic policies, which struggle to meet the demand for behavioral diversity in tasks such as language model fine-tuning or scientific discovery. This work proposes a novel paradigm based on distributions over reward functions, introducing nonlinear objectives over action sets and leveraging a principled gradient estimator derived from contextual bandits. The approach enables controllable induction of policy diversity without compromising expected return. By unifying classical policy gradient methods with action-set optimization frameworks, the proposed method demonstrates robust generation of diverse behaviors in complex tasks, significantly outperforming conventional approaches in empirical evaluations.

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, 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. Our empirical results 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.

Problem

• reinforcement learning
• behavioral diversity
• reward uncertainty
• policy optimization
• contextual bandits

Innovation

• reward uncertainty
• behavioral diversity
• non-linear objective
• contextual bandits
• policy gradient