Using Reward Uncertainty to Induce Diverse Behaviour in Reinforcement Learning

TL;DR

ROSA (Randomized Objective, Set Actions) is a reformulation of the RL objective that replaces the scalar reward with a distribution over reward functions and applies a nonlinear set function (e.g. max, softmax) over the rewards of $n$ i.i.d.-sampled actions. Unlike vanilla policy gradient (GRPO, RLOO, REINFORCE) — which collapses to a deterministic policy even when multiple actions are equally optimal — ROSA provably admits an optimal stochastic policy whose action probabilities are calibrated to the reward-function distribution. In the contextual-bandit setting the change is essentially a one-line advantage swap (group-mean → group-max/softmax over per-reward-function outcomes), it drops directly into standard RLHF-style post-training pipelines, and empirically it (a) covers all four length preferences in a MATH fine-tune of Gemma 2 2B at no cost to pass@k, and (b) is more robust than PG and pessimistic-ensemble baselines under noisy-judge rewards, with better zero-shot AIME 2025 pass@k.

Key claims

Standard policy gradient over a distribution of reward functions collapses to the most probable reward function’s optimum, because linear expectation collapses the reward distribution to its mean [§3, Fig. 2(a)]; adding entropy regularization or diversity bonuses fixes this only by introducing new failures — putting mass on bad actions, reordering suboptimal policies against reward, or lacking unbiased sample-based gradients [§3, Table 1].
The ROSA+Max objective $\mathbb{E}_{R \sim \rho}\, \mathbb{E}_{Y_1,\ldots,Y_n \sim \pi}\, \max_i R(Y_i)$ has an unbiased variance-reduced policy-gradient estimator obtained by taking, per sampled reward function $R$ , the max-argmax action’s log-prob score minus a leave-one-out baseline — a single-line modification of GRPO/RLOO advantage computation [§3.1, Eq. 4].
Under uniformly-weighted binary reward functions with distinct optimal actions, the ROSA+Max objective has a unique global optimum: the maximally-diverse policy that samples each of the $m$ optimal actions with probability $1/m$ , for any action-set size $n \ge 2$ [Prop. 3.2].
The optimality result generalizes: any strictly increasing and strictly concave function of the per-reward success count (e.g. max, softmax) yields a ROSA+ $f$ objective whose optimum is the maximally-diverse policy — max and softmax are the same optimum with different optimization landscapes [Thm. 4.1].
Under non-uniform reward-function weights $\rho(R_k)$ , ROSA+Max has an optimal policy that samples the top- $k$ optimal actions with a predictable, closed-form dependence on the weights — the reward distribution is a controllable diversity knob [Prop. 4.2, Fig. 4].
On an adversarial two-reward parity task where PG’s linear scalarization exactly cancels the opposing rewards, ROSA is the only method that makes any progress at all [§5.1, Fig. 5].
Fine-tuning Gemma 2 2B on MATH with four length-preference reward functions, PG concentrates on a single dominant length while ROSA covers all four length preferences at no cost to overall pass@ $k$ [§5.2, Fig. 6].
Under an “ensemble of judges” reward with 6/8, 3/8, or 1/8 correct judges, ROSA holds up as judge accuracy degrades while PG and pessimistic ensemble baselines (min, mean−stddev) collapse; the effect transfers — ROSA-trained MATH checkpoints achieve higher zero-shot AIME 2025 pass@ $k$ [§5.3, Fig. 8].
Curvature of the ROSA+Max landscape at the optimum decays with $n$ (Hessian diagonal $\propto (1-1/m)^{n-1}$ ), suggesting small-to-moderate action-set sizes ( $m \le n \le 3m$ ) as the sensible hyperparameter range for optimization stability [Lemma 3.3, Appendix B.1].

Method

ROSA is defined by three ingredients on top of standard policy gradient: (i) a distribution $\rho$ over reward functions rather than a single scalar reward, (ii) an action multiset $(Y_1, \ldots, Y_n)$ sampled i.i.d. from the policy $\pi$ , and (iii) a nonlinear permutation-invariant set function $f$ that aggregates the $n$ rewards under each sampled reward function. The primary objective is ROSA+Max: $\mathcal{J}(\pi) = \mathbb{E}_{R \sim \rho} \mathbb{E}_{(Y_i) \sim \pi}[\max_i R(Y_i)]$ .

The practical gradient estimator (Eq. 4) is a single change to the standard REINFORCE/GRPO advantage: instead of the group-mean baseline, take a per-reward-function group max, weight by $\rho$ , and subtract a leave-one-out (LOO) baseline for variance reduction. Only the maximal-reward action per sampled reward function contributes a non-zero score term. Softmax is a smooth alternative that removes the estimator’s sparsity while preserving the maximally-diverse optimum guarantee.

The recipe generalizes GRPO/RLOO in the single-reward case: mean-aggregation over the group recovers vanilla PG, and max-aggregation recovers best-of- $n$ (pass@ $k$ ) training. What’s new is that with multiple reward functions plus a non-linear set aggregator, the optimum is provably stochastic and maximally diverse, and controllable through $\rho$ — properties no single-reward RL objective can express.

Results

Adversarial parity task (contradictory reward functions): ROSA is the only method (vs. PG, PG+entropy, GRPO-style, diversity-bonus) that achieves non-zero reward — PG collapses to the mean, which is exactly zero by construction [Fig. 5].
MATH fine-tune, Gemma 2 2B, four length-preference reward functions: PG concentrates on one dominant length; ROSA samples all four length modes with no pass@ $k$ regression [Fig. 6, Fig. 17].
Ensemble-of-judges MATH task with 8 judges (75% / 37.5% / 12.5% correct-judge fraction): ROSA remains robust as judge accuracy drops while PG and pessimistic ensemble baselines (min, mean-stddev) degrade [Fig. 8(a)].
Zero-shot transfer of noisy-MATH-trained checkpoints to AIME 2025: ROSA achieves higher out-of-distribution pass@ $k$ than PG at every noise level [Fig. 8(b)].

Why it’s interesting

ROSA sits at a novel spot in the Reasoning RL design space. The concept page currently tracks a growing catalog of reward-shape interventions on top of GRPO (hierarchical distance, correctness+format, perception-stack, PARL’s parallelism-incentive) and training-mechanic interventions (Training-Free Group Relative Policy Optimization frozen-model context-space updates, On-Policy Self-Distillation in Diffusion Models on-policy self-distillation, AgentOPSD: Recursive Self-Distillation for Agentic Reinforcement Learning turn-level Bayesian credit). ROSA is neither — it’s a change to the objective itself, keeping GRPO’s group-relative-advantage skeleton but replacing scalar reward with a reward distribution and mean-aggregation with max/softmax. The direct implication is a principled treatment of the mode-collapse-under-noisy-reward problem that Natural Emergent Misalignment from Reward Hacking in Production RL documented as a source of emergent misalignment and that LoRA Without Regret showed to be information-theoretically inevitable under scalar-reward PG. It also generalizes best-of- $n$ / pass@ $k$ training (Flow Matching Policy Gradients et al.) from single-reward to reward-distribution settings, giving those objectives their first provably-stochastic optimum.