A Distributional Perspective on Reinforcement Learning — part 2/6

Cleaned-text code points [15049, 29161). Bibliography omitted here; source.md includes it.

We view the reward function as a random vector R ∈ 𝒵 R\in\mathcal{Z} , and define the transition operator P π : 𝒵 → 𝒵 P^{\pi}:\mathcal{Z}\to\mathcal{Z} P π ​ Z ​ ( x , a ) \displaystyle P^{\pi}Z(x,a) : ⁣ = 𝐷 Z ( X ′ , A ′ ) \displaystyle\overset{D}{:=}Z(X^{\prime},A^{\prime}) (4) X ′ \displaystyle X^{\prime} ∼ P ( ⋅ | x , a ) , A ′ ∼ π ( ⋅ | X ′ ) , \displaystyle\sim P(\cdot\,|\,x,a),\,A^{\prime}\sim\pi(\cdot\,|\,X^{\prime}), where we use capital letters to emphasize the random nature of the next state-action pair ( X ′ , A ′ ) (X^{\prime},A^{\prime}) . We define the distributional Bellman operator 𝒯 π : 𝒵 → 𝒵 \mathcal{T}^{\pi}:\mathcal{Z}\to\mathcal{Z} as 𝒯 π Z ( x , a ) : ⁣ = 𝐷 R ( x , a ) + γ P π Z ( x , a ) . \mathcal{T}^{\pi}Z(x,a)\overset{D}{:=}R(x,a)+\gamma P^{\pi}Z(x,a). (5) While 𝒯 π \mathcal{T}^{\pi} bears a surface resemblance to the usual Bellman operator ( 2 ), it is fundamentally different. In particular, three sources of randomness define the compound distribution 𝒯 π ​ Z \mathcal{T}^{\pi}Z : a) The randomness in the reward R R ,

b) The randomness in the transition P π P^{\pi} , and

c) The next-state value distribution Z ⁡ ( X ′ , A ′ ) Z(X^{\prime},A^{\prime}) .

In particular, we make the usual assumption that these three quantities are independent. In this section we will show that ( 5 ) is a contraction mapping whose unique fixed point is the random return Z π Z^{\pi} .

#### 3.3.1 Contraction in d ¯ p \bar{d}_{p}

Consider the process Z k + 1 := 𝒯 π ​ Z k Z_{k+1}:=\mathcal{T}^{\pi}Z_{k} , starting with some Z 0 ∈ 𝒵 Z_{0}\in\mathcal{Z} . We may expect the limiting expectation of { Z k } \{Z_{k}\} to converge exponentially quickly, as usual, to Q π Q^{\pi} . As we now show, the process converges in a stronger sense: 𝒯 π \mathcal{T}^{\pi} is a contraction in d ¯ p \bar{d}_{p} , which implies that all moments also converge exponentially quickly.

###### Lemma 3 .

𝒯 π : 𝒵 → 𝒵 \mathcal{T}^{\pi}:\mathcal{Z}\to\mathcal{Z} is a γ \gamma -contraction in d ¯ p \bar{d}_{p} .

Using Lemma 3 , we conclude using Banach’s fixed point theorem that 𝒯 π \mathcal{T}^{\pi} has a unique fixed point. By inspection, this fixed point must be Z π Z^{\pi} as defined in ( 1 ). As we assume all moments are bounded, this is sufficient to conclude that the sequence { Z k } \{Z_{k}\} converges to Z π Z^{\pi} in d ¯ p \bar{d}_{p} for 1 ≤ p ≤ ∞ 1\leq p\leq\infty .

To conclude, we remark that not all distributional metrics are equal; for example, Chung & Sobel (1987) have shown that 𝒯 π \mathcal{T}^{\pi} is not a contraction in total variation distance. Similar results can be derived for the Kullback-Leibler divergence and the Kolmogorov distance.

#### 3.3.2 Contraction in Centered Moments

Observe that d 2 ​ ( U , V ) d_{2}(U,V) (and more generally, d p d_{p} ) relates to a coupling C ⁡ ( ω ) := U ⁡ ( ω ) − V ⁡ ( ω ) C(\omega):=U(\omega)-V(\omega) , in the sense that d 2 2 ​ ( U , V ) \displaystyle d_{2}^{2}(U,V) ≤ 𝔼 [ ( U − V ) 2 ] = 𝕍 ⁡ ( C ) + ( 𝔼 C ) 2 . \displaystyle\leq\expect[(U-V)^{2}]=\mathbb{V}\big(C\big)+\big(\expect C\big)^{2}. As a result, we cannot directly use d 2 d_{2} to bound the variance difference | 𝕍 ⁡ ( 𝒯 π ​ Z ​ ( x , a ) ) − 𝕍 ⁡ ( Z π ​ ( x , a ) ) | |\mathbb{V}(\mathcal{T}^{\pi}Z(x,a))-\mathbb{V}(Z^{\pi}(x,a))| . However, 𝒯 π \mathcal{T}^{\pi} is in fact a contraction in variance ( Sobel, 1982 , see also appendix) . In general, 𝒯 π \mathcal{T}^{\pi} is not a contraction in the p t ​ h p^{th} centered moment, p > 2 p>2 , but the centered moments of the iterates { Z k } \{Z_{k}\} still converge exponentially quickly to those of Z π Z^{\pi} ; the proof extends the result of Rösler (1992) .

### 3.4 Control

Thus far we have considered a fixed policy π \pi , and studied the behaviour of its associated operator 𝒯 π \mathcal{T}^{\pi} . We now set out to understand the distributional operators of the control setting – where we seek a policy π \pi that maximizes value – and the corresponding notion of an optimal value distribution. As with the optimal value function, this notion is intimately tied to that of an optimal policy. However, while all optimal policies attain the same value Q ∗ Q^{*} , in our case a difficulty arises: in general there are many optimal value distributions.

In this section we show that the distributional analogue of the Bellman optimality operator converges, in a weak sense, to the set of optimal value distributions. However, this operator is not a contraction in any metric between distributions, and is in general much more temperamental than the policy evaluation operators. We believe the convergence issues we outline here are a symptom of the inherent instability of greedy updates, as highlighted by e.g. Tsitsiklis (2002) and most recently Harutyunyan et al. (2016) .

Let Π ∗ \Pi^{*} be the set of optimal policies. We begin by characterizing what we mean by an optimal value distribution .

###### Definition 1 (Optimal value distribution) .

An optimal value distribution is the v.d. of an optimal policy. The set of optimal value distributions is 𝒵 ∗ := { Z π ∗ : π ∗ ∈ Π ∗ } \mathcal{Z}^{*}:=\{Z^{\pi^{*}}:\pi^{*}\in\Pi^{*}\} .

We emphasize that not all value distributions with expectation Q ∗ Q^{*} are optimal: they must match the full distribution of the return under some optimal policy.

###### Definition 2 .

A greedy policy π \pi for Z ∈ 𝒵 Z\in\mathcal{Z} maximizes the expectation of Z Z . The set of greedy policies for Z Z is 𝒢 Z := { π : ∑ a π ( a | x ) 𝔼 Z ( x , a ) = max a ′ ∈ 𝒜 𝔼 Z ( x , a ′ ) } . \mathcal{G}_{Z}:=\{\pi:\sum\nolimits_{a}\pi(a\,|\,x)\expect Z(x,a)=\max_{a^{\prime}\in\mathcal{A}}\expect Z(x,a^{\prime})\}.

Recall that the expected Bellman optimality operator 𝒯 \mathcal{T} is 𝒯 ​ Q ​ ( x , a ) = 𝔼 R ​ ( x , a ) + γ ​ 𝔼 P ​ max a ′ ∈ 𝒜 ⁡ Q ⁡ ( x ′ , a ′ ) . \mathcal{T}Q(x,a)=\expect R(x,a)+\gamma\expect\nolimits_{P}\max_{a^{\prime}\in\mathcal{A}}Q(x^{\prime},a^{\prime}). (6) The maximization at x ′ x^{\prime} corresponds to some greedy policy. Although this policy is implicit in ( 6 ), we cannot ignore it in the distributional setting. We will call a distributional Bellman optimality operator any operator 𝒯 \mathcal{T} which implements a greedy selection rule, i.e. 𝒯 ​ Z = 𝒯 π ​ Z ​ for some ​ π ∈ 𝒢 Z . \mathcal{T}Z=\mathcal{T}^{\pi}Z\text{ for some }\pi\in\mathcal{G}_{Z}. As in the policy evaluation setting, we are interested in the behaviour of the iterates Z k + 1 := 𝒯 ​ Z k Z_{k+1}:=\mathcal{T}Z_{k} , Z 0 ∈ 𝒵 Z_{0}\in\mathcal{Z} . Our first result is to assert that 𝔼 Z k \expect Z_{k} behaves as expected.

###### Lemma 4 .

Let Z 1 , Z 2 ∈ 𝒵 Z_{1},Z_{2}\in\mathcal{Z} . Then ‖ 𝔼 𝒯 Z 1 − 𝔼 𝒯 Z 2 ‖ ∞ ≤ γ ‖ 𝔼 Z 1 − 𝔼 Z 2 ‖ ∞ , \left\|\expect\mathcal{T}Z_{1}-\expect\mathcal{T}Z_{2}\right\|_{\infty}\leq\gamma\left\|\expect Z_{1}-\expect Z_{2}\right\|_{\infty}, and in particular 𝔼 Z k → Q ∗ \expect Z_{k}\to Q^{*} exponentially quickly.

By inspecting Lemma 4 , we might expect that Z k Z_{k} converges quickly in d ¯ p \bar{d}_{p} to some fixed point in 𝒵 ∗ \mathcal{Z}^{*} . Unfortunately, convergence is neither quick nor assured to reach a fixed point. In fact, the best we can hope for is pointwise convergence, not even to the set 𝒵 ∗ \mathcal{Z}^{*} but to the larger set of nonstationary optimal value distributions .

###### Definition 3 .

A nonstationary optimal value distribution Z ∗ ⁣ ∗ Z^{**} is the value distribution corresponding to a sequence of optimal policies. The set of n.o.v.d. is 𝒵 ∗ ⁣ ∗ \mathcal{Z}^{**} .

###### Theorem 1 (Convergence in the control setting) .

Let 𝒳 \mathcal{X} be measurable and suppose that 𝒜 \mathcal{A} is finite. Then lim k → ∞ inf Z ∗ ⁣ ∗ ∈ 𝒵 ∗ ⁣ ∗ d p ​ ( Z k ​ ( x , a ) , Z ∗ ⁣ ∗ ​ ( x , a ) ) = 0 ∀ x , a . \lim_{k\to\infty}\inf_{Z^{**}\in\mathcal{Z}^{**}}d_{p}(Z_{k}(x,a),Z^{**}(x,a))=0\quad\forall x,a. If 𝒳 \mathcal{X} is finite, then Z k Z_{k} converges to 𝒵 ∗ ⁣ ∗ \mathcal{Z}^{**} uniformly. Furthermore, if there is a total ordering ≺ \prec on Π ∗ \Pi^{*} , such that for any Z ∗ ∈ 𝒵 ∗ Z^{*}\in\mathcal{Z}^{*} , 𝒯 ​ Z ∗ = 𝒯 π ​ Z ∗ ​ with ​ π ∈ 𝒢 Z ∗ , π ≺ π ′ ​ ∀ π ′ ∈ 𝒢 Z ∗ ∖ { π } . \mathcal{T}Z^{*}=\mathcal{T}^{\pi}Z^{*}\;\text{with}\;\pi\in\mathcal{G}_{Z^{*}},\;\pi\prec\pi^{\prime}\;\;\forall\pi^{\prime}\in\mathcal{G}_{Z^{*}}\setminus\{\pi\}. Then 𝒯 \mathcal{T} has a unique fixed point Z ∗ ∈ 𝒵 ∗ Z^{*}\in\mathcal{Z}^{*} .

Comparing Theorem 1 to Lemma 4 reveals a significant difference between the distributional framework and the usual setting of expected return. While the mean of Z k Z_{k} converges exponentially quickly to Q ∗ Q^{*} , its distribution need not be as well-behaved! To emphasize this difference, we now provide a number of negative results concerning 𝒯 \mathcal{T} .

###### Proposition 1 .

The operator 𝒯 \mathcal{T} is not a contraction.

Consider the following example (Figure 2 , left). There are two states, x 1 x_{1} and x 2 x_{2} ; a unique transition from x 1 x_{1} to x 2 x_{2} ; from x 2 x_{2} , action a 1 a_{1} yields no reward, while the optimal action a 2 a_{2} yields 1 + ϵ 1+\epsilon or − 1 + ϵ -1+\epsilon with equal probability. Both actions are terminal. There is a unique optimal policy and therefore a unique fixed point Z ∗ Z^{*} . Now consider Z Z as given in Figure 2 (right), and its distance to Z ∗ Z^{*} : d ¯ 1 ​ ( Z , Z ∗ ) = d 1 ​ ( Z ⁡ ( x 2 , a 2 ) , Z ∗ ​ ( x 2 , a 2 ) ) = 2 ​ ϵ , \bar{d}_{1}(Z,Z^{*})=d_{1}(Z(x_{2},a_{2}),Z^{*}(x_{2},a_{2}))=2\epsilon,\\ where we made use of the fact that Z = Z ∗ Z=Z^{*} everywhere except at ( x 2 , a 2 ) (x_{2},a_{2}) . When we apply 𝒯 \mathcal{T} to Z Z , however, the greedy action a 1 a_{1} is selected and 𝒯 ​ Z ​ ( x 1 ) = Z ⁡ ( x 2 , a 1 ) \mathcal{T}Z(x_{1})=Z(x_{2},a_{1}) . But d 1 ​ ( 𝒯 ​ Z , 𝒯 ​ Z ∗ ) \displaystyle d_{1}(\mathcal{T}Z,\mathcal{T}Z^{*}) = d 1 ​ ( 𝒯 ​ Z ​ ( x 1 ) , Z ∗ ​ ( x 1 ) ) \displaystyle=d_{1}(\mathcal{T}Z(x_{1}),Z^{*}(x_{1})) = 1 2 ​ | 1 − ϵ | + 1 2 ​ | 1 + ϵ | > 2 ​ ϵ \displaystyle=\tfrac{1}{2}|1-\epsilon|+\tfrac{1}{2}|1+\epsilon|>2\epsilon for a sufficiently small ϵ \epsilon . This shows that the undiscounted update is not a nonexpansion: d ¯ 1 ​ ( 𝒯 ​ Z , 𝒯 ​ Z ∗ ) > d ¯ 1 ​ ( Z , Z ∗ ) \bar{d}_{1}(\mathcal{T}Z,\mathcal{T}Z^{*})>\bar{d}_{1}(Z,Z^{*}) . With γ < 1 \gamma<1 , the same proof shows it is not a contraction. Using a more technically involved argument, we can extend this result to any metric which separates Z Z and 𝒯 ​ Z \mathcal{T}Z .

x 1 x 2 , a 1 x 2 , a 2 Z ∗ ϵ ± 1 0 ϵ ± 1 Z ϵ ± 1 0 − ϵ ± 1 𝒯 ​ Z 0 0 ϵ ± 1 \begin{array}[c]{lccc}\hfil\hfil\lx@intercol&\lx@intercol\hfil x_{1}\hfil\lx@intercol&\lx@intercol\hfil x_{2},a_{1}\hfil\lx@intercol&\lx@intercol\hfil x_{2},a_{2}\hfil\lx@intercol\\ \hline\cr Z^{*}&\pagecolor{lightsunflower}\epsilon\pm 1&0&\pagecolor{lightsunflower}\epsilon\pm 1\\ \hline\cr\hline\cr Z&\epsilon\pm 1&0&\pagecolor{lightsunflower}-\epsilon\pm 1\\ \hline\cr\mathcal{T}Z&\pagecolor{lightsunflower}{0}&0&\epsilon\pm 1\\ \hline\cr\hline\cr\end{array}

###### Proposition 2 .

Not all optimality operators have a fixed point Z ∗ = 𝒯 ​ Z ∗ Z^{*}=\mathcal{T}Z^{*} .

To see this, consider the same example, now with ϵ = 0 \epsilon=0 , and a greedy operator 𝒯 \mathcal{T} which breaks ties by picking a 2 a_{2} if Z ⁡ ( x 1 ) = 0 Z(x_{1})=0 , and a 1 a_{1} otherwise. Then the sequence 𝒯 ​ Z ∗ ​ ( x 1 ) , ( 𝒯 ) 2 ​ Z ∗ ​ ( x 1 ) , … \mathcal{T}Z^{*}(x_{1}),(\mathcal{T})^{2}Z^{*}(x_{1}),\,\dots alternates between Z ∗ ​ ( x 2 , a 1 ) Z^{*}(x_{2},a_{1}) and Z ∗ ​ ( x 2 , a 2 ) Z^{*}(x_{2},a_{2}) .

###### Proposition 3 .

That 𝒯 \mathcal{T} has a fixed point Z ∗ = 𝒯 ​ Z ∗ Z^{*}=\mathcal{T}Z^{*} is insufficient to guarantee the convergence of { Z k } \{Z_{k}\} to 𝒵 ∗ \mathcal{Z}^{*} .

Theorem 1 paints a rather bleak picture of the control setting. It remains to be seen whether the dynamical eccentricies highlighted here actually arise in practice. One open question is whether theoretically more stable behaviour can be derived using stochastic policies, for example from conservative policy iteration ( Kakade & Langford, 2002 ) .

## 4 Approximate Distributional Learning

In this section we propose an algorithm based on the distributional Bellman optimality operator. In particular, this will require choosing an approximating distribution. Although the Gaussian case has previously been considered ( Morimura et al., 2010a ; Tamar et al., 2016 ) , to the best of our knowledge we are the first to use a rich class of parametric distributions.

### 4.1 Parametric Distribution

We will model the value distribution using a discrete distribution parametrized by N ∈ ℕ N\in\mathbb{N} and V min , V max ∈ ℝ V_{\textsc{min}},V_{\textsc{max}}\in\mathbb{R} , and whose support is the set of atoms { z i = V min + i ​ △ ​ z : 0 ≤ i < N } \{z_{i}=V_{\textsc{min}}+i\triangle z:0\leq i<N\} , △ ​ z := V max − V min N − 1 \triangle z:=\tfrac{V_{\textsc{max}}-V_{\textsc{min}}}{N-1} . In a sense, these atoms are the “canonical returns” of our distribution. The atom probabilities are given by a parametric model θ : 𝒳 × 𝒜 → ℝ N \theta:\mathcal{X}\times\mathcal{A}\to\mathbb{R}^{N} Z θ ​ ( x , a ) = z i w.p. ​ p i ​ ( x , a ) := e θ i ​ ( x , a ) ∑ j e θ j ​ ( x , a ) . Z_{\theta}(x,a)=z_{i}\quad\text{w.p. }\;p_{i}(x,a):=\frac{e^{\theta_{i}(x,a)}}{\sum\nolimits_{j}e^{\theta_{j}(x,a)}}. The discrete distribution has the advantages of being highly expressive and computationally friendly ( Van den Oord et al., 2016 , see e.g.) .

### 4.2 Projected Bellman Update

Using a discrete distribution poses a problem: the Bellman update 𝒯 ​ Z θ \mathcal{T}Z_{\theta} and our parametrization Z θ Z_{\theta} almost always have disjoint supports. From the analysis of Section 3 it would seem natural to minimize the Wasserstein metric (viewed as a loss) between 𝒯 ​ Z θ \mathcal{T}Z_{\theta} and Z θ Z_{\theta} , which is also conveniently robust to discrepancies in support. However, a second issue prevents this: in practice we are typically restricted to learning from sample transitions, which is not possible under the Wasserstein loss (see Prop. 5 and toy results in the appendix).