Using Reward Uncertainty to Induce Diverse Behaviour in Reinforcement Learning

Abstract: Classical reinforcement learning (RL) typically seeks a deterministic policy that maximizes the expected sum of a scalar reward. Yet, modern applications such as LLM 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.

• The paper introduces ROSA, a novel RL framework that models reward functions as distributions to induce behavioral diversity.
• It derives unbiased policy gradients via a leave-one-out advantage, ensuring robustness against reward model uncertainty.
• Empirical results demonstrate ROSA's effectiveness in covering multiple correct actions and outperforming traditional methods.

Using Reward Uncertainty to Induce Diverse Behaviour in Reinforcement Learning

Motivation and Problem Statement

Traditional RL objectives presume a known scalar reward function, leading to optimal deterministic policies. However, many practical applications—such as LLM fine-tuning, creative code synthesis, and scientific discovery—require persistent behavioral diversity: multiple valid actions or solutions should be covered by the trained policy. The dominant methodological responses, entropy regularization and diversity bonuses, each induce diversity only through explicit, often heuristic modifications of the reward, which deteriorate expected return or distort the ordering of policies with respect to their quality. Moreover, they fail to meaningfully express uncertainty in user preferences or epistemic uncertainty about the reward function itself.

This work posits that meaningful diversity should emerge as a rational policy response to uncertainty in the reward. When reward models are imperfect or user preferences are ambiguous, forcing a deterministic solution is inherently sub-optimal. The authors introduce a new RL formulation based on modeling the reward function as a random variable—a distribution rather than a point estimate—then aggregating over sets of actions. The resulting “Randomized Objectives, Set Actions” (ROSA) framework generalizes classic single reward RL as well as best-of-$n$ and set-based RL methods, while naturally supporting stochastic policies and diversity.

Figure 1: Examples of distributions of reward functions and the desired policy. Left: a binary reward setting with multiple correct actions; Right: modeling epistemic uncertainty over the reward function, motivating stochastic policies.

The ROSA Objective

The ROSA objective replaces the standard expected reward maximization with an expectation over both a distribution of reward functions and an i.i.d. set of actions sampled from the policy. The primary instance, ROSA+Max, optimizes

$$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$

where $\rho$ is a (potentially continuous) distribution over reward functions, and $n$ is the action set size. The critical insight is that maximizing the set-aggregated reward over a reward function distribution incentivizes the policy to put probability mass over the set of optimal actions for each possible $R$. This resolves a fundamental limitation: vanilla policy gradients with reward averaging collapse onto behaviors that perform adequately under the mean reward but lose diversity and coverage over alternative optima.

The authors derive an unbiased, low-variance policy gradient estimator for these objectives via the leave-one-out advantage:

$$\hat{g}(\theta) = \frac{1}{m} \sum_{k=1}^m \sum_{i=1}^n \left( \max_j R_k(X, Y_j) - \max_{j \neq i} R_k(X, Y_j) \right) \nabla_\theta \log \pi_\theta(Y_i|X).$$

Figure 2: Schematic computation for ROSA+Max: given sampled reward functions and actions from the policy, aggregate maximum reward across the set, then optimize the weighted sum across reward distribution.

Theoretical Properties

Optimal Policy Characterization

Under separated, distinct reward functions (e.g., one per unique correct action), and a uniform reward distribution $\rho$, the unique optimizer of ROSA+Max is the uniform mixture over all actions optimal under some $R$:

$\pi^* = \sum_{i=1}^{m} \frac{1}{m} \delta_{y^*_i},$

where $y^*_i$ is the optimal action for $$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$0. This yields maximum support over all optimal answers and removes mass from globally sub-optimal actions.

Figure 3: Objective landscape simplex: Vanilla PG (leftmost) collapses to one mode, whereas ROSA+Max (rightmost) maximizes at the uniform distribution over all optimal actions.

By contrast, entropy regularized or diversity-bonus-augmented RL objectives either assign mass to all actions, including poor ones, or invert the ordering of suboptimal policies, favoring lower-quality but diverse policies. ROSA+Max ensures reward maximization and diversity correctly align in the policy improvement trajectory.

Controllable Diversity via Reward Distribution

For non-uniform $$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$1, policy probability mass on each action is directly determined by the weights in $$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$2, offering precise, interpretable control of diversity and allowing smooth interpolation between deterministic and maximally stochastic policies.

Figure 4: Landscape for non-uniform reward distributions: ROSA+Max finds a stochastic optimum weighted according to the probability of each reward in $$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$3—unlike PG, which remains deterministic.

Extension to Other Set Functions and Risk Sensitivity

The framework encompasses arbitrary permutation-invariant set functions beyond max, such as softmax-based aggregations, establishing a formal correspondence with distortion risk measures and utility uncertainty modeling. This flexibility admits landscapes with differentiated curvature, sparsity, and convergence properties, and connects conceptually to risk-sensitive RL and scalarization methods in multi-objective RL. However, unlike linear scalarization, ROSA's stochastic policies arise due to the nonlinearity of the set function over randomized rewards.

Figure 5: Set function choice controls the optimization behaviour: Softmax (middle) interpolates between mean and max, adjusting the diversity-strength-reward trade-off, whereas multiplicative diversity baselines yield distorted suboptimal policy orderings.

Empirical Evaluation

Competing or Uncertain Preferences

In tasks where reward functions are contradictory (e.g., parity/anti-parity scoring for sequence generation), vanilla PG is ineffective—rewards cancel, preventing meaningful learning. ROSA methods, in contrast, can simultaneously balance and represent both preferences, producing a policy that probabilistically covers both.

Figure 6: Sequence generation under competing reward functions—ROSA-trained policies achieve robust coverage of both optimal solutions; PG and diversity-based baselines fail in the presence of reward cancellation.

Multiple Correct Answers

ROSA supports tasks natively involving multiple valid answers, e.g., different correct program lengths or equivalent code variants. Experiments on mathematical reasoning demonstrate that while vanilla PG collapses to the most-commonly-rewarded length, ROSA-trained policies cover all specified preferred lengths without accuracy loss.

Figure 7: Training with length-based reward functions on a code/data generation task—ROSA produces outputs matching all length preferences proportional to their specified importance, while maximizing pass@$$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$4 accuracy.

Robustness Under Reward Model Uncertainty

ROSA provides robustness under epistemic uncertainty about the reward, as occurs with ensemble or noisy LLM reward judges. The trained policy matches coverage to reward function likelihood, and is empirically more robust to noisy or misspecified judges than pessimistic ensemble baselines or entropy-regularized methods.

Figure 8: Coverage in policy space under uncertain reward for 2D toy tasks: ROSA-trained policy matches the reward function distribution, whereas entropy or PG baselines either collapse or assign support to low-quality actions.

Figure 9: Large-scale evaluation on LLM mathematical reasoning: Under increasing judge noise, ROSA methods maintain higher in- and out-of-distribution pass@$$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$5 rates compared to ensemble or regularized alternatives.

Implications and Future Directions

The ROSA framework reframes diversity as an inevitable outcome of rational policy optimization in face of explicit reward uncertainty, rather than an objective bolted onto deterministic RL. This delivers several advantages: stochastic policies that maximize expected reward, direct modeling of human and epistemic ambiguities, and robustness to scaling errors in reward models or conflicting preferences. Practically, the formalism enables:

• Fine-grained control of behavioral diversity via reward distributions: This is critical for applications in alignment, creative generation, and RL from human feedback, where reward is rarely perfectly specified.
• Unbiased, computationally efficient policy gradients: The framework avoids the variance and bias pitfalls of two-stage or heuristic diversity bonuses.
• Generalization and compatibility: ROSA subsumes policy gradient, best-of-$$\mathcal{J}_{\text{ROSA+Max}}(\pi) = \mathbb{E}_{X, Y_{1:n} \sim \pi, R \sim \rho} \left[ \max_{1 \leq i \leq n} R(X, Y_i) \right],$$6, and scalarization-based multi-objective RL as special cases, and future work can exploit this connection for Pareto frontier exploration and sophisticated utility combination.

Notably, these results have implications for the design of RL pipelines in LLM alignment (to reduce mode collapse), cross-domain generalization, and robustness to distributional shift in the reward. The framework admits extensions to non-set action selection, sequential dependencies, non-convex reward uncertainty (as in program synthesis or combinatorial search), and gradient-compatible forms of oracle or learned reward modeling. For risk-sensitive settings, the architecture bridges to distortion risk measures.

Conclusion

ROSA constitutes a general-purpose RL framework that natively encodes reward function uncertainty and thus elicits calibrated, controllable behavioral diversity without distorting incentive alignment or computational tractability. Theoretical properties guarantee maximally diverse, reward-maximizing stochastic policies are optimal, and empirical evidence substantiates effectiveness in tasks where standard RL fails to maintain necessary solution breadth. This work lays a foundation for robust, diversity-conscious RL in settings where ambiguity and multiplicity of correct outcomes are intrinsic.