Implicit Quantile Networks for Distributional Reinforcement Learning

Introduction

Distributional reinforcement learning (Jaquette, 1973; Sobel, 1982; White, 1988; Morimura et al., 2010b; Bellemare et al., 2017) focuses on the intrinsic randomness of returns within the reinforcement learning (RL) framework. As the agent interacts with the environment, irreducible randomness seeps in through the stochasticity of these interactions, the approximations in the agent’s representation, and even the inherently chaotic nature of physical interaction (Yu et al., 2016). Distributional RL aims to model the distribution over returns, whose mean is the traditional value function, and to use these distributions to evaluate and optimize a policy.

Any distributional RL algorithm is characterized by two aspects: the parameterization of the return distribution, and the distance metric or loss function being optimized. Together, these choices control assumptions about the random returns and how approximations will be traded off. Categorical DQN (Bellemare et al., 2017, C51) combines a categorical distribution and the cross-entropy loss with the Cramér-minimizing projection (Rowland et al., 2018). For this, it assumes returns are bounded in a known range and trades off mean-preservation at the cost of overestimating variance.

C51 outperformed all previous improvements to DQN on a set of 57 Atari 2600 games in the Arcade Learning Environment (Bellemare et al., 2013), which we refer to as the Atari-57 benchmark. Subsequently, several papers have built upon this successful combination to achieve significant improvements to the state-of-the-art in Atari-57 (Hessel et al., 2018; Gruslys et al., 2018), and challenging continuous control tasks (Barth-Maron et al., 2018).

These algorithms are restricted to assigning probabilities to an a priori fixed, discrete set of possible returns. Dabney et al. (2018) propose an alternate pair of choices, parameterizing the distribution by a uniform mixture of Diracs whose locations are adjusted using quantile regression. Their algorithm, QR-DQN, while restricted to a discrete set of quantiles, automatically adapts return quantiles to minimize the Wasserstein distance between the Bellman updated and current return distributions. This flexibility allows QR-DQN to significantly improve on C51’s Atari-57 performance.

In this paper, we extend the approach of Dabney et al. (2018), from learning a discrete set of quantiles to learning the full quantile function, a continuous map from probabilities to returns. When combined with a base distribution, such as $U()$, this forms an implicit distribution capable of approximating any distribution over returns given sufficient network capacity. Our approach, implicit quantile networks (IQN), is best viewed as a simple distributional generalization of the DQN algorithm (Mnih et al., 2015), and provides several benefits over QR-DQN.

First, the approximation error for the distribution is no longer controlled by the number of quantiles output by the network, but by the size of the network itself, and the amount of training. Second, IQN can be used with as few, or as many, samples per update as desired, providing improved data efficiency with increasing number of samples per training update. Third, the implicit representation of the return distribution allows us to expand the class of policies to more fully take advantage of the learned distribution. Specifically, by taking the base distribution to be non-uniform, we expand the class of policies to $\epsilon$-greedy policies on arbitrary distortion risk measures (Yaari, 1987; Wang, 1996).

We begin by reviewing distributional reinforcement learning, related work, and introducing the concepts surrounding risk-sensitive RL. In subsequent sections, we introduce our proposed algorithm, IQN, and present a series of experiments using the Atari-57 benchmark, investigating the robustness and performance of IQN. Despite being a simple distributional extension to DQN, and forgoing any other improvements, IQN significantly outperforms QR-DQN and nearly matches the performance of Rainbow, which combines many orthogonal advances. In fact, in human-starts as well as in the hardest Atari games (where current RL agents still underperform human players) IQN improves over Rainbow.

Background / Related Work

We consider the standard RL setting, in which the interaction of an agent and an environment is modeled as a Markov Decision Process $(\mathcal{X},\mathcal{A},R,P,\gamma)$ (Puterman, 1994), where $\mathcal{X}$ and $\mathcal{A}$ denote the state and action spaces, $R$ the (state- and action-dependent) reward function, $P(\cdot|x,a)$ the transition kernel, and $\gamma\in(0,1)$ a discount factor. A policy $\pi(\cdot|x)$ maps a state to a distribution over actions.

To this end, Q-learning (Watkins, 1989) iteratively improves an estimate, $Q_{\theta}$, of the optimal action-value function, $Q^{*}$, by repeatedly applying the Bellman update:

The action-value function can be approximated by a parameterized function $Q_{\theta}$ (e.g. a neural network), and trained by minimizing the squared temporal difference (TD) error,

over samples $(x_{t},a_{t},r_{t},x_{t+1})$ observed while following an $\epsilon$-greedy policy over $Q_{\theta}$. This policy acts greedily with respect to $Q_{\theta}$ with probability $1-\epsilon$ and uniformly at random otherwise. DQN (Mnih et al., 2015) uses a convolutional neural network to parameterize $Q_{\theta}$ and the Q-learning algorithm to achieve human-level play on the Atari-57 benchmark.

In distributional RL, the distribution over returns (the law of $Z^{\pi}$) is considered instead of the scalar value function $Q^{\pi}$ that is its expectation. This change in perspective has yielded new insights into the dynamics of RL (Azar et al., 2012), and been a useful tool for analysis (Lattimore & Hutter, 2012). Empirically, distributional RL algorithms show improved sample complexity and final performance, as well as increased robustness to hyperparameter variation (Barth-Maron et al., 2018).

An analogous distributional Bellman equation of the form

can be derived, where $A\stackrel{{\scriptstyle D}}{{=}}B$ denotes that two random variables $A$ and $B$ have equal probability laws, and the random variables $X^{\prime}$ and $A^{\prime}$ are distributed according to $P(\cdot|x,a)$ and $\pi(\cdot|x^{\prime})$, respectively.

Morimura et al. (2010a) defined the distributional Bellman operator explicitly in terms of conditional probabilities, parameterized by the mean and scale of a Gaussian or Laplace distribution, and minimized the Kullback-Leibler (KL) divergence between the Bellman target and the current estimated return distribution. However, the distributional Bellman operator is not a contraction in the KL.

As with the scalar setting, a distributional Bellman optimality operator can be defined by

with $X^{\prime}$ distributed according to $P(\cdot|x,a)$. While the distributional Bellman operator for policy evaluation is a contraction in the $p$-Wasserstein distance (Bellemare et al., 2017), this no longer holds for the control case. Convergence to the optimal policy can still be established, but requires a more involved argument.

Bellemare et al. (2017) parameterize the return distribution as a categorical distribution over a fixed set of equidistant points and minimize the KL divergence to the projected distributional Bellman target. Their algorithm, C51, outperformed previous DQN variants on the Atari-57 benchmark. Subsequently, Hessel et al. (2018) combined C51 with enhancements such as prioritized experience replay (Schaul et al., 2016), $n$-step updates (Sutton, 1988), and the dueling architecture (Wang et al., 2016), leading to the Rainbow agent, current state-of-the-art in Atari-57.

The categorical parameterization, using the projected KL loss, has also been used in recent work to improve the critic of a policy gradient algorithm, D4PG, achieving significantly improved robustness and state-of-the-art performance across a variety of continuous control tasks (Barth-Maron et al., 2018).

2 p𝑝p-Wasserstein Metric

The $p$-Wasserstein metric, for $p\in[1,\infty]$, plays a key role in recent results in distributional RL (Bellemare et al., 2017; Dabney et al., 2018). It has also been a topic of increasing interest in generative modeling (Arjovsky et al., 2017; Bousquet et al., 2017; Tolstikhin et al., 2017), because unlike the KL divergence, the Wasserstein metric inherently trades off approximate solutions with likelihoods.

The $p$-Wasserstein distance is the $L_{p}$ metric on inverse cumulative distribution functions (c.d.f.), also known as quantile functions (Müller, 1997). For random variables $U$ and $V$ with quantile functions $F_{U}^{-1}$ and $F_{V}^{-1}$, respectively, the $p$-Wasserstein distance is given by

3 Quantile Regression for Distributional RL

Bellemare et al. (2017) showed that the distributional Bellman operator is a contraction in the $p$-Wasserstein metric, but as the proposed algorithm did not itself minimize the Wasserstein metric, this left a theory-practice gap for distributional RL. Recently, this gap was closed, in both directions. First and most relevant to this work, Dabney et al. (2018) proposed the use of quantile regression for distributional RL and showed that by choosing the quantile targets suitably the resulting projected distributional Bellman operator is a contraction in the $\infty$-Wasserstein metric. Concurrently, Rowland et al. (2018) showed the original class of categorical algorithms are a contraction in the Cramér distance, the $L_{2}$ metric on cumulative distribution functions.

By estimating the quantile function at precisely chosen points, QR-DQN minimizes the Wasserstein distance to the distributional Bellman target (Dabney et al., 2018). This estimation uses quantile regression, which has been shown to converge to the true quantile function value when minimized using stochastic approximation (Koenker, 2005).

In QR-DQN, the random return is approximated by a uniform mixture of $N$ Diracs,

with each $\theta_{i}$ assigned a fixed quantile target, $\hat{\tau}_{i}=\frac{\tau_{i-1}+\tau_{i}}{2}$ for $1\leq i\leq N$, where $\tau_{i}=i/N$. These quantile estimates are trained using the Huber (1964) quantile regression loss, with threshold $\kappa$,

At the time of this writing, QR-DQN achieves the best performance on Atari-57, human-normalized mean and median, of all agents that do not combine distributional RL, prioritized replay, and $n$-step updates (Dabney et al., 2018; Hessel et al., 2018; Gruslys et al., 2018).

4 Risk in Reinforcement Learning

Distributional RL algorithms have been theoretically justified for the Wasserstein and Cramér metrics (Bellemare et al., 2017; Rowland et al., 2018), and learning the distribution over returns, in and of itself, empirically results in significant improvements to data efficiency, final performance, and stability (Bellemare et al., 2017; Dabney et al., 2018; Gruslys et al., 2018; Barth-Maron et al., 2018). However, in each of these recent works the policy used was based entirely on the mean of the return distribution, just as in standard reinforcement learning. A natural question arises: can we expand the class of policies using information provided by the distribution over returns (i.e. to the class of risk-sensitive policies)? Furthermore, when would this larger policy class be beneficial?

Here, ‘risk’ refers to the uncertainty over possible outcomes, and risk-sensitive policies are those which depend upon more than the mean of the outcomes. At this point, it is important to highlight the difference between intrinsic uncertainty, captured by the distribution over returns, and parametric uncertainty, the uncertainty over the value estimate typically associated with Bayesian approaches such as PSRL (Osband et al., 2013) and Kalman TD (Geist & Pietquin, 2010). Distributional RL seeks to capture the former, which classic approaches to risk are built uponOne exception is the recent work (Moerland et al., 2017) towards combining both forms of uncertainty to improve exploration..

Expected utility theory states that if a decision policy is consistent with a particular set of four axioms regarding its choices then the decision policy behaves as though it is maximizing the expected value of some utility function $U$ (von Neumann & Morgenstern, 1947),

This is perhaps the most pervasive notion of risk-sensitivity. A policy maximizing a linear utility function is called risk-neutral, whereas concave or convex utility functions give rise to risk-averse or risk-seeking policies, respectively. Many previous studies on risk-sensitive RL adopt the utility function approach (Howard & Matheson, 1972; Marcus et al., 1997; Maddison et al., 2017).

A crucial axiom of expected utility is independence: given random variables $X$, $Y$ and $Z$, such that $X\succ Y$ ($X$ preferred over $Y$), any mixture between $X$ and $Z$ is preferred to the same mixture between $Y$ and $Z$ (von Neumann & Morgenstern, 1947). Stated in terms of the cumulative probability functions, $\alpha F_{X}+(1-\alpha)F_{Z}\geq\alpha F_{Y}+(1-\alpha)F_{Z},\ \forall\alpha\in$. This axiom in particular has troubled many researchers because it is consistently violated by human behavior (Tversky & Kahneman, 1992). The Allais paradox is a frequently used example of a decision problem where people violate the independence axiom of expected utility theory (Allais, 1990).

However, as Yaari (1987) showed, this axiom can be replaced by one in terms of convex combinations of outcome values, instead of mixtures of distributions. Specifically, if as before $X\succ Y$, then for any $\alpha\in$ and random variable $Z$, $\alpha F_{X}^{-1}+(1-\alpha)F_{Z}^{-1}\geq\alpha F_{Y}^{-1}+(1-\alpha)F_{Z}^{-1}$. This leads to an alternate, dual, theory of choice than that of expected utility. Under these axioms the decision policy behaves as though it is maximizing a distorted expectation, for some continuous monotonic function $h$:

Such a function $h$ is known as a distortion risk measure, as it distorts the cumulative probabilities of the random variable (Wang, 1996). That is, we have two fundamentally equivalent approaches to risk-sensitivity. Either, we choose a utility function and follow the expectation of this utility. Or, we choose a reweighting of the distribution and compute expectation under this distortion measure. Indeed, Yaari (1987) further showed that these two functions are inverses of each other. The choice between them amounts to a choice over whether the behavior should be invariant to mixing with random events or to convex combinations of outcomes.

Distortion risk measures include, as special cases, cumulative probability weighting used in cumulative prospect theory (Tversky & Kahneman, 1992), conditional value at risk (Chow & Ghavamzadeh, 2014), and many other methods (Morimura et al., 2010b). Recently Majumdar & Pavone (2017) argued for the use of distortion risk measures in robotics.

Implicit Quantile Networks

We now introduce the implicit quantile network (IQN), a deterministic parametric function trained to reparameterize samples from a base distribution, e.g. $\tau\sim U()$, to the respective quantile values of a target distribution. IQN provides an effective way to learn an implicit representation of the return distribution, yielding a powerful function approximator for a new DQN-like agent.

Let $F^{-1}_{Z}(\tau)$ be the quantile function at $\tau\in$ for the random variable $Z$. For notational simplicity we write $Z_{\tau}:=F^{-1}_{Z}(\tau)$, thus for $\tau\sim U()$ the resulting state-action return distribution sample is $Z_{\tau}(x,a)\sim Z(x,a)$.

We propose to model the state-action quantile function as a mapping from state-actions and samples from some base distribution, typically $\tau\sim U()$, to $Z_{\tau}(x,a)$, viewed as samples from the implicitly defined return distribution.

Let $\beta\colon\to$ be a distortion risk measure, with identity corresponding to risk-neutrality. Then, the distorted expectation of $Z(x,a)$ under $\beta$ is given by

Notice that the distorted expectation is equal to the expected value of $F^{-1}_{Z(x,a)}$ weighted by $\beta$, that is, $Q_{\beta}=\int_{0}^{1}F^{-1}_{Z}(\tau)d\beta(\tau)$. The immediate implication of this is that for any $\beta$, there exists a sampling distribution for $\tau$ such that the mean of $Z_{\tau}$ is equal to the distorted expectation of $Z$ under $\beta$, that is, any distorted expectation can be represented as a weighted sum over the quantiles (Dhaene et al., 2012). Denote by $\pi_{\beta}$ the risk-sensitive greedy policy

For two samples $\tau,\tau^{\prime}\sim U()$, and policy $\pi_{\beta}$, the sampled temporal difference (TD) error at step $t$ is

Implicit quantile networks differ from the approach of Dabney et al. (2018) in two ways. First, instead of approximating the quantile function at $n$ fixed values of $\tau$ we approximate it with $Z_{\tau}(x,a)\approx f(\psi(x),\phi(\tau))_{a}$ for some differentiable functions $f$, $\psi$, and $\phi$. If we ignore the distributional interpretation for a moment and view each $Z_{\tau}(x,a)$ as a separate action-value function, this highlights that implicit quantile networks are a type of universal value function approximator (UVFA) (Schaul et al., 2015). There may be additional benefits to implicit quantile networks beyond the obvious increase in representational fidelity. As with UVFAs, we might hope that training over many different $\tau$’s (goals in the case of the UVFA) leads to better generalization between values and improved sample complexity than attempting to train each separately.

As the network for $f$ is not particularly deep, we use the multiplicative form, $\psi\odot\phi$, to force interaction between the convolutional features and the sample embedding. Alternative functional forms, e.g. concatenation or a ‘residual’ function $\psi\odot(1+\phi)$, are conceivable, and $\phi(\tau)$ can be parameterized in different ways. To investigate these, we compared performance across a number of architectural variants on six Atari 2600 games (Asterix, Assault, Breakout, Ms.Pacman, QBert, Space Invaders). Full results are given in the Appendix. Despite minor variation in performance, we found the general approach to be robust to the various choices. Based upon the results we used the following function in our later experiments, for embedding dimension $n=64$:

After settling on a network architecture, we study the effect of the number of samples, $N$ and $N^{\prime}$, used in the estimate terms of Equation 3.

We hypothesized that $N$, the number of samples of $\tau\sim U()$, would affect the sample complexity of IQN, with larger values leading to faster learning, and that with $N=1$ one would potentially approach the performance of DQN. This would support the hypothesis that the improved performance of many distributional RL algorithms rests on their effect as auxiliary loss functions, which would vanish in the case of $N=1$. Furthermore, we believed that $N^{\prime}$, the number of samples of $\tau^{\prime}\sim U()$, would affect the variance of the gradient estimates much like a mini-batch size hyperparameter. Our prediction was that $N^{\prime}$ would have the greatest effect on variance of the long-term performance of the agent.

We used the same set of six games as before, with our chosen architecture, and varied $N,N^{\prime}\in\{1,8,32,64\}$. In Figure 2 we report the average human-normalized scores on the six games for each configuration. Figure 2 (left) shows the average performance over the first ten million frames, while (right) shows the average performance over the last ten million (from 190M to 200M).

As expected, we found that $N$ has a dramatic effect on early performance, shown by the continual improvement in score as the value increases. Additionally, we observed that $N^{\prime}$ affected performance very differently than expected: it had a strong effect on early performance, but minimal impact on long-term performance past $N^{\prime}=8$.

Overall, while using more samples for both distributions is generally favorable, $N=N^{\prime}=8$ appears to be sufficient to achieve the majority of improvements offered by IQN for long-term performance, with variation past this point largely insignificant. To our surprise we found that even for $N=N^{\prime}=1$, which is comparable to DQN in the number of loss components, the longer term performance is still quite strong ($\approx 3\times$ DQN).

In an informal evaluation, we did not find IQN to be sensitive to $K$, the number of samples used for the policy, and have fixed it at $K=32$ for all experiments.

Risk-Sensitive Reinforcement Learning

In this section, we explore the effects of varying the distortion risk measure, $\beta$, away from identity. This only affects the policy, $\pi_{\beta}$, used both in Equation 2 and for acting in the environment. As we have argued, evaluating under different distortion risk measures is equivalent to changing the sampling distribution for $\tau$, allowing us to achieve various forms of risk-sensitive policies. We focus on a handful of sampling distributions and their corresponding distortion measures. The first one is the cumulative probability weighting parameterization proposed in cumulative prospect theory (Tversky & Kahneman, 1992; Gonzalez & Wu, 1999):

In particular, we use the parameter value $\eta=0.71$ found by Wu & Gonzalez (1996) to most closely match human subjects. This choice is interesting as, unlike the others we consider, it is neither globally convex nor concave. For small values of $\tau$ it is locally concave and for larger values of $\tau$ it becomes locally convex. Recall that concavity corresponds to risk-averse and convexity to risk-seeking policies.

Second, we consider the distortion risk measure proposed by Wang (2000), where $\Phi$ and $\Phi^{-1}$ are taken to be the standard Normal cumulative distribution function and its inverse:

For $\eta<0$, this produces risk-averse policies and we include it due to its simple interpretation and ability to switch between risk-averse and risk-seeking distortions.

Third, we consider a simple power formula for risk-averse ($\eta<0$) or risk-seeking ($\eta>0$) policies:

Finally, we consider conditional value-at-risk (CVaR):

CVaR has been widely studied in and out of reinforcement learning (Chow & Ghavamzadeh, 2014). Its implementation as a modification to the sampling distribution of $\tau$ is particularly simple, as it changes $\tau\sim U()$ to $\tau\sim U([0,\eta])$. Another interesting sampling distribution, not included in our experiments, is denoted $\operatorname{Norm}(\eta)$ and corresponds to $\tau$ sampled by averaging $\eta$ samples from $U()$.

In Figure 3 (right) we give an example of a distribution (Neutral) and how each of these distortion measures affects the implied distribution due to changing the sampling distribution of $\tau$. $\operatorname{Norm}(3)$ and $\operatorname{CPW}(.71)$ reduce the impact of the tails of the distribution, while $\operatorname{Wang}$ and $\operatorname{CVaR}$ heavily shift the distribution mass towards the tails, creating a risk-averse or risk-seeking preference. Additionally, while CVaR entirely ignores all values corresponding to $\tau>\eta$, $\operatorname{Wang}$ gives these non-zero, but vanishingly small, probability.

By using these sampling distributions we can induce various risk-sensitive policies in IQN. We evaluate these on the same set of six Atari 2600 games previously used. Our algorithm simply changes the policy to maximize the distorted expectations instead of the usual sample mean. Figure 3 (left) shows our results in this experiment, with average scores reported under the usual, risk-neutral, evaluation criterion.

Intuitively, we expected to see a qualitative effect from risk-sensitive training, e.g. strengthened exploration from a risk-seeking objective. Although we did see qualitative differences, these did not always match our expectations. For two of the games, Asterix and Assault, there is a very significant advantage to the risk-averse policies. Although $\operatorname{CPW}$ tends to perform almost identically to the standard risk-neutral policy, and the risk-seeking $\operatorname{Wang}(1.5)$ performs as well or worse than risk-neutral, we find that both risk-averse policies improve performance over standard IQN. However, we also observe that the more risk-averse of the two, $\operatorname{CVaR}(0.1)$, suffers some loss in performance on two other games (QBert and Space Invaders).

Additionally, we note that the risk-seeking policy significantly underperforms the risk-neutral policy on three of the six games. It remains an open question as to exactly why we see improved performance for risk-averse policies. There are many possible explanations for this phenomenon, e.g. that risk-aversion encodes a heuristic to stay alive longer, which in many games is correlated with increased rewards.

Full Atari-57 Results

Finally, we evaluate IQN on the full Atari-57 benchmark, comparing with the state-of-the-art performance of Rainbow, a distributional RL agent that combines several advances in deep RL (Hessel et al., 2018), the closely related algorithm QR-DQN (Dabney et al., 2018), prioritized experience replay DQN (Schaul et al., 2016), and the original DQN agent (Mnih et al., 2015). Note that in this section we use the risk-neutral variant of the IQN, that is, the policy of the IQN agent is the regular $\epsilon$-greedy policy with respect to the mean of the state-action return distribution.

It is important to remember that Rainbow builds upon the distributional RL algorithm C51 (Bellemare et al., 2017), but also includes prioritized experience replay (Schaul et al., 2016), Double DQN (van Hasselt et al., 2016), Dueling Network architecture (Wang et al., 2016), Noisy Networks (Fortunato et al., 2017), and multi-step updates (Sutton, 1988). In particular, besides the distributional update, $n$-step updates and prioritized experience replay were found to have significant impact on the performance of Rainbow. Our other competitive baseline is QR-DQN, which is currently state-of-the-art for agents that do not combine distributional updates, $n$-step updates, and prioritized replay.

Thus, between QR-DQN and the much more complex Rainbow we compare to the two most closely related, and best performing, agents in published work. In particular, we would expect that IQN would benefit from the additional enhancements in Rainbow, just as Rainbow improved significantly over C51.

Figure 4 shows the mean (left) and median (right) human-normalized scores during training over the Atari-57 benchmark. IQN dramatically improves over QR-DQN, which itself improves on many previously published results. At 100 million frames IQN has reached the same level of performance as QR-DQN at 200 million frames. Table 1 gives a comparison between the same methods in terms of their best, human-normalized, scores per game under the 30 random no-op start condition. These are averages over the given number of seeds. Additionally, using human-starts, IQN achieves $162\%$ median human-normalized score, whereas Rainbow reaches $153\%$ (Hessel et al., 2018), see Table 2.

Finally, we took a closer look at the games in which each algorithm continues to underperform humans, and computed, on average, how far below human-level they performDetails of how this is computed can be found in the Appendix.. We refer to this value as the human-gapThanks to Joseph Modayil for proposing this metric. metric and give results in Table 1. Interestingly, C51 outperforms QR-DQN in this metric, and IQN outperforms all others. This shows that the remaining gap between Rainbow and IQN is entirely from games on which both algorithms are already super-human. The games where the most progress in RL is needed happen to be the games where IQN shows the greatest improvement over QR-DQN and Rainbow.

Discussion and Conclusions

We have proposed a generalization of recent work based around using quantile regression to learn the distribution over returns of the current policy. Our generalization leads to a simple change to the DQN agent to enable distributional RL, the natural integration of risk-sensitive policies, and significantly improved performance over existing methods. The IQN algorithm provides, for the first time, a fully integrated distributional RL agent without prior assumptions on the parameterization of the return distribution.

IQN can be trained with as little as a single sample from each state-action value distribution, or as many as computational limits allow to improve the algorithm’s data efficiency. Furthermore, IQN allows us to expand the class of control policies to a large class of risk-sensitive policies connected to distortion risk measures. Finally, we show substantial gains on the Atari-57 benchmark over QR-DQN, and even halving the distance between QR-DQN and Rainbow.

Despite the significant empirical successes in this paper there are many areas in need of additional theoretical analysis. We highlight a few particularly relevant open questions we were unable to address in the present work. First, sample-based convergence results have been recently shown for a class of categorical distributional RL algorithms (Rowland et al., 2018). Could existing sample-based RL convergence results be extended to the QR-based algorithms?

Second, can the contraction mapping results for a fixed grid of quantiles given by Dabney et al. (2018) be extended to the more general class of approximate quantile functions studied in this work? Finally, and particularly salient to our experiments with distortion risk measures, theoretical guarantees for risk-sensitive RL have been building over recent years, but have been largely limited to special cases and restricted classes of risk-sensitive policies. Can the convergence of the distribution of returns under the Bellman operator be leveraged to show convergence to a fixed-point in distorted expectations? In particular, can the control results of Bellemare et al. (2017) be expanded to cover some class of risk-sensitive policies?

There remain many intriguing directions for future research into distributional RL, even on purely empirical fronts. Hessel et al. (2018) recently showed that distributional RL agents can be significantly improved, when combined with other techniques. Creating a Rainbow-IQN agent could yield even greater improvements on Atari-57. We also recall the surprisingly rich return distributions found by Barth-Maron et al. (2018), and hypothesize that the continuous control setting may be a particularly fruitful area for the application of distributional RL in general, and IQN in particular.

References

Appendix

For the embedding $\phi$, we considered a number of variants: a learned linear embedding, a learned MLP embedding with a single hidden layer of size $n$, and a learned linear function of $n$ cosine basis functions of the form $\cos(\pi i\tau),i=1,\dots,n$. Each of those was followed by either a ReLU or sigmoid nonlinearity.

For the merging function $m$, the simplest choice would be a simple vector concatenation of $\psi(x)$ and $\phi(\tau)$. Note however, that the MLP $f$ which takes in the output of $m$ and outputs the action-value quantiles, only has a single hidden layer in the DQN network. Therefore, to force a sufficiently early interaction between the two representations, we also considered a multiplicative function $m(\psi,\phi)=\psi\odot\phi$, where $\odot$ denotes the element-wise (Hadamard) product of two vectors, as well as a ‘residual’ function $m(\psi,\phi)=\psi\odot(1+\phi)$.

Early experiments showed that a simple linear embedding of $\tau$ was insufficient to achieve good performance, and the residual version of $m$ didn’t show any marked difference to the multiplicative variant, so we do not include results for these here. For the other configurations, Figure 5 shows pairwise comparisons between 1) a cosine basis function embedding and a completely learned MLP embedding, 2) an embedding size (hidden layer size or number of cosine basis elements) 32 and 64, 3) ReLU and sigmoid nonlinearity following the embedding, and 4) concatenation and a multiplicative interaction between $\psi(x)$ and $\phi(\tau)$.

Each comparison ‘violin plot’ can be understood as a marginalization over the other variants of the architecture, with the human-normalized performance at the end of training, averaged across six Atari 2600 games, on the y-axis. Each white dot corresponds to a configuration (each represented by two seeds), the black dots show the position of our preferred configuration. The width of the colored regions corresponds to a kernel density estimate of the number of configurations at each performance level.

Our final choice is a multiplicative interaction with a linear function of a cosine embedding, with $n=64$ and a ReLU nonlinearity (see Equation 4), as this configuration yielded the highest performance consistently over multiple seeds. Also noteworthy is the overall robustness of the approach to these variations: most of the configurations consistently outperform the QR-DQN baseline shown as a grey horizontal line for comparison.

We give pseudo-code for the IQN loss in Algorithm 1. All other hyperparameters for this agent correspond to the ones used by Dabney et al. (2018). In particular, the Bellman target is computed using a target network. Notice that IQN will generally be more computationally expensive per-sample than QR-DQN. However, in practice IQN requires many fewer samples per update than QR-DQN so that the actual running times are comparable.

Evaluation

The human-normalized scores reported in this paper are given by the formula (van Hasselt et al., 2016; Dabney et al., 2018)

where $agent$, $human$ and $random$ are the per-game raw scores (undiscounted returns) for the given agent, a reference human player, and random agent baseline (Mnih et al., 2015).

The ‘human-gap’ metric referred to at the end of Section 5 builds on the human-normalized score, but emphasizes the remaining improvement for the agent to reach super-human performance. It is given by $gap=\max(1-score,0)$, with a value of $1$ corresponding to random play, and a value of corresponding to super-human level of performance. To avoid degeneracies in the case of $human<random$, the quantity is being clipped above at $1$.