Source-linked AI summary

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

arXiv:2608.16072v1cs.LGcs.AI

TL;DR

Fixed scalarization can obscure distinct reward profiles and keep optimizing saturated objectives in multi-reward language-model reinforcement learning. SA-MRPO reweights independently normalized objectives by saturation, improving harder objectives across reasoning and coding tasks while largely preserving already-satisfied objectives.

  • Problem

    Fixed reward scalarization can collapse distinct reward profiles and ignore how much improvement remains for each objective.

  • Method

    SA-MRPO independently normalizes reward objectives and adaptively downweights those nearing saturation to emphasize objectives with greater remaining headroom.

  • Results

    SA-MRPO improves the harder objective in most mathematical, adaptive-reasoning, and coding comparisons while largely preserving already-satisfied objectives.

  • Takeaways & Limitations

    Saturation-aware reweighting shifts optimization toward objectives with greater remaining headroom across mathematical reasoning, adaptive reasoning, and coding tasks.

  • Takeaways & Limitations

    When objectives have conflicting gradients, reallocating effort toward a less saturated objective can reduce performance on a previously well-optimized objective.

Abstract

from arXiv · show

Reinforcement learning (RL) with group-relative advantages has become the de facto standard for post-training language model reasoners. However, when optimizing multiple reward objectives, existing methods typically scalarize the reward vector with a fixed weighted sum before group-wise standardization. We show that this design leads to two fundamental problems: rollouts with distinct reward profiles can receive identical advantages, and all objectives are optimized with fixed relative weights regardless of their current level of saturation. As a result, training continues to allocate gradient budget to already-solved objectives instead of focusing on those with greater remaining headroom. We introduce \textbf{Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization} (SA-MRPO), which standardizes each reward objective independently and adaptively discounts its contribution according to a batch-level estimate of objective saturation. This dynamically reallocates optimization effort toward under-optimized objectives while empirically maintaining performance on those that are already well satisfied. We further show that saturation-aware reweighting can reverse the sign of an update, rather than merely rescale its magnitude. Across mathematical reasoning with two- and three-objective reward combinations, SA-MRPO improves the harder correctness objective over GDPO in 12 of 15 benchmark comparisons, with gains of up to $5\%$ on AIME24. On adaptive reasoning it improves accuracy on all five benchmarks, by $3.8\%$ on average and up to $9.2 \%$ on AMC23, and on coding benchmarks it improves pass rate by up to $2.3\%$, while in all settings maintaining the easier objectives near their already satisfied levels.

1 Introduction

The introduction identifies reward-resolution loss and fixed-weight saturation blindness in scalarized multi-reward policy optimization. It proposes SA-MRPO, which independently normalizes objectives and adaptively downweights saturated rewards to redirect optimization toward under-optimized objectives.

  • Motivation: Scalarization can assign identical advantages to rollouts with different reward profiles, losing resolution across objectives.Different weighted sums of objective rewards may produce the same scalar value.
  • Motivation: Fixed objective weights can keep prioritizing well-optimized rewards while overlooking harder objectives that remain under-optimized.The problem persists even after reward decoupling.
  • Method: SA-MRPO independently normalizes reward objectives and progressively downweights each one as its batch-mean reward approaches its attainable maximum.Its saturation estimate uses the batch-mean reward relative to each objective’s attainable range.
  • Experiments: 12 of 15 benchmark comparisons show higher mathematical-reasoning accuracy than GDPO, while adaptive reasoning improves on all five benchmarks with 3.8% average improvement.The mathematical results largely maintain performance on the saturated objective.
  • Method: SA-MRPO can alter the direction of policy updates while dynamically reallocating optimization emphasis toward under-optimized objectives.The method is described as strictly generalizing GDPO and GRPO as special cases.

2 Related Work

Related work spans critic-free RL with verifiable rewards, multi-objective language-model alignment, and methods for combining or dynamically weighting rewards in group-relative policy optimization. Recent approaches adapt weighting using reward variance, objective conflicts, informativeness, or saturation.

  • RLVR and multi-objective alignment: RL with verifiable rewards has become standard for improving language-model reasoning, motivating increasingly effective critic-free policy-optimization schemes.Representative methods include GRPO, DeepSeek R1, DAPO, and REINFORCE++.
  • RLVR and multi-objective alignment: Multi-objective language-model alignment addresses correctness, efficiency, helpfulness, harmlessness, and other preference dimensions through preference control, preference optimization, Pareto optimization, and adaptive reward weighting.
  • Multi-reward group-relative policy optimization: GDPO independently normalizes reward dimensions before aggregation, while DVAO adapts objective weights to reward variance and GD2PO addresses conflicts among rewards.These methods directly study multi-reward combination within group-relative policy optimization.
  • Dynamic objective allocation: Dynamic Reward Weighting, SAW, and Focal Reward recognize that objectives progress at different rates and optimization effort should evolve accordingly.SAW uses reward variability to measure objective informativeness, whereas Focal Reward estimates saturation for rubric-based reward criteria.

3 Preliminaries

The paper formulates multi-objective reinforcement learning for an autoregressive language model using grouped rollouts sampled under a frozen behavior policy. Its GRPO baseline scalarizes multiple rewards before standardizing them within each query’s rollout group.

  • Problem setup: The policy πθ generates autoregressive output sequences for queries sampled in batches from a data distribution.Each output sequence is modeled token by token conditioned on the query and preceding prefix.
  • Problem setup: For each query, a frozen behavior policy generates a group of G ≥2 rollouts, each represented by its token sequence and prefixes.The training setup uses grouped comparisons among rollouts associated with the same query.
  • Multi-reward objectives: The policy is trained with n reward objectives, where objective k assigns reward R(k)(qi, oi,j) and has prescribed weight wk ≥0.Rewards are defined separately for each query–rollout pair and objective.
  • Group Relative Policy Optimization: GRPO compares each rollout’s reward with the other rollouts for the same query, but standard multi-reward GRPO first combines objectives into a scalar score and then group-standardizes it.This scalarization precedes the GRPO advantage calculation.
  • Group Relative Policy Optimization: The clipped GRPO surrogate uses the policy ratio ρi,j,t(θ), with an optional KL penalty against a fixed reference policy omitted because it is independent of reward construction.The clipping threshold satisfies ϵ > 0, and the KL coefficient is β ≥0.

4 Saturation Aware Advantage Reweighting

SA-MRPO addresses scalarization’s collapse of distinct reward profiles and fixed allocation across objectives by independently standardizing rewards and reweighting them according to remaining saturation. The method preserves the underlying clipped group-relative objective, but can reduce already optimized objectives when gradients conflict or nominal reward headroom is not optimizable.

  • Motivation: Scalarization can assign identical advantages to distinct reward profiles and continue favoring saturated objectives instead of directing updates toward objectives with more remaining improvement.Under equal weights, (1, 0) and (0, 1) become indistinguishable.
  • Method: SA-MRPO combines each objective’s within-group rollout quality with a batch-level saturation estimate, reducing an objective’s contribution as its realized reward fraction increases.The factor (1 − s(k))^γ encodes remaining reward headroom alongside prescribed objective importance.
  • Method: SA-MRPO changes only rollout-advantage construction and then uses the resulting normalized advantage in the standard clipped group-relative surrogate objective.When γ = 0, the saturation factors equal one and recover the corresponding GDPO advantage under the same weights and normalization.
  • Adaptive allocation: Increasing γ shifts relative allocation toward objectives with larger remaining reward headroom and reduces the relative contribution of more saturated objectives.For objectives with positive weights and saturation values below one, the effective weight ratio decreases with γ when one objective is more saturated.
  • Limitations: Conflicting objective gradients can cause saturation-aware reweighting to reduce an already optimized objective because its reduced weight weakens the self-improvement term protecting it.The aggregate direction cannot decrease an objective to first order when all positively weighted cross-objective gradient inner products are nonnegative; otherwise conflict can dominate.
  • Limitations: Nominal reward headroom may not equal optimizable headroom because a large unrealized reward range does not ensure that the current policy class can realize it.When reward bounds are specified directly by the reward function, the remaining fraction is observable without estimating the reward ceiling.

5 Experiments

Across mathematical, adaptive, and code reasoning, SA-MRPO reallocates optimization toward harder or under-satisfied objectives while preserving already-satisfied auxiliary objectives. It improves correctness or pass rates over GDPO across most evaluated settings, with the strongest mathematical gains reaching 5.0 percentage points on AIME24.

  • Mathematical reasoning: 12 of 15 benchmark comparisons show higher mathematical accuracy for SA-MRPO than GDPO across two- and three-objective configurations.For Qwen2.5-7B-Instruct with three objectives, gains include 5.0 percentage points on AIME24 and 3.5 percentage points on MATH500.
  • Adaptive reasoning: 3.8 percentage points average accuracy gain over GDPO occurs across all five adaptive-reasoning benchmarks, with AMC23 improving from 28.3% to 37.5%.SA-MRPO also increases average response length from 333 to 459 tokens, while remaining below the saturation threshold Bmin = 1024.
  • Code reasoning: 2.3 percentage points is SA-MRPO’s largest code pass-rate improvement over GDPO, with higher pass rates on three of four benchmarks and comparable bug rates.The pass-rate improvements are 0.6, 1.4, and 2.3 percentage points on APPS, CodeContests, and Codeforces; TACO is 0.4 percentage points lower.
  • Sensitivity to γ: Every positive γ improves average mathematical accuracy relative to γ = 0, while larger γ values generally reduce length reward and increase EXCEED.γ = 0.5 achieves the highest average accuracy and strongest AIME24 and AMC23 performance, whereas larger values remain competitive but generally increase EXCEED.

6 Conclusion

The conclusion frames multi-reward policy optimization as an allocation problem across objectives at different training stages. It introduces SA-MRPO to adapt each objective’s contribution according to its current saturation, addressing a limitation of reward decoupling.

  • Multi-reward policy optimization should be viewed through how optimization effort is allocated across objectives at different training stages.
  • Reward decoupling preserves information from individual reward dimensions but does not distinguish difficult objectives from those near saturation.
  • SA-MRPO adapts each objective’s contribution according to its current saturation to address this limitation.
Loading 2608.16072v1…