A Distributional Perspective on Reinforcement Learning — part 3/6

Cleaned-text code points [29161, 45137). Bibliography omitted here; source.md includes it.
Instead, we project the sample Bellman update 𝒯 ^  Z θ \hat{\mathcal{T}}Z_{\theta} onto the support of Z θ Z_{\theta} (Figure 1 , Algorithm 1 ), effectively reducing the Bellman update to multiclass classification. Let π \pi be the greedy policy w.r.t. 𝔼 Z θ \expect Z_{\theta} . Given a sample transition ( x , a , r , x ′ ) (x,a,r,x^{\prime}) , we compute the Bellman update 𝒯 ^  z j := r + γ  z j \hat{\mathcal{T}}z_{j}:=r+\gamma z_{j} for each atom z j z_{j} , then distribute its probability p j  ( x ′ , π ⁡ ( x ′ ) ) p_{j}(x^{\prime},\pi(x^{\prime})) to the immediate neighbours of 𝒯 ^  z j \hat{\mathcal{T}}z_{j} . The i t  h i^{th} component of the projected update Φ  𝒯 ^  Z θ  ( x , a ) \Phi\hat{\mathcal{T}}Z_{\theta}(x,a) is ( Φ  𝒯 ^  Z θ  ( x , a ) ) i = ∑ j = 0 N − 1 [ 1 − | [ 𝒯 ^  z j ] V min V max − z i | △  z ] 0 1  p j  ( x ′ , π ⁡ ( x ′ ) ) , (\Phi\hat{\mathcal{T}}Z_{\theta}(x,a))_{i}=\sum_{j=0}^{N-1}\left[1-\frac{|[\hat{\mathcal{T}}z_{j}]^{V_{\textsc{max}}}_{V_{\textsc{min}}}-z_{i}|}{\triangle z}\right]^{1}_{0}p_{j}(x^{\prime},\pi(x^{\prime})), (7) where [ ⋅ ] a b [\cdot]^{b}_{a} bounds its argument in the range [ a , b ] [a,b] . 1 1 1 Algorithm 1 computes this projection in time linear in N N . As is usual, we view the next-state distribution as parametrized by a fixed parameter θ ~ \tilde{\theta} . The sample loss ℒ x , a  ( θ ) \mathcal{L}_{x,a}(\theta) is the cross-entropy term of the KL divergence D kl ( Φ 𝒯 ^ Z θ ~ ( x , a ) ∥ Z θ ( x , a ) ) , D_{\textsc{kl}}(\Phi\hat{\mathcal{T}}Z_{\tilde{\theta}}(x,a)\,\|\,Z_{\theta}(x,a)), which is readily minimized e.g. using gradient descent. We call this choice of distribution and loss the categorical algorithm . When N = 2 N=2 , a simple one-parameter alternative is Φ 𝒯 ^ Z θ ( x , a ) := [ 𝔼 [ 𝒯 ^ Z θ ( x , a ) ] − V min ) / △ z ] 0 1 ; \Phi\hat{\mathcal{T}}Z_{\theta}(x,a):=[\expect[\hat{\mathcal{T}}Z_{\theta}(x,a)]-V_{\textsc{min}})/\triangle z]^{1}_{0}; we call this the Bernoulli algorithm . We note that, while these algorithms appear unrelated to the Wasserstein metric, recent work ( Bellemare et al., 2017 ) hints at a deeper connection.

## 5 Evaluation on Atari 2600 Games

To understand the approach in a complex setting, we applied the categorical algorithm to games from the Arcade Learning Environment ( Bellemare et al., 2013 , ALE;) . While the ALE is deterministic, stochasticity does occur in a number of guises: 1) from state aliasing, 2) learning from a nonstationary policy, and 3) from approximation errors. We used five training games (Fig 3 ) and 52 testing games.

For our study, we use the DQN architecture ( Mnih et al., 2015 ) , but output the atom probabilities p i  ( x , a ) p_{i}(x,a) instead of action-values, and chose V max = − V min = 10 V_{\textsc{max}}=-V_{\textsc{min}}=10 from preliminary experiments over the training games. We call the resulting architecture Categorical DQN . We replace the squared loss ( r + γ  Q  ( x ′ , π ⁡ ( x ′ ) ) − Q ⁡ ( x , a ) ) 2 (r+\gamma Q(x^{\prime},\pi(x^{\prime}))-Q(x,a))^{2} by ℒ x , a  ( θ ) \mathcal{L}_{x,a}(\theta) and train the network to minimize this loss. 2 2 2 For N = 51 N=51 , our TensorFlow implementation trains at roughly 75% of DQN’s speed. As in DQN, we use a simple ϵ \epsilon -greedy policy over the expected action-values; we leave as future work the many ways in which an agent could select actions on the basis of the full distribution. The rest of our training regime matches Mnih et al. ’s, including the use of a target network for θ ~ \tilde{\theta} .

Figure 4 illustrates the typical value distributions we observed in our experiments. In this example, three actions (those including the button press) lead to the agent releasing its laser too early and eventually losing the game. The corresponding distributions reflect this: they assign a significant probability to 0 (the terminal value). The safe actions have similar distributions ( left , which tracks the invaders’ movement, is slightly favoured). This example helps explain why our approach is so successful: the distributional update keeps separated the low-value, ‘‘losing’’ event from the high-value, ‘‘survival’’ event, rather than average them into one (unrealizable) expectation. 3 3 3 Video: http://youtu.be/yFBwyPuO2Vg .

One surprising fact is that the distributions are not concentrated on one or two values, in spite of the ALE’s determinism, but are often close to Gaussians. We believe this is due to our discretizing the diffusion process induced by γ \gamma .

### 5.1 Varying the Number of Atoms

We began by studying our algorithm’s performance on the training games in relation to the number of atoms (Figure 3 ). For this experiment, we set ϵ = 0.05 \epsilon=0.05 . From the data, it is clear that using too few atoms can lead to poor behaviour, and that more always increases performance; this is not immediately obvious as we may have expected to saturate the network capacity. The difference in performance between the 51-atom version and DQN is particularly striking: the latter is outperformed in all five games, and in Seaquest we attain state-of-the-art performance. As an additional point of the comparison, the single-parameter Bernoulli algorithm performs better than DQN in 3 games out of 5, and is most notably more robust in Asterix .

One interesting outcome of this experiment was to find out that our method does pick up on stochasticity. Pong exhibits intrinsic randomness: the exact timing of the reward depends on internal registers and is truly unobservable. We see this clearly reflected in the agent’s prediction (Figure 5 ): over five consecutive frames, the value distribution shows two modes indicating the agent’s belief that it has yet to receive a reward. Interestingly, since the agent’s state does not include past rewards, it cannot even extinguish the prediction after receiving the reward, explaining the relative proportions of the modes.

### 5.2 State-of-the-Art Results

The performance of the 51-atom agent (from here onwards, C51) on the training games, presented in the last section, is particularly remarkable given that it involved none of the other algorithmic ideas present in state-of-the-art agents. We next asked whether incorporating the most common hyperparameter choice, namely a smaller training ϵ \epsilon , could lead to even better results. Specifically, we set ϵ = 0.01 \epsilon=0.01 (instead of 0.05 0.05 ); furthermore, every 1 million frames, we evaluate our agent’s performance with ϵ = 0.001 \epsilon=0.001 .

We compare our algorithm to DQN ( ϵ = 0.01 \epsilon=0.01 ), Double DQN ( van Hasselt et al., 2016 ) , the Dueling architecture ( Wang et al., 2016 ) , and Prioritized Replay ( Schaul et al., 2016 ) , comparing the best evaluation score achieved during training. We see that C51 significantly outperforms these other algorithms (Figures 6 and 7 ). In fact, C51 surpasses the current state-of-the-art by a large margin in a number of games, most notably Seaquest . One particularly striking fact is the algorithm’s good performance on sparse reward games, for example Venture and Private Eye . This suggests that value distributions are better able to propagate rarely occurring events. Full results are provided in the appendix.

We also include in the appendix (Figure 12 ) a comparison, averaged over 3 seeds, showing the number of games in which C51’s training performance outperforms fully-trained DQN and human players. These results continue to show dramatic improvements, and are more representative of an agent’s average performance. Within 50 million frames, C51 has outperformed a fully trained DQN agent on 45 out of 57 games. This suggests that the full 200 million training frames, and its ensuing computational cost, are unnecessary for evaluating reinforcement learning algorithms within the ALE.

The most recent version of the ALE contains a stochastic execution mechanism designed to ward against trajectory overfitting.Specifically, on each frame the environment rejects the agent’s selected action with probability p = 0.25 p=0.25 . Although DQN is mostly robust to stochastic execution, there are a few games in which its performance is reduced. On a score scale normalized with respect to the random and DQN agents, C51 obtains mean and median score improvements of 126 % 126\% and 21.5 % 21.5\% respectively, confirming the benefits of C51 beyond the deterministic setting.

## 6 Discussion

In this work we sought a more complete picture of reinforcement learning, one that involves value distributions. We found that learning value distributions is a powerful notion that allows us to surpass most gains previously made on Atari 2600, without further algorithmic adjustments.

### 6.1 Why does learning a distribution matter?

It is surprising that, when we use a policy which aims to maximize expected return, we should see any difference in performance. The distinction we wish to make is that learning distributions matters in the presence of approximation . We now outline some possible reasons.

Reduced chattering. Our results from Section 3.4 highlighted a significant instability in the Bellman optimality operator. When combined with function approximation, this instability may prevent the policy from converging, what Gordon (1995) called chattering . We believe the gradient-based categorical algorithm is able to mitigate these effects by effectively averaging the different distributions, similar to conservative policy iteration ( Kakade & Langford, 2002 ) . While the chattering persists, it is integrated to the approximate solution.

State aliasing. Even in a deterministic environment, state aliasing may result in effective stochasticity. McCallum (1995) , for example, showed the importance of coupling representation learning with policy learning in partially observable domains. We saw an example of state aliasing in Pong , where the agent could not exactly predict the reward 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 multitude of predictions ( Caruana 1997 ; Utgoff & Stracuzzi 2002 ; Sutton et al. 2011 ; Jaderberg et al. 2017 ). The distributional 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 proposed approaches, however, the accuracy of these predictions is tightly coupled with the agent’s performance.

Framework for inductive bias. The distributional perspective on reinforcement learning allows a more natural framework within which we can impose assumptions about the domain or the learning problem itself. In this work we used distributions with support bounded in [ V min , V max ] [V_{\textsc{min}},V_{\textsc{max}}] . Treating this support as a hyperparameter allows us to change the optimization problem by treating all extremal returns (e.g. greater than V max V_{\textsc{max}} ) as equivalent. Surprisingly, a similar value clipping in DQN significantly degrades performance in most games. To take another example: interpreting the discount factor γ \gamma as a proper probability, as some authors have argued, leads to a different algorithm.

Well-behaved optimization. It is well-accepted that the KL divergence between categorical distributions is a reasonably easy loss to minimize. This may explain some of our empirical performance. Yet early experiments with alternative losses, such as KL divergence between continuous densities, were not fruitful, in part because the KL divergence is insensitive to the values of its outcomes. A closer minimization of the Wasserstein metric should yield even better results than what we presented here.

In closing, we believe our results highlight the need to account for distribution in the design, theoretical or otherwise, of algorithms.

## Acknowledgements

The authors acknowledge the important role played by their colleagues at DeepMind throughout the development of this work. Special thanks to Yee Whye Teh, Alex Graves, Joel Veness, Guillaume Desjardins, Tom Schaul, David Silver, Andre Barreto, Max Jaderberg, Mohammad Azar, Georg Ostrovski, Bernardo Avila Pires, Olivier Pietquin, Audrunas Gruslys, Tom Stepleton, Aaron van den Oord; and particularly Chris Maddison for his comprehensive review of an earlier draft. Thanks also to Marek Petrik for pointers to the relevant literature, and Mark Rowland for fine-tuning details in the final version.

## Erratum

The camera-ready copy of this paper incorrectly reported a mean score of 1010% for C51. The corrected figure stands at 701%, which remains higher than the other comparable baselines. The median score remains unchanged at 178%.

The error was due to evaluation episodes in one game (Atlantis) lasting over 30 minutes; in comparison, the other results presented here cap episodes at 30 minutes, as is standard. The previously reported score on Atlantis was 3.7 million; our 30-minute score is 841,075, which we believe is close to the achievable maximum in this time frame. Capping at 30 minutes brings our human-normalized score on Atlantis from 22824% to a mere (!) 5199%, unfortunately enough to noticeably affect the mean score, whose sensitivity to outliers is well-documented.

## Appendix A Related Work

To the best of our knowledge, the work closest to ours are two papers ( Morimura et al., 2010b ; Morimura et al., 2010a ) studying the distributional Bellman equation from the perspective of its cumulative distribution functions. The authors propose both parametric and nonparametric solutions to learn distributions for risk-sensitive reinforcement learning. They also provide some theoretical analysis for the policy evaluation setting, including a consistency result in the nonparametric case. By contrast, we also analyze the control setting, and emphasize the use of the distributional equations to improve approximate reinforcement learning.

The variance of the return has been extensively studied in the risk-sensitive setting. Of note, Tamar et al. (2016) analyze the use of linear function approximation to learn this variance for policy evaluation, and Prashanth & Ghavamzadeh (2013) estimate the return variance in the design of a risk-sensitive actor-critic algorithm. Mannor & Tsitsiklis (2011) provides negative results regarding the computation of a variance-constrained solution to the optimal control problem.

The distributional formulation also arises when modelling uncertainty. Dearden et al. (1998) considered a Gaussian approximation to the value distribution, and modelled the uncertainty over the parameters of this approximation using a Normal-Gamma prior. Engel et al. (2005) leveraged the distributional Bellman equation to define a Gaussian process over the unknown value function. More recently, Geist & Pietquin (2010) proposed an alternative solution to the same problem based on unscented Kalman filters. We believe much of the analysis we provide here, which deals with the intrinsic randomness of the environment, can also be applied to modelling uncertainty.

Our work here is based on a number of foundational results, in particular concerning alternative optimality criteria. Early on, Jaquette (1973) showed that a moment optimality criterion, which imposes a total ordering on distributions, is achievable and defines a stationary optimal policy, echoing the second part of Theorem 1 . Sobel (1982) is usually cited as the first reference to Bellman equations for the higher moments (but not the distribution) of the return. Chung & Sobel (1987) provides results concerning the convergence of the distributional Bellman operator in total variation distance. White (1988) studies “nonstandard MDP criteria” from the perspective of optimizing the state-action pair occupancy.