# Parameter-Efficient Distributional RL via Normalizing Flows and a Geometry-Aware Cram'er Surrogate

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

## Key facts

- No citing papers in the map yet (data through 7 Oct 2026).
- Builds on 11 of its 19 references that are in the map.
- First posted on arXiv 7 May 2025.

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

- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [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)
- [Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents](https://mlsanity.com/paper/1709.06009/) (2017)
- [Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Design](https://mlsanity.com/paper/1902.00275/) (2019)
- [Search on the Replay Buffer: Bridging Planning and Reinforcement Learning](https://mlsanity.com/paper/1906.05253/) (2019)
- [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)
- [Reinforcement Learning in Economics and Finance](https://mlsanity.com/paper/2003.10014/) (2020)
- [Distributional Reinforcement Learning with Unconstrained Monotonic Neural Networks](https://mlsanity.com/paper/2106.03228/) (2021)

## Abstract (arXiv)

Distributional Reinforcement Learning (DistRL) improves upon expectation-based methods by modeling full return distributions, but standard approaches often remain far from parsimonious. Categorical methods (e.g., C51) rely on fixed supports where parameter counts scale linearly with resolution, while quantile methods approximate distributions as discrete mixtures whose piecewise-constant densities can be wasteful when modeling complex multi-modal or heavy-tailed returns. We introduce NFDRL, a parsimonious architecture that models return distributions using continuous normalizing flows. Unlike categorical baselines, our flow-based model maintains a compact parameter footprint that does not grow with the effective resolution of the distribution, while providing a dynamic, adaptive support for returns. To train this continuous representation, we propose a Cram'er-inspired, geometry-aware distance defined over probability masses obtained from the flow. We show that this distance is a true probability metric, that the associated distributional Bellman operator is a sqrt(gamma)-contraction, and that the resulting objective admits unbiased sample gradients, properties that are typically not simultaneously guaranteed in prior PDF-based DistRL methods. Empirically, NFDRL recovers rich, multi-modal return landscapes on toy MDPs and achieves performance competitive with categorical baselines on the Atari-5 benchmark, while offering substantially better parameter efficiency.