Technical Summary: Using Reward Uncertainty to Induce Diverse Behaviour in Reinforcement Learning

Problem Statement

Classical reinforcement learning (RL) typically seeks a deterministic policy that maximizes the expected sum of a scalar reward. While this assumption guarantees the existence of an optimal deterministic policy, it presents significant limitations in modern applications such as large language model (LLM) fine-tuning and scientific discovery. In these domains, convergence to a single behavior is often undesirable. For instance, preserving the diversity of pre-trained policies is crucial for creative tasks and chain-of-thought reasoning. Furthermore, when reward models are imperfect proxies for true underlying preferences, over-optimization toward a single deterministic policy can lead to the deterioration of generation quality (reward hacking).

Existing remedies, such as entropy regularization or diversity bonuses, attempt to enforce stochasticity but suffer from critical drawbacks:

1. Fragile Trade-offs: Entropy regularization forces a trade-off where increased stochasticity comes at the cost of reduced expected rewards.
2. Incorrect Policy Ordering: Diversity bonuses often result in undesirable rankings where sub-optimal policies with high diversity are preferred over policies with higher average rewards.
3. Inexpressivity: Many objectives, including vanilla policy gradient and expected scalarized returns (ESR), inherently possess deterministic optimal policies, failing to induce necessary diversity.
4. Optimization Challenges: Scalarized expected return (SER) formulations often lack unbiased, sample-based gradient estimators for non-linear objectives.

The authors argue that diversity should be understood not as an artificial constraint, but as the rational response to uncertainty in the reward function. When the reward function is not perfectly known (due to ambiguous preferences or imperfect reward models), committing to a single action can be sub-optimal.

Methodology: ROSA

The paper proposes 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. This framework is termed Randomized Objectives, Set Actions (ROSA).

The ROSA+Max Criterion

Focusing on the contextual bandit setting, the authors define the primary objective, ROSA+Max. Let $\rho$ be a distribution over reward functions $R$, and let $\mathbf{Y} = (Y_i)_{i=1}^n$ be a multiset of $n$ actions sampled i.i.d. from a policy $\pi_\theta$. The objective is:

$$J_{\text{ROSA+Max}}(\pi_\theta) = \mathbb{E}_{X \sim \mu, \mathbf{Y} \sim \pi_\theta} \left[ \mathbb{E}_{R \sim \rho} \left[ \max_{1 \le i \le n} R(X, Y_i) \right] \right]$$

Intuitively, this criterion aggregates performance by taking the maximum reward achieved within a set of sampled actions for each sampled reward function. This encourages the policy to ensure that at least one action in the set performs well for every possible reward function realization, rather than optimizing for the average reward.

Gradient Estimation

The authors derive a principled, unbiased, and variance-reduced policy gradient estimator for this objective. Given $m$ sampled reward functions and $n$ sampled actions, the gradient estimator is:

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

This estimator effectively assigns credit only to the action that achieves the maximum reward for a specific reward function sample, with the second term serving as a control variate for variance reduction.

Generalization

The framework generalizes to a family of set functions $f$. The authors prove that for binary reward functions, any set function $f$ that is strictly increasing and strictly concave with respect to the sum of rewards (e.g., Softmax) yields a maximally diverse optimal policy.

Key Contributions

1. ROSA Framework: Introduction of a new family of RL objectives that provide calibrated diversity control through a distribution over reward functions.
2. Theoretical Analysis:
   • Optimal Policy: Proved that for a uniform distribution over binary reward functions, the unique optimal policy of ROSA+Max is the one that assigns equal probability mass to all distinct optimal actions.
   • Controllable Diversity: Derived a formula showing how non-uniform reward function distributions induce specific, controllable stochastic policies, allowing the probability of sampling specific correct answers to be tuned via the reward distribution weights.
   • Landscape Analysis: Demonstrated that ROSA+Max avoids the "flat" landscapes of entropy regularization and the "re-ordering" issues of diversity bonuses, providing a gradient direction that simultaneously increases reward and diversity.
3. Empirical Validation: Extensive experiments in toy domains and large-scale settings (including LLM post-training) demonstrating that ROSA induces diverse behaviors without sacrificing expected reward, outperforming standard policy gradient and diversity-bonus baselines in scenarios with conflicting preferences and unknown ground-truth rewards.

Experimental Results

The paper evaluates ROSA across three primary scenarios:

1. Conflicting Preferences: In a task with two contradictory reward functions (e.g., parity vs. anti-parity), vanilla policy gradient fails due to reward cancellation. ROSA successfully learns a policy that balances both preferences, generating diverse valid outputs.
2. Multiple Correct Answers: In mathematical reasoning tasks where multiple answer lengths are valid, ROSA-trained policies sample from all specified length preferences, whereas standard policy gradient collapses to a single length range.
3. Unknown Ground-Truth Rewards: In settings with noisy reward models (simulated by an ensemble of judges with varying accuracy), ROSA demonstrates superior robustness. It captures the uncertainty in the reward function by distributing probability mass across reward modes in proportion to their weights, whereas entropy-regularized PG either overfits to the most likely reward or explores too broadly into sub-optimal actions.

In out-of-distribution evaluations (AIME 2025), models trained with ROSA on noisy MATH data achieved higher pass@k rates compared to standard PG and pessimistic ensemble baselines.

Significance and Claims

The paper claims that ROSA offers a robust and theoretically grounded alternative for complex RL tasks where traditional formulations fail to induce the desired breadth of agent behavior. By reframing diversity as a rational response to reward uncertainty, the framework resolves the tension between reward maximization and policy entropy.

The authors emphasize that ROSA:

• Naturally generalizes both vanilla policy gradient and recent action-set approaches (like best-of-$n$).
• Provides a mechanism to control diversity through the reward function distribution rather than fragile hyperparameter tuning of entropy coefficients.
• Ensures that the optimal solution is a stochastic policy that reflects the distribution of possible reward functions, preventing the collapse to a single deterministic behavior even when the reward model is imperfect.

Ultimately, the work provides a principled foundation for training stochastic policies that are as diverse and adaptable as the environments they inhabit, particularly in domains like LLM fine-tuning and scientific discovery where multiple valid solutions exist.