# Distributional Reinforcement Learning with Quantile Regression

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.

## Authors

• 19 publications
• 19 publications
• 33 publications
• 62 publications
• ### QR-MIX: Distributional Value Function Factorisation for Cooperative Multi-Agent Reinforcement Learning

In Cooperative Multi-Agent Reinforcement Learning (MARL) and under the s...
09/09/2020 ∙ by Jian Hu, et al. ∙ 7

• ### A Distributional Perspective on Reinforcement Learning

In this paper we argue for the fundamental importance of the value distr...
07/21/2017 ∙ by Marc G. Bellemare, et al. ∙ 0

The goal of reinforcement learning algorithms is to estimate and/or opti...
05/24/2018 ∙ by Zhongwen Xu, et al. ∙ 0

• ### Distributional reinforcement learning with linear function approximation

Despite many algorithmic advances, our theoretical understanding of prac...
02/08/2019 ∙ by Marc G. Bellemare, et al. ∙ 18

• ### Distributional Multivariate Policy Evaluation and Exploration with the Bellman GAN

The recently proposed distributional approach to reinforcement learning ...
08/06/2018 ∙ by Dror Freirich, et al. ∙ 0

• ### A Distributional Analysis of Sampling-Based Reinforcement Learning Algorithms

We present a distributional approach to theoretical analyses of reinforc...
03/27/2020 ∙ by Philip Amortila, et al. ∙ 6

• ### A neurally plausible model learns successor representations in partially observable environments

Animals need to devise strategies to maximize returns while interacting ...
06/22/2019 ∙ by Eszter Vertes, et al. ∙ 0

##### This week in AI

Get the week's most popular data science and artificial intelligence research sent straight to your inbox every Saturday.

## Introduction

In reinforcement learning, the value of an action in state describes the expected return, or discounted sum of rewards, obtained from beginning in that state, choosing action , and subsequently following a prescribed policy. Because knowing this value for the optimal policy is sufficient to act optimally, it is the object modelled by classic value-based methods such as SARSA [Rummery and Niranjan1994] and Q-Learning [Watkins and Dayan1992], which use Bellman’s equation [Bellman1957] to efficiently reason about value.

Recently, c51 (c51) showed that the distribution of the random returns, whose expectation constitutes the aforementioned value, can be described by the distributional analogue of Bellman’s equation, echoing previous results in risk-sensitive reinforcement learning [Heger1994, Morimura et al.2010, Chow et al.2015]. In this previous work, however, the authors argued for the usefulness in modeling this value distribution in and of itself. Their claim was asserted by exhibiting a distributional reinforcement learning algorithm, c51, which achieved state-of-the-art on the suite of benchmark Atari 2600 games [Bellemare et al.2013].

One of the theoretical contributions of the c51

work was a proof that the distributional Bellman operator is a contraction in a maximal form of the Wasserstein metric between probability distributions. In this context, the Wasserstein metric is particularly interesting because it does not suffer from disjoint-support issues

[Arjovsky, Chintala, and Bottou2017] which arise when performing Bellman updates. Unfortunately, this result does not directly lead to a practical algorithm: as noted by the authors, and further developed by bellemare17cramer (bellemare17cramer), the Wasserstein metric, viewed as a loss, cannot generally be minimized using stochastic gradient methods.

This negative result left open the question as to whether it is possible to devise an online distributional reinforcement learning algorithm which takes advantage of the contraction result. Instead, the c51

algorithm first performs a heuristic projection step, followed by the minimization of a KL divergence between projected Bellman update and prediction. The work therefore leaves a theory-practice gap in our understanding of distributional reinforcement learning, which makes it difficult to explain the good performance of

c51. Thus, the existence of a distributional algorithm that operates end-to-end on the Wasserstein metric remains an open question.

In this paper, we answer this question affirmatively. By appealing to the theory of quantile regression [Koenker2005], we show that there exists an algorithm, applicable in a stochastic approximation setting, which can perform distributional reinforcement learning over the Wasserstein metric. Our method relies on the following techniques:

• We “transpose” the parametrization from c51: whereas the former uses fixed locations for its approximation distribution and adjusts their probabilities, we assign fixed, uniform probabilities to adjustable locations;

• We show that quantile regression may be used to stochastically adjust the distributions’ locations so as to minimize the Wasserstein distance to a target distribution.

• We formally prove contraction mapping results for our overall algorithm, and use these results to conclude that our method performs distributional RL end-to-end under the Wasserstein metric, as desired.

The main interest of the original distributional algorithm was its state-of-the-art performance, despite still acting by maximizing expectations. One might naturally expect that a direct minimization of the Wasserstein metric, rather than its heuristic approximation, may yield even better results. We derive the Q-Learning analogue for our method (qr-dqn), apply it to the same suite of Atari 2600 games, and find that it achieves even better performance. By using a smoothed version of quantile regression, Huber quantile regression, we gain an impressive median score increment over the already state-of-the-art c51.

## Distributional RL

We model the agent-environment interactions by a Markov decision process (MDP) (

) [Puterman1994], with and the state and action spaces,

the random variable reward function,

the probability of transitioning from state to state after taking action , and the discount factor. A policy maps each state to a distribution over .

For a fixed policy , the return, , is a random variable representing the sum of discounted rewards observed along one trajectory of states while following . Standard RL algorithms estimate the expected value of , the value function,

 Vπ(x):=E[Zπ(x)]=E[∞∑t=0γtR(xt,at) | x0=x]. (1)

Similarly, many RL algorithms estimate the action-value function,

 Qπ(x,a):=E[Zπ(x,a)]=E[∞∑t=0γtR(xt,at)], (2) xt∼P(⋅|xt−1,at−1),at∼π(⋅|xt),x0=x,a0=a.

The -greedy policy on chooses actions uniformly at random with probability and otherwise according to .

In distributional RL the distribution over returns (i.e. the probability law of ), plays the central role and replaces the value function. We will refer to the value distribution by its random variable. When we say that the value function is the mean of the value distribution we are saying that the value function is the expected value, taken over all sources of intrinsic randomness [Goldstein, Misra, and Courtage1981], of the value distribution. This should highlight that the value distribution is not designed to capture the uncertainty in the estimate of the value function [Dearden, Friedman, and Russell1998, Engel, Mannor, and Meir2005], that is the parametric uncertainty, but rather the randomness in the returns intrinsic to the MDP.

Temporal difference (TD) methods significantly speed up the learning process by incrementally improving an estimate of using dynamic programming through the Bellman operator [Bellman1957],

 TπQ(x,a)=E[R(x,a)]+γEP,π[Q(x′,a′)].

Similarly, the value distribution can be computed through dynamic programming using a distributional Bellman operator [Bellemare, Dabney, and Munos2017],

 TπZ(x,a):D=R(x,a)+γZ(x′,a′), (3) x′∼P(⋅|x,a),a′∼π(⋅|x′),

where denotes equality of probability laws, that is the random variable is distributed according to the same law as .

The c51 algorithm models using a discrete distribution supported on a “comb” of fixed locations uniformly spaced over a predetermined interval. The parameters of that distribution are the probabilities

, expressed as logits, associated with each location

. Given a current value distribution, the c51 algorithm applies a projection step to map the target onto its finite element support, followed by a Kullback-Leibler (KL) minimization step (see Figure 1). c51 achieved state-of-the-art performance on Atari 2600 games, but did so with a clear disconnect with the theoretical results of c51 (c51). We now review these results before extending them to the case of approximate distributions.

### The Wasserstein Metric

The -Wasserstein metric , for , also known as the Mallows metric [Bickel and Freedman1981] or the Earth Mover’s Distance (EMD) when [Levina and Bickel2001], is an integral probability metric between distributions. The -Wasserstein distance is characterized as the

metric on inverse cumulative distribution functions (inverse CDFs)

[Müller1997]. That is, the -Wasserstein metric between distributions and is given by,111For , .

 Wp(U,Y)=(∫10|F−1Y(ω)−F−1U(ω)|pdω)1/p, (4)

where for a random variable , the inverse CDF of is defined by

 F−1Y(ω):=inf{y∈R:ω≤FY(y)}, (5)

where is the CDF of . Figure 2 illustrates the 1-Wasserstein distance as the area between two CDFs.

Recently, the Wasserstein metric has been the focus of increased research due to its appealing properties of respecting the underlying metric distances between outcomes [Arjovsky, Chintala, and Bottou2017, Bellemare et al.2017]

. Unlike the Kullback-Leibler divergence, the Wasserstein metric is a true probability metric and considers both the probability of and the distance between various outcome events. These properties make the Wasserstein well-suited to domains where an underlying similarity in outcome is more important than exactly matching likelihoods.

### Convergence of Distributional Bellman Operator

In the context of distributional RL, let

be the space of action-value distributions with finite moments:

 Z={ Z:X×A→P(R)| E[|Z(x,a)|p]<∞, ∀(x,a),p≥1}.

Then, for two action-value distributions , we will use the maximal form of the Wasserstein metric introduced by [Bellemare, Dabney, and Munos2017],

 ¯dp(Z1,Z2):=supx,aWp(Z1(x,a),Z2(x,a)). (6)

It was shown that is a metric over value distributions. Furthermore, the distributional Bellman operator is a contraction in , a result that we now recall.

###### Lemma 1 (Lemma 3, c51 c51).

is a -contraction: for any two ,

 ¯dp(TπZ1,TπZ2)≤γ¯dp(Z1,Z2).

Lemma 1 tells us that is a useful metric for studying the behaviour of distributional reinforcement learning algorithms, in particular to show their convergence to the fixed point . Moreover, the lemma suggests that an effective way in practice to learn value distributions is to attempt to minimize the Wasserstein distance between a distribution and its Bellman update , analogous to the way that TD-learning attempts to iteratively minimize the distance between and .

Unfortunately, another result shows that we cannot in general minimize the Wasserstein metric (viewed as a loss) using stochastic gradient descent.

###### Theorem 1 (Theorem 1, bellemare17cramer bellemare17cramer).

Let be the empirical distribution derived from samples

drawn from a Bernoulli distribution

. Let be a Bernoulli distribution parametrized by , the probability of the variable taking the value . Then the minimum of the expected sample loss is in general different from the minimum of the true Wasserstein loss; that is,

This issue becomes salient in a practical context, where the value distribution must be approximated. Crucially, the c51 algorithm is not guaranteed to minimize any -Wasserstein metric. This gap between theory and practice in distributional RL is not restricted to c51. morimura10parametric (morimura10parametric) parameterize a value distribution with the mean and scale of a Gaussian or Laplace distribution, and minimize the KL divergence between the target and the prediction . They demonstrate that value distributions learned in this way are sufficient to perform risk-sensitive Q-Learning. However, any theoretical guarantees derived from their method can only be asymptotic; the Bellman operator is at best a non-expansion in KL divergence.

## Approximately Minimizing Wasserstein

Recall that c51 approximates the distribution at each state by attaching variable (parametrized) probabilities to fixed locations . Our approach is to “transpose” this parametrization by considering fixed probabilities but variable locations. Specifically, we take uniform weights, so that for each .

Effectively, our new approximation aims to estimate quantiles of the target distribution. Accordingly, we will call it a quantile distribution, and let be the space of quantile distributions for fixed . We will denote the cumulative probabilities associated with such a distribution (that is, the discrete values taken on by the CDF) by , so that for . We will also write to simplify notation.

Formally, let

be some parametric model. A quantile distribution

maps each state-action pair to a uniform probability distribution supported on . That is,

 Zθ(x,a):=1NN∑i=1δθi(x,a), (7)

where denotes a Dirac at .

Compared to the original parametrization, the benefits of a parameterized quantile distribution are threefold. First, (1) we are not restricted to prespecified bounds on the support, or a uniform resolution, potentially leading to significantly more accurate predictions when the range of returns vary greatly across states. This also (2) lets us do away with the unwieldy projection step present in c51, as there are no issues of disjoint supports. Together, these obviate the need for domain knowledge about the bounds of the return distribution when applying the algorithm to new tasks. Finally, (3) this reparametrization allows us to minimize the Wasserstein loss, without suffering from biased gradients, specifically, using quantile regression.

### The Quantile Approximation

It is well-known that in reinforcement learning, the use of function approximation may result in instabilities in the learning process [Tsitsiklis and Van Roy1997]. Specifically, the Bellman update projected onto the approximation space may no longer be a contraction. In our case, we analyze the distributional Bellman update, projected onto a parameterized quantile distribution, and prove that the combined operator is a contraction.

#### Quantile Projection

We are interested in quantifying the projection of an arbitrary value distribution onto , that is

 ΠW1Z:=argminZθ∈ZQW1(Z,Zθ),

Let be a distribution with bounded first moment and

a uniform distribution over

Diracs as in (7), with support . Then

 W1(Y,U)=N∑i=1∫τiτi−1|F−1Y(ω)−θi|dω.
###### Lemma 2.

For any with and cumulative distribution function with inverse , the set of minimizing

 ∫τ′τ|F−1(ω)−θ|dω,

is given by

 {θ∈R∣∣∣F(θ)=(τ+τ′2)}.

In particular, if is the inverse CDF, then is always a valid minimizer, and if is continuous at , then is the unique minimizer.

These quantile midpoints will be denoted by for . Therefore, by Lemma 2, the values for that minimize are given by . Figure 2 shows an example of the quantile projection minimizing the -Wasserstein distance to .222We save proofs for the appendix due to space limitations.

### Quantile Regression

The original proof of Theorem 1 only states the existence of a distribution whose gradients are biased. As a result, we might hope that our quantile parametrization leads to unbiased gradients. Unfortunately, this is not true.

###### Proposition 1.

Let be a quantile distribution, and the empirical distribution composed of samples from . Then for all , there exists a such that

 argminE[Wp(^Zm,Zθ)]≠argminWp(Z,Zθ).

However, there is a method, more widely used in economics than machine learning, for unbiased stochastic approximation of the quantile function.

Quantile regression, and conditional quantile regression, are methods for approximating the quantile functions of distributions and conditional distributions respectively [Koenker2005]. These methods have been used in a variety of settings where outcomes have intrinsic randomness [Koenker and Hallock2001]; from food expenditure as a function of household income [Engel1857], to studying value-at-risk in economic models [Taylor1999].

The quantile regression loss, for quantile

, is an asymmetric convex loss function that penalizes overestimation errors with weight

and underestimation errors with weight . For a distribution , and a given quantile , the value of the quantile function may be characterized as the minimizer of the quantile regression loss

 Lτ\textscqr(θ):=E^Z∼Z[ρτ(^Z−θ)], where ρτ(u)=u(τ−δ{u<0}), ∀u∈R. (8)

More generally, by Lemma 2 we have that the minimizing values of for are those that minimize the following objective:

 N∑i=1E^Z∼Z[ρ^τi(^Z−θi)]

In particular, this loss gives unbiased sample gradients. As a result, we can find the minimizing by stochastic gradient descent.

#### Quantile Huber Loss

The quantile regression loss is not smooth at zero; as , the gradient of Equation 8 stays constant. We hypothesized that this could limit performance when using non-linear function approximation. To this end, we also consider a modified quantile loss, called the quantile Huber loss.333Our quantile Huber loss is related to, but distinct from that of aravkin2014sparse (aravkin2014sparse). This quantile regression loss acts as an asymmetric squared loss in an interval around zero and reverts to a standard quantile loss outside this interval.

The Huber loss is given by [Huber1964],

 Lκ(u)={12u2,if |u|≤κκ(|u|−12κ),otherwise. (9)

The quantile Huber loss is then simply the asymmetric variant of the Huber loss,

 ρκτ(u)=|τ−δ{u<0}|Lκ(u). (10)

For notational simplicity we will denote , that is, it will revert to the standard quantile regression loss.

### Combining Projection and Bellman Update

We are now in a position to prove our main result, which states that the combination of the projection implied by quantile regression with the Bellman operator is a contraction. The result is in -Wasserstein metric, i.e. the size of the largest gap between the two CDFs.

###### Proposition 2.

Let be the quantile projection defined as above, and when applied to value distributions gives the projection for each state-value distribution. For any two value distributions for an MDP with countable state and action spaces,

 ¯d∞(ΠW1TπZ1,ΠW1TπZ2)≤γ¯d∞(Z1,Z2). (11)

We therefore conclude that the combined operator has a unique fixed point , and the repeated application of this operator, or its stochastic approximation, converges to . Because , we conclude that convergence occurs for all . Interestingly, the contraction property does not directly hold for ; see Lemma 5 in the appendix.

## Distributional RL using Quantile Regression

We can now form a complete algorithmic approach to distributional RL consistent with our theoretical results. That is, approximating the value distribution with a parameterized quantile distribution over the set of quantile midpoints, defined by Lemma 2. Then, training the location parameters using quantile regression (Equation 8).

### Quantile Regression Temporal Difference Learning

Recall the standard TD update for evaluating a policy ,

 V(x)←V(x)+α(r+γV(x′)−V(x)), a∼π(⋅|x),r∼R(x,a),x′∼P(⋅|x,a).

TD allows us to update the estimated value function with a single unbiased sample following . Quantile regression also allows us to improve the estimate of the quantile function for some target distribution, , by observing samples and minimizing Equation 8.

Furthermore, we have shown that by estimating the quantile function for well-chosen values of we can obtain an approximation with minimal 1-Wasserstein distance from the original (Lemma 2). Finally, we can combine this with the distributional Bellman operator to give a target distribution for quantile regression. This gives us the quantile regression temporal difference learning (qrtd) algorithm, summarized simply by the update,

 θi(x)←θi(x)+α(^τi−δ{r+γz′<θi(x))}), (12) a∼π(⋅|x),r∼R(x,a),x′∼P(⋅|x,a),z′∼Zθ(x′),

where is a quantile distribution as in (7), and is the estimated value of in state . It is important to note that this update is for each value of and is defined for a single sample from the next state value distribution. In general it is better to draw many samples of and minimize the expected update. A natural approach in this case, which we use in practice, is to compute the update for all pairs of (). Next, we turn to a control algorithm and the use of non-linear function approximation.

### Quantile Regression dqn

Q-Learning is an off-policy reinforcement learning algorithm built around directly learning the optimal action-value function using the Bellman optimality operator [Watkins and Dayan1992],

 TQ(x,a)=E[R(x,a)]+γEx′∼P[maxa′Q(x′,a′)].

The distributional variant of this is to estimate a state-action value distribution and apply a distributional Bellman optimality operator,

 TZ(x,a)=R(x,a)+γZ(x′,a′), (13) x′∼P(⋅|x,a),a′=argmaxa′Ez∼Z(x′,a′)[z].

Notice in particular that the action used for the next state is the greedy action with respect to the mean of the next state-action value distribution.

For a concrete algorithm we will build on the dqn architecture [Mnih et al.2015]. We focus on the minimal changes necessary to form a distributional version of dqn. Specifically, we require three modifications to dqn

. First, we use a nearly identical neural network architecture as

dqn, only changing the output layer to be of size , where is a hyper-parameter giving the number of quantile targets. Second, we replace the Huber loss used by dqn444dqn

uses gradient clipping of the squared error that makes it equivalent to a Huber loss with

.
, with , with a quantile Huber loss (full loss given by Algorithm 1

). Finally, we replace RMSProp

[Tieleman and Hinton2012] with Adam [Kingma and Ba2015]. We call this new algorithm quantile regression dqn (qr-dqn).

Unlike c51, qr-dqn does not require projection onto the approximating distribution’s support, instead it is able to expand or contract the values arbitrarily to cover the true range of return values. As an additional advantage, this means that qr-dqn does not require the additional hyper-parameter giving the bounds of the support required by c51. The only additional hyper-parameter of qr-dqn not shared by dqn is the number of quantiles , which controls with what resolution we approximate the value distribution. As we increase , qr-dqn goes from dqn to increasingly able to estimate the upper and lower quantiles of the value distribution. It becomes increasingly capable of distinguishing low probability events at either end of the cumulative distribution over returns.

## Experimental Results

In the introduction we claimed that learning the distribution over returns had distinct advantages over learning the value function alone. We have now given theoretically justified algorithms for performing distributional reinforcement learning, qrtd for policy evaluation and qr-dqn for control. In this section we will empirically validate that the proposed distributional reinforcement learning algorithms: (1) learn the true distribution over returns, (2) show increased robustness during training, and (3) significantly improve sample complexity and final performance over baseline algorithms.

#### Value Distribution Approximation Error

We begin our experimental results by demonstrating that qrtd actually learns an approximate value distribution that minimizes the -Wasserstein to the ground truth distribution over returns. Although our theoretical results already establish convergence of the former to the latter, the empirical performance helps to round out our understanding.

We use a variant of the classic windy gridworld domain [Sutton and Barto1998], modified to have two rooms and randomness in the transitions. Figure 3(a) shows our version of the domain, where we have combined the transition stochasticity, wind, and the doorway to produce a multi-modal distribution over returns when anywhere in the first room. Each state transition has probability of moving in a random direction, otherwise the transition is affected by wind moving the agent northward. The reward function is zero until reaching the goal state , which terminates the episode and gives a reward of . The discount factor is .

We compute the ground truth value distribution for optimal policy , learned by policy iteration, at each state by performing Monte-Carlo (MC) rollouts and recording the observed returns as an empirical distribution, shown in Figure 3(b). Next, we ran both 0 and qrtd while following for episodes. Each episode begins in the designated start state (). Both algorithms started with a learning rate of . For qrtd we used and drop by half every episodes.

Let be the MC estimated distribution over returns from the start state , similarly its mean. In Figure 3 we show the approximation errors at for both algorithms with respect to the number of episodes. In (d) we evaluated, for both 0 and qrtd, the squared error, , and in (e) we show the -Wasserstein metric for qrtd, , where and are the expected returns and value distribution at state estimated by the algorithm. As expected both algorithms converge correctly in mean, and qrtd minimizes the -Wasserstein distance to .

### Evaluation on Atari 2600

We now provide experimental results that demonstrate the practical advantages of minimizing the Wasserstein metric end-to-end, in contrast to the c51 approach. We use the 57 Atari 2600 games from the Arcade Learning Environment (ALE) [Bellemare et al.2013]. Both c51 and qr-dqn build on the standard dqn architecture, and we expect both to benefit from recent improvements to dqn such as the dueling architectures [Wang et al.2016] and prioritized replay [Schaul et al.2016]. However, in our evaluations we compare the pure versions of c51 and qr-dqn without these additions. We present results for both a strict quantile loss, (qr-dqn-), and with a Huber quantile loss with (qr-dqn-).

We performed hyper-parameter tuning over a set of five training games and evaluated on the full set of 57 games using these best settings (, , and ).555We swept over in (); in (); () As with dqn we use a target network when computing the distributional Bellman update. We also allow to decay at the same rate as in dqn, but to a lower value of , as is common in recent work [Bellemare, Dabney, and Munos2017, Wang et al.2016, van Hasselt, Guez, and Silver2016].

Out training procedure follows that of mnih15nature (mnih15nature)’s, and we present results under two evaluation protocols: best agent performance and online performance. In both evaluation protocols we consider performance over all 57 Atari 2600 games, and transform raw scores into human-normalized scores [van Hasselt, Guez, and Silver2016].

#### Best agent performance

To provide comparable results with existing work we report test evaluation results under the best agent protocol. Every one million training frames, learning is frozen and the agent is evaluated for frames while recording the average return. Evaluation episodes begin with up to random no-ops [Mnih et al.2015], and the agent uses a lower exploration rate (). As training progresses we keep track of the best agent performance achieved thus far.

Table 1 gives the best agent performance, at million frames trained, for qr-dqn, c51, dqn, Double dqn [van Hasselt, Guez, and Silver2016], Prioritized replay [Schaul et al.2016], and Dueling architecture [Wang et al.2016]. We see that qr-dqn outperforms all previous agents in mean and median human-normalized score.

#### Online performance

In this evaluation protocol (Figure 4) we track the average return attained during each testing (left) and training (right) iteration. For the testing performance we use a single seed for each algorithm, but show online performance with no form of early stopping. For training performance, values are averages over three seeds. Instead of reporting only median performance, we look at the distribution of human-normalized scores over the full set of games. Each bar represents the score distribution at a fixed percentile (th, th, th, th, and th). The upper percentiles show a similar trend but are omitted here for visual clarity, as their scale dwarfs the more informative lower half.

From this, we can infer a few interesting results. (1) Early in learning, most algorithms perform worse than random for at least of games. (2) qrtd gives similar improvements to sample complexity as prioritized replay, while also improving final performance. (3) Even at million frames, there are of games where all algorithms reach less than of human. This final point in particular shows us that all of our recent advances continue to be severely limited on a small subset of the Atari 2600 games.

## Conclusions

The importance of the distribution over returns in reinforcement learning has been (re)discovered and highlighted many times by now. In c51 (c51) the idea was taken a step further, and argued to be a central part of approximate reinforcement learning. However, the paper left open the question of whether there exists an algorithm which could bridge the gap between Wasserstein-metric theory and practical concerns.

In this paper we have closed this gap with both theoretical contributions and a new algorithm which achieves state-of-the-art performance in Atari 2600. There remain many promising directions for future work. Most exciting will be to expand on the promise of a richer policy class, made possible by action-value distributions. We have mentioned a few examples of such policies, often used for risk-sensitive decision making. However, there are many more possible decision policies that consider the action-value distributions as a whole.

Additionally, qr-dqn is likely to benefit from the improvements on dqn made in recent years. For instance, due to the similarity in loss functions and Bellman operators we might expect that qr-dqn suffers from similar over-estimation biases to those that Double dqn was designed to address [van Hasselt, Guez, and Silver2016]. A natural next step would be to combine qr-dqn with the non-distributional methods found in Table 1.

## Acknowledgements

The authors acknowledge the vital contributions of their colleagues at DeepMind. Special thanks to Tom Schaul, Audrunas Gruslys, Charles Blundell, and Benigno Uria for their early suggestions and discussions on the topic of quantile regression. Additionally, we are grateful for feedback from David Silver, Yee Whye Teh, Georg Ostrovski, Joseph Modayil, Matt Hoffman, Hado van Hasselt, Ian Osband, Mohammad Azar, Tom Stepleton, Olivier Pietquin, Bilal Piot; and a second acknowledgement in particular of Tom Schaul for his detailed review of an previous draft.

## References

• [Aravkin et al.2014] Aravkin, A. Y.; Kambadur, A.; Lozano, A. C.; and Luss, R. 2014. Sparse Quantile Huber Regression for Efficient and Robust Estimation. arXiv.
• [Arjovsky, Chintala, and Bottou2017] Arjovsky, M.; Chintala, S.; and Bottou, L. 2017. Wasserstein Generative Adversarial Networks. In Proceedings of the 34th International Conference on Machine Learning (ICML).
• [Bellemare et al.2013] Bellemare, M. G.; Naddaf, Y.; Veness, J.; and Bowling, M. 2013. The Arcade Learning Environment: An Evaluation Platform for General Agents.

Journal of Artificial Intelligence Research

47:253–279.
• [Bellemare et al.2017] Bellemare, M. G.; Danihelka, I.; Dabney, W.; Mohamed, S.; Lakshminarayanan, B.; Hoyer, S.; and Munos, R. 2017. The Cramer Distance as a Solution to Biased Wasserstein Gradients. arXiv.
• [Bellemare, Dabney, and Munos2017] Bellemare, M. G.; Dabney, W.; and Munos, R. 2017. A Distributional Perspective on Reinforcement Learning. Proceedings of the 34th International Conference on Machine Learning (ICML).
• [Bellman1957] Bellman, R. E. 1957. Dynamic Programming. Princeton, NJ: Princeton University Press.
• [Bickel and Freedman1981] Bickel, P. J., and Freedman, D. A. 1981. Some Asymptotic Theory for the Bootstrap. The Annals of Statistics 1196–1217.
• [Chow et al.2015] Chow, Y.; Tamar, A.; Mannor, S.; and Pavone, M. 2015. Risk-Sensitive and Robust Decision-Making: a CVaR Optimization Approach. In Advances in Neural Information Processing Systems (NIPS), 1522–1530.
• [Dearden, Friedman, and Russell1998] Dearden, R.; Friedman, N.; and Russell, S. 1998. Bayesian Q-learning. In Proceedings of the National Conference on Artificial Intelligence.
• [Engel, Mannor, and Meir2005] Engel, Y.; Mannor, S.; and Meir, R. 2005. Reinforcement Learning with Gaussian Processes. In Proceedings of the International Conference on Machine Learning (ICML).
• [Engel1857] Engel, E. 1857. Die Productions-und Consumtionsverhältnisse des Königreichs Sachsen. Zeitschrift des Statistischen Bureaus des Königlich Sächsischen Ministeriums des Innern 8:1–54.
• [Goldstein, Misra, and Courtage1981] Goldstein, S.; Misra, B.; and Courtage, M. 1981. On Intrinsic Randomness of Dynamical Systems. Journal of Statistical Physics 25(1):111–126.
• [Heger1994] Heger, M. 1994. Consideration of Risk in Reinforcement Learning. In Proceedings of the 11th International Conference on Machine Learning, 105–111.
• [Huber1964] Huber, P. J. 1964. Robust Estimation of a Location Parameter. The Annals of Mathematical Statistics 35(1):73–101.
• [Kingma and Ba2015] Kingma, D., and Ba, J. 2015. Adam: A Method for Stochastic Optimization. Proceedings of the International Conference on Learning Representations.
• [Koenker and Hallock2001] Koenker, R., and Hallock, K. 2001. Quantile Regression: An Introduction. Journal of Economic Perspectives 15(4):43–56.
• [Koenker2005] Koenker, R. 2005. Quantile Regression. Cambridge University Press.
• [Levina and Bickel2001] Levina, E., and Bickel, P. 2001. The Earth Mover’s Distance is the Mallows Distance: Some Insights from Statistics. In

The 8th IEEE International Conference on Computer Vision (ICCV)

.
IEEE.
• [Mnih et al.2015] Mnih, V.; Kavukcuoglu, K.; Silver, D.; Rusu, A. A.; Veness, J.; Bellemare, M. G.; Graves, A.; Riedmiller, M.; Fidjeland, A. K.; Ostrovski, G.; et al. 2015. Human-level Control through Deep Reinforcement Learning. Nature 518(7540):529–533.
• [Morimura et al.2010] Morimura, T.; Hachiya, H.; Sugiyama, M.; Tanaka, T.; and Kashima, H. 2010. Parametric Return Density Estimation for Reinforcement Learning. In Proceedings of the Conference on Uncertainty in Artificial Intelligence (UAI).
• [Müller1997] Müller, A. 1997. Integral Probability Metrics and their Generating Classes of Functions. Advances in Applied Probability 29(2):429–443.
• [Nair et al.2015] Nair, A.; Srinivasan, P.; Blackwell, S.; Alcicek, C.; Fearon, R.; De Maria, A.; Panneershelvam, V.; Suleyman, M.; Beattie, C.; and Petersen, S. e. a. 2015. Massively Parallel Methods for Deep Reinforcement Learning. In

ICML Workshop on Deep Learning

.
• [Puterman1994] Puterman, M. L. 1994. Markov Decision Processes: Discrete stochastic dynamic programming. John Wiley & Sons, Inc.
• [Rummery and Niranjan1994] Rummery, G. A., and Niranjan, M. 1994. On-line Q-learning using Connectionist Systems. Technical report, Cambridge University Engineering Department.
• [Schaul et al.2016] Schaul, T.; Quan, J.; Antonoglou, I.; and Silver, D. 2016. Prioritized Experience Replay. In Proceedings of the International Conference on Learning Representations (ICLR).
• [Sutton and Barto1998] Sutton, R. S., and Barto, A. G. 1998. Reinforcement Learning: An Introduction. MIT Press.
• [Taylor1999] Taylor, J. W. 1999. A Quantile Regression Approach to Estimating the Distribution of Multiperiod Returns. The Journal of Derivatives 7(1):64–78.
• [Tieleman and Hinton2012] Tieleman, T., and Hinton, G. 2012. Lecture 6.5: Rmsprop. COURSERA: Neural Networks for Machine Learning 4(2).
• [Tsitsiklis and Van Roy1997] Tsitsiklis, J. N., and Van Roy, B. 1997. An Analysis of Temporal-Difference Learning with Function Approximation. IEEE Transactions on Automatic Control 42(5):674–690.
• [van Hasselt, Guez, and Silver2016] van Hasselt, H.; Guez, A.; and Silver, D. 2016. Deep Reinforcement Learning with Double Q-Learning. In Proceedings of the AAAI Conference on Artificial Intelligence.
• [Wang et al.2016] Wang, Z.; Schaul, T.; Hessel, M.; Hasselt, H. v.; Lanctot, M.; and de Freitas, N. 2016. Dueling Network Architectures for Deep Reinforcement Learning. In Proceedings of the International Conference on Machine Learning (ICML).
• [Watkins and Dayan1992] Watkins, C. J., and Dayan, P. 1992. Q-learning. Machine Learning 8(3):279–292.

## Appendix

### Proofs

See 2

###### Proof.

For any , the function is convex, and has subgradient given by

 θ↦⎧⎪⎨⎪⎩1\ if\ θF−1(ω),

so the function is also convex, and has subgradient given by

 θ↦∫F(θ)τ−1dω+∫τ′F(θ)1dω.

Setting this subgradient equal to yields

 (τ+τ′)−2F(θ)=0, (14)

and since is the identity map on , it is clear that satisfies Equation 14. Note that in fact any such that yields a subgradient of , which leads to a multitude of minimizers if is not continuous at . ∎

See 1

###### Proof.

Write , with . We take to be of the same form as . Specifically, consider given by

 Z=N∑i=11Nδi,

supported on the set , and take . Then clearly the unique minimizing for is given by taking . However, consider the gradient with respect to for the objective

 E[Wp(^ZN,Zθ)].

We have

 ∇θ1E[Wp(^ZN,Zθ)]|θ1=1=E[∇θ1Wp(^ZN,Zθ)|θ1=1].

In the event that the sample distribution has an atom at , then the optimal transport plan pairs the atom of at with this atom of , and gradient with respect to of is . If the sample distribution does not contain an atom at , then the left-most atom of is greater than (since is supported on . In this case, the gradient on is negative. Since this happens with non-zero probability, we conclude that

 ∇θ1E[Wp(^ZN,Zθ)]|θ1=1<0,

and therefore cannot be the minimizer of . ∎

See 2

###### Proof.

We assume that instantaneous rewards given a state-action pair are deterministic; the general case is a straightforward generalization. Further, since the operator is a -contraction in , it is sufficient to prove the claim in the case . In addition, since Wasserstein distances are invariant under translation of the support of distributions, it is sufficient to deal with the case where for all . The proof then proceeds by first reducing to the case where every value distribution consists only of single Diracs, and then dealing with this reduced case using Lemma 3.

We write and , for some functions . Let be a state-action pair, and let be all the state-action pairs that are accessible from in a single transition, where is a (finite or countable) indexing set. Write for the probability of transitioning from to , for each . We now construct a new MDP and new value distributions for this MDP in which all distributions are given by single Diracs, with a view to applying Lemma 3. The new MDP is of the following form. We take the state-action pair , and define new states, actions, transitions, and a policy , so that the state-action pairs accessible from in this new MDP are given by , and the probability of reaching the state-action pair is . Further, we define new value distributions as follows. For each and , we set:

 ˜Z(˜xji,˜aji)=δθj(xi,ai) ˜Y(˜xji,˜aji)=δψj(xi,ai).

The construction is illustrated in Figure 5.

Since, by Lemma 4, the distance between the 1-Wasserstein projections of two real-valued distributions is the max over the difference of a certain set of quantiles, we may appeal to Lemma 3 to obtain the following:

 d∞(ΠW1(T˜π˜Z)(x′,a′),ΠW1(T˜π˜Y)(x′,a′)) ≤ supi=1∈Ij=1,…,N|θj(xi,ai)−ψj(xi,ai)| = supi=1∈Id∞(Z(xi,ai),Y(xi,ai)) (15)

Now note that by construction, (respectively, ) has the same distribution as (respectively, ), and so

 d∞(ΠW1(T˜π˜Z)(x′,a′),ΠW1(T˜π˜Y)(x′,a′)) =d∞(ΠW1(TπZ)(x′,a′),ΠW1(TπY)(x′,a′)).

Therefore, substituting this into the Inequality 15, we obtain

 d∞(ΠW1(TπZ)(x′,a′),ΠW1(TπY)(x′,a′)) ≤ supi∈Id∞(Z(xi,ai),Y(xi,ai)).

Taking suprema over the initial state then yields the result. ∎

### Supporting results

###### Lemma 3.

Consider an MDP with countable state and action spaces. Let be value distributions such that each state-action distribution , is given by a single Dirac. Consider the particular case where rewards are identically and , and let . Denote by the projection operator that maps a probability distribution onto a Dirac delta located at its th quantile. Then

 ¯¯¯d∞(ΠτTπZ,ΠτTπY)≤¯¯¯d∞(Z,Y)
###### Proof.

Let and for each state-action pair , for some functions . Let be a state-action pair, and let be all the state-action pairs that are accessible from in a single transition, with a (finite or countably infinite) indexing set. To lighten notation, we write for and for . Further, let the probability of transitioning from to be , for all .

Then we have

 (TπZ)(x′,a′)=∑i∈Ipiδθi (16) (TπY)(x′,a′)=∑i∈Ipiδψi. (17)

Now consider the th quantile of each of these distributions, for arbitrary. Let be such that is equal to this quantile of , and let such that is equal to this quantile of . Now note that

 d∞(ΠτTπZ(x′,a′),ΠτTπY(x′,a′))=|θu−ψv|

We now show that

 |θu−ψv|>|θi−ψi|   ∀i∈I (18)

is impossible, from which it will follow that

 d∞(ΠτTπZ(x′,a′),ΠτTπY(x′,a′))≤¯¯¯d∞(Z,Y),

and the result then follows by taking maxima over state-action pairs . To demonstrate the impossibility of (18), without loss of generality we take .

We now introduce the following partitions of the indexing set . Define:

 I≤θu={i∈I|θi≤θu}, I>θu={i∈I|θi>θu}, I<ψv={i∈I|ψi<ψv}, I≥ψv={i∈I|ψi≥ψv},

and observe that we clearly have the following disjoint unions:

 I=I≤θu∪I>θu, I=I<ψv∪I≥ψv.

If (18) is to hold, then we must have . Therefore, we must have . But if this is the case, then since is the th quantile of , we must have

 ∑i∈I≤θupi≥τ,

and so consequently

 ∑i∈I<ψvpi≥τ,

from which we conclude that the th quantile of is less than , a contradiction. Therefore (18) cannot hold, completing the proof. ∎

###### Lemma 4.

For any two probability distributions over the real numbers, and the Wasserstein projection operator that projects distributions onto support of size , we have that

 d∞(ΠW1ν1,ΠW1ν2) = maxi=1,…,n∣∣∣F−1ν1(2i−12n)−F−1ν2(2i−12n)∣∣∣.
###### Proof.

By the discussion surrounding Lemma 2, we have that for . Therefore, the optimal coupling between and must be given by for each . This immediately leads to the expression of the lemma. ∎

### Further theoretical results

###### Lemma 5.

The projected Bellman operator is in general not a non-expansion in , for .

###### Proof.

Consider the case where the number of Dirac deltas in each distribution, , is equal to , and let . We consider an MDP with a single initial state, , and two terminal states, and . We take the action space of the MDP to be trivial, and therefore omit it in the notation that follows. Let the MDP have a probability of transitioning from to , and probability of transitioning from to . We take all rewards in the MDP to be identically . Further, consider two value distributions, and , given by:

 Z(x1)=12δ0+12δ2, Y(x1)=12δ1+12δ2, Z(x2)=12δ