# How Does Return Distribution in Distributional Reinforcement Learning Help Optimization?

arXiv: https://arxiv.org/abs/2209.14513 · HTML: https://mlsanity.com/paper/2209.14513/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 34 references that are in the map.
- First posted on arXiv 29 Sep 2022.

## Builds on (most cited first, 25 of 25)

- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Distributional Reinforcement Learning with Quantile Regression](https://mlsanity.com/paper/1710.10044/) (2017)
- [Implicit Quantile Networks for Distributional Reinforcement Learning](https://mlsanity.com/paper/1806.06923/) (2018)
- [How Does Batch Normalization Help Optimization?](https://mlsanity.com/paper/1805.11604/) (2018)
- [On the Global Convergence Rates of Softmax Policy Gradient Methods](https://mlsanity.com/paper/2005.06392/) (2020)
- [A Theoretical Analysis of Deep Q-Learning](https://mlsanity.com/paper/1901.00137/) (2019)
- [On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift](https://mlsanity.com/paper/1908.00261/) (2019)
- [Understanding the impact of entropy on policy optimization](https://mlsanity.com/paper/1811.11214/) (2018)
- [The Cramer Distance as a Solution to Biased Wasserstein Gradients](https://mlsanity.com/paper/1705.10743/) (2017)
- [Fully Parameterized Quantile Function for Distributional Reinforcement Learning](https://mlsanity.com/paper/1911.02140/) (2019)
- [Conservative Offline Distributional Reinforcement Learning](https://mlsanity.com/paper/2107.06106/) (2021)
- [Revisiting Rainbow: Promoting more Insightful and Inclusive Deep Reinforcement Learning Research](https://mlsanity.com/paper/2011.14826/) (2020)
- [Improving Regression Performance with Distributional Losses](https://mlsanity.com/paper/1806.04613/) (2018)
- [Spectral Normalisation for Deep Reinforcement Learning: an Optimisation Perspective](https://mlsanity.com/paper/2105.05246/) (2021)
- [Distributional Reinforcement Learning for Efficient Exploration](https://mlsanity.com/paper/1905.06125/) (2019)
- [A Comparative Analysis of Expected and Distributional Reinforcement Learning](https://mlsanity.com/paper/1901.11084/) (2019)
- [Towards Deeper Deep Reinforcement Learning with Spectral Normalization](https://mlsanity.com/paper/2106.01151/) (2021)
- [Towards Understanding Label Smoothing](https://mlsanity.com/paper/2006.11653/) (2020)
- [Functional Regularization for Reinforcement Learning via Learned Fourier Features](https://mlsanity.com/paper/2112.03257/) (2021)
- [Distributional Reinforcement Learning via Moment Matching](https://mlsanity.com/paper/2007.12354/) (2020)

## Built on by (most cited first, 4 of 4)

- [A Survey of Temporal Credit Assignment in Deep Reinforcement Learning](https://mlsanity.com/paper/2312.01072/) (2023)
- [Qsharp: Provably Optimal Distributional RL for LLM Post-Training](https://mlsanity.com/paper/2502.20548/) (2025)
- [FlowCritic: Bridging Value Estimation with Flow Matching in Reinforcement Learning](https://mlsanity.com/paper/2510.22686/) (2025)
- [On Distributional Reinforcement Learning in Chaotic Dynamical Systems](https://mlsanity.com/paper/2605.30160/) (2026)

## Abstract (arXiv)

Distributional reinforcement learning, which focuses on learning the entire return distribution instead of only its expectation in standard RL, has demonstrated remarkable success in enhancing performance. Despite these advancements, our comprehension of how the return distribution within distributional RL still remains limited. In this study, we investigate the optimization advantages of distributional RL by utilizing its extra return distribution knowledge over classical RL within the Neural Fitted Z-Iteration~(Neural FZI) framework. To begin with, we demonstrate that the distribution loss of distributional RL has desirable smoothness characteristics and hence enjoys stable gradients, which is in line with its tendency to promote optimization stability. Furthermore, the acceleration effect of distributional RL is revealed by decomposing the return distribution. It shows that distributional RL can perform favorably if the return distribution approximation is appropriate, measured by the variance of gradient estimates in each environment. Rigorous experiments validate the stable optimization behaviors of distributional RL and its acceleration effects compared to classical RL. Our research findings illuminate how the return distribution in distributional RL algorithms helps the optimization.