# Distributional Reinforcement Learning via Moment Matching

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 24 Jul 2020.

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

- [Exploring the Training Robustness of Distributional Reinforcement Learning against Noisy State Observations](https://mlsanity.com/paper/2109.08776/) (2021)
- [Generalization Bounds with Minimal Dependency on Hypothesis Class via Distributionally Robust Optimization](https://mlsanity.com/paper/2106.11180/) (2021)
- [Distributional Reinforcement Learning with Regularized Wasserstein Loss](https://mlsanity.com/paper/2202.00769/) (2022)
- [MMD-MIX: Value Function Factorisation with Maximum Mean Discrepancy for Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2106.11652/) (2021)
- [Intrinsic Benefits of Categorical Distributional Loss: Uncertainty-aware Regularized Exploration in Reinforcement Learning](https://mlsanity.com/paper/2110.03155/) (2021)
- [How Does Return Distribution in Distributional Reinforcement Learning Help Optimization?](https://mlsanity.com/paper/2209.14513/) (2022)
- [Cooperative Multi-Agent Reinforcement Learning with Hypergraph Convolution](https://mlsanity.com/paper/2112.06771/) (2021)
- [PACER: A Fully Push-forward-based Distributional Reinforcement Learning Algorithm](https://mlsanity.com/paper/2306.06637/) (2023)
- [Risk-Aware Continuous Control with Neural Contextual Bandits](https://mlsanity.com/paper/2312.09961/) (2023)
- [Quantile Geometry Regularization for Distributional Reinforcement Learning](https://mlsanity.com/paper/2605.08182/) (2026)

## Abstract (arXiv)

We consider the problem of learning a set of probability distributions from the empirical Bellman dynamics in distributional reinforcement learning (RL), a class of state-of-the-art methods that estimate the distribution, as opposed to only the expectation, of the total return. We formulate a method that learns a finite set of statistics from each return distribution via neural networks, as in (Bellemare, Dabney, and Munos 2017; Dabney et al. 2018b). Existing distributional RL methods however constrain the learned statistics to predefined functional forms of the return distribution which is both restrictive in representation and difficult in maintaining the predefined statistics. Instead, we learn unrestricted statistics, i.e., deterministic (pseudo-)samples, of the return distribution by leveraging a technique from hypothesis testing known as maximum mean discrepancy (MMD), which leads to a simpler objective amenable to backpropagation. Our method can be interpreted as implicitly matching all orders of moments between a return distribution and its Bellman target. We establish sufficient conditions for the contraction of the distributional Bellman operator and provide finite-sample analysis for the deterministic samples in distribution approximation. Experiments on the suite of Atari games show that our method outperforms the standard distributional RL baselines and sets a new record in the Atari games for non-distributed agents.