A Distributional Perspective on Reinforcement Learning — part 5/6
Cleaned-text code points [61045, 77024). Bibliography omitted here; source.md includes it.
Because 𝒜 \mathcal{A} is finite, the gap Δ ( x ) := Q ∗ ( x , π ∗ ( x ) ) − max a ≠ π ∗ ( x ) Q ∗ ( x , a ) \Delta(x):=Q^{*}(x,\pi^{*}(x))-\max_{a\neq\pi^{*}(x)}Q^{*}(x,a) is attained for some strictly positive Δ ( x ) > 0 \Delta(x)>0 . By definition, there exists a k k such that ϵ k = γ k B < Δ ( x ) 2 , \epsilon_{k}=\gamma^{k}B<\frac{\Delta(x)}{2}, and hence every x ∈ 𝒳 x\in\mathcal{X} must eventually be in 𝒳 k \mathcal{X}_{k} . ∎
This lemma allows us to guarantee the existence of an iteration k k after which sufficiently many states are well-behaved, in the sense that the greedy policy at those states chooses the optimal action. We will call these states “solved”. We in fact require not only these states to be solved, but also most of their successors, and most of the successors of those, and so on. We formalize this notion as follows: fix some δ > 0 \delta>0 , let 𝒳 k , 0 := 𝒳 k \mathcal{X}_{k,0}:=\mathcal{X}_{k} , and define for i > 0 i>0 the set 𝒳 k , i := { x : x ∈ 𝒳 k , P ( 𝒳 k − 1 , i − 1 | x , π ∗ ( x ) ) ≥ 1 − δ } , \mathcal{X}_{k,i}:=\big\{x:x\in\mathcal{X}_{k},P(\mathcal{X}_{k-1,i-1}\,|\,x,\pi^{*}(x))\geq 1-\delta\big\}, As the following lemma shows, any x x is eventually contained in the recursively-defined sets 𝒳 k , i \mathcal{X}_{k,i} , for any i i .
###### Lemma 6 .
For any i ∈ ℕ i\in\mathbb{N} and any x ∈ 𝒳 x\in\mathcal{X} , there exists a k k such that for all k ′ ≥ k k^{\prime}\geq k , x ∈ 𝒳 k ′ , i x\in\mathcal{X}_{k^{\prime},i} .
###### Proof.
Fix i i and let us suppose that 𝒳 k , i ↑ 𝒳 \mathcal{X}_{k,i}\uparrow\mathcal{X} . By Lemma 5 , this is true for i = 0 i=0 . We infer that for any probability measure P P on 𝒳 \mathcal{X} , P ( 𝒳 k , i ) → P ( 𝒳 ) = 1 P(\mathcal{X}_{k,i})\to P(\mathcal{X})=1 . In particular, for a given x ∈ 𝒳 k x\in\mathcal{X}_{k} , this implies that P ( 𝒳 k , i | x , π ∗ ( x ) ) → P ( 𝒳 | x , π ∗ ( x ) ) = 1 . P(\mathcal{X}_{k,i}\,|\,x,\pi^{*}(x))\to P(\mathcal{X}\,|\,x,\pi^{*}(x))=1. Therefore, for any x x , there exists a time after which it is and remains a member of 𝒳 k , i + 1 \mathcal{X}_{k,i+1} , the set of states for which P ( 𝒳 k − 1 , i | x , π ∗ ( x ) ) ≥ 1 − δ P(\mathcal{X}_{k-1,i}\,|\,x,\pi^{*}(x))\geq 1-\delta . We conclude that 𝒳 k , i + 1 ↑ 𝒳 \mathcal{X}_{k,i+1}\uparrow\mathcal{X} also. The statement follows by induction. ∎
###### Proof of Theorem 1 .
The proof is similar to policy iteration-type results, but requires more care in dealing with the metric and the possibly infinite state space. We will write W k ( x ) := Z k ( x , π k ( x ) ) W_{k}(x):=Z_{k}(x,\pi_{k}(x)) , define W ∗ W^{*} similarly and with some overload of notation write 𝒯 W k ( x ) := W k + 1 ( x ) = 𝒯 Z k ( x , π k + 1 ( x ) ) \mathcal{T}W_{k}(x):=W_{k+1}(x)=\mathcal{T}Z_{k}(x,\pi_{k+1}(x)) . Finally, let S i k ( x ) := 𝕀 [ x ∈ 𝒳 k , i ] S^{k}_{i}(x):=\mathbb{I}\left[x\in\mathcal{X}_{k,i}\right] and S ¯ i k ( x ) = 1 − S i k ( x ) \bar{S}^{k}_{i}(x)=1-S^{k}_{i}(x) .
Fix i > 0 i>0 and x ∈ 𝒳 k + 1 , i + 1 ⊆ 𝒳 k x\in\mathcal{X}_{k+1,i+1}\subseteq\mathcal{X}_{k} . We begin by using Lemma 1 to separate the transition from x x into a solved term and an unsolved term: P π k W k ( x ) = S i k W k ( X ′ ) + S ¯ i k W k ( X ′ ) , P^{\pi_{k}}W_{k}(x)=S^{k}_{i}W_{k}(X^{\prime})+\bar{S}^{k}_{i}W_{k}(X^{\prime}), where X ′ X^{\prime} is the random successor from taking action π k ( x ) := π ∗ ( x ) \pi_{k}(x):=\pi^{*}(x) , and we write S i k = S i k ( X ′ ) , S ¯ i k = S ¯ i k ( X ′ ) S^{k}_{i}=S^{k}_{i}(X^{\prime}),\bar{S}^{k}_{i}=\bar{S}^{k}_{i}(X^{\prime}) to ease the notation. Similarly, P π k W ∗ ( x ) = S i k W ∗ ( X ′ ) + S ¯ i k W ∗ ( X ′ ) . P^{\pi_{k}}W^{*}(x)=S^{k}_{i}W^{*}(X^{\prime})+\bar{S}^{k}_{i}W^{*}(X^{\prime}). Now d p ( W k + 1 ( x ) , W ∗ ( x ) ) \displaystyle d_{p}(W_{k+1}(x),W^{*}(x)) = d p ( 𝒯 W k ( x ) , 𝒯 W ∗ ( x ) ) \displaystyle=d_{p}(\mathcal{T}W_{k}(x),\mathcal{T}W^{*}(x)) ≤ ( a ) γ d p ( P π k W k ( x ) , P π ∗ W ∗ ( x ) ) \displaystyle\hskip-70.0001pt\overset{(a)}{\leq}\gamma d_{p}(P^{\pi_{k}}W_{k}(x),P^{\pi^{*}}W^{*}(x)) ≤ ( b ) γ d p ( S i k W k ( X ′ ) , S i k W ∗ ( X ′ ) ) \displaystyle\hskip-70.0001pt\overset{(b)}{\leq}\gamma d_{p}(S^{k}_{i}W_{k}(X^{\prime}),S^{k}_{i}W^{*}(X^{\prime})) + γ d p ( S ¯ i k W k ( X ′ ) , S ¯ i k W ∗ ( X ′ ) ) , \displaystyle\hskip-60.00009pt+\gamma d_{p}(\bar{S}^{k}_{i}W_{k}(X^{\prime}),\bar{S}^{k}_{i}W^{*}(X^{\prime})), (12) where in ( a ) (a) we used Properties P1 and P2 of the Wasserstein metric, and in (b) we separate states for which π k = π ∗ \pi_{k}=\pi^{*} from the rest using Lemma 1 ( { S i k , S ¯ i k } \{S^{k}_{i},\bar{S}^{k}_{i}\} form a partition of Ω \Omega ). Let δ i := Pr { X ′ ∉ 𝒳 k , i } = 𝔼 { S ¯ i k ( X ′ ) } = ∥ S ¯ i k ( X ′ ) ∥ p \delta_{i}:=\Pr\{X^{\prime}\notin\mathcal{X}_{k,i}\}=\expect\{\bar{S}^{k}_{i}(X^{\prime})\}=\|\bar{S}_{i}^{k}(X^{\prime})\|_{p} . From property P3 of the Wasserstein metric, we have d p ( S ¯ i k W k ( X ′ ) , S ¯ i k W ∗ ( X ′ ) ) \displaystyle d_{p}(\bar{S}^{k}_{i}W_{k}(X^{\prime}),\bar{S}^{k}_{i}W^{*}(X^{\prime})) ≤ sup x ′ d p ( S ¯ i k ( X ′ ) W k ( x ′ ) , S ¯ i k ( X ′ ) W ∗ ( x ′ ) ) \displaystyle\hskip-70.0001pt\leq\sup_{x^{\prime}}d_{p}(\bar{S}^{k}_{i}(X^{\prime})W_{k}(x^{\prime}),\bar{S}^{k}_{i}(X^{\prime})W^{*}(x^{\prime})) ≤ ‖ S ¯ i k ( X ′ ) ‖ p sup x ′ d p ( W k ( x ′ ) , W ∗ ( x ′ ) ) \displaystyle\hskip-70.0001pt\leq\|\bar{S}_{i}^{k}(X^{\prime})\|_{p}\sup_{x^{\prime}}d_{p}(W_{k}(x^{\prime}),W^{*}(x^{\prime})) ≤ δ i sup x ′ d p ( W k ( x ′ ) , W ∗ ( x ′ ) ) \displaystyle\hskip-70.0001pt\leq\delta_{i}\sup_{x^{\prime}}d_{p}(W_{k}(x^{\prime}),W^{*}(x^{\prime})) ≤ δ i B . \displaystyle\hskip-70.0001pt\leq\delta_{i}B. Recall that B < ∞ B<\infty is the largest attainable ‖ Z ‖ ∞ \left\|Z\right\|_{\infty} . Since also δ i < δ \delta_{i}<\delta by our choice of x ∈ 𝒳 k + 1 , i + 1 x\in\mathcal{X}_{k+1,i+1} , we can upper bound the second term in ( 12 ) by γ δ B \gamma\delta B . This yields d p ( W k + 1 ( x ) , W ∗ ( x ) ) \displaystyle d_{p}(W_{k+1}(x),W^{*}(x)) ≤ \displaystyle\leq γ d p ( S i k W k ( X ′ ) , S i k W ∗ ( X ′ ) ) + γ δ B . \displaystyle\hskip-60.00009pt\gamma d_{p}(S^{k}_{i}W_{k}(X^{\prime}),S^{k}_{i}W^{*}(X^{\prime}))+\gamma\delta B. By induction on i > 0 i>0 , we conclude that for x ∈ 𝒳 k + i , i x\in\mathcal{X}_{k+i,i} and some random state X ′′ X^{\prime\prime} i i steps forward, d p ( W k + i ( x ) , W ∗ ( x ) ) \displaystyle d_{p}(W_{k+i}(x),W^{*}(x)) ≤ \displaystyle\leq γ i d p ( S 0 k W k ( X ′′ ) , S 0 k W ∗ ( X ′′ ) ) + δ B 1 − γ \displaystyle\hskip-60.00009pt\gamma^{i}d_{p}(S^{k}_{0}W_{k}(X^{\prime\prime}),S^{k}_{0}W^{*}(X^{\prime\prime}))+\frac{\delta B}{1-\gamma} ≤ γ i B + δ B 1 − γ . \displaystyle\hskip-60.00009pt\leq\gamma^{i}B+\frac{\delta B}{1-\gamma}. Hence for any x ∈ 𝒳 x\in\mathcal{X} , ϵ > 0 \epsilon>0 , we can take δ \delta , i i , and finally k k large enough to make d p ( W k ( x ) , W ∗ ( x ) ) < ϵ d_{p}(W_{k}(x),W^{*}(x))<\epsilon . The proof then extends to Z k ( x , a ) Z_{k}(x,a) by considering one additional application of 𝒯 \mathcal{T} .
We now consider the more general case where there are multiple optimal policies. We expand the definition of 𝒳 k , i \mathcal{X}_{k,i} as follows: 𝒳 k , i := { x ∈ 𝒳 k : ∀ π ∗ ∈ Π ∗ , 𝔼 a ∗ ∼ π ∗ ( x ) P ( 𝒳 k − 1 , i − 1 | x , a ∗ ) ≥ 1 − δ } , \mathcal{X}_{k,i}:=\big\{x\in\mathcal{X}_{k}:\forall\pi^{*}\in\Pi^{*},\;\;\mathclap{\expect_{a^{*}\sim\pi^{*}(x)}}\;\;P(\mathcal{X}_{k-1,i-1}\,|\,x,a^{*})\geq 1-\delta\big\}, Because there are finitely many actions, Lemma 6 also holds for this new definition. As before, take x ∈ 𝒳 k , i x\in\mathcal{X}_{k,i} , but now consider the sequence of greedy policies π k , π k − 1 , … \pi_{k},\pi_{k-1},\dots selected by successive applications of 𝒯 \mathcal{T} , and write 𝒯 π ¯ k := 𝒯 π k 𝒯 π k − 1 ⋯ 𝒯 π k − i + 1 , \mathcal{T}^{\bar{\pi}_{k}}:=\mathcal{T}^{\pi_{k}}\mathcal{T}^{\pi_{k-1}}\cdots\mathcal{T}^{\pi_{k-i+1}}, such that Z k + 1 = 𝒯 π ¯ k Z k − i + 1 . Z_{k+1}=\mathcal{T}^{\bar{\pi}_{k}}Z_{k-i+1}. Now denote by 𝒵 ∗ ∗ \mathcal{Z}^{**} the set of nonstationary optimal policies. If we take any Z ∗ ∈ 𝒵 ∗ Z^{*}\in\mathcal{Z}^{*} , we deduce that inf Z ∗ ∗ ∈ 𝒵 ∗ ∗ d p ( 𝒯 π ¯ k Z ∗ ( x , a ) , Z ∗ ∗ ( x , a ) ) ≤ δ B 1 − γ , \inf_{Z^{**}\in\mathcal{Z}^{**}}d_{p}(\mathcal{T}^{\bar{\pi}_{k}}Z^{*}(x,a),Z^{**}(x,a))\leq\frac{\delta B}{1-\gamma}, since Z ∗ Z^{*} corresponds to some optimal policy π ∗ \pi^{*} and π ¯ k \bar{\pi}_{k} is optimal along most of the trajectories from ( x , a ) (x,a) . In effect, 𝒯 π ¯ k Z ∗ \mathcal{T}^{\bar{\pi}_{k}}Z^{*} is close to the value distribution of the nonstationary optimal policy π ¯ k π ∗ \bar{\pi}_{k}\pi^{*} . Now for this Z ∗ Z^{*} , inf Z ∗ ∗ d p ( Z k ( x , a ) , Z ∗ ∗ ( x , a ) ) \displaystyle\inf_{Z^{**}}d_{p}(Z_{k}(x,a),Z^{**}(x,a)) ≤ d p ( Z k ( x , a ) , 𝒯 π ¯ k Z ∗ ( x , a ) ) \displaystyle\hskip-80.00012pt\leq d_{p}(Z_{k}(x,a),\mathcal{T}^{\bar{\pi}_{k}}Z^{*}(x,a)) + inf Z ∗ ∗ d p ( 𝒯 π ¯ k Z ∗ ( x , a ) , Z ∗ ∗ ( x , a ) ) \displaystyle\hskip-70.0001pt+\inf_{Z^{**}}d_{p}(\mathcal{T}^{\bar{\pi}_{k}}Z^{*}(x,a),Z^{**}(x,a)) ≤ d p ( 𝒯 π ¯ k Z k − i + 1 ( x , a ) , 𝒯 π ¯ k Z ∗ ( x , a ) ) + δ B 1 − γ \displaystyle\hskip-80.00012pt\leq d_{p}(\mathcal{T}^{\bar{\pi}_{k}}Z_{k-i+1}(x,a),\mathcal{T}^{\bar{\pi}_{k}}Z^{*}(x,a))+\frac{\delta B}{1-\gamma} ≤ γ i B + 2 δ B 1 − γ , \displaystyle\hskip-80.00012pt\leq\gamma^{i}B+\frac{2\delta B}{1-\gamma}, using the same argument as before with the newly-defined 𝒳 k , i \mathcal{X}_{k,i} . It follows that inf Z ∗ ∗ ∈ 𝒵 ∗ ∗ d p ( Z k ( x , a ) , Z ∗ ∗ ( x , a ) ) → 0 . \inf_{Z^{**}\in\mathcal{Z}^{**}}d_{p}(Z_{k}(x,a),Z^{**}(x,a))\to 0. When 𝒳 \mathcal{X} is finite, there exists a fixed k k after which 𝒳 k = 𝒳 \mathcal{X}_{k}=\mathcal{X} . The uniform convergence result then follows.
To prove the uniqueness of the fixed point Z ∗ Z^{*} when 𝒯 \mathcal{T} selects its actions according to the ordering ≺ \prec , we note that for any optimal value distribution Z ∗ Z^{*} , its set of greedy policies is Π ∗ \Pi^{*} . Denote by π ∗ \pi^{*} the policy coming first in the ordering over Π ∗ \Pi^{*} . Then 𝒯 = 𝒯 π ∗ \mathcal{T}=\mathcal{T}^{\pi^{*}} , which has a unique fixed point (Section 3.3 ). ∎
###### Proposition 4 .
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}^{*} .
We provide here a sketch of the result. Consider a single state x 1 x_{1} with two actions, a 1 a_{1} and a 2 a_{2} (Figure 8 ). The first action yields a reward of 1 / 2 1/2 , while the other either yields 0 0 or 1 1 with equal probability, and both actions are optimal. Now take γ = 1 / 2 \gamma=1/2 and write R 0 , R 1 , … R_{0},R_{1},\dots for the received rewards. Consider a stochastic policy that takes action a 2 a_{2} with probability p p . For p = 0 p=0 , the return is Z p = 0 = 1 1 − γ 1 2 = 1 . Z_{p=0}=\frac{1}{1-\gamma}\frac{1}{2}=1. For p = 1 p=1 , on the other hand, the return is random and is given by the following fractional number (in binary): Z p = 1 = R 0 . R 1 R 2 R 3 ⋯ . Z_{p=1}=R_{0}.R_{1}R_{2}R_{3}\cdots. As a result, Z p = 1 Z_{p=1} is uniformly distributed between 0 0 and 2 2 ! In fact, note that Z p = 0 = 0.11111 ⋯ = 1 . Z_{p=0}=0.11111\cdots=1. For some intermediary value of p p , we obtain a different probability of the different digits, but always putting some probability mass on all returns in [ 0 , 2 ] [0,2] .
Now suppose we follow the nonstationary policy that takes a 1 a_{1} on the first step, then a 2 a_{2} from there on. By inspection, the return will be uniformly distributed on the interval [ 1 / 2 , 3 / 2 ] [1/2,3/2] , which does not correspond to the return under any value of p p . But now we may imagine an operator 𝒯 \mathcal{T} which alternates between a 1 a_{1} and a 2 a_{2} depending on the exact value distribution it is applied to, which would in turn converge to a nonstationary optimal value distribution.
###### Lemma 7 (Sample Wasserstein distance) .
Let { P i } \{P_{i}\} be a collection of random variables, I ∈ ℕ I\in\mathbb{N} a random index independent from { P i } \{P_{i}\} , and consider the mixture random variable P = P I P=P_{I} . For any random variable Q Q independent of I I , d p ( P , Q ) ≤ 𝔼 i ∼ I d p ( P i , Q ) , d_{p}(P,Q)\leq\expect_{i\sim I}d_{p}(P_{i},Q), and in general the inequality is strict and ∇ Q d p ( P I , Q ) ≠ 𝔼 i ∼ I ∇ Q d p ( P i , Q ) . \nabla_{Q}d_{p}(P_{I},Q)\neq\expect_{i\sim I}\nabla_{Q}d_{p}(P_{i},Q).
###### Proof.
We prove this using Lemma 1 . Let A i := 𝕀 [ I = i ] A_{i}:=\mathbb{I}\left[I=i\right] . We write d p ( P , Q ) \displaystyle d_{p}(P,Q) = d p ( P I , Q ) \displaystyle=d_{p}(P_{I},Q) = d p ( ∑ i A i P i , ∑ i A i Q ) \displaystyle=d_{p}\Big(\sum\nolimits_{i}A_{i}P_{i},\sum\nolimits_{i}A_{i}Q\Big) ≤ ∑ i d p ( A i P i , A i Q ) \displaystyle\leq\sum\nolimits_{i}d_{p}(A_{i}P_{i},A_{i}Q) ≤ ∑ i Pr { I = i } d p ( P i , Q ) \displaystyle\leq\sum\nolimits_{i}\Pr\{I=i\}d_{p}(P_{i},Q) = 𝔼 I d P ( P i , Q ) . \displaystyle=\expect\nolimits_{I}d_{P}(P_{i},Q). where in the penultimate line we used the independence of I I from P i P_{i} and Q Q to appeal to property P3 of the Wasserstein metric.
To show that the bound is in general strict, consider the mixture distribution depicted in Figure 9 . We will simply consider the d 1 d_{1} metric between this distribution P P and another distribution Q Q . The first distribution is P = { 0 w.p. 1 / 2 1 w.p. 1 / 2 . P=\left\{\begin{array}[]{ll}0&\text{w.p. }1/2\\ 1&\text{w.p. }1/2.\end{array}\right. In this example, i ∈ { 1 , 2 } i\in\{1,2\} , P 1 = 0 P_{1}=0 , and P 2 = 1 P_{2}=1 . Now consider the distribution with the same support but that puts probability p p on 0 0 : Q = { 0 w.p. p 1 w.p. 1 − p . Q=\left\{\begin{array}[]{ll}0&\text{w.p. }p\\ 1&\text{w.p. }1-p.\end{array}\right. The distance between P P and Q Q is d 1 ( P , Q ) = | p − 1 2 | . d_{1}(P,Q)=|p-\tfrac{1}{2}|. This is d 1 ( P , Q ) = 1 2 d_{1}(P,Q)=\frac{1}{2} for p ∈ { 0 , 1 } p\in\{0,1\} , and strictly less than 1 2 \frac{1}{2} for any other values of p p . On the other hand, the corresponding expected distance (after sampling an outcome x 1 x_{1} or x 2 x_{2} with equal probability) is 𝔼 I d 1 ( P i , Q ) = 1 2 p + 1 2 ( 1 − p ) = 1 2 . \expect\nolimits_{I}d_{1}(P_{i},Q)=\tfrac{1}{2}p+\tfrac{1}{2}(1-p)=\tfrac{1}{2}. Hence d 1 ( P , Q ) < 𝔼 I d 1 ( P i , Q ) d_{1}(P,Q)<\expect\nolimits_{I}d_{1}(P_{i},Q) for p ∈ ( 0 , 1 ) p\in(0,1) . This shows that the bound is in general strict. By inspection, it is clear that the two gradients are different. ∎
###### Proposition 5 .
Fix some next-state distribution Z Z and policy π \pi . Consider a parametric value distribution Z θ Z_{\theta} , and and define the Wasserstein loss ℒ W ( θ ) := d p ( Z θ ( x , a ) , R ( x , a ) + γ Z ( X ′ , π ( X ′ ) ) ) . \mathcal{L}_{W}(\theta):=d_{p}(Z_{\theta}(x,a),R(x,a)+\gamma Z(X^{\prime},\pi(X^{\prime}))). Let r ∼ R ( x , a ) r\sim R(x,a) and x ′ ∼ P ( ⋅ | x , a ) x^{\prime}\sim P(\cdot\,|\,x,a) and consider the sample loss L W ( θ , r , x ′ ) := d p ( Z θ ( x , a ) , r + γ Z ( x ′ , π ( x ′ ) ) CLOSE . L_{W}(\theta,r,x^{\prime}):=d_{p}(Z_{\theta}(x,a),r+\gamma Z(x^{\prime},\pi(x^{\prime})). Its expectation is an upper bound on the loss ℒ W \mathcal{L}_{W} : ℒ W ( θ ) ≤ 𝔼 R , P L W ( θ , r , x ′ ) , \mathcal{L}_{W}(\theta)\leq\expect_{R,P}L_{W}(\theta,r,x^{\prime}), in general with strict inequality.
The result follows directly from the previous lemma.
## Appendix C Algorithmic Details