Distributional Policy Iteration

Definition

Policy Evaluation

Compute $Z^{\pi }$ from the current policy $\pi$.

Distributional Bellman operator ${\mathcal{T}}^{\pi }$ for a fixed policy $\pi$ is a Contraction Map in a maximal form of Wasserstein Metric. ${\overline{W}}_p(Z_1,Z_2)={\sup }_{x,a}{W}_p(Z_1(x,a),Z_2(x,a))$

Thus, the iteration $Z\leftarrow {\mathcal{T}}^{\pi }Z$ converges to a unique fixed point $Z^{\pi }$ (by Contraction Lemma).

Control (Policy Improvement)

Seek a greedy policy $\pi$ based on the $Z^{\ast }$, which maximizes the expectation of $Z^{\ast }$.

Distributional Bellman optimality operator $\mathcal{T}$ is not a Contraction Map, not even continuous in any Metric, causing the distributional instability. To overcome the problem, discrete action-value distributions are used to approximate the true distribution $Z$.