Learn What’s Left, Not What’s Mastered: Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization

Yixuan Wang, Yifei Chen, Haichao Zhang, Haozheng Luo, Xander Wu, Jie Ni, Yun Fu, Nuno Vasconcelos, Yijiang Li

SA-MRPO dynamically reweights multi-objective RL updates based on objective saturation to prioritize under-optimized tasks.

How can we optimize language models for multiple reward objectives without letting "easy" or saturated rewards drown out the learning signal for harder, unsaturated objectives?

Multi-objective reinforcement learning for language models often uses fixed reward weights, which causes training to waste gradient budget on objectives that are already mastered while neglecting harder, under-optimized tasks. The authors introduce Saturation-Aware Advantage Reweighting for Multi-Reward Policy Optimization (SA-MRPO), which independently normalizes each reward and adaptively discounts its contribution based on how close the current batch mean is to the objective's maximum possible reward. This mechanism shifts optimization effort toward objectives with the most remaining headroom, improving performance on correctness and reasoning benchmarks by up to 9.2% while maintaining performance on already-satisfied constraints.

Paper Primer

Standard multi-reward optimization scalarizes heterogeneous rewards into a single value before group-relative normalization. This approach loses resolution—where different reward profiles collapse into identical advantages—and ignores the "saturation" of individual objectives, leading the model to continue optimizing tasks that have already reached their performance ceiling.

SA-MRPO is an adaptive objective-allocation rule: it calculates a saturation ratio for each objective based on the batch mean relative to its attainable range, then uses this ratio to exponentiate and down-weight the influence of saturated objectives. This acts like a dynamic throttle on a multi-lane highway, where the system automatically reduces the speed limit in lanes that are already congested to prioritize traffic flow in lanes with more open space.

SA-MRPO consistently outperforms fixed-weight baselines (GDPO) on harder reasoning objectives.

Across 15 mathematical reasoning benchmarks, SA-MRPO improved accuracy in 12 comparisons, with gains of up to 5% on AIME24 and 9.2% on AMC23.

The method preserves performance on easier, already-satisfied objectives.

In multi-objective settings involving correctness, length, and format constraints, the model maintains performance on the easier constraints while reallocating gradient pressure to the correctness objective.

Why is this approach better than simply adjusting the fixed weights manually?

Manual weight tuning is static and cannot account for the fact that different objectives reach saturation at different stages of training. SA-MRPO automates this by using the batch-level saturation ratio to adjust weights dynamically throughout the training process.

Does this method guarantee that performance on a saturated objective will never drop?

No. SA-MRPO is an adaptive allocation rule, not a constrained optimization method. If a less-saturated objective has a sufficiently conflicting gradient, the reallocation of effort can lead to a first-order decrease in the performance of a previously well-optimized objective.

The saturation ratio $s(k)$ measures nominal headroom within a prescribed reward range, not necessarily the absolute limit of the model's capacity. A high saturation ratio indicates the objective is well-learned relative to the reward function, but it does not guarantee that further improvement is impossible if the model's architecture is the bottleneck.

Researchers can now optimize multiple, heterogeneous reward objectives without manual weight scheduling, as SA-MRPO automatically shifts the model's focus toward the most difficult remaining tasks.

Motivation and Problem Framing

Identifies why scalarizing multi‑reward RL fails and motivates a saturation‑aware reweighting.

Standard multi‑reward RL scalarizes the reward vector before computing advantages, which leads to two fundamental problems: (1) distinct reward profiles can receive identical advantages, and (2) all objectives are optimized with fixed relative weights regardless of their current level of saturation.

When multiple objectives are collapsed into a single scalar, the optimizer cannot tell which objectives still have room for improvement, so it wastes effort on already‑saturated goals while neglecting harder ones.

The failure of scalarization in multi‑objective RL is that it masks which objectives are saturated, causing optimization to waste effort on already‑solved goals.

Foundations of Multi-Reward RL

Defines the RL setting, reward handling, and the GRPO baseline used later.

We consider reinforcement learning for an auto‑regressive language model that must satisfy n different reward objectives. For each query $q$, the policy $\pi$$\theta$ defines a distribution over output sequences $o$ by factorizing the probability of each token conditioned on the prefix. Training proceeds by sampling B queries, generating G \geq 2 rollouts per query with a frozen behavior policy $\pi$$\theta_{o}$ld, and evaluating each rollout under the n reward functions.

Scalarization collapses several reward signals into a single score by taking a weighted sum, letting the optimizer treat the multi‑objective problem as a standard single‑reward RL task.

GRPO measures how a rollout’s scalar reward compares to its peers for the same query, normalizing by the group’s mean and variance to obtain a relative advantage.

The SA-MRPO Mechanism

Reweights reward components by their remaining headroom to focus learning on unsaturated objectives.

Scalarizing multiple rewards hides distinct reward profiles and keeps the optimizer stuck on objectives that have already hit their ceiling, as illustrated in Figure 1.

Think of the optimizer as a thermostat that turns down heating for rooms that are already warm—here each “room” is a reward objective, and the thermostat reduces its weight once the objective is close to its reward ceiling.

How does this differ from the usual scalar weighting of rewards?

Standard scalar weighting applies a fixed coefficient $w_k$ regardless of how close the policy is to the objective’s maximum reward. Saturation‑aware reweighting multiplies $w_k$ by $(1 - s(k))^{\gamma}$, so the coefficient shrinks as the objective approaches its ceiling, dynamically reallocating learning effort.

Compute means: $\mu_1^{1}=1$, $\mu_1^{2}=0.6$.

Compute std devs: $\sigma_1^{1}=0$ (treated as 1 to avoid division by zero), $\sigma_1^{2}=0.2$.

Batch‑average rewards: $\bar r(1)=1$, $\bar r(2)=0.6$.

Compute saturation ratios: $s(1)=(1-0)/(1-0)=1$, $s(2)=(0.6-0)/(1-0)=0.6$.

Effective weights: $w_1^{*}=0.5\,(1-1)^{0.5}=0$, $w_2^{*}=0.5\,(1-0.6)^{0.5}=0.5\sqrt{0.4}\approx0.316$.

Relative advantages: $A_{1,1}^{1}=0$ (zero std), $A_{1,1}^{2}=(0.4-0.6)/0.2=-1$; $A_{1,2}^{2}=(0.8-0.6)/0.2=1$.

Aggregate raw advantage for rollout 1: $A_{1,1}=w_2^{*}\,A_{1,1}^{2}=0.316\times(-1)=-0.316$; for rollout 2: $A_{1,2}=0.316\times1=0.316$.

Batch‑normalize: mean $=0$, std $=0.447$, yielding $A_{1,1}^{\text{SA}}=-0.707$, $A_{1,2}^{\text{SA}}=0.707$.

When an objective is fully saturated its weight drops to zero, so the policy’s update is driven entirely by the remaining unsaturated objective.

SA‑MRPO policy update (Algorithm 1)

Because the effective weights shrink for saturated objectives, the update can inadvertently reduce performance on those objectives if a less‑saturated objective has a strongly conflicting gradient.

Empirical Results and Ablations

SA‑MRPO consistently outperforms GDPO across diverse reward configurations.

Recall that standard multi‑reward RL scalarizes rewards, causing saturated objectives to be ignored. SA‑MRPO instead reweights each reward component based on its saturation level.

SA‑MRPO improves accuracy over GDPO by up to 5.0 percentage points on the AIME24 benchmark.

Table 1 shows a 5.0 pp gain for Qwen2.5‑7B‑Instruct when adding the saturation‑aware reweighting.

GDPO treats each reward group as a distribution and optimizes the policy to improve expected returns across those groups.

How does GDPO differ from the earlier GRPO baseline?

GRPO uses a single group‑relative advantage estimate, while GDPO forms a full distribution per group and optimizes a surrogate that respects each group’s variance, enabling finer‑grained reweighting.

**Table 1.** Compares GDPO and SA-MRPO across model scales and reward configurations. Across the three configurations, SA-MRPO achieves higher accuracy than GDPO in 12 of the 15 benchmark comparisons. For Qwen2.5-7B-Instruct with three reward objectives, SA-MRPO improves four of the five benchmarks, including gains of 5.0 percentage points on AIME24 and 3.5 percentage points on MATH500. The same pattern is observed for Qwen2.5-3B-Instruct, where SA-MRPO improves four.

**Table 2.** Adaptive reasoning with an explicitly saturated length objective.

**Table 3.** Code reasoning results for Qwen2.5-7B-Instruct.

**Table.** The table evaluates the performance of Qwen2.5-3B-Instruct using different values of $\gamma$ (ranging from 0 to 1.0) across five benchmarks: AIME24, Minerva, AMC23, MATH500, and Olympiad. Metrics include Accuracy (Acc $\uparrow$) and Exceed ($\downarrow$).

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers