Reward Decomposition & Distributional RL

• Reward Decomposition and Distributional RL are techniques that decompose reward signals into interpretable channels and model full return distributions to enhance multi-objective decision making.
• These methods employ novel architectures like MD3QN and DRDRL to jointly optimize for diverse objectives while capturing cross-channel dependencies and uncertainty.
• Practical applications in Atari gaming, multi-agent coordination, and sparse-reward tasks demonstrate improved sample efficiency, risk-sensitive control, and interpretable credit assignment.

Reward decomposition and distributional reinforcement learning (RL) constitute a rapidly converging synthesis within contemporary RL. Reward decomposition refers to factorizing the immediate reward signal into multiple interpretable sources or channels, while distributional RL focuses on learning full return distributions rather than scalar expected values. Their intersection has led to novel theoretical and algorithmic breakthroughs, enabling agents to jointly optimize for multiple objectives, capture uncertainty and cross-channel dependencies, and improve sample efficiency, credit assignment, and risk-sensitive control.

1. Formalism: Multi-Dimensional Returns and Bellman Operators

A Markov decision process (MDP) equipped with $N$ reward sources yields, under policy $\pi$, a multi-dimensional return vector:

$$G_t = (G^1_t, G^2_t, ..., G^N_t)^\top, \qquad G^n_t = \sum_{k=0}^\infty \gamma^k r_{t+k,n},$$

where $r_{t,n}$ indexes each reward channel and $\gamma$ is the discount factor (Zhang et al., 2021, Wiltzer et al., 2024).

The joint return distribution at $(s,a)$ is represented by a probability measure $\mu(s,a)\in\mathcal{P}(\mathbb{R}^N)$. The core operator is the joint distributional Bellman operator:

$$(\mathcal{T}^\pi\mu)(s,a) = \text{Law}[r + \gamma\,Z'], \quad r \sim R(\cdot|s,a), \ s'\sim P(\cdot|s,a),\ a'\sim\pi(\cdot|s'),\ Z'\sim\mu(s',a').$$

This generalizes the scalar case and captures not only marginal randomness in each channel but also their correlations.

A key result is that this operator is a $\gamma$-contraction under the $p$-Wasserstein metric:

$$\bar{d}_p(\mathcal{T}^\pi\mu_1, \mathcal{T}^\pi\mu_2) \leq \gamma\,\bar{d}_p(\mu_1,\mu_2),$$

guaranteeing a unique fixed point for the joint return law (Zhang et al., 2021, Lee et al., 2024, Wiltzer et al., 2024).

2. Algorithmic Architectures for Reward Decomposition and Distributional RL

Several principled architectural strategies have emerged:

• Multi-dimensional Particle Networks (MD3QN): For each action, the network outputs $M$ deterministic samples ("particles") in $\mathbb{R}^N$, forming an empirical joint law. The reward decomposition is encoded as $N$ output dimensions per particle; distributional approximation employs Gaussian kernel maximum mean discrepancy (MMD) objectives (Zhang et al., 2021).
• Factorized Categorical Head Models (DRDRL): The agent maintains $N$ categorical distribution heads, each approximating the law of sub-return $Z_i$. The joint return is recovered via discrete convolution. A KL-based disentanglement bonus regularizes heads to specialize (Lin et al., 2019).
• Oracle-Free Particle/Categorical Dynamic Programming (MV-DRL): Equally weighted particle mixtures (EWP) or categorical tabulation over grids in $\mathbb{R}^d$ allow for scalable, provably convergent algorithms, with signed-measure projections ensuring contraction and convergence for temporal difference learning in $d>1$ (Wiltzer et al., 2024).
• Likelihood-based Reward Redistribution: In environments with delayed, sparse scalar feedback, per-step rewards are modeled as conditional distributions, and dense surrogate signals are generated by maximizing the likelihood of the observed trajectory sum via leave-one-out objectives (Xiao et al., 20 Mar 2025).
• GMM-Based Multi-Agent Decomposition (NDD): A global noisy reward is fit as a Gaussian mixture, with learned weights assigning local probabilistic reward channels. Each agent updates locally using distributional RL, and regularization terms address ambiguity in the decomposition (Geng et al., 2023).

3. Theoretical Properties and Contraction Results

Distributional RL with reward decomposition inherits and generalizes foundational contraction properties:

• Unique Fixed Points: Joint Bellman operators in $\mathcal{P}(\mathbb{R}^N)$ (via $W_p$ or MMD metrics) are strict contractions ($\gamma$-factor), resulting in unique fixed points for both multi-dimensional and infinite-dimensional Banach reward spaces (Lee et al., 2024, Wiltzer et al., 2024, Zhang et al., 2021).
• Convergence of Algorithms: Both particle-based (with $O(m^{-1/2})$ error for $m$ particles and polynomial dependence on $N$) and categorical signed-measure TD (error $O(m^{-\alpha/(2N)})$, i.e., exponential in dimension) provably approach the optimal distributional solution (Wiltzer et al., 2024).
• Consistency in Value Aggregation and Maximization: In multi-agent reward decomposition via GMM, monotonicity is proven: locally greedy actions align with global maximization, provided mixture weights are nonnegative (Geng et al., 2023).

4. Loss Functions and Objective Design

Typical training losses combine distributional divergence and decomposition-specific regularization:

• MMD Loss (MD3QN): For online and target networks, the minimization of squared maximum mean discrepancy between empirical joint particles and Bellman target samples drives convergence (Zhang et al., 2021, Wiltzer et al., 2024).
• KL Divergence (DRDRL): The per-sample loss is KL-divergence between the projected Bellman target (via convolution across heads) and the network's joint output, penalized by pairwise disentanglement (Lin et al., 2019).
• Likelihood Objective (LRR): The surrogate reward model parameters are optimized via log-likelihood over leave-one-out "virtual" observations, with an inherent uncertainty regularization term $\log\sigma_\theta(s,a)$ (Xiao et al., 20 Mar 2025).
• PDF Matching with Regularizers (NDD): Decomposition is trained by minimizing squared difference to the global reward PDF, plus penalty terms aligning component means and mixture weights to avoid ambiguity (Geng et al., 2023).

5. Practical Considerations: Scalability, Representation Fidelity, and Sample Efficiency

Reward decomposition and distributional RL scale with several considerations:

Method	Dimensional scalability	Representation fidelity	Sample efficiency
Particles (MD3QN, MV-DRL)	$O(N)$ per particle, polynomial	Good for moderate $N$	High for small-medium $N$
Categorical (DRDRL, MV-DRL)	$O(m^{N})$ for grid	Decays as $O(m^{-1/2N})$	Expensive for $N>3$
Likelihood-based (LRR)	Per-step, independent	Dense signal via variance	High for sparse reward
GMM-based (NDD)	Number of agents $N$	Provable consistency	Data augmentation improves

Particle and signed-measure categorical projections mitigate the curse of dimensionality and maintain contraction structure (Wiltzer et al., 2024). For high-dimensional rewards, Euclidean embeddings and quantization reduce approximation error to acceptable levels (Lee et al., 2024). Choice of $N$ (number of sources), $m$ (particle/categorical atoms), and kernel bandwidths are critical for performance (Zhang et al., 2021). Diffusion-based data augmentation can further reduce sample complexity (Geng et al., 2023).

6. Empirical Evaluation and Qualitative Analysis

A variety of benchmarks demonstrate tangible benefits:

• Atari Games (MD3QN, DRDRL): Agents that model joint return distributions outperform scalar and hybrid architectures, especially when reward channels exhibit strong correlation (e.g., combat vs. survival rewards) (Zhang et al., 2021, Lin et al., 2019).
• Multi-Agent Environments (NDD): Decomposition accuracy is maintained under complex non-Gaussian noise; risk-sensitive policies optimized through distortion mappings yield superior tradeoffs in multi-agent coordination (Geng et al., 2023).
• Sparse-Reward Control (LRR): Likelihood-based reward redistribution leads to faster learning and higher return in delayed feedback tasks by providing dense, uncertainty-informed reward signals (Xiao et al., 20 Mar 2025).
• High-Dimensional MDPs (MV-DRL, Off-policy DRL): Approximating high-dimensional return with grid-based or particle-based models achieves rapid convergence and enables user-specified nonlinear utilities in policy optimization (Wiltzer et al., 2024, Lee et al., 2024).

Qualitative analysis reveals channel-specific behavioral differentiation, interpretable credit assignment, and accurate cross-channel correlation modeling.

7. Challenges and Trade-offs in Representation Choice

The rapid scaling of categorical models with the number of reward channels/dimensions ($O(m^{1/2N})$ error) presents a practical bottleneck for large $N$ (Wiltzer et al., 2024). Particle-based and signed-measure categorical algorithms offer scalable alternatives, though may sacrifice some pointwise convergence guarantees. Regularization is essential to avoid trivial decompositions and ambiguity (Lin et al., 2019, Geng et al., 2023). In multi-agent and noisy environments, integrating distributional RL with robust decomposition can offer risk-sensitive policy constraints and credible value factorization (Geng et al., 2023).

A plausible implication is that matching the decomposition granularity to the number of meaningful reward channels—ideally $N\leq5$–$10$—maximizes sample efficiency without overwhelming representation complexity (Zhang et al., 2021, Wiltzer et al., 2024).

---

In sum, reward decomposition and distributional RL provide a unified framework for multi-objective, uncertainty-aware sequential decision making. Recent research achieves provably convergent, scalable algorithms by leveraging contraction properties, efficient projections, and statistically robust reward splitting, advancing both practical policy performance and theoretical understanding across single-agent, multi-agent, and high-dimensional tasks (Zhang et al., 2021, Lin et al., 2019, Lee et al., 2024, Xiao et al., 20 Mar 2025, Geng et al., 2023, Wiltzer et al., 2024).