Distributional Reinforcement Learning with Quantile Regression

Will Dabney; Mark Rowland; Marc G. Bellemare; Rémi Munos

Source-linked AI summary

Distributional Reinforcement Learning with Quantile Regression

Will Dabney, Mark Rowland, Marc G. Bellemare, Rémi Munos

arXiv:1710.10044v1Oct 27, 2017cs.AIcs.LGstat.ML

TL;DR

Traditional reinforcement learning estimates expected returns, leaving open whether distributional learning can use the Wasserstein contraction theory end-to-end. This paper develops a quantile-regression distributional algorithm, proves its contraction properties, and reports state-of-the-art Atari 2600 performance. The authors conclude that the theory–practice gap is closed, while identifying possible benefits from combining QR-DQN with later DQN improvements.

Problem
Existing distributional reinforcement learning theory established a Wasserstein contraction, but C51’s heuristic projection and KL-divergence optimization left end-to-end practical use of that result open.
Method
The paper uses quantile regression with fixed uniform probabilities and adjustable distributional locations to approximate return distributions over the Wasserstein metric.
Results
The combined projected Bellman operator has a unique fixed point and converges under the Wasserstein-based formulation, while the resulting method achieves state-of-the-art performance on Atari 2600.
Takeaways & Limitations
The paper closes the theory–practice gap and supports richer policies based on action-value distributions, including policies relevant to risk-sensitive decision making.
Takeaways & Limitations
QR-DQN may retain overestimation biases similar to those addressed by Double DQN, motivating combination with nondistributional DQN improvements.

Abstract

from arXiv · show

In reinforcement learning an agent interacts with the environment by taking actions and observing the next state and reward. When sampled probabilistically, these state transitions, rewards, and actions can all induce randomness in the observed long-term return. Traditionally, reinforcement learning algorithms average over this randomness to estimate the value function. In this paper, we build on recent work advocating a distributional approach to reinforcement learning in which the distribution over returns is modeled explicitly instead of only estimating the mean. That is, we examine methods of learning the value distribution instead of the value function. We give results that close a number of gaps between the theoretical and algorithmic results given by Bellemare, Dabney, and Munos (2017). First, we extend existing results to the approximate distribution setting. Second, we present a novel distributional reinforcement learning algorithm consistent with our theoretical formulation. Finally, we evaluate this new algorithm on the Atari 2600 games, observing that it significantly outperforms many of the recent improvements on DQN, including the related distributional algorithm C51.

Introduction

The paper addresses a theory–practice gap in distributional reinforcement learning: C51 achieved strong Atari performance, but its heuristic projection did not directly implement the Wasserstein contraction theory. It introduces quantile-regression methods that provide an end-to-end Wasserstein approach and reports improved Atari performance.

Motivation: C51 established strong Atari performance while modeling return distributions, but its heuristic projection and KL-divergence optimization did not directly follow the Wasserstein contraction result.This left the existence of an end-to-end Wasserstein distributional algorithm unresolved.
Contribution: The paper answers this open question by using quantile regression to perform distributional reinforcement learning over the Wasserstein metric in stochastic approximation.The method is designed to connect the theoretical metric with a practical learning algorithm.
Method: The approximation transposes C51’s parameterization: it fixes uniform probabilities and adjusts N distributional locations.C51 instead uses fixed locations and adjusts their probabilities.
Method: Quantile regression stochastically adjusts distributional locations to minimize the Wasserstein distance to a target distribution.The approach estimates quantiles of the target distribution rather than only its expected return.
Theory: The paper formally proves contraction results for its overall algorithm, establishing end-to-end distributional reinforcement learning under the Wasserstein metric.This theoretical result closes part of the gap between the distributional Bellman theory and an implementable algorithm.
Results: QR-DQN achieves better performance than the related C51 algorithm on the same Atari 2600 game suite.The paper also reports a 33% median score increment from using Huber quantile regression, although the supplied passage is truncated after this comparison.

Distributional RL

Distributional reinforcement learning models the full return distribution rather than only its expected value, using Wasserstein-based theory to study learning and approximation. The section motivates this approach, describes the theory–practice gap in C51, and introduces the contraction framework underlying distributional updates.

Distributional RL: Distributional RL models the probability law of returns, capturing intrinsic randomness in the MDP rather than uncertainty in the value estimate.The conventional value function is the expected return, while distributional methods represent the return distribution itself.
Distributional RL: The distributional Bellman operator updates return distributions, analogously to how the standard Bellman operator updates action-value estimates.The update is defined by the random reward plus the discounted next-state return under the policy.
The Wasserstein Metric: The p-Wasserstein metric compares distributions through the Lp distance between their inverse cumulative distribution functions and accounts for distances between outcomes.For p = 1, it is also called the Earth Mover’s Distance; Figure 2 depicts the distance as area between CDFs.
Convergence of Distributional Bellman Operator: The maximal Wasserstein metric is a metric over value distributions, and the distributional Bellman operator is a γ-contraction in that metric.This contraction supports convergence analysis toward the fixed-point value distribution and motivates minimizing Wasserstein distance to Bellman updates.
Convergence of Distributional Bellman Operator: Stochastic gradient descent cannot generally minimize the Wasserstein loss directly because expected sample-loss minimizers may differ from true-loss minimizers.This mismatch becomes important when practical algorithms approximate value distributions; C51 instead projects targets and minimizes KL divergence, without a guarantee of minimizing a Wasserstein metric.

Approximately Minimizing Wasserstein

The paper transposes C51’s parametrization to fixed uniform probabilities over adjustable locations, enabling quantile regression to approximate distributions in Wasserstein distance. This projection, combined with the distributional Bellman operator, is theoretically contractive and supports convergence.

The quantile parametrization fixes uniform probabilities and learns adjustable locations, rather than learning probabilities at fixed support locations.This avoids prespecified support bounds and uniform resolution constraints.
Quantile projection uses inverse-CDF values at quantile midpoints to minimize the 1-Wasserstein distance to a target distribution.For continuous inverse CDFs, the midpoint inverse-CDF value is the unique minimizer.
Quantile regression provides unbiased sample gradients, allowing stochastic gradient descent to learn the minimizing quantile locations.The loss asymmetrically penalizes overestimation and underestimation according to the quantile level.
The quantile Huber loss smooths quantile regression near zero while retaining the asymmetric quantile loss outside a bounded interval.It was introduced because the nonsmooth standard loss might limit performance with nonlinear function approximation.
The projected distributional Bellman operator is a contraction in the ∞-Wasserstein metric, yielding a unique fixed point and convergence for all p ∈ [1, ∞].The contraction inequality is d̄∞(ΠW1TπZ1, ΠW1TπZ2) ≤ γ d̄∞(Z1, Z2).

Distributional RL using Quantile Regression

The algorithm represents value distributions with quantile locations and trains them using quantile regression temporal-difference updates. Its control version, QR-DQN, modifies DQN’s output layer, loss, and optimizer while avoiding C51’s support projection.

The complete distributional algorithm parameterizes value distributions over quantile midpoints and trains location parameters with quantile regression.The QRTD update samples transitions and next-state distribution values, then applies quantile regression to the locations.
Distributional Q-Learning estimates state-action value distributions and selects the next action greedily by the mean of those distributions.The distributional Bellman optimality operator uses a′ = arg maxa′ E[z].
QR-DQN changes DQN’s output to |A| × N quantile targets, replaces its Huber loss with quantile Huber loss, and replaces RMSProp with Adam.These are the three modifications required to form the distributional DQN variant.
QR-DQN avoids C51’s support projection and support-bound hyper-parameter, while N controls the resolution of its value-distribution approximation.Its learned locations can expand or contract to cover the true return range.
In the two-room windy gridworld, transition randomness, wind, and a doorway produce a multimodal return distribution that the algorithms approximate from sampled episodes.The experimental setup uses Monte Carlo rollouts as ground truth and compares TD(0) with QRTD.

Experimental Results

The experiments evaluate QRTD and QR-DQN on stochastic gridworld and 57 Atari 2600 games, measuring distributional accuracy, sample complexity, robustness, and final performance. QR-DQN outperforms prior agents in mean and median human-normalized score, while online results reveal persistent difficulty on a subset of games.

Distributional policy evaluation: QRTD minimizes the 1-Wasserstein distance to the ground-truth return distribution while both QRTD and TD(0) converge correctly in mean.The comparison uses a stochastic two-room windy gridworld and evaluates approximation error from the designated start state.
Atari evaluation: The Atari evaluation covers all 57 games using human-normalized scores and compares best-agent and online performance protocols.Best-agent evaluation freezes learning every one million frames for 500K frames, while online evaluation tracks testing and training returns during learning.
Best-agent performance: QR-DQN outperforms QR-DQN’s listed baselines, including C51, DQN, Double DQN, prioritized replay, and dueling architecture, in mean and median human-normalized score at 200 million frames.Table 1 reports best-agent performance across the 57 Atari 2600 games.
Online performance: QRTD provides sample-complexity improvements similar to prioritized replay while also improving final performance.The online results summarize performance distributions across games rather than only median performance.
Online performance: At 200 million frames, 10% of games remain below 10% of human performance for all algorithms.The result identifies a persistent limitation affecting a small subset of Atari 2600 games.

Conclusions

The paper closes the gap between Wasserstein-metric theory and practical distributional reinforcement learning with theoretical results and a new Atari-performing algorithm. It also identifies richer policies based on action-value distributions as a direction for future work, while noting that QR-DQN may benefit from existing DQN improvements.

Contributions: The paper claims to close the gap between Wasserstein-metric theory and practical distributional reinforcement learning through theoretical contributions and a new algorithm.The conclusion states that the new algorithm achieves state-of-the-art performance in Atari 2600.
Future work: Action-value distributions may support richer decision policies, including policies that consider the distributions as a whole.The paper specifically highlights risk-sensitive decision making as one source of such policies.
Future work: QR-DQN is likely to benefit from later DQN improvements, including methods addressing overestimation bias such as Double DQN.The authors suggest combining QR-DQN with the nondistributional methods listed in Table 1.

Proofs

The proofs establish quantile-based projection properties and use them to show that the projected distributional Bellman operator contracts and converges under the maximal Wasserstein metric. The construction transforms the MDP into an equivalent setting with Dirac value distributions to apply supporting lemmas.

Quantile projection: For quantile intervals [τ, τ′], minimizing the corresponding objective selects the inverse CDF at the midpoint F^-1((τ + τ′)/2).If the inverse CDF is continuous there, the minimizer is unique; otherwise multiple minimizers may exist.
Approximate distributions: The empirical-sample analysis shows that minimizing expected Wasserstein distance to sampled distributions need not recover the original distribution.When the sample distribution lacks an atom at the boundary, the gradient can be nonzero, ruling out the original distribution as the minimizer.
Contraction proof: The projection proposition establishes a Wasserstein contraction relationship for projected value distributions under the distributional Bellman operator.The proof assumes deterministic instantaneous rewards and reduces the argument using contraction and translation invariance properties.
Contraction proof: The proof constructs a transformed MDP whose value distributions are Dirac distributions while preserving the relevant Bellman-update distributions.The construction assigns transition probabilities and transformed distributions so the original and transformed updates have matching distributions.

Supporting results

Supporting lemmas connect quantile projections to Wasserstein distances and prove that projected Dirac-valued Bellman updates preserve the needed quantile ordering. These results provide the local inequalities used in the broader contraction argument.

Quantile projections: The projection onto a Dirac delta at quantile τ is the operator used to compare corresponding quantiles of Bellman-updated distributions.The lemma considers countable state and action spaces with zero rewards and γ = 1.
Quantile ordering: For Dirac-valued distributions, the proof compares the τth quantiles of two Bellman updates through accessible successor state-action pairs.The argument partitions successors according to whether their values fall below or above the selected quantiles.
Quantile ordering: The assumed ordering of the selected quantiles leads to disjointness and containment conditions that contradict the quantile definition.The contradiction shows that the relevant inequality cannot fail, completing the supporting lemma.
Wasserstein projection: The Wasserstein projection lemma expresses the distance between projected distributions through the differences between their selected quantiles.The proof identifies the optimal coupling by pairing corresponding quantile locations.

Further theoretical results

The projected Bellman operator is not generally non-expansive in p-Wasserstein distance, and a concrete two-Dirac example demonstrates this limitation.

The projected Bellman operator ΠW1T π is generally not a non-expansion in dp for p ∈ [1, ∞).
For a two-Dirac construction with γ = 1, the initial value distributions are specified over one initial and two terminal states.The transition probabilities from the initial state are 2/3 and 1/3, with all rewards equal to zero.
The Wasserstein distance between the constructed distributions is dp(Z, Y) = 1/2^1/p.
The projected backup is obtained by first computing the full backup distribution.
With two equally weighted Diracs, the projected locations correspond to the 25% and 75% quantiles of the original distributions.

Notation

The paper defines notation for human-normalized Atari scores and for its distributional reinforcement-learning framework, while Figure 6 reports online training curves and Figure 7 reports raw scores across games.

Human-normalized scores use per-game agent, human-baseline, and random-agent raw scores.
Table 2 lists notation for value distributions, Wasserstein metrics, projections, quantile losses, and parameterized distributions.
Figure 6 shows online training curves for DQN, C51, and QR-DQN across 57 Atari 2600 games.The curves average three seeds, use a five-iteration sliding-window smoother, and display standard-deviation error bands.
Figure 7 presents raw scores across all games after 30 no-op actions, alongside reference values from Wang et al. (2016) and Bellemare, Dabney, and Munos (2017).