A Distributional Perspective on Reinforcement Learning
Marc G. Bellemare * 1 Will Dabney * 1 Remi´ Munos 1
Abstract ment learning. Specifically, the main object of our study is In this paper we argue for the fundamental impor- the random return Z whose expectation is the value Q. This tance of the value distribution: the distribution random return is also described by a recursive equation, but of the random return received by a reinforcement one of a distributional nature: learning agent. This is in contrast to the com- D Z(x, a) = R(x, a) + γZ(X0,A0). mon approach to reinforcement learning which models the expectation of this return, or value. The distributional Bellman equation states that the distribu- Although there is an established body of liter- tion of Z is characterized by the interaction of three random ature studying the value distribution, thus far it variables: the reward R, the next state-action (X0,A0), and has always been used for a specific purpose such its random return Z(X0,A0). By analogy with the well- as implementing risk-aware behaviour. We begin known case, we call this quantity the value distribution. with theoretical results in both the policy eval- uation and control settings, exposing a signifi- Although the distributional perspective is almost as old cant distributional instability in the latter. We as Bellman’s equation itself (Jaquette, 1973; Sobel, 1982; then use the distributional perspective to design White, 1988), in reinforcement learning it has thus far been a new algorithm which applies Bellman’s equa- subordinated to specific purposes: to model parametric un- tion to the learning of approximate value distri- certainty (Dearden et al., 1998), to design risk-sensitive al- butions. We evaluate our algorithm using the gorithms (Morimura et al., 2010b;a), or for theoretical anal- suite of games from the Arcade Learning En- ysis (Azar et al., 2012; Lattimore & Hutter, 2012). By con- vironment. We obtain both state-of-the-art re- trast, we believe the value distribution has a central role to sults and anecdotal evidence demonstrating the play in reinforcement learning. importance of the value distribution in approxi- Contraction of the policy evaluation Bellman operator. mate reinforcement learning. Finally, we com- Basing ourselves on results byR osler¨ (1992) we show that, bine theoretical and empirical evidence to high- for a fixed policy, the Bellman operator over value distribu- light the ways in which the value distribution im- tions is a contraction in a maximal form of the Wasserstein pacts learning in the approximate setting. (also called Kantorovich or Mallows) metric. Our partic- ular choice of metric matters: the same operator is not a contraction in total variation, Kullback-Leibler divergence, 1. Introduction or Kolmogorov distance. One of the major tenets of reinforcement learning states Instability in the control setting. We will demonstrate an that, when not otherwise constrained in its behaviour, an instability in the distributional version of Bellman’s opti- arXiv:1707.06887v1 [cs.LG] 21 Jul 2017 agent should aim to maximize its expected utility Q, or mality equation, in contrast to the policy evaluation case. value (Sutton & Barto, 1998). Bellman’s equation succintly Specifically, although the optimality operator is a contrac- describes this value in terms of the expected reward and ex- tion in expected value (matching the usual optimality re- 0 0 pected outcome of the random transition (x, a) (X ,A ): sult), it is not a contraction in any metric over distributions. → 0 0 These results provide evidence in favour of learning algo- Q(x, a) = E R(x, a) + γ E Q(X ,A ). rithms that model the effects of nonstationary policies. In this paper, we aim to go beyond the notion of value and argue in favour of a distributional perspective on reinforce- Better approximations. From an algorithmic standpoint, there are many benefits to learning an approximate distribu- *Equal contribution 1DeepMind, London, UK. Correspon- tion rather than its approximate expectation. The distribu- dence to: Marc G. Bellemare
π π 0 0 Q (x, a) = E R(x, a) + γ E Q (x , a ). P,π 3.1. Distributional Equations In reinforcement learning we are typically interested in act- It will sometimes be convenient to make use of the proba- ing so as to maximize the return. The most common ap- bility space (Ω, , Pr). The reader unfamiliar with mea- F A Distributional Perspective on Reinforcement Learning sure theory may think of Ω as the space of all possible of U, V . The metric dp has the following properties: outcomes of an experiment (Billingsley, 1995). We will X write u p to denote the Lp norm of a vector u R for dp(aU, aV ) a dp(U, V ) (P1) k k X∈ ×A ≤ | | 1 p ; the same applies to vectors in R . The dp(A + U, A + V ) dp(U, V ) (P2) ≤ ≤ ∞ X X ×A ≤ Lp norm of a random vector U :Ω R (or R ) is → dp(AU, AV ) A pdp(U, V ). (P3) p1/p ≤ k k then U p := E U(ω) , and for p = we have k k k kp ∞ U ∞ = ess sup U(ω) ∞ (we will omit the dependency We will need the following additional property, which onk kω Ω wheneverk unambiguous).k We will denote the makes no independence assumptions on its variables. Its ∈ c.d.f. of a random variable U by FU (y) := Pr U y , proof, and that of later results, is given in the appendix. −1 { ≤ } and its inverse c.d.f. by F (q) := inf y : FU (y) q . U { ≥ } Lemma 1 (Partition lemma). Let A1,A2,... be a set of D random variables describing a partition of Ω, i.e. A (ω) A distributional equation U := V indicates that the ran- i 0, 1 and for any ω there is exactly one A with A (ω) =∈ dom variable U is distributed according to the same law i i {1. Let} U, V be two random variables. Then as V . Without loss of generality, the reader can understand the two sides of a distributional equation as relating the dis- X dp U, V dp(AiU, AiV ). tributions of two independent random variables. Distribu- ≤ i tional equations have been used in reinforcement learning by Engel et al.(2005); Morimura et al.(2010a) among oth- Let denote the space of value distributions with bounded Z ers, and in operations research by White(1988). moments. For two value distributions Z1,Z2 we will make use of a maximal form of the Wasserstein∈ metric: Z 3.2. The Wasserstein Metric d¯p(Z1,Z2) := sup dp(Z1(x, a),Z2(x, a)). The main tool for our analysis is the Wasserstein metric dp x,a between cumulative distribution functions (see e.g. Bickel ¯ & Freedman, 1981, where it is called the Mallows metric). We will use dp to establish the convergence of the distribu- For F , G two c.d.fs over the reals, it is defined as tional Bellman operators.
Lemma 2. d¯p is a metric over value distributions. dp(F,G) := inf U V p, U,V k − k 3.3. Policy Evaluation where the infimum is taken over all pairs of random vari- In the policy evaluation setting (Sutton & Barto, 1998) we ables with respective cumulative distributions (U, V ) F are interested in the value function V π associated with a and . The infimum is attained by the inverse c.d.f. trans- G given policy π. The analogue here is the value distribu- form of a random variable uniformly distributed on [0, 1]: π π U tion Z . In this section we characterize Z and study the π −1 −1 behaviour of the policy evaluation operator . We em- d (F,G) = F ( ) G ( ) . π T p k U − U kp phasize that Z describes the intrinsic randomness of the agent’s interactions with its environment, rather than some For p < this is more explicitly written as ∞ measure of uncertainty about the environment itself. Z 1 1/p We view the reward function as a random vector R , −1 −1 p π ∈ Z dp(F,G) = F (u) G (u) du . and define the transition operator P : 0 − Z → Z D P πZ(x, a) := Z(X0,A0) (4) Given two random variables U, V with c.d.fs FU , FV , we 0 0 0 will write dp(U, V ) := dp(FU ,FV ). We will find it conve- X P ( x, a),A π( X ), nient to conflate the random variables under consideration ∼ · | ∼ · | with their versions under the inf, writing where we use capital letters to emphasize the random na- ture of the next state-action pair (X0,A0). We define the π dp(U, V ) = inf U V p. distributional Bellman operator : as U,V k − k T Z → Z D πZ(x, a) := R(x, a) + γP πZ(x, a). (5) whenever unambiguous; we believe the greater legibility T justifies the technical inaccuracy. Finally, we extend this While π bears a surface resemblance to the usual Bell- metric to vectors of random variables, such as value distri- man operatorT (2), it is fundamentally different. In particu- butions, using the corresponding L norm. p lar, three sources of randomness define the compound dis- Consider a scalar a and a random variable A independent tribution πZ: T A Distributional Perspective on Reinforcement Learning
a) The randomness in the reward R, a difficulty arises: in general there are many optimal value b) The randomness in the transition P π, and distributions. c) The next-state value distribution Z(X0,A0). In this section we show that the distributional analogue of the Bellman optimality operator converges, in a weak In particular, we make the usual assumption that these three sense, to the set of optimal value distributions. However, quantities are independent. In this section we will show this operator is not a contraction in any metric between dis- that (5) is a contraction mapping whose unique fixed point tributions, and is in general much more temperamental than π is the random return Z . the policy evaluation operators. We believe the conver- gence issues we outline here are a symptom of the inherent ¯ 3.3.1. CONTRACTIONIN dp instability of greedy updates, as highlighted by e.g. Tsitsik- Consider the process Z := πZ , starting with some lis(2002) and most recently Harutyunyan et al.(2016). k+1 T k Z0 . We may expect the limiting expectation of Zk Let Π∗ be the set of optimal policies. We begin by charac- ∈ Z π { } to converge exponentially quickly, as usual, to Q . As we terizing what we mean by an optimal value distribution. now show, the process converges in a stronger sense: π T Definition 1 (Optimal value distribution). An optimal is a contraction in d¯p, which implies that all moments also value distribution is the v.d. of an optimal policy. The set converge exponentially quickly. ∗ of optimal value distributions is ∗ := Zπ : π∗ Π∗ . Lemma 3. π : is a γ-contraction in d¯ . Z { ∈ } T Z → Z p We emphasize that not all value distributions with expecta- Using Lemma3, we conclude using Banach’s fixed point ∗ π tion Q are optimal: they must match the full distribution theorem that has a unique fixed point. By inspection, of the return under some optimal policy. this fixed pointT must be Zπ as defined in (1). As we assume Definition 2. A greedy policy π for Z maximizes the all moments are bounded, this is sufficient to conclude that ∈ Z the sequence Z converges to Zπ in d¯ for 1 p . expectation of Z. The set of greedy policies for Z is { k} p ≤ ≤ ∞ To conclude, we remark that not all distributional metrics X 0 Z := π : π(a x) E Z(x, a) = max E Z(x, a ) . are equal; for example, Chung & Sobel(1987) have shown G { a | a0∈A } that π is not a contraction in total variation distance. Sim- ilar resultsT can be derived for the Kullback-Leibler diver- Recall that the expected Bellman optimality operator is gence and the Kolmogorov distance. T 0 0 Q(x, a) = E R(x, a) + γ EP max Q(x , a ). (6) T a0∈A 3.3.2. CONTRACTIONIN CENTERED MOMENTS The maximization at x0 corresponds to some greedy policy. Observe that d2(U, V ) (and more generally, dp) relates to a coupling C(ω) := U(ω) V (ω), in the sense that Although this policy is implicit in (6), we cannot ignore it − in the distributional setting. We will call a distributional 2 2 2 Bellman optimality operator any operator which imple- d (U, V ) E[(U V ) ] = V C + E C . 2 ≤ − ments a greedy selection rule, i.e. T As a result, we cannot directly use d to bound the variance 2 π π π π Z = Z for some π . difference V( Z(x, a)) V(Z (x, a)) . However, T T ∈ GZ is in fact a| contractionT in variance− (Sobel|, 1982, see alsoT appendix). In general, π is not a contraction in the pth As in the policy evaluation setting, we are interested in the T behaviour of the iterates Z := Z , Z . Our first centered moment, p > 2, but the centered moments of the k+1 T k 0 ∈ Z result is to assert that E Zk behaves as expected. iterates Zk still converge exponentially quickly to those of Zπ; the{ proof} extends the result ofR osler¨ (1992). Lemma 4. Let Z ,Z . Then 1 2 ∈ Z
3.4. Control E Z1 E Z2 ∞ γ E Z1 E Z2 ∞ , k T − T k ≤ k − k ∗ Thus far we have considered a fixed policy π, and studied and in particular E Zk Q exponentially quickly. the behaviour of its associated operator π. We now set → T out to understand the distributional operators of the control By inspecting Lemma4, we might expect that Zk con- ∗ setting – where we seek a policy π that maximizes value verges quickly in d¯p to some fixed point in . Unfor- – and the corresponding notion of an optimal value distri- tunately, convergence is neither quick nor assuredZ to reach bution. As with the optimal value function, this notion is a fixed point. In fact, the best we can hope for is pointwise intimately tied to that of an optimal policy. However, while convergence, not even to the set ∗ but to the larger set of all optimal policies attain the same value Q∗, in our case nonstationary optimal value distributionsZ . A Distributional Perspective on Reinforcement Learning
Definition 3. A nonstationary optimal value distribution R = 0 R = ! ± 1 x x , a x , a Z∗∗ is the value distribution corresponding to a sequence 1 2 1 2 2 ∗∗ ∗ of optimal policies. The set of n.o.v.d. is . a1 a2 Z 1 0 1 ± ± Z x Theorem 1 (Convergence in the control setting). Let be 2 Z 1 0 1 measurable and suppose that is finite. Then X ± − ± A Z 0 0 1 x T ± ∗∗ 1 lim inf dp(Zk(x, a),Z (x, a)) = 0 x, a. k→∞ Z∗∗∈Z∗∗ ∀
∗∗ Figure 2. Undiscounted two-state MDP for which the optimality If is finite, then Zk converges to uniformly. Further- operator T is not a contraction, with example. The entries that X Z ∗ ∗ ∗ more, if there is a total ordering on Π , such that for any contribute to d¯1(Z,Z ) and d¯1(T Z,Z ) are highlighted. Z∗ ∗, ≺ ∈ Z ∗ π ∗ 0 0 ∗ ∗ Z = Z with π Z∗ , π π π Z∗ π . Proposition 3. That has a fixed point Z = Z is T T ∈ G ≺ ∀ ∈ G \{ } T T ∗ insufficient to guarantee the convergence of Zk to . Then has a unique fixed point Z∗ ∗. { } Z T ∈ Z Theorem1 paints a rather bleak picture of the control set- Comparing Theorem1 to Lemma4 reveals a significant ting. It remains to be seen whether the dynamical eccen- difference between the distributional framework and the tricies highlighted here actually arise in practice. One open usual setting of expected return. While the mean of Zk question is whether theoretically more stable behaviour can converges exponentially quickly to Q∗, its distribution need be derived using stochastic policies, for example from con- not be as well-behaved! To emphasize this difference, we servative policy iteration (Kakade & Langford, 2002). now provide a number of negative results concerning . T Proposition 1. The operator is not a contraction. 4. Approximate Distributional Learning T Consider the following example (Figure2, left). There are In this section we propose an algorithm based on the dis- two states, x1 and x2; a unique transition from x1 to x2; tributional Bellman optimality operator. In particular, this from x2, action a1 yields no reward, while the optimal ac- will require choosing an approximating distribution. Al- tion a2 yields 1 + or 1 + with equal probability. Both though the Gaussian case has previously been considered actions are terminal. There− is a unique optimal policy and (Morimura et al., 2010a; Tamar et al., 2016), to the best of therefore a unique fixed point Z∗. Now consider Z as given our knowledge we are the first to use a rich class of para- in Figure2 (right), and its distance to Z∗: metric distributions.
¯ ∗ ∗ d1(Z,Z ) = d1(Z(x2, a2),Z (x2, a2)) = 2, 4.1. Parametric Distribution where we made use of the fact that Z = Z∗ everywhere We will model the value distribution using a discrete distri- except at (x , a ). When we apply to Z, however, the bution parametrized by N N and VMIN,VMAX R, and 2 2 T ∈ ∈ greedy action a1 is selected and Z(x1) = Z(x2, a1). But whose support is the set of atoms zi = VMIN + i z : 0 T VMAX −VMIN { 4 ≤ i < N , z := N−1 . In a sense, these atoms are the ∗ ∗ } 4 d1( Z, Z ) = d1( Z(x1),Z (x1)) “canonical returns” of our distribution. The atom probabil- T T T ities are given by a parametric model N = 1 1 + 1 1 + > 2 θ : R 2 | − | 2 | | X × A → eθi(x,a) for a sufficiently small . This shows that the undiscounted Zθ(x, a) = zi w.p. pi(x, a) := . ∗ ∗ P eθj (x,a) update is not a nonexpansion: d¯1( Z, Z ) > d¯1(Z,Z ). j With γ < 1, the same proof showsT itT is not a contraction. The discrete distribution has the advantages of being highly Using a more technically involved argument, we can extend expressive and computationally friendly (see e.g. Van den this result to any metric which separates Z and Z. T Oord et al., 2016). Proposition 2. Not all optimality operators have a fixed point Z∗ = Z∗. T 4.2. Projected Bellman Update To see this, consider the same example, now with = 0, Using a discrete distribution poses a problem: the Bell- and a greedy operator which breaks ties by picking a man update Z and our parametrization Z almost al- T 2 T θ θ if Z(x1) = 0, and a1 otherwise. Then the sequence ways have disjoint supports. From the analysis of Section ∗ 2 ∗ ∗ Z (x1), ( ) Z (x1),... alternates between Z (x2, a1) 3 it would seem natural to minimize the Wasserstein met- Tand Z∗(x ,T a ). ric (viewed as a loss) between Z and Z , which is also 2 2 T θ θ A Distributional Perspective on Reinforcement Learning conveniently robust to discrepancies in support. However, cade Learning Environment (ALE; Bellemare et al., 2013). a second issue prevents this: in practice we are typically While the ALE is deterministic, stochasticity does occur in restricted to learning from sample transitions, which is not a number of guises: 1) from state aliasing, 2) learning from possible under the Wasserstein loss (see Prop.5 and toy a nonstationary policy, and 3) from approximation errors. results in the appendix). We used five training games (Fig3) and 52 testing games. Instead, we project the sample Bellman update ˆZ onto For our study, we use the DQN architecture (Mnih et al., T θ the support of Zθ (Figure1, Algorithm1), effectively re- 2015), but output the atom probabilities pi(x, a) instead ducing the Bellman update to multiclass classification. Let of action-values, and chose V = V = 10 from π be the greedy policy w.r.t. Z . Given a sample transi- MAX MIN E θ preliminary experiments over the training− games. We call 0 ˆ tion (x, a, r, x ), we compute the Bellman update zj := the resulting architecture Categorical DQN. We replace the r + γz for each atom z , then distribute its probabilityT j j squared loss (r + γQ(x0, π(x0)) Q(x, a))2 by (θ) p (x0, π(x0)) to the immediate neighbours of ˆz . The ith − Lx,a j T j and train the network to minimize this loss.2 As in DQN, component of the projected update Φ ˆZ (x, a) is T θ we use a simple -greedy policy over the expected action- 1 N−1 " VMAX # values; we leave as future work the many ways in which an |[Tˆzj ] − zi| ˆ X VMIN 0 0 (ΦT Zθ(x, a))i = 1 − pj (x , π(x )), agent could select actions on the basis of the full distribu- 4z j=0 0 tion. The rest of our training regime matches Mnih et al.’s, (7) ˜ b 1 including the use of a target network for θ. where [ ]a bounds its argument in the range [a, b]. As is usual, we· view the next-state distribution as parametrized Figure4 illustrates the typical value distributions we ob- by a fixed parameter θ˜. The sample loss x,a(θ) is the served in our experiments. In this example, three actions cross-entropy term of the KL divergence L (those including the button press) lead to the agent releas- ing its laser too early and eventually losing the game. The ˆ DKL(Φ Z˜(x, a) Z (x, a)), T θ k θ corresponding distributions reflect this: they assign a sig- which is readily minimized e.g. using gradient descent. We nificant probability to 0 (the terminal value). The safe call this choice of distribution and loss the categorical al- actions have similar distributions (LEFT, which tracks the gorithm. When N = 2, a simple one-parameter alternative invaders’ movement, is slightly favoured). This example ˆ ˆ 1 helps explain why our approach is so successful: the dis- is Φ Zθ(x, a) := [E[ Zθ(x, a)] VMIN)/ z]0; we call this theT Bernoulli algorithmT . We note− that, while4 these al- tributional update keeps separated the low-value, “losing” gorithms appear unrelated to the Wasserstein metric, recent event from the high-value, “survival” event, rather than av- 3 work (Bellemare et al., 2017) hints at a deeper connection. erage them into one (unrealizable) expectation. One surprising fact is that the distributions are not concen- Algorithm 1 Categorical Algorithm trated on one or two values, in spite of the ALE’s determin- input A transition xt, at, rt, xt+1, γt [0, 1] ism, but are often close to Gaussians. We believe this is due P ∈ Q(xt+1, a) := i zipi(xt+1, a) to our discretizing the diffusion process induced by γ. a∗ arg max Q(x , a) ← a t+1 mi = 0, i 0,...,N 1 5.1. Varying the Number of Atoms for j 0,...,N∈ 1 do− ∈ − We began by studying our algorithm’s performance on the # Compute the projection of ˆzj onto the support zi VMAX T { } training games in relation to the number of atoms (Figure ˆzj [rt + γtzj] T ← VMIN 3). For this experiment, we set = 0.05. From the data, it b ( ˆz V )/∆z # b [0,N 1] j ← T j − MIN j ∈ − is clear that using too few atoms can lead to poor behaviour, l bj , u bj and that more always increases performance; this is not im- ← b c ← d e ˆ # Distribute probability of zj mediately obvious as we may have expected to saturate the ∗ T ml ml + pj(xt+1, a )(u bj) network capacity. The difference in performance between m ← m + p (x , a∗)(b− l) u ← u j t+1 j − the 51-atom version and DQN is particularly striking: the end for latter is outperformed in all five games, and in SEAQUEST output P m log p (x , a ) # Cross-entropy loss − i i i t t we attain state-of-the-art performance. As an additional point of the comparison, the single-parameter Bernoulli al- gorithm performs better than DQN in 3 games out of 5, and 5. Evaluation on Atari 2600 Games is most notably more robust in ASTERIX.
To understand the approach in a complex setting, we ap- 2For N = 51, our TensorFlow implementation trains at plied the categorical algorithm to games from the Ar- roughly 75% of DQN’s speed. 3Video: http://youtu.be/yFBwyPuO2Vg. 1Algorithm1 computes this projection in time linear in N. A Distributional Perspective on Reinforcement Learning
ASTERIX
BREAKOUT PONG
SEAQUEST
Average Score Categorical DQN Dueling Arch. 5 returns 21 returns 11 r e t u r n s 51 returns Q*BERT Bernoulli DQN Training Frames (millions)
Figure 3. Categorical DQN: Varying number of atoms in the discrete distribution. Scores are moving averages over 5 million frames.
evaluate our agent’s performance with = 0.001. Laser Left+Laser We compare our algorithm to DQN ( = 0.01), Double Right+Laser DQN (van Hasselt et al., 2016), the Dueling architecture Right (Wang et al., 2016), and Prioritized Replay (Schaul et al., Probability Left 2016), comparing the best evaluation score achieved during Noop training. We see that C51 significantly outperforms these Return other algorithms (Figures6 and7). In fact, C51 surpasses the current state-of-the-art by a large margin in a number of EAQUEST Figure 4. Learned value distribution during an episode of SPACE games, most notably S . One particularly striking INVADERS. Different actions are shaded different colours. Re- fact is the algorithm’s good performance on sparse reward turns below 0 (which do not occur in SPACE INVADERS) are not games, for example VENTURE and PRIVATE EYE. This shown here as the agent assigns virtually no probability to them. suggests that value distributions are better able to propa- gate rarely occurring events. Full results are provided in the appendix. One interesting outcome of this experiment was to find We also include in the appendix (Figure 12) a compari- out that our method does pick up on stochasticity. PONG son, averaged over 3 seeds, showing the number of games exhibits intrinsic randomness: the exact timing of the re- in which C51’s training performance outperforms fully- ward depends on internal registers and is truly unobserv- trained DQN and human players. These results continue able. We see this clearly reflected in the agent’s prediction to show dramatic improvements, and are more representa- (Figure5): over five consecutive frames, the value distribu- tive of an agent’s average performance. Within 50 million tion shows two modes indicating the agent’s belief that it frames, C51 has outperformed a fully trained DQN agent has yet to receive a reward. Interestingly, since the agent’s on 45 out of 57 games. This suggests that the full 200 mil- state does not include past rewards, it cannot even extin- lion training frames, and its ensuing computational cost, guish the prediction after receiving the reward, explaining are unnecessary for evaluating reinforcement learning al- the relative proportions of the modes. gorithms within the ALE. The most recent version of the ALE contains a stochastic 5.2. State-of-the-Art Results execution mechanism designed to ward against trajectory The performance of the 51-atom agent (from here onwards, overfitting.Specifically, on each frame the environment re- C51) on the training games, presented in the last section, is jects the agent’s selected action with probability p = 0.25. particularly remarkable given that it involved none of the Although DQN is mostly robust to stochastic execution, other algorithmic ideas present in state-of-the-art agents. there are a few games in which its performance is reduced. We next asked whether incorporating the most common On a score scale normalized with respect to the random hyperparameter choice, namely a smaller training , could and DQN agents, C51 obtains mean and median score im- lead to even better results. Specifically, we set = 0.01 provements of 126% and 21.5% respectively, confirming (instead of 0.05); furthermore, every 1 million frames, we the benefits of C51 beyond the deterministic setting. A Distributional Perspective on Reinforcement Learning
Figure 5. Intrinsic stochasticity in PONG.
Mean Median > H.B. > DQN butions, similar to conservative policy iteration (Kakade & DQN 228% 79% 24 0 Langford, 2002). While the chattering persists, it is inte- DDQN 307% 118% 33 43 grated to the approximate solution. DUEL. 373% 151% 37 50 State aliasing. Even in a deterministic environment, state PRIOR. 434% 124% 39 48 aliasing may result in effective stochasticity. McCallum PR.DUEL. 592% 172% 39 44 (1995), for example, showed the importance of coupling C51 701% 178% 40 50 representation learning with policy learning in partially ob- † UNREAL 880% 250% - - servable domains. We saw an example of state aliasing in PONG, where the agent could not exactly predict the re- Figure 6. Mean and median scores across 57 Atari games, mea- sured as percentages of human baseline (H.B., Nair et al., 2015). ward timing. Again, by explicitly modelling the resulting distribution we provide a more stable learning target. A richer set of predictions. A recurring theme in artificial intelligence is the idea of an agent learning from a mul- titude of predictions (Caruana 1997; Utgoff & Stracuzzi 2002; Sutton et al. 2011; Jaderberg et al. 2017). The dis- tributional approach naturally provides us with a rich set of auxiliary predictions, namely: the probability that the return will take on a particular value. Unlike previously Figure 7. Percentage improvement, per-game, of C51 over Dou- proposed approaches, however, the accuracy of these pre- ble DQN, computed using van Hasselt et al.’s method. dictions is tightly coupled with the agent’s performance. Framework for inductive bias. The distributional per- 6. Discussion spective on reinforcement learning allows a more natural framework within which we can impose assumptions about In this work we sought a more complete picture of rein- the domain or the learning problem itself. In this work we forcement learning, one that involves value distributions. used distributions with support bounded in [V ,V ]. We found that learning value distributions is a powerful no- MIN MAX Treating this support as a hyperparameter allows us to tion that allows us to surpass most gains previously made change the optimization problem by treating all extremal on Atari 2600, without further algorithmic adjustments. returns (e.g. greater than VMAX) as equivalent. Surprisingly, a similar value clipping in DQN significantly degrades per- 6.1. Why does learning a distribution matter? formance in most games. To take another example: in- It is surprising that, when we use a policy which aims to terpreting the discount factor γ as a proper probability, as maximize expected return, we should see any difference some authors have argued, leads to a different algorithm. in performance. The distinction we wish to make is that Well-behaved optimization. It is well-accepted that the learning distributions matters in the presence of approxi- KL divergence between categorical distributions is a rea- mation. We now outline some possible reasons. sonably easy loss to minimize. This may explain some of Reduced chattering. Our results from Section 3.4 high- our empirical performance. Yet early experiments with al- lighted a significant instability in the Bellman optimal- ternative losses, such as KL divergence between continu- ity operator. When combined with function approxima- ous densities, were not fruitful, in part because the KL di- tion, this instability may prevent the policy from converg- vergence is insensitive to the values of its outcomes. A ing, what Gordon(1995) called chattering. We believe closer minimization of the Wasserstein metric should yield the gradient-based categorical algorithm is able to mitigate even better results than what we presented here. these effects by effectively averaging the different distri- In closing, we believe our results highlight the need to ac- † The UNREAL results are not altogether comparable, as count for distribution in the design, theoretical or other- they were generated in the asynchronous setting with per-game wise, of algorithms. hyperparameter tuning (Jaderberg et al., 2017). A Distributional Perspective on Reinforcement Learning
Acknowledgements Caruana, Rich. Multitask learning. Machine Learning, 28(1): 41–75, 1997. The authors acknowledge the important role played by their colleagues at DeepMind throughout the development of Chung, Kun-Jen and Sobel, Matthew J. Discounted mdps: Distri- this work. Special thanks to Yee Whye Teh, Alex Graves, bution functions and exponential utility maximization. SIAM Journal on Control and Optimization, 25(1):49–62, 1987. Joel Veness, Guillaume Desjardins, Tom Schaul, David Silver, Andre Barreto, Max Jaderberg, Mohammad Azar, Dearden, Richard, Friedman, Nir, and Russell, Stuart. Bayesian Georg Ostrovski, Bernardo Avila Pires, Olivier Pietquin, Q-learning. In Proceedings of the National Conference on Ar- Audrunas Gruslys, Tom Stepleton, Aaron van den Oord; tificial Intelligence, 1998. and particularly Chris Maddison for his comprehensive re- Engel, Yaakov, Mannor, Shie, and Meir, Ron. Reinforcement view of an earlier draft. Thanks also to Marek Petrik for learning with gaussian processes. In Proceedings of the In- pointers to the relevant literature, and Mark Rowland for ternational Conference on Machine Learning, 2005. fine-tuning details in the final version. Geist, Matthieu and Pietquin, Olivier. Kalman temporal differ- ences. Journal of Artificial Intelligence Research, 39:483–532, Erratum 2010. The camera-ready copy of this paper incorrectly reported a Gordon, Geoffrey. Stable function approximation in dynamic pro- gramming. In Proceedings of the Twelfth International Confer- mean score of 1010% for C51. The corrected figure stands ence on Machine Learning, 1995. at 701%, which remains higher than the other comparable baselines. The median score remains unchanged at 178%. Harutyunyan, Anna, Bellemare, Marc G., Stepleton, Tom, and Munos, Remi.´ Q(λ) with off-policy corrections. In Proceed- The error was due to evaluation episodes in one game (At- ings of the Conference on Algorithmic Learning Theory, 2016. lantis) lasting over 30 minutes; in comparison, the other results presented here cap episodes at 30 minutes, as is Hoffman, Matthew D., de Freitas, Nando, Doucet, Arnaud, and Peters, Jan. An expectation maximization algorithm for con- standard. The previously reported score on Atlantis was tinuous markov decision processes with arbitrary reward. In 3.7 million; our 30-minute score is 841,075, which we be- Proceedings of the International Conference on Artificial In- lieve is close to the achievable maximum in this time frame. telligence and Statistics, 2009. Capping at 30 minutes brings our human-normalized score Jaderberg, Max, Mnih, Volodymyr, Czarnecki, Wojciech Marian, on Atlantis from 22824% to a mere (!) 5199%, unfortu- Schaul, Tom, Leibo, Joel Z, Silver, David, and Kavukcuoglu, nately enough to noticeably affect the mean score, whose Koray. Reinforcement learning with unsupervised auxiliary sensitivity to outliers is well-documented. tasks. Proceedings of the International Conference on Learn- ing Representations, 2017.
References Jaquette, Stratton C. Markov decision processes with a new opti- mality criterion: Discrete time. The Annals of Statistics, 1(3): Azar, Mohammad Gheshlaghi, Munos, Remi,´ and Kappen, 496–505, 1973. Hilbert. On the sample complexity of reinforcement learning with a generative model. In Proceedings of the International Kakade, Sham and Langford, John. Approximately optimal ap- Conference on Machine Learning, 2012. proximate reinforcement learning. In Proceedings of the Inter- national Conference on Machine Learning, 2002. Bellemare, Marc G, Naddaf, Yavar, Veness, Joel, and Bowling, Michael. The arcade learning environment: An evaluation plat- Kingma, Diederik and Ba, Jimmy. Adam: A method for stochas- form for general agents. Journal of Artificial Intelligence Re- tic optimization. Proceedings of the International Conference search, 47:253–279, 2013. on Learning Representations, 2015. Bellemare, Marc G., Danihelka, Ivo, Dabney, Will, Mo- Lattimore, Tor and Hutter, Marcus. PAC bounds for discounted hamed, Shakir, Lakshminarayanan, Balaji, Hoyer, Stephan, MDPs. In Proceedings of the Conference on Algorithmic and Munos, Remi.´ The cramer distance as a solution to biased Learning Theory, 2012. wasserstein gradients. arXiv, 2017. Bellman, Richard E. Dynamic programming. Princeton Univer- Mannor, Shie and Tsitsiklis, John N. Mean-variance optimization sity Press, Princeton, NJ, 1957. in markov decision processes. 2011. Bertsekas, Dimitri P. and Tsitsiklis, John N. Neuro-Dynamic Pro- McCallum, Andrew K. Reinforcement learning with selective per- gramming. Athena Scientific, 1996. ception and hidden state. PhD thesis, University of Rochester, 1995. Bickel, Peter J. and Freedman, David A. Some asymptotic the- ory for the bootstrap. The Annals of Statistics, pp. 1196–1217, Mnih, Volodymyr, Kavukcuoglu, Koray, Silver, David, Rusu, An- 1981. drei A, Veness, Joel, Bellemare, Marc G, Graves, Alex, Ried- miller, Martin, Fidjeland, Andreas K, Ostrovski, Georg, et al. Billingsley, Patrick. Probability and measure. John Wiley & Human-level control through deep reinforcement learning. Na- Sons, 1995. ture, 518(7540):529–533, 2015. A Distributional Perspective on Reinforcement Learning
Morimura, Tetsuro, Hachiya, Hirotaka, Sugiyama, Masashi, van Hasselt, Hado, Guez, Arthur, and Silver, David. Deep rein- Tanaka, Toshiyuki, and Kashima, Hisashi. Parametric return forcement learning with double Q-learning. In Proceedings of density estimation for reinforcement learning. In Proceed- the AAAI Conference on Artificial Intelligence, 2016. ings of the Conference on Uncertainty in Artificial Intelligence, 2010a. Veness, Joel, Bellemare, Marc G., Hutter, Marcus, Chua, Alvin, and Desjardins, Guillaume. Compress and control. In Proceed- Morimura, Tetsuro, Sugiyama, Masashi, Kashima, Hisashi, ings of the AAAI Conference on Artificial Intelligence, 2015. Hachiya, Hirotaka, and Tanaka, Toshiyuki. Nonparametric re- turn distribution approximation for reinforcement learning. In Wang, Tao, Lizotte, Daniel, Bowling, Michael, and Schuurmans, Proceedings of the 27th International Conference on Machine Dale. Dual representations for dynamic programming. Journal Learning (ICML-10), pp. 799–806, 2010b. of Machine Learning Research, pp. 1–29, 2008. Nair, Arun, Srinivasan, Praveen, Blackwell, Sam, Alcicek, Wang, Ziyu, Schaul, Tom, Hessel, Matteo, Hasselt, Hado van, Cagdas, Fearon, Rory, De Maria, Alessandro, Panneershelvam, Lanctot, Marc, and de Freitas, Nando. Dueling network archi- Vedavyas, Suleyman, Mustafa, Beattie, Charles, and Petersen, tectures for deep reinforcement learning. In Proceedings of the Stig et al. Massively parallel methods for deep reinforcement International Conference on Machine Learning, 2016. learning. In ICML Workshop on Deep Learning, 2015. White, D. J. Mean, variance, and probabilistic criteria in finite Prashanth, LA and Ghavamzadeh, Mohammad. Actor-critic algo- markov decision processes: a review. Journal of Optimization rithms for risk-sensitive mdps. In Advances in Neural Informa- Theory and Applications, 56(1):1–29, 1988. tion Processing Systems, 2013. Puterman, Martin L. Markov Decision Processes: Discrete stochastic dynamic programming. John Wiley & Sons, Inc., 1994. Rosler,¨ Uwe. A fixed point theorem for distributions. Stochastic Processes and their Applications, 42(2):195–214, 1992. Schaul, Tom, Quan, John, Antonoglou, Ioannis, and Silver, David. Prioritized experience replay. In Proceedings of the International Conference on Learning Representations, 2016. Sobel, Matthew J. The variance of discounted markov decision processes. Journal of Applied Probability, 19(04):794–802, 1982. Sutton, Richard S. and Barto, Andrew G. Reinforcement learning: An introduction. MIT Press, 1998. Sutton, R.S., Modayil, J., Delp, M., Degris, T., Pilarski, P.M., White, A., and Precup, D. Horde: A scalable real-time archi- tecture for learning knowledge from unsupervised sensorimo- tor interaction. In Proceedings of the International Conference on Autonomous Agents and Multiagents Systems, 2011. Tamar, Aviv, Di Castro, Dotan, and Mannor, Shie. Learning the variance of the reward-to-go. Journal of Machine Learning Research, 17(13):1–36, 2016. Tieleman, Tijmen and Hinton, Geoffrey. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magni- tude. COURSERA: Neural networks for machine learning, 4 (2), 2012. Toussaint, Marc and Storkey, Amos. Probabilistic inference for solving discrete and continuous state markov decision pro- cesses. In Proceedings of the International Conference on Ma- chine Learning, 2006. Tsitsiklis, John N. On the convergence of optimistic policy itera- tion. Journal of Machine Learning Research, 3:59–72, 2002. Utgoff, Paul E. and Stracuzzi, David J. Many-layered learning. Neural Computation, 14(10):2497–2529, 2002. Van den Oord, Aaron, Kalchbrenner, Nal, and Kavukcuoglu, Ko- ray. Pixel recurrent neural networks. In Proceedings of the International Conference on Machine Learning, 2016. A Distributional Perspective on Reinforcement Learning
A. Related Work sequence of actions leading to maximal expected reward. Wang et al.(2008) considered the dual formulation of re- To the best of our knowledge, the work closest to ours are inforcement learning, where one optimizes the stationary two papers (Morimura et al., 2010b;a) studying the distri- distribution subject to constraints given by the transition butional Bellman equation from the perspective of its cu- function (Puterman, 1994), in particular its relationship to mulative distribution functions. The authors propose both linear approximation. Related to this dual is the Compress parametric and nonparametric solutions to learn distribu- and Control algorithm Veness et al.(2015), which describes tions for risk-sensitive reinforcement learning. They also a value function by learning a return distribution using den- provide some theoretical analysis for the policy evaluation sity models. One of the aims of this work was to address setting, including a consistency result in the nonparamet- the question left open by their work of whether one could ric case. By contrast, we also analyze the control setting, be design a practical distributional algorithm based on the and emphasize the use of the distributional equations to im- Bellman equation, rather than Monte Carlo estimation. prove approximate reinforcement learning. The variance of the return has been extensively stud- B. Proofs ied in the risk-sensitive setting. Of note, Tamar et al. (2016) analyze the use of linear function approximation Lemma 1 (Partition lemma). Let A1,A2,... be a set of to learn this variance for policy evaluation, and Prashanth random variables describing a partition of Ω, i.e. Ai(ω) ∈ & Ghavamzadeh(2013) estimate the return variance in the 0, 1 and for any ω there is exactly one Ai with Ai(ω) = { } design of a risk-sensitive actor-critic algorithm. Mannor 1. Let U, V be two random variables. Then & Tsitsiklis(2011) provides negative results regarding the X dp U, V dp(AiU, AiV ). computation of a variance-constrained solution to the opti- ≤ i mal control problem. Proof. We will give the proof for p < , noting that the The distributional formulation also arises when modelling D ∞ D same applies to p = . Let Y := A U and Z := A V , uncertainty. Dearden et al.(1998) considered a Gaussian i i i i respectively. First note∞ that approximation to the value distribution, and modelled the p p uncertainty over the parameters of this approximation us- dp(AiU, AiV ) = inf E Yi Zi Yi,Zi | − | ing a Normal-Gamma prior. Engel et al.(2005) leveraged h p i the distributional Bellman equation to define a Gaussian = inf E E Yi Zi Ai . Yi,Zi | − | | process over the unknown value function. More recently, p Geist & Pietquin(2010) proposed an alternative solution to Now, AiU AiV = 0 whenever Ai = 0. It follows that | − | p the same problem based on unscented Kalman filters. We we can choose Yi,Zi so that also Yi Zi = 0 whenever | − | believe much of the analysis we provide here, which deals Ai = 0, without increasing the expected norm. Hence with the intrinsic randomness of the environment, can also p dp(AiU, AiV ) = be applied to modelling uncertainty. p inf Pr Ai = 1 E Yi Zi Ai = 1 . (8) Our work here is based on a number of foundational re- Yi,Zi { } | − | | sults, in particular concerning alternative optimality crite- Next, we claim that ria. Early on, Jaquette(1973) showed that a moment opti- X h p i inf Pr Ai = 1 E AiU AiV Ai = 1 (9) mality criterion, which imposes a total ordering on distri- U,V i { } − | butions, is achievable and defines a stationary optimal pol- X h p i inf Pr Ai = 1 E Yi Zi Ai = 1 . icy, echoing the second part of Theorem1. Sobel(1982) ≤ Y1,Y2,... i { } | − | is usually cited as the first reference to Bellman equations Z1,Z2,... for the higher moments (but not the distribution) of the re- Specifically, the left-hand side of the equation is an infi- turn. Chung & Sobel(1987) provides results concerning mum over all r.v.’s whose cumulative distributions are FU the convergence of the distributional Bellman operator in and FV , respectively, while the right-hand side is an in- total variation distance. White(1988) studies “nonstandard fimum over sequences of r.v’s Y1,Y2,... and Z1,Z2,...
MDP criteria” from the perspective of optimizing the state- whose cumulative distributions are FAiU ,FAiV , respec- action pair occupancy. tively. To prove this upper bound, consider the c.d.f. of : A number of probabilistic frameworks for reinforcement U learning have been proposed in recent years. The plan- F (y) = Pr U y U { ≤ } ning as inference approach (Toussaint & Storkey, 2006; X = Pr Ai = 1 Pr U y Ai = 1 Hoffman et al., 2009) embeds the return into a graphical i { } { ≤ | } model, and applies probabilistic inference to determine the X = Pr Ai = 1 Pr AiU y Ai = 1 . i { } { ≤ | } A Distributional Perspective on Reinforcement Learning
Hence the distribution F is equivalent, in an almost sure Lemma 3. π : is a γ-contraction in d¯ . U T Z → Z p sense, to one that first picks an element Ai of the partition, then picks a value for U conditional on the choice Ai. On Proof. Consider Z1,Z2 . By definition, D ∈ Z the other hand, the c.d.f. of Yi = AiU is d¯ ( πZ , πZ ) = sup d ( πZ (x, a), πZ (x, a)). p T 1 T 2 p T 1 T 2 F (y) = Pr A = 1 Pr A U y A = 1 x,a AiU { i } { i ≤ | i } (10) + Pr A = 0 Pr A U y A = 0 By the properties of d , we have { i } { i ≤ | i } p = Pr Ai = 1 Pr AiU y Ai = 1 π π { } { ≤ | } dp(T Z1(x, a), T Z2(x, a)) π π + Pr Ai = 0 I [y 0] . = dp(R(x, a) + γP Z1(x, a),R(x, a) + γP Z2(x, a)) { } ≥ π π Thus the right-hand side infimum in (9) has the additional ≤ γdp(P Z1(x, a),P Z2(x, a)) 0 0 0 0 constraint that it must preserve the conditional c.d.fs, in ≤ γ sup dp(Z1(x , a ),Z2(x , a )), x0,a0 particular when y 0. Put another way, instead of hav- ≥ ing the freedom to completely reorder the mapping U : where the last line follows from the definition of P π (see Ω R, we can only reorder it within each element of the (4)). Combining with (10) we obtain partition.→ We now write π π π π d¯p( Z1, Z2) = sup dp( Z1(x, a), Z2(x, a)) p T T x,a T T dp(U, V ) = inf U V p U,V k − k 0 0 0 0 γ sup dp(Z1(x , a ),Z2(x , a )) p = inf E U V ≤ x0,a0 U,V | − | = γd¯p(Z1,Z2). (a) X p = inf Pr Ai = 1 E U V Ai = 1 U,V i { } | − | | Proposition 1 (Sobel, 1982). Consider two value distri- X p butions , and write to be the vector of = inf Pr Ai = 1 E AiU AiV Ai = 1 , Z1,Z2 V(Zi) U,V i { } | − | | ∈ Z variances of Zi. Then where (a) follows because A1,A2,... is a partition. Using π π E Z1 E Z2 γ E Z1 E Z2 , and (9), this implies k T − T k∞ ≤ k − k∞ π π 2 V( Z1) V( Z2) γ VZ1 VZ2 . p k T − T k∞ ≤ k − k∞ dp(U, V ) X h p i Proof. The first statement is standard, and its proof follows = inf Pr Ai = 1 E AiU AiV Ai = 1 U,V i { } − | from E πZ = π E Z, where the second π denotes the T T T X h p i usual operator over value functions. Now, by independence inf Pr Ai = 1 E Yi Zi Ai = 1 π ≤ Y1,Y2,... i { } − | of R and P Zi: Z1,Z2,... (b) X h p i π π = inf Pr Ai = 1 E Yi Zi Ai = 1 V( Zi(x, a)) = V R(x, a) + γP Zi(x, a) i Yi,Zi { } − | T 2 π (c) X = V(R(x, a)) + γ V(P Zi(x, a)). = dp(AiU, AiV ), i And now because in (b) the individual components of the sum are π π independently minimized; and (c) from (8). V( Z1) V( Z2) k T − T k∞ π π Lemma 2. d¯ is a metric over value distributions. = sup V( Z1(x, a)) V( Z2(x, a)) p x,a T − T 2 π π = sup γ [V(P Z1(x, a)) V(P Z2(x, a))] Proof. The only nontrivial property is the triangle inequal- x,a − ity. For any value distribution Y , write 2 0 0 0 0 ∈ Z = sup γ E [V(Z1(X ,A )) V(Z2(X ,A ))] x,a − d¯ (Z ,Z ) = sup d (Z (x, a),Z (x, a)) p 1 2 p 1 2 2 0 0 0 0 x,a sup γ V(Z1(x , a )) V(Z2(x , a )) ≤ x0,a0 − (a) 2 sup [dp(Z1(x, a),Y (x, a)) + dp(Y (x, a),Z2(x, a))] γ VZ1 VZ2 . ≤ x,a ≤ k − k∞
sup dp(Z1(x, a),Y (x, a)) + sup dp(Y (x, a),Z2(x, a)) Lemma 4. Let Z1,Z2 . Then ≤ x,a x,a ∈ Z
= d¯p(Z1,Y ) + d¯p(Y,Z2), E Z1 E Z2 γ E Z1 E Z2 , k T − T k∞ ≤ k − k∞ ∗ where in (a) we used the triangle inequality for dp. and in particular E Zk Q exponentially quickly. → A Distributional Perspective on Reinforcement Learning
Proof. The proof follows by linearity of expectation. Write is attained for some strictly positive ∆(x) > 0. By defini- D for the distributional operator and E for the usual op- tion, there exists a k such that erator.T Then T k ∆(x) k = γ B < , E DZ1 E DZ2 = E E Z1 E E Z2 2 k T − T k∞ kT − T k∞ γ Z1 Z2 ∞ . and hence every x must eventually be in . ≤ k − k ∈ X Xk Theorem 1 (Convergence in the control setting). Let This lemma allows us to guarantee the existence of an Z := Z with Z . Let be measurable and k k−1 0 iteration k after which sufficiently many states are well- supposeT that is finite. Then∈ Z X A behaved, in the sense that the greedy policy at those states ∗∗ chooses the optimal action. We will call these states lim inf dp(Zk(x, a),Z (x, a)) = 0 x, a. k→∞ Z∗∗∈Z∗∗ ∀ “solved”. We in fact require not only these states to be
∗∗ solved, but also most of their successors, and most of the If is finite, then Zk converges to uniformly. Further- successors of those, and so on. We formalize this notion as more,X if there is a total ordering Zon Π∗, such that for any follows: fix some δ > 0, let k,0 := k, and define for Z∗ ∗, ≺ X X ∈ Z i > 0 the set ∗ π ∗ 0 0 Z = Z with π Z∗ , π π π Z∗ π , ∗ := x : x ,P ( − − x, π (x)) 1 δ , T T ∈ G ≺ ∀ ∈ G \{ } Xk,i ∈ Xk Xk 1,i 1 | ≥ − then has a unique fixed point Z∗ ∗. T ∈ Z As the following lemma shows, any x is eventually con- tained in the recursively-defined sets , for any i. The gist of the proof of Theorem1 consists in showing that Xk,i for every state x, there is a time k after which the greedy Lemma 6. For any i N and any x , there exists a k 0 ∈ ∈ X policy w.r.t. Q is mostly optimal. To clearly expose the such that for all k k, x k0,i. k ≥ ∈ X steps involved, we will first assume a unique (and there- fore deterministic) optimal policy π∗, and later return to Proof. Fix i and let us suppose that . By Lemma Xk,i ↑ X the general case; we will denote the optimal action at x by 5, this is true for i = 0. We infer that for any probability π∗(x). For notational convenience, we will write Q := measure P on , P ( ) P ( ) = 1. In particular, for k X Xk,i → X E Zk and k := Zk . Let B := 2 supZ∈Z Z ∞ < a given x k, this implies that G k G k k ∞ ∈ X and let k := γ B. We first define the set of states k ∗ X ⊆ X P ( x, π∗(x)) P ( x, π∗(x)) = 1. whose values must be sufficiently close to Q at time k: Xk,i | → X |
n ∗ ∗ ∗ o Therefore, for any x, there exists a time after which it is k := x : Q (x, π (x)) max Q (x, a) > 2k . X − a=6 π∗(x) and remains a member of k,i+1, the set of states for which ∗ X (11) P ( k−1,i x, π (x)) 1 δ. We conclude that k,i+1 Indeed, by Lemma4, we know that after k iterations also.X The| statement≥ follows− by induction. X ↑ X Q (x, a) Q∗(x, a) γk Q (x, a) Q∗(x, a) . | k − | ≤ | 0 − | ≤ k Proof of Theorem1. The proof is similar to policy iteration-type results, but requires more care in dealing For x , write a∗ := π∗(x). For any a , we deduce ∈ X ∈ A with the metric and the possibly infinite state space. that ∗ We will write Wk(x) := Zk(x, πk(x)), define W ∗ ∗ ∗ ∗ similarly and with some overload of notation write Qk(x, a ) Qk(x, a) Q (x, a ) Q (x, a) 2k. − ≥ − − W (x) := W (x) = Z (x, π (x)). Finally, let T k k+1 T k k+1 ∗ 0 Sk(x) := [x ] and S¯k(x) = 1 Sk(x). It follows that if x k, then also Qk(x, a ) > Qk(x, a ) i I k,i i i 0 ∗ ∈ X ∈ X − for all a = π (x): for these states, the greedy policy Fix i > 0 and x . We begin by using 6 ∈ Xk+1,i+1 ⊆ Xk πk(x) := arg maxa Qk(x, a) corresponds to the optimal Lemma1 to separate the transition from x into a solved ∗ policy π . term and an unsolved term: Lemma 5. For each x there exists a k such 0 0 0 ∈ X P πk W (x) = SkW (X ) + S¯kW (X ), that, for all k k, x k0 , and in particular k i k i k arg max Q (x, a)≥ = π∗(x). ∈ X a k where X0 is the random successor from taking action ∗ k k 0 ¯k πk(x) := π (x), and we write Si = Si (X ), Si = Proof. Because is finite, the gap ¯k 0 A Si (X ) to ease the notation. Similarly, ∗ ∗ ∗ ∆(x) := Q (x, π (x)) max Q (x, a) πk ∗ k ∗ 0 ¯k ∗ 0 − a=6 π∗(x) P W (x) = Si W (X ) + Si W (X ). A Distributional Perspective on Reinforcement Learning
Now such that π¯k ∗ ∗ Zk+1 = Zk−i+1. dp(Wk+1(x),W (x)) = dp( Wk(x), W (x)) T T T ∗∗ (a) ∗ Now denote by the set of nonstationary optimal poli- πk π ∗ Z ∗ ∗ γdp(P Wk(x),P W (x)) cies. If we take any Z , we deduce that ≤ ∈ Z (b) k 0 k ∗ 0 γdp(Si Wk(X ),Si W (X )) π¯k ∗ ∗∗ δB ≤ inf dp( Z (x, a),Z (x, a)) , ¯k 0 ¯k ∗ 0 Z∗∗∈Z∗∗ T ≤ 1 γ + γdp(Si Wk(X ), Si W (X )), (12) − ∗ ∗ where in (a) we used Properties P1 and P2 of the Wasser- since Z corresponds to some optimal policy π and π¯k is optimal along most of the trajectories from (x, a). In effect, stein metric, and in (b) we separate states for which πk = π¯ ∗ π∗ from the rest using Lemma1( Sk, S¯k form a parti- k Z is close to the value distribution of the nonstation- i i T ∗ ∗ 0 { } ¯k 0 ary optimal policy π¯kπ . Now for this Z , tion of Ω). Let δi := Pr X / k,i = E Si (X ) = ¯k 0 { ∈ X } { } Si (X ) p. From property P3 of the Wasserstein metric, ∗∗ k k inf dp(Zk(x, a),Z (x, a)) we have Z∗∗ π¯k ∗ ¯k 0 ¯k ∗ 0 dp(Zk(x, a), Z (x, a)) dp(Si Wk(X ), Si W (X )) ≤ T π¯k ∗ ∗∗ + inf dp( Z (x, a),Z (x, a)) ¯k 0 0 ¯k 0 ∗ 0 ∗∗ sup dp(Si (X )Wk(x ), Si (X )W (x )) Z T ≤ x0 π¯k π¯k ∗ δB ¯k 0 0 ∗ 0 dp( Zk−i+1(x, a), Z (x, a)) + Si (X ) p sup dp(Wk(x ),W (x )) ≤ T T 1 γ ≤ k k x0 − 0 ∗ 0 i 2δB δi sup dp(Wk(x ),W (x )) γ B + , ≤ x0 ≤ 1 γ − δ B. ≤ i using the same argument as before with the newly-defined Recall that B < is the largest attainable Z . Since k,i. It follows that ∞ k k∞ X also δi < δ by our choice of x k+1,i+1, we can upper ∈ X ∗∗ bound the second term in (12) by γδB. This yields inf dp(Zk(x, a),Z (x, a)) 0. Z∗∗∈Z∗∗ → ∗ dp(Wk+1(x),W (x)) ≤ When is finite, there exists a fixed k after which k = 0 ∗ 0 X X γd (SkW (X ),SkW (X )) + γδB. . The uniform convergence result then follows. p i k i X ∗ By induction on , we conclude that for To prove the uniqueness of the fixed point Z when se- i > 0 x k+i,i T and some random state X00 i steps forward, ∈ X lects its actions according to the ordering , we note that for any optimal value distribution Z∗, its set≺ of greedy poli- ∗ ∗ ∗ dp(Wk+i(x),W (x)) cies is Π . Denote by π the policy coming first in the or- ≤ ∗ δB dering over Π∗. Then = π , which has a unique fixed γid (SkW (X00),SkW ∗(X00)) + T T p 0 k 0 1 γ point (Section 3.3). − δB γiB + . Proposition 4. That has a fixed point Z∗ = Z∗ is ≤ 1 γ insufficient to guaranteeT the convergence of Z toT ∗. − { k} Z Hence for any x , > 0, we can take δ, i, and finally k ∈ X ∗ We provide here a sketch of the result. Consider a single large enough to make dp(Wk(x),W (x)) < . The proof state x with two actions, a and a (Figure8). The first then extends to Zk(x, a) by considering one additional ap- 1 1 2 plication of . action yields a reward of 1/2, while the other either yields T 0 or 1 with equal probability, and both actions are optimal. We now consider the more general case where there are Now take γ = 1/2 and write R0,R1,... for the received multiple optimal policies. We expand the definition of k,i as follows: X rewards. Consider a stochastic policy that takes action a2 with probability p. For p = 0, the return is ∗ ∗ ∗ Xk,i := x ∈ Xk : ∀π ∈ Π , P (Xk−1,i−1 | x, a ) ≥ 1−δ , ∗ E∗ a ∼π (x) 1 1 Z = = 1. Because there are finitely many actions, Lemma6 also p=0 1 γ 2 − holds for this new definition. As before, take x k,i, but ∈ X For p = 1, on the other hand, the return is random and is now consider the sequence of greedy policies πk, πk−1,... selected by successive applications of , and write given by the following fractional number (in binary): T π¯k := πk πk−1 πk−i+1 , Z = R .R R R . T T T ···T p=1 0 1 2 3 ··· A Distributional Perspective on Reinforcement Learning
R = 1/2 R = 0 or 1 consider the d1 metric between this distribution P and an- other distribution Q. The first distribution is a a1 2 0 w.p. 1/2 P = 1 w.p. 1/2. x1 In this example, i 1, 2 , P1 = 0, and P2 = 1. Now consider the distribution∈ { with} the same support but that puts Figure 8. A simple example illustrating the effect of a nonstation- probability p on 0: ary policy on the value distribution. 0 w.p. p Q = 1 w.p. 1 p. As a result, Zp=1 is uniformly distributed between 0 and 2! − In fact, note that The distance between P and Q is
Zp=0 = 0.11111 = 1. ··· d (P,Q) = p 1 . 1 | − 2 | For some intermediary value of p, we obtain a different 1 probability of the different digits, but always putting some This is d1(P,Q) = 2 for p 0, 1 , and strictly less than 1 ∈ { } probability mass on all returns in [0, 2]. 2 for any other values of p. On the other hand, the corre- sponding expected distance (after sampling an outcome x1 Now suppose we follow the nonstationary policy that takes or x2 with equal probability) is a1 on the first step, then a2 from there on. By inspec- tion, the return will be uniformly distributed on the interval 1 1 1 EI d1(Pi,Q) = 2 p + 2 (1 p) = 2 . [1/2, 3/2], which does not correspond to the return under − any value of p. But now we may imagine an operator Hence d1(P,Q) < EI d1(Pi,Q) for p (0, 1). This shows T ∈ which alternates between a1 and a2 depending on the ex- that the bound is in general strict. By inspection, it is clear act value distribution it is applied to, which would in turn that the two gradients are different. converge to a nonstationary optimal value distribution. Lemma 7 (Sample Wasserstein distance). Let P be a { i} R = 0 R = 1 collection of random variables, I N a random index x x independent from P , and consider∈ the mixture random 1 2 { i} variable P = PI . For any random variable Q independent ½ ½ of I, x dp(P,Q) E dp(Pi,Q), ≤ i∼I and in general the inequality is strict and Figure 9. Example MDP in which the expected sample Wasser- stein distance is greater than the Wasserstein distance. Qdp(PI ,Q) = E Qdp(Pi,Q). ∇ 6 i∼I ∇ Proposition 5. Fix some next-state distribution Z and pol- Proof. We prove this using Lemma1. Let A := [I = i]. i I icy π. Consider a parametric value distribution Zθ, and We write and define the Wasserstein loss
dp(P,Q) = dp(PI ,Q) 0 0 W (θ) := dp(Zθ(x, a),R(x, a) + γZ(X , π(X ))). X X L = dp AiPi, AiQ i i Let r R(x, a) and x0 P ( x, a) and consider the X ∼ ∼ · | dp(AiPi,AiQ) sample loss ≤ i X 0 0 0 Pr I = i dp(Pi,Q) LW (θ, r, x ) := dp(Zθ(x, a), r + γZ(x , π(x )). ≤ i { } = EI dP (Pi,Q). Its expectation is an upper bound on the loss : LW where in the penultimate line we used the independence of 0 W (θ) E LW (θ, r, x ), I from Pi and Q to appeal to property P3 of the Wasserstein L ≤ R,P metric. in general with strict inequality. To show that the bound is in general strict, consider the mixture distribution depicted in Figure9. We will simply The result follows directly from the previous lemma. A Distributional Perspective on Reinforcement Learning
(a) Stochastic Bellman Target (b) Wasserstein Categorical ) ✓ Monte-Carlo Target ,Z Z ⇡ Wasserstein F Z
( Categorical 1 d
# Atoms Return
π Figure 10. (a) Wasserstein distance between ground truth distribution Z and approximating distributions Zθ. Varying number of atoms in approximation, training target, and loss function. (b) Approximate cumulative distributions for five representative states in CliffWalk.
C. Algorithmic Details r = -1 safe path While our training regime closely follows that of DQN (Mnih et al., 2015), we use Adam (Kingma & Ba, 2015) optimal path instead of RMSProp (Tieleman & Hinton, 2012) for gra- S The Cliff G dient rescaling. We also performed some hyperparam- eter tuning for our final results. Specifically, we eval- r = -100 uated two hyperparameters over our five training games and choose the values that performed best. The hyperpa- Figure 11. CliffWalk Environment (Sutton & Barto, 1998). rameter values we considered were VMAX 3, 10, 100 and 1/L, 0.1/L, 0.01/L, 0.001/L,∈ {0.0001/L }, adam ∈ { } where L = 32 is the minibatch size. We found VMAX = 10 and adam = 0.01/L performed best. We used the same lowing a policy that tries to take the safe path but has a 10% step-size value as DQN (α = 0.00025). chance of taking another action uniformly at random. Pseudo-code for the categorical algorithm is given in Algo- We compute a ground-truth distribution of returns Zπ using rithm1. We apply the Bellman update to each atom sepa- 10000 Monte-Carlo (MC) rollouts from each state. We then rately, and then project it into the two nearest atoms in the perform two experiments, approximating the value distri- original support. Transitions to a terminal state are handled bution at each state with our discrete distributions. with γ = 0. t In the first experiment, we perform supervised learning us- ing either the Wasserstein loss or categorical projection D. Comparison of Sampled Wasserstein Loss (Equation7) with cross-entropy loss. We use Zπ as the and Categorical Projection supervised target and perform 5000 sweeps over all states to ensure both approaches have converged. In the second Lemma3 proves that for a fixed policy the distributional π experiment, we use the same loss functions, but the training Bellman operator is a -contraction in ¯ , and therefore γ dp target comes from the one-step distributional Bellman op- that π will converge in distribution to the true distribution erator with sampled transitions. We use V = 100 and of returnsT Zπ. In this section, we empirically validate these MIN V = 1.4 For the sample updates we perform− 10 times results on the CliffWalk domain shown in Figure 11. The MAX as many− sweeps over the state space. Fundamentally, these dynamics of the problem match those given by Sutton & experiments investigate how well the two training regimes Barto(1998). We also study the convergence of the distri- butional Bellman operator under the sampled Wasserstein 4Because there is a small probability of larger negative returns, loss and the categorical projection (Equation7) while fol- some approximation error is unavoidable. However, this effect is relatively negligible in our experiments. A Distributional Perspective on Reinforcement Learning
(minimizing the Wasserstein or categorical loss) minimize the Wasserstein metric under both ideal (supervised target) and practical (sampled one-step Bellman target) conditions. In Figure 10a we show the final Wasserstein distance π d1(Z ,Zθ) between the learned distributions and the ground-truth distribution as we vary the number of atoms. The graph shows that the categorical algorithm does indeed minimize the Wasserstein metric in both the supervised and sample Bellman setting. It also highlights that minimizing the Wasserstein loss with stochastic gradient descent is in general flawed, confirming the intuition given by Propo- sition5. In repeat experiments the process converged to π different values of d1(Z ,Zθ), suggesting the presence of local minima (more prevalent with fewer atoms). Figure 10 provides additional insight into why the sampled Wasserstein distance may perform poorly. Here, we see the cumulative densities for the approximations learned under these two losses for five different states along the safe path in CliffWalk. The Wasserstein has converged to a fixed- point distribution, but not one that captures the true (Monte Carlo) distribution very well. By comparison, the categor- ical algorithm captures the variance of the true distribution much more accurately. GAMES VIDEOURL Freeway http://youtu.be/97578n9kFIk E. Supplemental Videos and Results Pong http://youtu.be/vIz5P6s80qA Q*Bert http://youtu.be/v-RbNX4uETw In Figure 13 we provide links to supplemental videos show- Seaquest http://youtu.be/d1yz4PNFUjI ing the C51 agent during training on various Atari 2600 Space Invaders http://youtu.be/yFBwyPuO2Vg games. Figure 12 shows the relative performance of C51 over the course of training. Figure 14 provides a table Figure 13. Supplemental videos of C51 during training. of evaluation results, comparing C51 to other state-of-the- art agents. Figures 15–18 depict particularly interesting frames.
C51 vs. DQN
C51 vs. HUMAN
DQN vs. HUMAN # # Games Superior
Training Frames (millions)
Figure 12. Number of Atari games where an agent’s training per- formance is greater than a baseline (fully trained DQN & human). Error bands give standard deviations, and averages are over num- ber of games. A Distributional Perspective on Reinforcement Learning
GAMES RANDOM HUMAN DQN DDQN DUEL PRIOR. DUEL. C51 Alien 227.8 7,127.7 1,620.0 3,747.7 4,461.4 3,941.0 3,166 Amidar 5.8 1,719.5 978.0 1,793.3 2,354.5 2,296.8 1,735 Assault 222.4 742.0 4,280.4 5,393.2 4,621.0 11,477.0 7,203 Asterix 210.0 8,503.3 4,359.0 17,356.5 28,188.0 375,080.0 406,211 Asteroids 719.1 47,388.7 1,364.5 734.7 2,837.7 1,192.7 1,516 Atlantis 12,850.0 29,028.1 279,987.0 106,056.0 382,572.0 395,762.0 841,075 Bank Heist 14.2 753.1 455.0 1,030.6 1,611.9 1,503.1 976 Battle Zone 2,360.0 37,187.5 29,900.0 31,700.0 37,150.0 35,520.0 28,742 Beam Rider 363.9 16,926.5 8,627.5 13,772.8 12,164.0 30,276.5 14,074 Berzerk 123.7 2,630.4 585.6 1,225.4 1,472.6 3,409.0 1,645 Bowling 23.1 160.7 50.4 68.1 65.5 46.7 81.8 Boxing 0.1 12.1 88.0 91.6 99.4 98.9 97.8 Breakout 1.7 30.5 385.5 418.5 345.3 366.0 748 Centipede 2,090.9 12,017.0 4,657.7 5,409.4 7,561.4 7,687.5 9,646 Chopper Command 811.0 7,387.8 6,126.0 5,809.0 11,215.0 13,185.0 15,600 Crazy Climber 10,780.5 35,829.4 110,763.0 117,282.0 143,570.0 162,224.0 179,877 Defender 2,874.5 18,688.9 23,633.0 35,338.5 42,214.0 41,324.5 47,092 Demon Attack 152.1 1,971.0 12,149.4 58,044.2 60,813.3 72,878.6 130,955 Double Dunk -18.6 -16.4 -6.6 -5.5 0.1 -12.5 2.5 Enduro 0.0 860.5 729.0 1,211.8 2,258.2 2,306.4 3,454 Fishing Derby -91.7 -38.7 -4.9 15.5 46.4 41.3 8.9 Freeway 0.0 29.6 30.8 33.3 0.0 33.0 33.9 Frostbite 65.2 4,334.7 797.4 1,683.3 4,672.8 7,413.0 3,965 Gopher 257.6 2,412.5 8,777.4 14,840.8 15,718.4 104,368.2 33,641 Gravitar 173.0 3,351.4 473.0 412.0 588.0 238.0 440 H.E.R.O. 1,027.0 30,826.4 20,437.8 20,130.2 20,818.2 21,036.5 38,874 Ice Hockey -11.2 0.9 -1.9 -2.7 0.5 -0.4 -3.5 James Bond 29.0 302.8 768.5 1,358.0 1,312.5 812.0 1,909 Kangaroo 52.0 3,035.0 7,259.0 12,992.0 14,854.0 1,792.0 12,853 Krull 1,598.0 2,665.5 8,422.3 7,920.5 11,451.9 10,374.4 9,735 Kung-Fu Master 258.5 22,736.3 26,059.0 29,710.0 34,294.0 48,375.0 48,192 Montezuma’s Revenge 0.0 4,753.3 0.0 0.0 0.0 0.0 0.0 Ms. Pac-Man 307.3 6,951.6 3,085.6 2,711.4 6,283.5 3,327.3 3,415 Name This Game 2,292.3 8,049.0 8,207.8 10,616.0 11,971.1 15,572.5 12,542 Phoenix 761.4 7,242.6 8,485.2 12,252.5 23,092.2 70,324.3 17,490 Pitfall! -229.4 6,463.7 -286.1 -29.9 0.0 0.0 0.0 Pong -20.7 14.6 19.5 20.9 21.0 20.9 20.9 Private Eye 24.9 69,571.3 146.7 129.7 103.0 206.0 15,095 Q*Bert 163.9 13,455.0 13,117.3 15,088.5 19,220.3 18,760.3 23,784 River Raid 1,338.5 17,118.0 7,377.6 14,884.5 21,162.6 20,607.6 17,322 Road Runner 11.5 7,845.0 39,544.0 44,127.0 69,524.0 62,151.0 55,839 Robotank 2.2 11.9 63.9 65.1 65.3 27.5 52.3 Seaquest 68.4 42,054.7 5,860.6 16,452.7 50,254.2 931.6 266,434 Skiing -17,098.1 -4,336.9 -13,062.3 -9,021.8 -8,857.4 -19,949.9 -13,901 Solaris 1,236.3 12,326.7 3,482.8 3,067.8 2,250.8 133.4 8,342 Space Invaders 148.0 1,668.7 1,692.3 2,525.5 6,427.3 15,311.5 5,747 Star Gunner 664.0 10,250.0 54,282.0 60,142.0 89,238.0 125,117.0 49,095 Surround -10.0 6.5 -5.6 -2.9 4.4 1.2 6.8 Tennis -23.8 -8.3 12.2 -22.8 5.1 0.0 23.1 Time Pilot 3,568.0 5,229.2 4,870.0 8,339.0 11,666.0 7,553.0 8,329 Tutankham 11.4 167.6 68.1 218.4 211.4 245.9 280 Up and Down 533.4 11,693.2 9,989.9 22,972.2 44,939.6 33,879.1 15,612 Venture 0.0 1,187.5 163.0 98.0 497.0 48.0 1,520 Video Pinball 16,256.9 17,667.9 196,760.4 309,941.9 98,209.5 479,197.0 949,604 Wizard Of Wor 563.5 4,756.5 2,704.0 7,492.0 7,855.0 12,352.0 9,300 Yars’ Revenge 3,092.9 54,576.9 18,098.9 11,712.6 49,622.1 69,618.1 35,050 Zaxxon 32.5 9,173.3 5,363.0 10,163.0 12,944.0 13,886.0 10,513
Figure 14. Raw scores across all games, starting with 30 no-op actions. Reference values from Wang et al.(2016). A Distributional Perspective on Reinforcement Learning
Figure 15. FREEWAY: Agent differentiates action-value distributions under pressure.
Figure 16. Q*BERT: Top, left and right: Predicting which actions are unrecoverably fatal. Bottom-Left: Value distribution shows steep consequences for wrong actions. Bottom-Right: The agent has made a huge mistake.
Figure 17. SEAQUEST: Left: Bimodal distribution. Middle: Might hit the fish. Right: Definitely going to hit the fish.
Figure 18. SPACE INVADERS: Top-Left: Multi-modal distribution with high uncertainty. Top-Right: Subsequent frame, a more certain demise. Bottom-Left: Clear difference between actions. Bottom-Middle: Uncertain survival. Bottom-Right: Certain success.