Per-Reward Advantage Decomposition

• Per-reward advantage decomposition is a framework that separates multi-dimensional reward signals to allow targeted optimization and precise credit assignment in complex RL tasks.
• Methodologies like MARBLE and SAC-D use per-reward normalization and gradient-space harmonization to mitigate interference from scalarization and improve learning efficiency.
• The approach offers empirical gains such as reduced gradient conflict, enhanced interpretability, and theoretical guarantees for convergence and generalization in multi-objective environments.

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 (Zhao et al., 7 May 2026). 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$0 for the NFT loss, enabling backward passes for each $R_k$1 and yielding individual gradients $R_k$2.
• Gradient-space harmonization via QP: Gradients $R_k$3 (renormalized to unit norm) are harmonized into a single update direction $R_k$4 by solving a simplex-constrained quadratic program: $R_k$5, minimizing inter-reward destructive interference.
• Amortized updates: To reduce computational cost, MARBLE leverages the fact that the NFT loss is affine in $R_k$6 and allows the combined gradient-update for a weighted sum of advantages to be computed via a single backward pass, updating $R_k$7 only every $R_k$8 steps and using an EMA-smoothed coefficient for stability.
• Update: The final update incorporates the harmonized direction with a KL-regularization gradient, proceeding as $R_k$9 (Zhao et al., 7 May 2026).

2.2. Value Function Decomposition for Actor-Critic

SAC-D decomposes the $w_k$0-function into $w_k$1 and analogously decomposes the advantage $w_k$2. By tracking $w_k$3 estimates, the algorithm enables per-reward diagnostics, influence quantification, and conflict-averse gradient adjustments (e.g., via CAGrad) (MacGlashan et al., 2022).

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 $w_k$4, with normalization performed either on the scalarized reward or each component individually. The decomposed approach forms $w_k$5 and composes policy updates as $w_k$6-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 (Adams et al., 21 Apr 2026).

2.4. Off-Policy Decomposition of Advantage and Return

In the return-decomposition paradigm, the overall return is recursively written as $w_k$7, with $w_k$8 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 (Pan et al., 2024).

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 $w_k$9 reward components and group size $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$0, the suboptimality gap between the joint and per-component approach is at most $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$1. This shows the loss from decomposition vanishes with large $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$2 and well-aligned (low-covariance) reward components (Adams et al., 21 Apr 2026).
• Bellman consistency: In value-decomposed actor-critic, linearity guarantees hold: $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$3 and $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$4 rigorously (MacGlashan et al., 2022).

These results establish principled foundations for the practical adoption of reward-wise decomposition in complex multi-aspect RL systems.

4. Empirical Benefits and Diagnostics

Empirical studies highlight several advantages of per-reward advantage decomposition:

• Conflict mitigation: MARBLE resolves destructive gradient alignment found in naïve weighted summation, improving the worst-case reward gradient cosine from negative to consistently positive in 80% of mini-batches.
• Balanced optimization: Unified training with MARBLE or decomposed GRPO simultaneously advances all reward channels, matching or exceeding sequential fine-tuning or specialist ensembles.
• Computational efficiency: Amortized harmonization in MARBLE achieves training speeds nearly on par with single-reward NFT, with modest GPU memory overhead (1.14x), outperforming naïve scalarization in both runtime and sample efficiency.
• Interpretability: Value/advantage decomposition in actor-critic variants like SAC-D enables influence tracing—diagnosing over- or under-fitting to particular reward channels, or evaluating the impact of each reward's removal via gradient difference metrics.
• Generalization: Decomposition is leveraged for deeper understanding of transfer in LVLMs, with bounds showing minimal penalty to out-of-distribution generalization given modest cross-covariances (Adams et al., 21 Apr 2026).

User studies and external metrics (Aesthetic, ImageReward, UniReward) confirm that per-reward harmonization yields higher human preference and robustness in generative settings (Zhao et al., 7 May 2026).

5. Extensions in Offline RL and Return Decomposition

In offline RL with trajectory-wise rewards, as in PARTED, advantage decomposition is realized by first deriving per-step proxy rewards through least-squares reward redistribution, then running pessimistic value iteration to form step-wise advantages. The resulting per-step advantage $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$5 decomposes as the sum of future proxy-reward contributions (offset by epistemic-uncertainty penalties), enabling fine-grained credit assignment in the absence of per-step supervision (Xu et al., 2022).

Similarly, Off-policy DAE explicitly constrains advantage and transition-noise terms to recover the correct return decomposition even under off-policy sampling, accelerating learning and reducing variance in both deterministic and stochastic environments (Pan et al., 2024).

6. Limitations and Open Challenges

Although per-reward advantage decomposition provides compelling benefits, several practical issues remain:

• Scalability to large $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$6: As the number of reward channels increases, harmonization may incur greater computational and memory costs despite amortization strategies.
• Correlation structure: When reward components are highly misaligned (large $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$7), decomposition alone cannot fully eliminate the risk of conflicting policy updates.
• Reward misspecification: Decomposed approaches are sensitive to the selection and definition of reward channels; missing or poorly specified components can confound diagnostics and optimization.
• Constraint enforcement: In the off-policy setting (e.g., DAE), estimation of transition-noise advantages sometimes requires generative models of the environment (e.g., CVAEs), introducing additional hyperparameters and potential estimation bias.

A plausible implication is that hybrid methods combining per-reward decomposition with joint scaling, regularization, or explicit conflict-resolution (as in QP harmonization) may yield the most robust solutions in high-dimensional or highly non-stationary reward settings.

7. Summary Table: Methodological Approaches

Method	Main Decomposition Mechanism	Distinctive Features
MARBLE (Zhao et al., 7 May 2026)	Per-reward normalized advantages; QP gradient harmonization	Affine loss structure for amortization; EMA smoothing
SAC-D (MacGlashan et al., 2022)	Reward-wise $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$8, $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$9 decomposition	Influence diagnostics; CAGrad
GRPO Decomposition (Adams et al., 21 Apr 2026)	Group-wise per-component normalization	Theoretical suboptimality bounds
PARTED (Xu et al., 2022)	Reward redistribution; proxy per-step advantages	Pessimistic value iteration; trajectory-wise rewards
Off-policy DAE (Pan et al., 2024)	Per-step advantage & transition-noise decomposition	Causal interpretation; off-policy correction

Each approach operationalizes per-reward advantage decomposition in a domain-specific manner while retaining the core benefit—sample-level preservation and coordinated optimization of individual objective signals.