--- title: Per-Reward Advantage Decomposition url: https://www.emergentmind.com/topics/per-reward-advantage-decomposition type: topic --- # Per-Reward Advantage Decomposition Per-reward advantage decomposition is a collection of methodologies and theoretical results enabling the disentanglement and coordinated optimization of multiple reward signals within reinforcement learning (RL) and reinforcement fine-tuning (RFT) frameworks. In settings where the target objective is inherently multi-aspect—such as aligning large diffusion models or multi-objective vision-language agents with human values—reward aggregation via scalarization often induces interference, masks specialist learning signals, and complicates credit assignment. Per-reward advantage decomposition addresses these issues by maintaining, optimizing, and harmonizing the gradient contributions of each reward channel, yielding both empirical improvements and new theoretical guarantees for convergence and generalization. ## 1. Formal Definition and Motivation In multi-objective RL, the scalar reward is typically composed as $R = \sum_{k=1}^K w_k R_k$, with $R_k$ denoting individual reward channels and $w_k$ positive weights. The canonical advantage function $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$ can be extended to each reward channel as $A_k^\pi(s,a) = Q_k^\pi(s,a) - V_k^\pi(s)$. Per-reward advantage decomposition refers to leveraging this finer structure either to form separate advantage estimates and policy gradients per reward or, more generally, to enable explicit tracking, diagnosis, and harmonization of each reward-specific update direction. Motivation for this approach arises from sample-level specialization: in many tasks, training rollouts are “specialists,” informative for a subset of reward dimensions only. Aggregating via weighted sums dilutes supervision and can drive updates that improve some objectives while harming others, especially when the reward components are misaligned or have disparate scales [2605.06507]. Direct per-reward decomposition preserves the informative directionality and granularity of such signals. ## 2. Algorithmic Realizations and Practical Implementations ### 2.1. MARBLE: Multi-Aspect Reward Balance for Diffusion RL The MARBLE framework embodies a state-of-the-art implementation of per-reward advantage decomposition in diffusion RL. Its workflow involves: - **Per-reward normalization**: For each reward $k$, a running-mean and std $(\mu_k, \sigma_k)$ yield a normalized advantage $A_k(x) = \frac{R_k(x) - \mu_k}{\sigma_k + \varepsilon}$ for sample $x$. - **Reward-specific gradient computation**: Each advantage $A_k$ produces an interpolation coefficient $r_k$ for the NFT loss, enabling backward passes for each $\ell_k(\theta;x,t)$ and yielding individual gradients $g_k$. - **Gradient-space harmonization via QP**: Gradients $g_k$ (renormalized to unit norm) are harmonized into a single update direction $d^* = \sum_k \alpha^*_k \hat{g}_k$ by solving a simplex-constrained quadratic program: $\alpha^* = \arg\min_{\alpha \in \Delta^K} \|\sum_{k} \alpha_k \hat{g}_k \|^2$, minimizing inter-reward destructive interference. - **Amortized updates**: To reduce computational cost, MARBLE leverages the fact that the NFT loss is affine in $A_k$ and allows the combined gradient-update for a weighted sum of advantages to be computed via a single backward pass, updating $\alpha^*$ only every $N$ steps and using an EMA-smoothed coefficient for stability. - **Update**: The final update incorporates the harmonized direction with a KL-regularization gradient, proceeding as $ \theta \leftarrow \theta - \eta(d_\text{final} + \beta_\text{KL} \nabla_\theta D_\text{KL} )$ [2605.06507]. ### 2.2. Value Function Decomposition for Actor-Critic SAC-D decomposes the $Q$-function into $Q^\pi(s,a) = \sum_{i=1}^m w_i Q_i^\pi(s,a)$ and analogously decomposes the advantage $A^\pi(s,a) = \sum_i w_i A_i^\pi(s,a)$. By tracking $A_i$ estimates, the algorithm enables per-reward diagnostics, influence quantification, and conflict-averse gradient adjustments (e.g., via CAGrad) [2206.13901]. ### 2.3. Group Relative Policy Optimization—Per-Reward Decomposition In settings such as GRPO for language-vision models with composite rewards, the group-relative advantage is $A_\text{GRPO}(o^{(i)}) = (R(x, o^{(i)}) - \bar{R}_G)/(\sigma_{R,G} + \varepsilon)$, with normalization performed either on the scalarized reward or each component individually. The decomposed approach forms $A_k(x, o^{(i)}) = (R_k(x,o^{(i)}) - \overline{R}_{k,G})/\sigma_{R_k,G}$ and composes policy updates as $\sum_k w_k A_k$-weighted gradients. Reward Decomposition Theorems explicitly quantify the (small) suboptimality gap between decomposed and joint updates under moderate covariance between rewards and sufficient group size [2604.19857]. ### 2.4. Off-Policy Decomposition of Advantage and Return In the return-decomposition paradigm, the overall return is recursively written as $G(\tau) = V^\pi(s_0) + \sum_{t} \gamma^t A^\pi(s_t,a_t) + (...\ \text{luck terms})$, with $A^\pi$ interpreted as the causal effect of actions. This underpins the Off-policy DAE approach: advantage and “transition-noise” (luck) terms are estimated separately under general behavior policies, enabling robust off-policy RL free from importance sampling [2402.12874]. ## 3. Theoretical Guarantees and Analysis Per-reward advantage decomposition admits rigorous theoretical analysis: - **Gradient-space harmonization**: The QP formulation used in MARBLE ensures the update direction is Pareto-improving in all reward dimensions up to the projection geometry of their gradients. - **Decomposition suboptimality**: In GRPO, Theorem 2 demonstrates that for $m$ reward components and group size $G$, the suboptimality gap between the joint and per-component approach is at most $ \frac{m(m-1)}{2G} \sum_{i