A Distributional Perspective on Reinforcement Learning — part 1/6
Cleaned-text code points [0, 15049). Bibliography omitted here; source.md includes it.
# A Distributional Perspective on Reinforcement Learning
###### Abstract
In this paper we argue for the fundamental importance of the value distribution : the distribution of the random return received by a reinforcement learning agent. This is in contrast to the common approach to reinforcement learning which models the expectation of this return, or value . Although there is an established body of literature studying the value distribution, thus far it has always been used for a specific purpose such as implementing risk-aware behaviour. We begin with theoretical results in both the policy evaluation and control settings, exposing a significant distributional instability in the latter. We then use the distributional perspective to design a new algorithm which applies Bellman’s equation to the learning of approximate value distributions. We evaluate our algorithm using the suite of games from the Arcade Learning Environment. We obtain both state-of-the-art results and anecdotal evidence demonstrating the importance of the value distribution in approximate reinforcement learning. Finally, we combine theoretical and empirical evidence to highlight the ways in which the value distribution impacts learning in the approximate setting.
## 1 Introduction
One of the major tenets of reinforcement learning states that, when not otherwise constrained in its behaviour, an agent should aim to maximize its expected utility Q Q , or value ( Sutton & Barto, 1998 ) . Bellman’s equation succintly describes this value in terms of the expected reward and expected outcome of the random transition ( x , a ) → ( X ′ , A ′ ) (x,a)\hskip-1.00006pt\to\hskip-1.00006pt(X^{\prime},A^{\prime}) : Q ( x , a ) = 𝔼 R ( x , a ) + γ 𝔼 Q ( X ′ , A ′ ) . Q(x,a)=\expect R(x,a)+\gamma\expect Q(X^{\prime},A^{\prime}). In this paper, we aim to go beyond the notion of value and argue in favour of a distributional perspective on reinforcement learning. Specifically, the main object of our study is the random return Z Z whose expectation is the value Q Q . This random return is also described by a recursive equation, but one of a distributional nature: Z ( x , a ) = 𝐷 R ( x , a ) + γ Z ( X ′ , A ′ ) . Z(x,a)\overset{D}{=}R(x,a)+\gamma Z(X^{\prime},A^{\prime}). The distributional Bellman equation states that the distribution of Z Z is characterized by the interaction of three random variables: the reward R R , the next state-action ( X ′ , A ′ ) (X^{\prime},A^{\prime}) , and its random return Z ( X ′ , A ′ ) Z(X^{\prime},A^{\prime}) . By analogy with the well-known case, we call this quantity the value distribution .
Although the distributional perspective is almost as old as Bellman’s equation itself ( Jaquette, 1973 ; Sobel, 1982 ; White, 1988 ) , in reinforcement learning it has thus far been subordinated to specific purposes: to model parametric uncertainty ( Dearden et al., 1998 ) , to design risk-sensitive algorithms ( Morimura et al., 2010b ; Morimura et al., 2010a ) , or for theoretical analysis ( Azar et al., 2012 ; Lattimore & Hutter, 2012 ) . By contrast, we believe the value distribution has a central role to play in reinforcement learning.
Contraction of the policy evaluation Bellman operator. Basing ourselves on results by Rösler (1992) we show that, for a fixed policy, the Bellman operator over value distributions is a contraction in a maximal form of the Wasserstein (also called Kantorovich or Mallows) metric. Our particular choice of metric matters: the same operator is not a contraction in total variation, Kullback-Leibler divergence, or Kolmogorov distance.
Instability in the control setting. We will demonstrate an instability in the distributional version of Bellman’s optimality equation, in contrast to the policy evaluation case. Specifically, although the optimality operator is a contraction in expected value (matching the usual optimality result), it is not a contraction in any metric over distributions. These results provide evidence in favour of learning algorithms that model the effects of nonstationary policies.
Better approximations. From an algorithmic standpoint, there are many benefits to learning an approximate distribution rather than its approximate expectation. The distributional Bellman operator preserves multimodality in value distributions, which we believe leads to more stable learning. Approximating the full distribution also mitigates the effects of learning from a nonstationary policy. As a whole, we argue that this approach makes approximate reinforcement learning significantly better behaved.
We will illustrate the practical benefits of the distributional perspective in the context of the Arcade Learning Environment ( Bellemare et al., 2013 ) . By modelling the value distribution within a DQN agent ( Mnih et al., 2015 ) , we obtain considerably increased performance across the gamut of benchmark Atari 2600 games, and in fact achieve state-of-the-art performance on a number of games. Our results echo those of Veness et al. (2015) , who obtained extremely fast learning by predicting Monte Carlo returns.
From a supervised learning perspective, learning the full value distribution might seem obvious: why restrict ourselves to the mean? The main distinction, of course, is that in our setting there are no given targets. Instead, we use Bellman’s equation to make the learning process tractable; we must, as Sutton & Barto (1998) put it, “learn a guess from a guess”. It is our belief that this guesswork ultimately carries more benefits than costs.
## 2 Setting
We consider an agent interacting with an environment in the standard fashion: at each step, the agent selects an action based on its current state, to which the environment responds with a reward and the next state. We model this interaction as a time-homogeneous Markov Decision Process ( 𝒳 , 𝒜 , R , P , γ ) (\mathcal{X},\mathcal{A},R,P,\gamma) . As usual, 𝒳 \mathcal{X} and 𝒜 \mathcal{A} are respectively the state and action spaces, P P is the transition kernel P ( ⋅ | x , a ) P(\cdot\,|\,x,a) , γ ∈ [ 0 , 1 ] \gamma\in[0,1] is the discount factor, and R R is the reward function, which in this work we explicitly treat as a random variable. A stationary policy π \pi maps each state x ∈ 𝒳 x\in\mathcal{X} to a probability distribution over the action space 𝒜 \mathcal{A} .
### 2.1 Bellman’s Equations
The return Z π Z^{\pi} is the sum of discounted rewards along the agent’s trajectory of interactions with the environment. The value function Q π Q^{\pi} of a policy π \pi describes the expected return from taking action a ∈ 𝒜 a\in\mathcal{A} from state x ∈ 𝒳 x\in\mathcal{X} , then acting according to π \pi : Q π ( x , a ) \displaystyle Q^{\pi}(x,a) : = 𝔼 Z π ( x , a ) = 𝔼 [ ∑ t = 0 ∞ γ t R ( x t , a t ) ] , \displaystyle:=\expect Z^{\pi}(x,a)=\expect\left[\sum_{t=0}^{\infty}\gamma^{t}R(x_{t},a_{t})\right], (1) x t ∼ P ( ⋅ CLOSE \displaystyle x_{t}\sim P(\cdot | x t − 1 , a t − 1 ) , a t ∼ π ( ⋅ | x t ) , x 0 = x , a 0 = a . \displaystyle\,|\,x_{t-1},a_{t-1}),a_{t}\sim\pi(\cdot\,|\,x_{t}),x_{0}=x,a_{0}=a. Fundamental to reinforcement learning is the use of Bellman’s equation ( Bellman, 1957 ) to describe the value function: Q π ( x , a ) = 𝔼 R ( x , a ) + γ 𝔼 P , π Q π ( x ′ , a ′ ) . Q^{\pi}(x,a)=\expect R(x,a)+\gamma\expect_{P,\pi}Q^{\pi}(x^{\prime},a^{\prime}). In reinforcement learning we are typically interested in acting so as to maximize the return. The most common approach for doing so involves the optimality equation Q ∗ ( x , a ) = 𝔼 R ( x , a ) + γ 𝔼 P max a ′ ∈ 𝒜 Q ∗ ( x ′ , a ′ ) . Q^{*}(x,a)=\expect R(x,a)+\gamma\expect\nolimits_{P}\max_{a^{\prime}\in\mathcal{A}}Q^{*}(x^{\prime},a^{\prime}). This equation has a unique fixed point Q ∗ Q^{*} , the optimal value function, corresponding to the set of optimal policies Π ∗ \Pi^{*} ( π ∗ \pi^{*} is optimal if 𝔼 a ∼ π ∗ Q ∗ ( x , a ) = max a Q ∗ ( x , a ) \expect_{a\sim\pi^{*}}Q^{*}(x,a)=\max_{a}Q^{*}(x,a) ).
We view value functions as vectors in ℝ 𝒳 × 𝒜 \mathbb{R}^{\mathcal{X}\times\mathcal{A}} , and the expected reward function as one such vector. In this context, the Bellman operator 𝒯 π \mathcal{T}^{\pi} and optimality operator 𝒯 \mathcal{T} are 𝒯 π Q ( x , a ) \displaystyle\mathcal{T}^{\pi}Q(x,a) : = 𝔼 R ( x , a ) + γ 𝔼 P , π Q ( x ′ , a ′ ) \displaystyle:=\expect R(x,a)+\gamma\expect_{P,\pi}Q(x^{\prime},a^{\prime}) (2) 𝒯 Q ( x , a ) \displaystyle\mathcal{T}Q(x,a) : = 𝔼 R ( x , a ) + γ 𝔼 P max a ′ ∈ 𝒜 Q ( x ′ , a ′ ) . \displaystyle:=\expect R(x,a)+\gamma\expect\nolimits_{P}\max_{a^{\prime}\in\mathcal{A}}Q(x^{\prime},a^{\prime}). (3) These operators are useful as they describe the expected behaviour of popular learning algorithms such as SARSA and Q-Learning. In particular they are both contraction mappings, and their repeated application to some initial Q 0 Q_{0} converges exponentially to Q π Q^{\pi} or Q ∗ Q^{*} , respectively ( Bertsekas & Tsitsiklis, 1996 ) .
## 3 The Distributional Bellman Operators
In this paper we take away the expectations inside Bellman’s equations and consider instead the full distribution of the random variable Z π Z^{\pi} . From here on, we will view Z π Z^{\pi} as a mapping from state-action pairs to distributions over returns, and call it the value distribution .
Our first aim is to gain an understanding of the theoretical behaviour of the distributional analogues of the Bellman operators, in particular in the less well-understood control setting. The reader strictly interested in the algorithmic contribution may choose to skip this section.
### 3.1 Distributional Equations
It will sometimes be convenient to make use of the probability space ( Ω , ℱ , Pr ) (\Omega,\mathcal{F},\Pr) . The reader unfamiliar with measure theory may think of Ω \Omega as the space of all possible outcomes of an experiment ( Billingsley, 1995 ) . We will write ‖ 𝐮 ‖ p \|\mathbf{u}\|_{p} to denote the L p L_{p} norm of a vector 𝐮 ∈ ℝ 𝒳 \mathbf{u}\in\mathbb{R}^{\mathcal{X}} for 1 ≤ p ≤ ∞ 1\leq p\leq\infty ; the same applies to vectors in ℝ 𝒳 × 𝒜 \mathbb{R}^{\mathcal{X}\times\mathcal{A}} . The L p L_{p} norm of a random vector U : Ω → ℝ 𝒳 U:\Omega\to\mathbb{R}^{\mathcal{X}} (or ℝ 𝒳 × 𝒜 \mathbb{R}^{\mathcal{X}\times\mathcal{A}} ) is then ‖ U ‖ p := [ 𝔼 [ ‖ U ( ω ) ‖ p p ] ] 1 / p \|U\|_{p}:=\left[\expect\left[\|U(\omega)\|_{p}^{p}\right]\right]^{1/p} , and for p = ∞ p=\infty we have ‖ U ‖ ∞ = ess sup ‖ U ( ω ) ‖ ∞ \|U\|_{\infty}=\textrm{ess}\sup{\|U(\omega)\|_{\infty}} (we will omit the dependency on ω ∈ Ω \omega\in\Omega whenever unambiguous). We will denote the c.d.f. of a random variable U U by F U ( y ) := Pr { U ≤ y } F_{U}(y):=\Pr\{U\leq y\} , and its inverse c.d.f. by F U − 1 ( q ) := inf { y : F U ( y ) ≥ q } F^{-1}_{U}(q):=\inf\{y:F_{U}(y)\geq q\} .
A distributional equation U : = 𝐷 V U\overset{D}{:=}V indicates that the random variable U U is distributed according to the same law as V V . Without loss of generality, the reader can understand the two sides of a distributional equation as relating the distributions of two independent random variables. Distributional equations have been used in reinforcement learning by Engel et al. (2005) ; Morimura et al. (2010a) among others, and in operations research by White (1988) .
### 3.2 The Wasserstein Metric
The main tool for our analysis is the Wasserstein metric d p d_{p} between cumulative distribution functions (see e.g. Bickel & Freedman, 1981 , where it is called the Mallows metric) . For F F , G G two c.d.fs over the reals, it is defined as d p ( F , G ) := inf U , V ‖ U − V ‖ p , d_{p}(F,G):=\inf_{U,V}\|U-V\|_{p}, where the infimum is taken over all pairs of random variables ( U , V ) (U,V) with respective cumulative distributions F F and G G . The infimum is attained by the inverse c.d.f. transform of a random variable 𝒰 \mathcal{U} uniformly distributed on [ 0 , 1 ] [0,1] : d p ( F , G ) = ‖ F − 1 ( 𝒰 ) − G − 1 ( 𝒰 ) ‖ p . d_{p}(F,G)=\|F^{-1}(\mathcal{U})-G^{-1}(\mathcal{U})\|_{p}. For p < ∞ p<\infty this is more explicitly written as d p ( F , G ) = ( ∫ 0 1 | F − 1 ( u ) − G − 1 ( u ) | p 𝑑 u ) 1 / p . d_{p}(F,G)=\left(\int_{0}^{1}\big|F^{-1}(u)-G^{-1}(u)\big|^{p}du\right)^{1/p}. Given two random variables U U , V V with c.d.fs F U F_{U} , F V F_{V} , we will write d p ( U , V ) := d p ( F U , F V ) d_{p}(U,V):=d_{p}(F_{U},F_{V}) . We will find it convenient to conflate the random variables under consideration with their versions under the inf \inf , writing d p ( U , V ) = inf U , V ‖ U − V ‖ p . d_{p}(U,V)=\inf_{U,V}\|U-V\|_{p}. whenever unambiguous; we believe the greater legibility justifies the technical inaccuracy. Finally, we extend this metric to vectors of random variables, such as value distributions, using the corresponding L p L_{p} norm.
Consider a scalar a a and a random variable A A independent of U , V U,V . The metric d p d_{p} has the following properties: d p ( a U , a V ) \displaystyle d_{p}(aU,aV) ≤ | a | d p ( U , V ) \displaystyle\leq|a|d_{p}(U,V) (P1) d p ( A + U , A + V ) \displaystyle d_{p}(A+U,A+V) ≤ d p ( U , V ) \displaystyle\leq d_{p}(U,V) (P2) d p ( A U , A V ) \displaystyle d_{p}(AU,AV) ≤ ‖ A ‖ p d p ( U , V ) . \displaystyle\leq\|A\|_{p}d_{p}(U,V). (P3) We will need the following additional property, which makes no independence assumptions on its variables. Its proof, and that of later results, is given in the appendix.
###### Lemma 1 (Partition lemma) .
Let A 1 , A 2 , … A_{1},A_{2},\dots be a set of random variables describing a partition of Ω \Omega , i.e. A i ( ω ) ∈ { 0 , 1 } A_{i}(\omega)\in\{0,1\} and for any ω \omega there is exactly one A i A_{i} with A i ( ω ) = 1 A_{i}(\omega)=1 . Let U , V U,V be two random variables. Then d p ( U , V ) ≤ ∑ i d p ( A i U , A i V ) . d_{p}\big(U,V\big)\leq\sum\nolimits_{i}d_{p}(A_{i}U,A_{i}V).
Let 𝒵 \mathcal{Z} denote the space of value distributions with bounded moments. For two value distributions Z 1 , Z 2 ∈ 𝒵 Z_{1},Z_{2}\in\mathcal{Z} we will make use of a maximal form of the Wasserstein metric: d ¯ p ( Z 1 , Z 2 ) := sup x , a d p ( Z 1 ( x , a ) , Z 2 ( x , a ) ) . \bar{d}_{p}(Z_{1},Z_{2}):=\sup_{x,a}d_{p}(Z_{1}(x,a),Z_{2}(x,a)). We will use d ¯ p \bar{d}_{p} to establish the convergence of the distributional Bellman operators.
###### Lemma 2 .
d ¯ p \bar{d}_{p} is a metric over value distributions.
### 3.3 Policy Evaluation
In the policy evaluation setting ( Sutton & Barto, 1998 ) we are interested in the value function V π V^{\pi} associated with a given policy π \pi . The analogue here is the value distribution Z π Z^{\pi} . In this section we characterize Z π Z^{\pi} and study the behaviour of the policy evaluation operator 𝒯 π \mathcal{T}^{\pi} . We emphasize that Z π Z^{\pi} describes the intrinsic randomness of the agent’s interactions with its environment, rather than some measure of uncertainty about the environment itself.