Distributional Reinforcement Learning with Quantile Regression

Introduction

In reinforcement learning, the value of an action $a$ in state $s$ describes the expected return, or discounted sum of rewards, obtained from beginning in that state, choosing action $a$, 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 (?) and Q-Learning (?), which use Bellman’s equation (?) to efficiently reason about value.

Recently, ? (?) 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 (?; ?; ?). 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 (?).

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 (?) 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 ? (?), 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 (?), 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 $N$ fixed locations for its approximation distribution and adjusts their probabilities, we assign fixed, uniform probabilities to $N$ 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 $33\%$ 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) ($\mathcal{X},\mathcal{A},R,P,\gamma$) (?), with $\mathcal{X}$ and $\mathcal{A}$ the state and action spaces, $R$ the random variable reward function, $P(x^{\prime}|x,a)$ the probability of transitioning from state $x$ to state $x^{\prime}$ after taking action $a$, and $\gamma\in[0,1)$ the discount factor. A policy $\pi(\cdot|x)$ maps each state $x\in\mathcal{X}$ to a distribution over $\mathcal{A}$.

For a fixed policy $\pi$, the return, $Z^{\pi}=\sum_{t=0}^{\infty}\gamma^{t}R_{t}$, is a random variable representing the sum of discounted rewards observed along one trajectory of states while following $\pi$. Standard RL algorithms estimate the expected value of $Z^{\pi}$, the value function,

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

The $\epsilon$-greedy policy on $Q^{\pi}$ chooses actions uniformly at random with probability $\epsilon$ and otherwise according to $\operatorname*{arg\,max}_{a}Q^{\pi}(x,a)$.

In distributional RL the distribution over returns (i.e. the probability law of $Z^{\pi}$), 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 (?), 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 (?; ?), 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 $Q^{\pi}$ using dynamic programming through the Bellman operator (?),

Similarly, the value distribution can be computed through dynamic programming using a distributional Bellman operator (?),

where $Y:\overset{D}{=}U$ denotes equality of probability laws, that is the random variable $Y$ is distributed according to the same law as $U$.

The c51 algorithm models $Z^{\pi}(x,a)$ using a discrete distribution supported on a “comb” of fixed locations $z_{1}\leq\cdots\leq z_{N}$ uniformly spaced over a predetermined interval. The parameters of that distribution are the probabilities $q_{i}$, expressed as logits, associated with each location $z_{i}$. Given a current value distribution, the c51 algorithm applies a projection step $\Phi$ to map the target $\mathcal{T}^{\pi}Z$ 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 ? (?). We now review these results before extending them to the case of approximate distributions.

The $p$-Wasserstein metric $W_{p}$, for $p\in[1,\infty]$, also known as the Mallows metric (?) or the Earth Mover’s Distance (EMD) when $p=1$ (?), is an integral probability metric between distributions. The $p$-Wasserstein distance is characterized as the $L^{p}$ metric on inverse cumulative distribution functions (inverse CDFs) (?). That is, the $p$-Wasserstein metric between distributions $U$ and $Y$ is given by,For $p=\infty$, $W_{\infty}(Y,U)=\sup_{\omega\in}|F_{Y}^{-1}(\omega)-F_{U}^{-1}(\omega)|$.

where for a random variable $Y$, the inverse CDF $F^{-1}_{Y}$ of $Y$ is defined by

where $F_{Y}(y)=Pr(Y\leq y)$ is the CDF of $Y$. 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 (?; ?). 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 $\mathcal{Z}$ be the space of action-value distributions with finite moments:

Then, for two action-value distributions $Z_{1},Z_{2}\in\mathcal{Z}$, we will use the maximal form of the Wasserstein metric introduced by (?),

It was shown that $\bar{d}_{p}$ is a metric over value distributions. Furthermore, the distributional Bellman operator $\mathcal{T}^{\pi}$ is a contraction in $\bar{d}_{p}$, a result that we now recall.

$\mathcal{T}^{\pi}$ is a $\gamma$-contraction: for any two $Z_{1},Z_{2}\in\mathcal{Z}$,

Lemma 1 tells us that $\bar{d}_{p}$ is a useful metric for studying the behaviour of distributional reinforcement learning algorithms, in particular to show their convergence to the fixed point $Z^{\pi}$. 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 $Z$ and its Bellman update $\mathcal{T}^{\pi}Z$, analogous to the way that TD-learning attempts to iteratively minimize the $L^{2}$ distance between $Q$ and $\mathcal{T}Q$.

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

Let $\hat{Y}_{m}:=\tfrac{1}{m}\sum_{i=1}^{m}\delta_{Y_{i}}$ be the empirical distribution derived from samples $Y_{1},\dots,Y_{m}$ drawn from a Bernoulli distribution $B$. Let $B_{\mu}$ be a Bernoulli distribution parametrized by $\mu$, the probability of the variable taking the value $1$. 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 $p$-Wasserstein metric. This gap between theory and practice in distributional RL is not restricted to c51. ? (?) parameterize a value distribution with the mean and scale of a Gaussian or Laplace distribution, and minimize the KL divergence between the target $\mathcal{T}^{\pi}Z$ and the prediction $Z$. 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 $q_{1},\dots,q_{N}$ to fixed locations $z_{1}\leq\cdots\leq z_{N}$. Our approach is to “transpose” this parametrization by considering fixed probabilities but variable locations. Specifically, we take uniform weights, so that $q_{i}=1/N$ for each $i=1,\ldots,N$.

Effectively, our new approximation aims to estimate quantiles of the target distribution. Accordingly, we will call it a quantile distribution, and let $\mathcal{Z}_{Q}$ be the space of quantile distributions for fixed $N$. We will denote the cumulative probabilities associated with such a distribution (that is, the discrete values taken on by the CDF) by $\tau_{1},\ldots,\tau_{N}$, so that $\tau_{i}=\frac{i}{N}$ for $i=1,\ldots,N$. We will also write $\tau_{0}=0$ to simplify notation.

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.

It is well-known that in reinforcement learning, the use of function approximation may result in instabilities in the learning process (?). 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.

We are interested in quantifying the projection of an arbitrary value distribution $Z\in\mathcal{Z}$ onto $\mathcal{Z}_{Q}$, that is

Let $Y$ be a distribution with bounded first moment and $U$ a uniform distribution over $N$ Diracs as in (7), with support $\{\theta_{1},\dots,\theta_{N}\}$. Then

In particular, if $F^{-1}$ is the inverse CDF, then $F^{-1}((\tau+\tau^{\prime})/2)$ is always a valid minimizer, and if $F^{-1}$ is continuous at $(\tau+\tau^{\prime})/2$, then $F^{-1}((\tau+\tau^{\prime})/2)$ is the unique minimizer.

These quantile midpoints will be denoted by $\hat{\tau}_{i}=\frac{\tau_{i-1}+\tau_{i}}{2}$ for $1\leq i\leq N$. Therefore, by Lemma 2, the values for $\{\theta_{1},\theta_{1},\ldots,\theta_{N}\}$ that minimize $W_{1}(Y,U)$ are given by $\theta_{i}=F_{Y}^{-1}(\hat{\tau}_{i})$. Figure 2 shows an example of the quantile projection $\Pi_{W_{1}}Z$ minimizing the $1$-Wasserstein distance to $Z$.We 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.

Let $Z_{\theta}$ be a quantile distribution, and $\hat{Z}_{m}$ the empirical distribution composed of $m$ samples from $Z$. Then for all $p\geq 1$, there exists a $Z$ such that

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 (?). These methods have been used in a variety of settings where outcomes have intrinsic randomness (?); from food expenditure as a function of household income (?), to studying value-at-risk in economic models (?).

The quantile regression loss, for quantile $\tau\in$, is an asymmetric convex loss function that penalizes overestimation errors with weight $\tau$ and underestimation errors with weight $1-\tau$. For a distribution $Z$, and a given quantile $\tau$, the value of the quantile function $F_{Z}^{-1}(\tau)$ may be characterized as the minimizer of the quantile regression loss

More generally, by Lemma 2 we have that the minimizing values of $\{\theta_{1},\ldots,\theta_{N}\}$ for $W_{1}(Z,Z_{\theta})$ are those that minimize the following objective:

In particular, this loss gives unbiased sample gradients. As a result, we can find the minimizing $\{\theta_{1},\ldots,\theta_{N}\}$ by stochastic gradient descent.

The quantile regression loss is not smooth at zero; as $u\to 0^{+}$, 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.Our quantile Huber loss is related to, but distinct from that of ? (?). This quantile regression loss acts as an asymmetric squared loss in an interval $[-\kappa,\kappa]$ around zero and reverts to a standard quantile loss outside this interval.

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

For notational simplicity we will denote $\rho^{0}_{\tau}=\rho_{\tau}$, 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 $\infty$-Wasserstein metric, i.e. the size of the largest gap between the two CDFs.

Let $\Pi_{W_{1}}$ 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 $Z_{1},Z_{2}\in\mathcal{Z}$ for an MDP with countable state and action spaces,

We therefore conclude that the combined operator $\Pi_{W_{1}}\mathcal{T}^{\pi}$ has a unique fixed point $\hat{Z}^{\pi}$, and the repeated application of this operator, or its stochastic approximation, converges to $\hat{Z}^{\pi}$. Because ${\bar{d}}_{p}\leq{\bar{d}}_{\infty}$, we conclude that convergence occurs for all $p\in[1,\infty]$. Interestingly, the contraction property does not directly hold for $p<\infty$; 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).

Recall the standard TD update for evaluating a policy $\pi$,

TD allows us to update the estimated value function with a single unbiased sample following $\pi$. Quantile regression also allows us to improve the estimate of the quantile function for some target distribution, $Y(x)$, by observing samples $y\sim Y(x)$ and minimizing Equation 8.

Furthermore, we have shown that by estimating the quantile function for well-chosen values of $\tau\in(0,1)$ 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,

where $Z_{\theta}$ is a quantile distribution as in (7), and $\theta_{i}(x)$ is the estimated value of $F^{-1}_{Z^{\pi}(x)}(\hat{\tau}_{i})$ in state $x$. It is important to note that this update is for each value of $\hat{\tau}_{i}$ and is defined for a single sample from the next state value distribution. In general it is better to draw many samples of $z^{\prime}\sim Z(x^{\prime})$ 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 ($\theta_{i}(x),\theta_{j}(x^{\prime})$). 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 (?),

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

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 (?). 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 $|\mathcal{A}|\times N$, where $N$ is a hyper-parameter giving the number of quantile targets. Second, we replace the Huber loss used by dqndqn uses gradient clipping of the squared error that makes it equivalent to a Huber loss with $\kappa=1$., $L_{\kappa}(r_{t}+\gamma\max_{a^{\prime}}Q(x_{t+1},a^{\prime})-Q(x_{t},a_{t}))$ with $\kappa=1$, with a quantile Huber loss (full loss given by Algorithm 1). Finally, we replace RMSProp (?) with Adam (?). 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 $N$, which controls with what resolution we approximate the value distribution. As we increase $N$, 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.

We begin our experimental results by demonstrating that qrtd actually learns an approximate value distribution that minimizes the $1$-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 (?), 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 $0.1$ 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 $x_{G}$, which terminates the episode and gives a reward of $1.0$. The discount factor is $\gamma=0.99$.

We compute the ground truth value distribution for optimal policy $\pi$, learned by policy iteration, at each state by performing $1K$ Monte-Carlo (MC) rollouts and recording the observed returns as an empirical distribution, shown in Figure 3(b). Next, we ran both $\textsc{td}($0$)$ and qrtd while following $\pi$ for $10K$ episodes. Each episode begins in the designated start state ($x_{S}$). Both algorithms started with a learning rate of $\alpha=0.1$. For qrtd we used $N=32$ and drop $\alpha$ by half every $2K$ episodes.

Let $Z^{\pi}_{MC}(x_{S})$ be the MC estimated distribution over returns from the start state $x_{S}$, similarly $V^{\pi}_{MC}(x_{S})$ its mean. In Figure 3 we show the approximation errors at $x_{S}$ for both algorithms with respect to the number of episodes. In (d) we evaluated, for both $\textsc{td}($0$)$ and qrtd, the squared error, $(V^{\pi}_{MC}-V(x_{S}))^{2}$, and in (e) we show the $1$-Wasserstein metric for qrtd, $W_{1}(Z^{\pi}_{MC}(x_{S}),Z(x_{S}))$, where $V(x_{S})$ and $Z(x_{S})$ are the expected returns and value distribution at state $x_{S}$ estimated by the algorithm. As expected both algorithms converge correctly in mean, and qrtd minimizes the $1$-Wasserstein distance to $Z^{\pi}_{MC}$.

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) (?). 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 (?) and prioritized replay (?). 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, $\kappa=0$ (qr-dqn-), and with a Huber quantile loss with $\kappa=1$ (qr-dqn-$1$).

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 ($\alpha=0.00005$, $\epsilon_{ADAM}=0.01/32$, and $N=200$).We swept over $\alpha$ in ($10^{-3},5\times 10^{-4},10^{-4},5\times 10^{-5},10^{-5}$); $\epsilon_{ADAM}$ in ($0.01/32,0.005/32,0.001/32$); $N$ ($10,50,100,200$) As with dqn we use a target network when computing the distributional Bellman update. We also allow $\epsilon$ to decay at the same rate as in dqn, but to a lower value of $0.01$, as is common in recent work (?; ?; ?).

Out training procedure follows that of ? (?)’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 (?).

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 $500K$ frames while recording the average return. Evaluation episodes begin with up to $30$ random no-ops (?), and the agent uses a lower exploration rate ($\epsilon=0.001$). As training progresses we keep track of the best agent performance achieved thus far.

Table 1 gives the best agent performance, at $200$ million frames trained, for qr-dqn, c51, dqn, Double dqn (?), Prioritized replay (?), and Dueling architecture (?). 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 ($10$th, $20$th, $30$th, $40$th, and $50$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 $10\%$ of games. (2) qrtd gives similar improvements to sample complexity as prioritized replay, while also improving final performance. (3) Even at $200$ million frames, there are $10\%$ of games where all algorithms reach less than $10\%$ 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 ? (?) 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 (?). 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

Appendix

For any $\omega\in$, the function $\theta\mapsto|F^{-1}(\omega)-\theta|$ is convex, and has subgradient given by

so the function $\theta\mapsto\int_{\tau}^{\tau^{\prime}}|F^{-1}(\omega)-\theta|d\omega$ is also convex, and has subgradient given by

and since $F\circ F^{-1}$ is the identity map on $$, it is clear that$\theta=F^{-1}((\tau+\tau^{\prime})/2)$satisfies Equation 14. Note that in fact any$\theta$such that$F(\theta)=(\tau+\tau^{\prime})/2$yields a subgradient of , which leads to a multitude of minimizers if$F^{-1}$is not continuous at$(\tau+\tau^{\prime})/2$. ∎

Write $Z_{\theta}=\sum_{i=1}^{N}\frac{1}{N}\delta_{\theta_{i}}$, with $\theta_{1}\leq\cdots\leq\theta_{N}$. We take $Z$ to be of the same form as $Z_{\theta}$. Specifically, consider $Z$ given by

supported on the set $\{1,\ldots,N\}$, and take $m=N$. Then clearly the unique minimizing $Z_{\theta}$ for $W_{p}(Z,Z_{\theta})$ is given by taking $Z_{\theta}=Z$. However, consider the gradient with respect to $\theta_{1}$ for the objective

In the event that the sample distribution $\hat{Z}_{N}$ has an atom at $1$, then the optimal transport plan pairs the atom of $Z_{\theta}$ at $\theta_{1}=1$ with this atom of $\hat{Z}_{N}$, and gradient with respect to $\theta_{1}$ of $W_{p}(\hat{Z}_{N},Z_{\theta})$ is . If the sample distribution $\hat{Z}_{N}$ does not contain an atom at $1$, then the left-most atom of $\hat{Z}_{N}$ is greater than $1$ (since $Z$ is supported on $\{1,\ldots,N\}$. In this case, the gradient on $\theta_{1}$ is negative. Since this happens with non-zero probability, we conclude that

We assume that instantaneous rewards given a state-action pair are deterministic; the general case is a straightforward generalization. Further, since the operator $\mathcal{T}^{\pi}$ is a $\gamma$-contraction in $\overline{d}_{\infty}$, it is sufficient to prove the claim in the case $\gamma=1$. In addition, since Wasserstein distances are invariant under translation of the support of distributions, it is sufficient to deal with the case where $r(x,a)\equiv 0$ for all $(x,a)\in\mathcal{X}\times\mathcal{A}$. 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.

The construction is illustrated in Figure 5.

Since, by Lemma 4, the $d_{\infty}$ 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:

Now note that by construction, $(\mathcal{T}^{\widetilde{\pi}}\widetilde{Z})(x^{\prime},a^{\prime})$ (respectively, $(\mathcal{T}^{\widetilde{\pi}}\widetilde{Y})(x^{\prime},a^{\prime})$) has the same distribution as $(\mathcal{T}^{\pi}Z)(x^{\prime},a^{\prime})$ (respectively, $(\mathcal{T}^{\pi}Y)(x^{\prime},a^{\prime})$), and so

Therefore, substituting this into the Inequality 15, we obtain

Taking suprema over the initial state $(x^{\prime},a^{\prime})$ then yields the result. ∎

Supporting results

Consider an MDP with countable state and action spaces. Let $Z,Y$ be value distributions such that each state-action distribution $Z(x,a)$, $Y(x,a)$ is given by a single Dirac. Consider the particular case where rewards are identically and $\gamma=1$, and let $\tau\in$. Denote by $\Pi_{\tau}$ the projection operator that maps a probability distribution onto a Dirac delta located at its $\tau$th quantile. Then

Now consider the $\tau$th quantile of each of these distributions, for $\tau\in$ arbitrary. Let $u\in I$ be such that $\theta_{u}$ is equal to this quantile of $(\mathcal{T}^{\pi}Z)(x^{\prime},a^{\prime})$, and let $v\in I$ such that $\psi_{v}$ is equal to this quantile of $(\mathcal{T}^{\pi}Y)(x^{\prime},a^{\prime})$. Now note that

is impossible, from which it will follow that

and the result then follows by taking maxima over state-action pairs $(x^{\prime},a^{\prime})$. To demonstrate the impossibility of (18), without loss of generality we take $\theta_{u}\leq\psi_{v}$.

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

and observe that we clearly have the following disjoint unions:

If (18) is to hold, then we must have $I_{\leq\theta_{u}}\cap I_{\geq\psi_{v}}=\emptyset$. Therefore, we must have $I_{\leq\theta_{u}}\subseteq I_{<\psi_{v}}$. But if this is the case, then since $\theta_{u}$ is the $\tau$th quantile of $(\mathcal{T}^{\pi}Z)(x^{\prime},a^{\prime})$, we must have

from which we conclude that the $\tau$th quantile of $(\mathcal{T}^{\pi}Y)(x^{\prime},a^{\prime})$ is less than $\psi_{v}$, a contradiction. Therefore (18) cannot hold, completing the proof. ∎

For any two probability distributions $\nu_{1},\nu_{2}$ over the real numbers, and the Wasserstein projection operator $\Pi_{W_{1}}$ that projects distributions onto support of size $n$, we have that

By the discussion surrounding Lemma 2, we have that $\Pi_{W_{1}}\nu_{k}=\sum_{i=1}^{n}\frac{1}{n}\delta_{F^{-1}_{\nu_{k}}(\frac{2i-1}{2n})}$ for $k=1,2$. Therefore, the optimal coupling between $\Pi_{W_{1}}\nu_{1}$ and $\Pi_{W_{1}}\nu_{2}$ must be given by $F^{-1}_{\nu_{1}}(\frac{2i-1}{2n})\mapsto F^{-1}_{\nu_{2}}(\frac{2i-1}{2n})$ for each $i=1,\ldots,n$. This immediately leads to the expression of the lemma. ∎

Further theoretical results

The projected Bellman operator $\Pi_{W_{1}}\mathcal{T}^{\pi}$ is in general not a non-expansion in $\overline{d}_{p}$, for $p\in[1,\infty)$.

Consider the case where the number of Dirac deltas in each distribution, $N$, is equal to $2$, and let $\gamma=1$. We consider an MDP with a single initial state, $x$, and two terminal states, $x_{1}$ and $x_{2}$. 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 $2/3$ probability of transitioning from $x$ to $x_{1}$, and $1/3$ probability of transitioning from $x$ to $x_{2}$. We take all rewards in the MDP to be identically . Further, consider two value distributions, $Z$ and $Y$, given by:

We now consider the projected backup for these two value distributions at the state $x$. We first compute the full backup:

Appealing to Lemma 2, we note that when projected these distributions onto two equally-weighted Diracs, the locations of these Diracs correspond to the 25% and 75% quantiles of the original distributions. We therefore have

Notation

Human-normalized scores are given by (?),

where $agent$, $human$ and $random$ represent the per-game raw scores for the agent, human baseline, and random agent baseline.