Policy optimization improves reward handling in multi-reward reinforcement learning
Whitened Advantage Summation Policy Optimization: Correlation-Aware Aggregation for Multi-Reward Reinforcement Learning
Machine LearningArtificial Intelligence
Summary
Reinforcement learning often uses multiple rewards added together, but when these rewards are related, this can cause overcounting of shared information. The authors propose a new method called WAS-PO that adjusts these rewards by factoring in their correlations before combining them, aiming for more balanced learning. They tested this method on language model tasks and found it improved training results compared to an existing method, though the gains did not extend clearly to accuracy on new data. Their work also reveals conditions where certain rewards might be downweighted or reversed during learning.
What this means in practice
- •For machine learning engineers: Improve training of language models that rely on multiple rewards by better handling correlations between these rewards during policy optimization.
- •For ai research teams: Design reinforcement learning experiments that involve several correlated objectives by applying whitening techniques to reduce overcounting of shared reward components.
Authors
Olaitan Olaleye, Jayson Meribe, Temidayo Olaleye
Abstract
Reinforcement learning for language models often combines several rewards by normalizing each one and adding them with fixed weights. When the rewards are correlated, the variation they share is counted once for every reward that carries it. In terms of the ZCA-decorrelated components of the reward signal, the weighted sum applies weights $Σ^{1/2} w$ rather than the chosen weights $w$, where $Σ$ is the reward correlation matrix. We study WAS-PO, which whitens the normalized advantages with $Σ^{-1/2}$ before adding them, so that the chosen weights apply to those components. We derive the scalarized objective whose gradient it follows when normalization statistics are held fixed, show that its weights on the raw rewards can be negative, give an exact two-reward condition under which the lighter-weighted reward is reversed, and give a loose bound on the error from estimating $Σ$ during training. On tool calling with Qwen2.5-1.5B-Instruct, WAS-PO reached a higher mean training reward than GDPO in all seven seeds with three rewards (3.8% higher) and all three seeds with four (3.9% higher). WAS-PO also replaces GDPO's per-prompt normalization with batch-level normalization, and we did not separate the two changes with three or four rewards. The training-reward gains did not carry over to held-out accuracy, which was slightly lower than GDPO's with three rewards (not significant) and 0.03 lower in absolute terms with four (two seeds evaluated). With two rewards, WAS-PO trailed GDPO in training reward, and a three-seed control weakly suggests that the gap comes from batch-level normalization rather than whitening. In a math setting where the two rewards were nearly identical by construction, WAS-PO drove down the lighter-weighted reward, as the sign condition predicts.