# Near-Minimax-Optimal Distributional Reinforcement Learning
with a Generative Model

Mark Rowland  Li Kevin Wenliang  Rémi Munos  Clare Lyle  Yunhao Tang  Will
Dabney

###### Abstract

We propose a new algorithm for model-based distributional reinforcement
learning (RL), and prove that it is minimax-optimal for approximating return
distributions with a generative model (up to logarithmic factors), resolving
an open question of Zhang et al. (2023). Our analysis provides new theoretical
results on categorical approaches to distributional RL, and also introduces a
new distributional Bellman equation, the stochastic categorical CDF Bellman
equation, which we expect to be of independent interest. We also provide an
experimental study comparing several model-based distributional RL algorithms,
with several takeaways for practitioners.

Machine Learning, Reinforcement Learning, Distributional Reinforcement
Learning, Sample Complexity

##  1 Introduction

In distributional reinforcement learning, the aim is to predict the full
probability distribution of possible returns at each state, rather than just
the mean return (Morimura et al., 2010a; Bellemare et al., 2017, 2023).
Applications of distributional reinforcement learning range from dopamine
response modelling in neuroscience (Dabney et al., 2020), to driving risk-
sensitive decision-making and exploration in domains such as robotics (Bodnar
et al., 2020), healthcare (Böck et al., 2022), and algorithm discovery (Fawzi
et al., 2022), as well as forming a core component of many deep reinforcement
learning architectures (Bellemare et al., 2017; Dabney et al., 2018b, a; Yang
et al., 2019; Bellemare et al., 2020; Shahriari et al., 2022; Wurman et al.,
2022).

Beyond asymptotic convergence (Rowland et al., 2018, 2023; Bellemare et al.,
2023), an emerging topic in the theory of distributional RL is that of finite-
sample bounds (Böck & Heirzinger, 2022; Wu et al., 2023): _How many samples
are required to learn accurate return-distribution estimates?_

Zhang et al. (2023) make a variety of contributions in the setting where
samples are obtained via a _generative model_ , whereby one can obtain sample
transitions at any desired state in the environment, a commonly studied regime
for mean-return RL (Kearns et al., 2002; Kakade, 2003; Azar et al., 2013).
Zhang et al. (2023) prove a lower bound of
$N=\Omega(\varepsilon^{-2}(1-\gamma)^{-3})$ samples required per state in
order to obtain $\varepsilon$-accurate return-distribution estimates with high
probability, as measured by Wasserstein distance (here, $\gamma\in[0,1)$ is
the discount factor). This matches the minimax sample complexity of learning
an accurate value function (Azar et al., 2013). Zhang et al. (2023) also
propose an algorithm with a sample complexity bound of
$N=\widetilde{O}(\varepsilon^{-2}(1-\gamma)^{-4})$ for the problem above.

This leaves two key questions open. First, there is gap of
$O((1-\gamma)^{-1})$ between the bounds above: Can the upper bound be
improved, or in other words, _is distributional RL statistically harder than
learning a value function with a generative model?_ Second, the algorithm
proposed by Zhang et al. (2023) is typically not implementable in practice, as
it involves solving the _unprojected_ distributional Bellman equation over an
infinite-dimensional space of probability distributions (the authors use
approximate methods in their experiments for this reason): _Is there a
tractable algorithm that achieves the lower bound?_ We resolve both these
questions with the following contributions:

  * •

We propose the _direct categorical fixed-point algorithm_ (DCFP), a new
distributional RL algorithm based on the categorical approach (Bellemare et
al., 2017).

  * •

We prove that with an appropriate number of categories, DCFP has a sample
complexity of $\widetilde{O}(\varepsilon^{-2}(1-\gamma)^{-3})$ (for attaining
$\varepsilon$-accurate Wasserstein-1 return-distribution approximations with
high probability), matching the minimax lower bound up to logarithmic factors.
This shows that in this context, distributional RL is no harder statistically
than learning an accurate value function.

  * •

Our analysis introduces a novel distributional RL concept which we expect to
be of independent interest: the _stochastic categorical CDF Bellman equation_.

  * •

We also study the empirical performance of several algorithms for
distributional RL with a generative model, and identify several key findings
for practitioners.

##  2 Background

Throughout the paper, we consider the problem of evaluation in an infinite-
horizon Markov reward process (MRP), with finite state space $\mathcal{X}$,
transition probabilities $P\in\mathbb{R}^{\mathcal{X}\times\mathcal{X}}$,
reward function $r:\mathcal{X}\rightarrow[0,1]$, and discount factor
$\gamma\in[0,1)$. A random trajectory $(X_{t},R_{t})_{t\geq 0}$ is generated
from an initial state $X_{0}=x$ according to the conditional distributions
$X_{t}\mid(X_{0},\ldots,X_{t-1})\sim P(\cdot|X_{t-1})$, and $R_{t}=r(X_{t})$.
The _return_ associated with the trajectory is given by the quantity
$\sum_{t\geq 0}\gamma^{t}R_{t}$. In RL, a central task is to estimate the
_value function_ $V^{*}:\mathcal{X}\rightarrow\mathbb{R}$, defined by

$\displaystyle\textstyle V^{*}(x)=\mathbb{E}[\sum_{t\geq 0}\gamma^{t}R_{t}\mid
X_{0}=x]\,,$ (1)

given some form of observations from the MRP. The value function defines the
expected return, conditional on each possible starting state in the MRP.

The value function satisfies the Bellman equation $V^{*}=TV^{*}$, where
$T:\mathbb{R}^{\mathcal{X}}\rightarrow\mathbb{R}^{\mathcal{X}}$ is defined by

$\displaystyle(TV)(x)=\mathbb{E}_{x}[R+\gamma V(X^{\prime})]\,,$ (2)

where $(X,R,X^{\prime})$ is a random transition in the environment,
distributed as described above. When the transition probabilities of the MRP
are known, the right-hand side can be evaluated as an affine transformation of
$V$. MRP theory (see, e.g., Puterman, 2014 for an overview) then provides two
methods for obtaining (an approximation to) $V^{*}$. First, _dynamic
programming_ (DP), in which an initial approximation
$V_{0}\in[0,(1-\gamma)^{-1}]^{\mathcal{X}}$ is taken, and the sequence
$(V_{k})_{k=0}^{\infty}$ is computed via the update $V_{k+1}=TV_{k}$; it is
guaranteed that $\|V_{k}-V^{*}\|_{\infty}\leq\gamma^{k}(1-\gamma)^{-1}$.
Second, _direct solution_ , in which the linear system $V=TV$ is solved
directly to obtain $V^{*}$ as its unique solution.

###  2.1 Reinforcement learning with a generative model

In many settings the transition probabilities of the MRP are unknown, and the
value function must be estimated based on data comprising sampled transitions,
introducing a statistical element to the problem. A commonly used model for
this data is a _generative model_ (Kearns et al., 2002; Kakade, 2003). In this
setting, for each state $x\in\mathcal{X}$, we observe $N$ i.i.d. samples
$(X^{x}_{i})_{i=1}^{N}$ from $P(\cdot|x)$, and this collection of
$N|\mathcal{X}|$ samples may then be used by an algorithm to estimate the
value function. Azar et al. (2013) showed that at least
$N=\Omega(\varepsilon^{-2}(1-\gamma)^{-3}\log(|\mathcal{X}|/\delta))$ samples
are required to obtain $\varepsilon$-accurate estimates of the value function
with high probability (measured in $L^{\infty}$ norm), and also showed that
this bound is attained (up to logarithmic factors) by a _certainty
equivalence_ algorithm, which treats the empirically observed transition
frequencies as the true ones, and solves for the value function of the
corresponding MRP.

Mathematically, the algorithm constructs the empirical transition probability
matrix $\hat{P}\in\mathbb{R}^{\mathcal{X}\times\mathcal{X}}$, given by
$\hat{P}(y|x)=N^{-1}\sum_{i=1}^{N}\mathbbm{1}\\{X^{x}_{i}=y\\}$. From this,
the empirical Bellman operator $\hat{T}$ can be defined as in Equation (2),
with the expectation now computed with respect to $\hat{P}$. One then solves
the empirical Bellman equation $V=\hat{T}V$ to obtain an estimated value
function $\hat{V}$. Azar et al. (2013) show that for sufficiently small
$\varepsilon>0$ and given $\delta\in(0,1)$, with
$N=\Omega(\varepsilon^{-2}(1-\gamma)^{-3}\log(|\mathcal{X}|/\delta))$ samples,
we have $\|\hat{V}-V^{*}\|_{\infty}\leq\varepsilon$ with probability at least
$1-\delta$, which matches the lower bound described above.

###  2.2 Distributional reinforcement learning

Distributional RL aims to capture the full probability distribution of the
random return at each state, not just its mean. Mathematically, the object of
interest is the return-distribution function (RDF)
$\eta^{*}:\mathcal{X}\rightarrow\mathscr{P}(\mathbb{R})$, defined by

$\displaystyle\textstyle\eta^{*}(x)=\mathcal{D}\big{(}\sum_{t=0}^{\infty}\gamma^{t}R_{t}\mid
X_{0}=x\big{)}\,,$ (3)

where $\mathcal{D}$ extracts the probability distribution of the input random
variable. The distributional perspective on reinforcement learning has proved
practically useful in a wide variety of applications, including healthcare
(Böck et al., 2022), navigation (Bellemare et al., 2020), and algorithm
discovery (Fawzi et al., 2022).

The central equation behind dynamic programming approaches to approximating
the return distribution function is the _distributional Bellman equation_
(Sobel, 1982; Morimura et al., 2010a; Bellemare et al., 2017), given by

$\displaystyle\eta^{*}=\mathcal{T}\eta^{*}\,,$

where
$\mathcal{T}:\mathscr{P}(\mathbb{R})^{\mathcal{X}}\rightarrow\mathscr{P}(\mathbb{R})^{\mathcal{X}}$
is the _distributional Bellman operator_ , defined by

$\displaystyle(\mathcal{T}\eta)(x)=\mathcal{D}\big{(}R+\gamma
G(X^{\prime})\mid X=x\big{)}\,,$

where independent from the random transition $(X,R,X^{\prime})$, we have
$G(x)\sim\eta(x)$ for each $x\in\mathcal{X}$.

###  2.3 Categorical dynamic programming

Given an initial RDF approximation
$\eta\in\mathscr{P}([0,(1-\gamma)^{-1}])^{\mathcal{X}}$, it also holds that
the update $\eta\leftarrow\mathcal{T}\eta$ converges to $\eta^{*}$ in an
appropriate sense (e.g., in Wasserstein distance; see Bellemare et al., 2017),
in analogy with mean-return DP. However, probability distributions are
infinite-dimensional objects, and as such it is rarely the case that this
update can actually be implemented and carried out in practice. Instead, the
use of some kind of approximate, tractable representation of probability
distributions is typically required.

Representations. In this paper, we focus on the categorical approach to
distributional reinforcement learning (Bellemare et al., 2017), in which
estimates of return distributions are represented as categorical distributions
over a finite number of outcomes $z_{1}<\cdots<z_{m}$. We will take
$\\{z_{1},\ldots,z_{m}\\}$ to be an equally spaced grid over the range of
possible returns $[0,(1-\gamma)^{-1}]$, so that
$z_{i}=\tfrac{i-1}{m-1}(1-\gamma)^{-1}$ for $i=1,\ldots,m$. Approximations of
the RDF are then represented as
$\eta_{p}:\mathcal{X}\rightarrow\mathscr{P}(\mathbb{R})$, where

$\displaystyle\textstyle\eta_{p}(x)=\sum_{i=1}^{m}p_{i}(x)\delta_{z_{i}}\,;$
(4)

here, $\delta_{z}$ is the Dirac distribution at the outcome $z$, and
$p=((p_{i}(x))_{i=1}^{m}:x\in\mathcal{X})$ are adjustable probability mass
parameters; see Figure 1(a). The number of categories $m$ can be interpreted
as controlling the _expressivity_ of the representation, and should be
carefully chosen in practice.

Dynamic programming. The iteration $\eta\leftarrow\mathcal{T}\eta$ _cannot_ be
used to update the parameters $p$ in Equation (4) directly, since the
distributions $(\mathcal{T}\eta)(x)$ are no longer supported on
$\\{z_{1},\ldots,z_{m}\\}$, and so cannot be expressed in the form given in
Equation (4); see Figure 1(c).

Bellemare et al. (2017) circumvent this issue by _projecting_ the resulting
distributions back onto the support set $\\{z_{1},\ldots,z_{m}\\}$ via a map
${\Pi_{m}}:\mathscr{P}([0,(1-\gamma)^{-1}])\rightarrow\mathscr{P}(\\{z_{1},\ldots,z_{m}\\}])$.
Intuitively, ${\Pi_{m}}$ can be thought of as allocating each outcome
$z\in[z_{i},z_{i+1}]$ to its neighbouring gridpoints $z_{i}$ and $z_{i+1}$, in
proportion to their proximity, so that the projection of the Dirac
distribution $\delta_{z}$, is defined by

$\displaystyle{\Pi_{m}}\delta_{z}=\frac{z_{i+1}-z}{z_{i+1}-z_{i}}\delta_{z_{i}}+\frac{z-z_{i}}{z_{i+1}-z_{i}}\delta_{z_{i+1}}\,.$

In this paper, we work with the equivalent definition of the projection
${\Pi_{m}}$ given by Rowland et al. (2018, Proposition 6), in which the
probability mass assigned to $z_{i}$ by ${\Pi_{m}}\nu$ is given by the
expectation $\mathbb{E}_{Z\sim\nu}[h_{i}(Z)]$, where
$h_{i}:[0,(1-\gamma)^{-1}]\rightarrow[0,1]$ is the “hat-function” at $z_{i}$,
which linearly interpolates between a value of $1$ at $z_{i}$, and 0 at
neighbouring gridpoints $z_{i-1}$, $z_{i+1}$, and is 0 outside this range.
Figure 1(b) illustrates $h_{i}$ and $h_{m}$; see Appendix B for a restatement
of the full definition given by Rowland et al. (2018).

Rowland et al. (2018) showed that _categorical dynamic programming_ (CDP), the
iterative computation $\eta\leftarrow{\Pi_{m}}\mathcal{T}\eta$ of _projected_
distributional Bellman updates (see Figure 1(c)), leads to convergence of
$\eta$ to a _categorical fixed point_ , which can be made arbitrarily close to
the true RDF by increasing $m$, as measured in Cramér distance (Cramér, 1928;
Székely, 2003; Székely & Rizzo, 2013).

Figure 1: (a) The density of a distribution $\nu$ (grey), and its categorical
projection ${\Pi_{m}}\nu\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})$ (blue). (b)
Hat functions $h_{i}$ (solid) and $h_{m}$ (dashed). (c) A categorical
distribution (blue); its update after being scaled by $\gamma$ and shifted by
$r$ by the distributional Bellman operator $\mathcal{T}$, moving its support
off the grid $\\{z_{1},\ldots,z_{m}\\}$ (pink); the resulting realigned
distribution supported on the grid $\\{z_{1},\ldots,z_{m}\\}$ after projection
via ${\Pi_{m}}$ (green).

######  Definition 2.1.

The _Cramér distance_
$\ell_{2}:\mathscr{P}(\mathbb{R})\times\mathscr{P}(\mathbb{R})\rightarrow\mathbb{R}$
is defined by

$\displaystyle\ell_{2}(\nu,\nu^{\prime})=\bigg{[}\int_{\mathbb{R}}(F_{\nu}(t)-F_{\nu^{\prime}}(t))^{2}\;\mathrm{d}t\bigg{]}^{1/2}\,,$

where $F_{\nu},F_{\nu^{\prime}}$ are the CDFs of $\nu,\nu^{\prime}$,
respectively. The _supremum-Cramér distance_ $\overline{\ell}_{2}$ on
$\mathscr{P}(\mathbb{R})^{\mathcal{X}}$ is defined by

$\displaystyle\overline{\ell}_{2}(\eta,\eta^{\prime})=\max_{x\in\mathcal{X}}\ell_{2}(\eta(x),\eta^{\prime}(x))\,.$

The central convergence results concerning CDP established by Rowland et al.
(2018) are summarised below.

######  Proposition 2.2.

(Rowland et al., 2018). The operator
${\Pi_{m}}\mathcal{T}:\mathscr{P}([0,(1-\gamma)^{-1}])^{\mathcal{X}}\rightarrow\mathscr{P}([0,(1-\gamma)^{-1}])^{\mathcal{X}}$
is a contraction mapping with respect to $\overline{\ell}_{2}$, with
contraction factor $\sqrt{\gamma}$, and has a unique fixed point,
$\eta_{\text{C}}\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})^{\mathcal{X}}$. As a
result, for any $\eta_{0}\in\mathscr{P}([0,(1-\gamma)^{-1}])^{\mathcal{X}}$,
with $\eta_{k+1}={\Pi_{m}}\mathcal{T}\eta_{k}$, we have
$\overline{\ell}_{2}(\eta_{k},\eta_{\text{C}})\leq(1-\gamma)^{-1}\gamma^{k}$.
Further, the distance between $\eta_{\text{C}}$ and the true RDF $\eta^{*}$
can be bounded as

$\displaystyle\overline{\ell}_{2}(\eta_{\text{C}},\eta^{*})\leq\frac{1}{(1-\gamma)\sqrt{m-1}}\,.$
(5)

This establishes CDP as a principled approach to approximating return
distributions, and also quantifies the accuracy achievable with CDP using $m$
categories, which will be central in informing our choice of $m$ to obtain a
sample-efficient, accurate algorithm below.

##  3 Distributional RL with a generative model

The central problem we study in this paper is how to do sample-efficient
_distributional_ RL with a generative model. That is, given the samples
$((X^{x}_{i})_{i=1}^{N}:x\in\mathcal{X})$ described in Section 2.1, how
accurate of an approximation to the return-distribution function in Equation
(3) can one compute?

This question was raised by Zhang et al. (2023), who proposed to perform
distributional dynamic programming updates
$\eta\leftarrow\hat{\mathcal{T}}\eta$ as described in Section 2.2, using the
_empirical_ distributional Bellman operator $\hat{\mathcal{T}}$ derived from
the empirical transition probabilities $\hat{P}$, producing an estimate
$\hat{\eta}\in\mathscr{P}([0,(1-\gamma)^{-1}])$ of the true RDF $\eta^{*}$
such that for any $\varepsilon>0$ and $\delta\in(0,1)$,

$\displaystyle w_{1}(\hat{\eta}(x),\eta^{*}(x))\leq\varepsilon\,,$

with probability at least $1-\delta$ for all $x\in\mathcal{X}$, whenever
$N=\widetilde{\Omega}(\varepsilon^{-2}(1-\gamma)^{-4}\text{polylog}(1/\delta))$.
Here, $w_{1}$ denotes the Wasserstein-1 distance between probability
distributions, defined for any $\nu,\nu^{\prime}\in\mathscr{P}(\mathbb{R})$
with CDFs $F_{\nu},F_{\nu^{\prime}}$ by

$\displaystyle
w_{1}(\nu,\nu^{\prime})=\int_{\mathbb{R}}|F_{\nu}(t)-F_{\nu^{\prime}}(t)|\;\mathrm{d}t\,.$

We focus on the Wasserstein-1 distance as the main metric of interest in this
paper as it is particularly compatible with existing methods for analysing
categorical approaches to distributional RL, and it provides upper bounds for
differences of many statistical functionals of interest, such as expectations
of Lipschitz functions (Villani, 2009; Peyré & Cuturi, 2019); and conditional-
value-at-risk (Rockafellar et al., 2000; Rockafellar & Uryasev, 2002; Bhat &
Prashanth, 2019, CVaR). Zhang et al. (2023) also prove a lower bound of
$N=\widetilde{\Omega}(\varepsilon^{-2}(1-\gamma)^{-3})$ samples required to
obtain such an accurate prediction with high probability, which follows from a
reduction to the mean-return case (Azar et al., 2013), as $w_{1}$-distance
upper-bounds difference of means.

There are two natural questions that the analysis of Zhang et al. (2023)
leaves open. Firstly, can the gap between the lower bound and upper bound as a
function of $(1-\gamma)^{-1}$ described above be closed? Zhang et al. (2023)
conjecture that their analysis is loose. Second, we also note that the
distributional dynamic programming procedure
$\eta\leftarrow\hat{\mathcal{T}}\eta$ proposed by Zhang et al. (2023), without
incorporating any restrictions on the representations of distributions, runs
into severe space and memory issues, and is not practical to run (and indeed
Zhang et al. (2023) introduce approximations to the algorithm when running
empirically for these reasons). A remaining question is then whether there are
tractable algorithms that can achieve the lower bound on sample complexity
described above. Our contributions below provide a new, tractable
distributional RL algorithm that attains (up to logarithmic factors) the lower
bound on sample complexity provided by Zhang et al. (2023), resolving these
questions.

##  4 Direct categorical fixed-point computation

Our approach to obtaining a sample-efficient algorithm that is near-minimax-
optimal in the sense described above begins with the categorical approach to
distributional dynamic programming described in Section 2.3. We begin first by
introducing a new algorithm for computing the categorical fixed point
$\eta_{\text{C}}$ referred to in Proposition 2.2 _directly_ , that avoids
computing an approximate solution via dynamic programming iterations
$\eta\leftarrow\Pi_{m}\mathcal{T}\eta$. We expect this algorithm to be of
independent interest within distributional reinforcement learning.

###  4.1 Direct categorical fixed-point computation

Our first contribution is to develop a new computational perspective on the
projected categorical Bellman operator ${\Pi_{m}}\mathcal{T}$, which results
in a new algorithm for computing the fixed point $\eta_{\text{C}}$ _exactly_ ,
without requiring the iterative categorical dynamic programming (CDP)
algorithm described in Section 2.3.

CDF operator and fixed-point equation. We first formulate the application of
the projected categorical Bellman operator ${\Pi_{m}}\mathcal{T}$ as a linear
map, and give an explicit expression for the matrix representing this linear
map when the input RDFs are represented with cumulative distribution functions
(CDFs). We consider the effect of applying ${\Pi_{m}}\mathcal{T}$ to an RDF
$\eta=\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})^{\mathcal{X}}$, with
$\eta(x)=\sum_{i=1}^{m}p_{i}(x)\delta_{z_{i}}$. By Rowland et al. (2018,
Proposition 6), the updated probabilities for ${\Pi_{m}}\mathcal{T}$ can be
expressed as

$\displaystyle
p^{\prime}_{i}(x)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m}P(y|x)h^{x}_{i,j}p_{j}(y)\,,$
(6)

where $h^{x}_{i,j}=h_{i}(r(x)+\gamma z_{j})$. We convert this into an
expression for cumulative probabilities, rather than individual probability
masses, to obtain a simpler analysis below. To do so, we introduce the
encoding of $\eta\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})^{\mathcal{X}}$ into
corresponding CDF values $F\in\mathbb{R}^{\mathcal{X}\times m}$, where
$F_{i}(x)=\eta(x)([z_{1},z_{i}])$ denotes the _cumulative_ mass for the
distribution at state $x$ over the set $\\{z_{1},\ldots,z_{i}\\}$.

######  Proposition 4.1.

If $\eta\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})^{\mathcal{X}}$ is an RDF with
corresponding CDF values $F\in\mathbb{R}^{\mathcal{X}\times m}$, then the
corresponding CDF values $F^{\prime}\in\mathbb{R}^{\mathcal{X}\times m}$ for
${\Pi_{m}}\mathcal{T}\eta$ satisfy

$\displaystyle
F^{\prime}_{i}(x)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m}P(y|x)(H^{x}_{i,j}-H^{x}_{i,j+1})F_{j}(y)\,,$
(7)

where

$\displaystyle\textstyle H^{x}_{i,j}=\sum_{l\leq i}h_{l}(r(x)+\gamma z_{j})$
(8)

for $j=1,\ldots,m$, and by convention we take $H^{x}_{i,m+1}=0$.

Under this notation, we can rewrite Equation (7) simply as a matrix-vector
multiplication in the space $\mathbb{R}^{\mathcal{X}\times[m]}$:

$\displaystyle F^{\prime}=T_{P}F\,,$

where $T_{P}$ is the $(\mathcal{X}\times[m])\times(\mathcal{X}\times[m])$
square matrix, with entries given by

$\displaystyle T_{P}(x,i;y,j)=P(y|x)(H^{x}_{i,j}-H^{x}_{i,j+1})\,,$ (9)

and $F,F^{\prime}\in\mathbb{R}^{\mathcal{X}\times m}$ above are interpreted in
vectorised form. We drop dependence on $m$ from the notation $T_{P}$ for
conciseness. Thus, CDP can be implemented via simple matrix multiplication on
CDF values, and the CDF values $F^{*}$ corresponding to the categorical fixed
point $\eta_{\text{C}}$ solve the equation

$\displaystyle F=T_{P}F\,,\text{ or equivalently }(I-T_{P})F=0\,.$ (10)

Obtaining a system with unique solution.  Equation (10) suggests that we can
directly solve a linear system to obtain the exact categorical fixed point,
rather than performing the iterative CDP algorithm to obtain an approximation.
Note, however, that $F^{*}$ is not the _unique_ solution of Equation (10); for
example, $F=0$ is also a solution. This arises because in Equation (10), the
_distribution masses_ at each state (that is, $F_{m}(x)$), are unconstrained.
By contrast, Proposition 2.2 establishes that $\eta_{\text{C}}$ is the unique
solution of $\eta={\Pi_{m}}\mathcal{T}\eta$ in the space
$\mathscr{P}([z_{1},z_{m}])^{\mathcal{X}}$, where each element $\eta(x)$ is
constrained to be a probability distribution _a priori_. Thus, Equation (10)
requires some modification to obtain a linear system with a unique solution.
This is obtained by removing $F_{m}(x)$ as a variable from the system (for
each $x\in\mathcal{X})$), replacing it by the constant 1, and removing
redundant rows from the resulting linear system, as the following proposition
describes; the “axis-aligned” nature of these constraints is the benefit of
working with CDF values.

######  Proposition 4.2.

The linear system in Equation (10), with the additional linear constraints
$F_{m}(x)=1$ for all $x\in\mathcal{X}$, is equivalent to the following linear
system in $\widetilde{F}\in\mathbb{R}^{\mathcal{X}\times[m-1]}$:

$\displaystyle(I-\widetilde{T}_{P})\widetilde{F}=\widetilde{H}\,,$ (11)

where the $(x,i;y,j)$ coordinate of $\widetilde{T}_{P}$ (for $1\leq i,j\leq
m-1$) is

$\displaystyle\widetilde{T}_{P}(x,i;y,j)=P(y|x)(H^{x}_{i,j}-H^{x}_{i,j+1})\,,$

and for each $x\in\mathcal{X}$, $1\leq i\leq m-1$:

$\displaystyle\widetilde{H}(x,i)=H^{x}_{i,m}\,.$

Having moved to the inhomogeneous system over
$\mathbb{R}^{\mathcal{X}\times[m-1]}$ in Equation (11), we can deduce the
following via the contraction theory in Proposition 2.2.

######  Proposition 4.3.

The linear system in Equation (11) has a unique solution, which is precisely
the CDF values $((F^{*}_{i}(x))_{i=1}^{m-1}:x\in\mathcal{X})$ of the
categorical fixed point.

The _direct categorical fixed-point algorithm_ (DCFP) consists of solving the
linear system in Equation (11) to obtain the exact categorical fixed point;
see Algorithm 1 for a summary.

Algorithm 1 DCFP

1: Calculate matrices $(H^{x}:x\in\mathcal{X})$ via Equation (8).

2: Calculate matrix $\widetilde{T}_{P}$ and vector $\widetilde{H}$ via
Equation (9).

3: Call linear system solver on Equation (11).

4: Obtain resulting solution
$\widetilde{F}^{*}\in\mathbb{R}^{\mathcal{X}\times[m-1]}$.

5: Return $F^{*}$, obtained by appending the values $F^{*}_{m}(x)=1$ to the
solution $\widetilde{F}^{*}$.

Complexity and implementation details. Representing $\widetilde{T}_{P}$ as a
dense matrix requires $O(|\mathcal{X}|^{2}m^{2})$ space, and solving the
corresponding linear system in Equation (11) with a standard linear solver
requires $O(|\mathcal{X}|^{3}m^{3})$ time. However, in many problems there are
cases where the DCFP algorithm can be implemented more efficiently. Crucially,
$\widetilde{T}_{P}$ often has sparse structure, and so sparse linear solvers
may afford an opportunity to make substantial improvements in computational
efficiency. We explore this point further empirically in Section 6, and
discuss the sparsity the linear system from a theoretical perspective in
Appendix F.

###  4.2 DCFP with a generative model

We now return to the setting where the Markov reward process in which we are
performing evaluation is unknown, and instead we have access to the random
next-state samples $((X^{x}_{i})_{i=1}^{N}:x\in\mathcal{X})$, as described in
Section 3.

Our model-based DCFP algorithm proceeds by first constructing the
corresponding empirical transition probabilities $\hat{P}$, so that
$\hat{P}(y|x)=N^{-1}\sum_{i=1}^{N}\mathbbm{1}\\{X^{x}_{i}=y\\}$, and then
calls the DCFP procedure outlined in Algorithm 1, treating $\hat{P}$ as the
true transition probabilities of the MRP when constructing the matrix in Line
2, which we denote here by $T_{\hat{P}}$, to reflect the fact that it is built
from $\hat{P}$. This produces the output CDF values $\hat{F}$, from which
estimated return distributions can be decoded as

$\displaystyle\hat{\eta}(x)=\sum_{i=1}^{m}(\hat{F}_{i}(x)-\hat{F}_{i-1}(x))\delta_{z_{i}}\,,$
(12)

where by convention $F_{0}(x)=0$.

##  5 Sample complexity analysis

Our goal now is to analyse the sample complexity of model-based DCFP, as
described in the previous section. We first introduce a notational shorthand
that will be of extensive use in the statement and proof of this result: when
writing distances between distributions, such as
$w_{1}(\eta^{*}(x),\hat{F}(x))$, we identify CDF vectors
$\hat{F}(x)\in\mathbb{R}^{m}$ with the distributions they represent as in
Equation (12). The core theoretical result of the paper is as follows.

######  Theorem 5.1.

Let $\varepsilon\in(0,(1-\gamma)^{-1/2})$ and $\delta\in(0,1)$, and suppose
the number of categories satisfies $m\geq 4(1-\gamma)^{-2}\varepsilon^{-2}+1$.
Then the output $\hat{F}$ of model-based DCFP with
$N=\Omega(\varepsilon^{-2}(1-\gamma)^{-3}\mathrm{polylog}(|\mathcal{X}|/\delta))$
samples satisfies

$\displaystyle\max_{x\in\mathcal{X}}w_{1}(\eta^{*}(x),\hat{F}(x))\leq\varepsilon\,,$

with probability at least $1-\delta$.

This result establishes that model-based DCFP attains the minimax lower-bound
(up to logarithmic factors) for high-probability return distribution
estimation in Wasserstein distance, resolving an open question raised by Zhang
et al. (2023); in a certain sense, _estimation of return distributions is no
more statistically difficult than that of mean returns with a generative
model_. We emphasise that there is no direct dependence of $N$ on $m$, so
there is no statistical penalty to using a large number of categories $m$.

###  5.1 Structure of the proof of Theorem 5.1

The remainder of this section provides a sketch proof of Theorem 5.1; a
complete proof is provided in the appendix. The proof is broadly motivated by
the approaches of Azar et al. (2013), Agarwal et al. (2020), and Pananjady &
Wainwright (2020), who analyse the mean-return case, and we highlight where
key ideas and new mathematical objects are required in this distributional
setting. In particular, we highlight the use of the _stochastic categorical
CDF Bellman equation_ , a new distributional Bellman equation that plays a key
role in our analysis, which we expect to be of independent interest.

Reduction to Cramér distance.  The first step of the analysis is to reduce
Theorem 5.1 to a statement about approximation in Cramér distance, which is
much better suited to the analysis of DCFP, owing to the results described in
Section 2.3. This can be done by upper-bounding Wasserstein distance by Cramér
distance using the following result, which is proven in the appendix via
Jensen’s inequality.

######  Lemma 5.2.

For any two distributions
$\nu,\nu^{\prime}\in\mathscr{P}([0,(1-\gamma)^{-1}])$, we have

$\displaystyle
w_{1}(\nu,\nu^{\prime})\leq(1-\gamma)^{-1/2}\ell_{2}(\nu,\nu^{\prime})\,.$

Theorem 5.1 now follows from the following result, which quantifies the sample
complexity of model-based DCFP when measuring distances using the Cramér
distance.

######  Theorem 5.3.

Let $\varepsilon\in(0,1)$ and $\delta\in(0,1)$, and suppose the number of
categories satisfies $m\geq 4(1-\gamma)^{-2}\varepsilon^{-2}+1$. Then the
output $\hat{F}$ of model-based DCFP with
$N=\Omega(\varepsilon^{-2}(1-\gamma)^{-2}\mathrm{polylog}(|\mathcal{X}|/\delta))$
samples satisfies

$\displaystyle\max_{x\in\mathcal{X}}\ell_{2}(\eta^{*}(x),\hat{F}(x))\leq\varepsilon\,,$
(13)

with probability at least $1-\delta$.

Reduction to categorical fixed-point error. Our first step in proving Theorem
5.3 is to use the triangle inequality to split the Cramér distance on the
left-hand side of Equation (13) into a representation approximation error, and
sample-based error:

$\displaystyle\overline{\ell}_{2}(\eta^{*},\hat{F})$
$\displaystyle\leq\overline{\ell}_{2}(\eta^{*},F^{*})+\overline{\ell}_{2}(F^{*},\hat{F})$
$\displaystyle\leq\frac{1}{(1-\gamma)\sqrt{m-1}}+\overline{\ell}_{2}(F^{*},\hat{F})\,,$

with the second inequality following from the fixed-point quality bound in
Equation (5). With $m$ as specified in the theorem statement, the first term
in the right-hand side above is bounded by $\varepsilon/2$. Thus, it suffices
to focus on the second term on the right-hand side, which quantifies the
sample-based error in estimating the categorical fixed point $F^{*}$.

Concentration. Through a combination of the use of a version of Bernstein’s
inequality in Hilbert space (Chatalic et al., 2022) and propagation of this
inequality across time steps in the MRP, we next arrive at the following
inequality with probability at least $1-\delta$:

$\displaystyle\ell_{2}(\hat{F}(x),F^{*}(x))\leq$ (14)
$\displaystyle\widetilde{O}\bigg{(}\\!\frac{1}{\sqrt{N(1-\gamma)}}\\!\sqrt{\|(I-\gamma\hat{P})^{-1}\sigma_{\hat{P}}\|_{\infty}}\\!+\\!\frac{1}{(1-\gamma)^{3/2}N^{3/4}}\\!\bigg{)}\,.$

Here, $\sigma_{\hat{P}}\in\mathbb{R}^{\mathcal{X}}$ is an instance of a new
class of distributional object, the _local squared-Cramér variation_. The
general definition of this object is given below, for the case of a general
transition matrix $Q\in\mathbb{R}^{\mathcal{X}\times\mathcal{X}}$, to avoid
conflation with the specific transition matrices $P$ and $\hat{P}$ that
Theorem 5.3 is concerned with.

######  Definition 5.4.

For a given transition matrix $Q$, the _single-sample operator_
$\hat{T}_{Q}:\mathbb{R}^{\mathcal{X}\times
m}\rightarrow\mathbb{R}^{\mathcal{X}\times m}$ is the random operator given
by: (i) constructing a _random transition matrix_ $\hat{Q}$ by, for each
$x\in\mathcal{X}$, sampling $X^{\prime}\sim Q(\cdot|x)$, and setting
$\hat{Q}(X^{\prime}|x)=1$; (ii) setting $\hat{T}_{Q}=T_{\hat{Q}}$.

######  Definition 5.5.

For a given transition matrix $Q$ with corresponding CDP fixed point
$F^{Q}\in\mathbb{R}^{\mathcal{X}\times m}$, the _local squared-Cramér
variation at $Q$_, $\sigma_{Q}\in\mathbb{R}^{\mathcal{X}}$, is defined by

$\displaystyle\sigma_{Q}(x)=\mathbb{E}[\ell_{2}^{2}((\hat{T}_{Q}F^{Q})(x),F^{Q}(x))]\,.$

Intuitively, the local squared-Cramér variation $\sigma_{Q}$ encodes the
variability of the fixed point $F^{Q}$ after a sample-based, rather than
exact, dynamic programming update. From this point of view, it is a natural
quantity to arise in Equation (14), and plays a similar role to the variance
in the classical Bernstein inequality (Bernstein, 1946).

In Corollary 5.12 below, we will deduce that under the conditions of Theorem
5.3, we have

$\displaystyle\|(I-\gamma\hat{P})^{-1}\sigma_{\hat{P}}\|_{\infty}\leq\frac{2}{1-\gamma}\,.$
(15)

Substituting this into Equation (14) gives

$\displaystyle\ell_{2}(\hat{F}(x),F^{*}(x))\leq\widetilde{O}\bigg{(}\frac{1}{(1-\gamma)\sqrt{N}}+\frac{1}{(1-\gamma)^{3/2}N^{3/4}}\bigg{)}$

with probability at least $1-\delta$, for all $x\in\mathcal{X}$. Now, taking
$N=\widetilde{\Omega}((1-\gamma)^{-2}\varepsilon^{-2})$ yields that this
expression is $O(\varepsilon)$, which completes the sketch proof of Theorem
5.3.

What remains to be described is how to arrive at the bound in Equation (15);
The section below provides a high-level overview of the technical details
involved in obtaining it.

###  5.2 The stochastic categorical CDF Bellman equation

The central idea is to relate the _local_ squared-Cramér variation to a
corresponding _global_ notion of variation, in analogy with the variance
Bellman equation (Sobel, 1982) used by Azar et al. (2013) in the mean-return
case. To define this corresponding global notion, we begin by defining a new
type of distributional Bellman equation, which can be intuitively thought of
as encoding the result of repeatedly applying a sequence of independent
single-sample operators. Again, we work with a general transition matrix $Q$.

######  Definition 5.6.

For a general transition matrix $Q$, the _stochastic categorical CDF (SC-CDF)
Bellman equation_ is given by

$\displaystyle\Phi(x)\overset{\mathcal{D}}{=}(\hat{T}_{Q}\Phi)(x)\,,$ (16)

where $\overset{\mathcal{D}}{=}$ denotes equality in distribution. Here,
$\hat{T}_{Q}$ is a _single-sample operator_ with respect to $Q$, as in
Definition 5.4. Each $\Phi(x)$ is a random variable taking values in the space
of valid CDF values

$\displaystyle\mathscr{F}=\\{F\in\mathbb{R}^{m}:0\leq F_{1}\leq\cdots\leq
F_{m-1}\leq F_{m}=1\\}$

for distributions in $\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})$, and importantly
is taken to be independent of the random operator $\hat{T}_{Q}$.

The intuition is that “unravelling” Equation (16) should lead to a solution of
the form
$\Phi(x)\overset{\mathcal{D}}{=}\lim_{k\rightarrow\infty}\hat{T}^{(k)}_{Q}\cdots\hat{T}^{(1)}_{Q}F$,
where $(\hat{T}^{(i)}_{Q})_{i=1}^{k}$ are independent single-sample operators,
so that $\Phi(x)$ encodes the fluctuations due to repeated CDP updates with
randomly-sampled transitions. To make this intuition precise, we first verify
that the SC-CDF Bellman equation has a unique solution.

######  Proposition 5.7.

The SC-CDF Bellman equation in Equation (16) has a unique solution, in the
sense that there is a unique distribution for each $\Phi(x)$ such that
Equation (16) holds for each $x\in\mathcal{X}$.

We write $\Phi^{Q}$ for the collection of random CDFs that satisfy the SC-CDF
Bellman equation in Equation (16), whose existence is guaranteed by
Proposition 5.7. We next relate $\Phi^{Q}$ to the standard categorical fixed
point $F^{Q}$.

######  Proposition 5.8.

We have $\mathbb{E}[\Phi^{Q}(x)]=F^{Q}(x)$ for all $x\in\mathcal{X}$.

Proposition 5.8 shows that $\Phi^{Q}$ can indeed thought of as encoding random
variation around the usual categorical fixed point $F^{Q}$; see Figure 2. This
motivates the following.

######  Definition 5.9.

For a given transition matrix $Q$ with corresponding CDP fixed point
$F^{Q}\in\mathbb{R}^{\mathcal{X}\times m}$, the _global squared-Cramér
variation at_ $Q$, $\Sigma_{Q}\in\mathbb{R}^{\mathcal{X}}$, is defined by

$\displaystyle\Sigma_{Q}(x)=\mathbb{E}[\ell_{2}^{2}(\Phi^{Q}(x),F^{Q}(x))]\,.$

######  Remark 5.10.

We pause to consider the nature of the new object defined by the SC-CDF
Bellman equation; $\Phi^{Q}(x)$ is a _doubly_ distributional object. It
represents a probability distribution around the object $F^{Q}(x)$, which
itself already summarises the distribution of the return. This reveals a dual
perspective on distributional RL itself: the randomness modelled can be: (i)
of intrinsic interest, and/or (ii) instrumental in analysing the estimation of
a non-random object from random data (c.f. the variance Bellman equation in
analysing the estimation of mean returns (Azar et al., 2013)). The object
$\Phi^{Q}$ is motivated by _both_ of these concerns.

Figure 2:  Left: Example MRP with $r(x_{0})=1,r(x_{1})=0$, $\gamma=0.9$.
Right: Categorical fixed point $F^{*}(x_{0})$ with $m=15$, and 5 independent
samples from the random CDF $\Phi^{*}(x_{0})$.

The following Bellman-like inequality draws a relationship between local and
global squared-Cramér variation, allowing us to make progress from Equation
(15).

######  Proposition 5.11.

We have

$\displaystyle\Sigma_{Q}\geq\sigma_{Q}+\gamma
Q\Sigma_{Q}-\bigg{(}\frac{2}{m\sqrt{1-\gamma}}+\frac{1}{m^{2}(1-\gamma)^{2}}\bigg{)}\mathbf{1}\,,$

where $\mathbf{1}\in\mathbb{R}^{\mathcal{X}}$ is a vector of ones, and the
inequality above is interpreted component-wise.

Rearrangement and bounding of the quantities in the inequality of Proposition
5.11, in the specific case $Q=\hat{P}$, then allows us to deduce the required
inequality in Equation (15) that completes the proof of Theorem 5.3.

######  Corollary 5.12.

We can bound the term $\|(I-\gamma\hat{P})^{-1}\sigma_{\hat{P}}\|_{\infty}$
appearing in Equation (14) under the assumptions of Theorem 5.3 as follows:

$\displaystyle\|(I-\gamma\hat{P})^{-1}\sigma_{\hat{P}}\|_{\infty}\leq\|\Sigma_{\hat{P}}\|_{\infty}+\frac{1}{1-\gamma}\leq\frac{2}{1-\gamma}\,.$

##  6 Empirical evaluation

To complement our theoretical analysis, which focuses on worst-case sample
complexity bounds, we report empirical findings for implementations of several
model-based distributional RL algorithms in the generative model setting. We
compare the new DCFP algorithm introduced in Section 4.1 with _quantile
dynamic programming_ (Dabney et al., 2018b; Bellemare et al., 2023, QDP) a
distinct approach to distributional RL in which return distributions are
approximated via dynamic programming with a finite number of quantiles.

In Figure 3 (left), we report results of running DCFP and QDP on a 5-state
environment with transition matrix randomly sampled from Dirichlet
distributions, and entries of the immediate reward function
$r\in\mathbb{R}^{\mathcal{X}}$ randomly sampled from $\text{Unif}([0,1])$,
with varying numbers of atoms $m$, environment samples per state $N$, and
discount $\gamma$. We report the maximum $w_{1}$-error against true return
distributions (estimated via Monte Carlo sampling). All runs are repeated 30
times, and error bars are 95% bootstrapped confidence intervals. Sufficient DP
iterations ensure approximate convergence to their fixed points. In Figure 3
(right), we plot estimation error against wallclock time for $m=100,N=10^{6}$,
including results for CDP (which approximates the solution of DCFP via dynamic
programming; Section 2.3); line plots indicate the estimation error/wallclock
time trade-off as we increase the number of DP iterations.

For low discount factors/atom counts QDP generally outperforms DCFP in terms
of asymptotic estimation error, due to QDP’s ability to modify its atom
support to regions of the interval $[0,(1-\gamma)^{-1}]$ where mass is
concentrated. However, we note that DCFP is generally faster than QDP.
Further, the gap between DCFP and QDP, in terms of estimation error, is vastly
reduced for larger discount and larger atom count. We also note that
particularly at high discounts, DCFP outperforms CDP and QDP in terms of
wallclock time, due to DP methods requiring many iterations to converge in
these cases (Rowland et al., 2018; Dabney et al., 2018b). Full results,
including on several further environments, are given in Appendix F: key
findings include that QDP works particularly well in near-deterministic
environments, and DCFP works particularly well in settings where there are
short high-probability paths from a state to itself. We also investigate
versions of DCFP and CDP that set atom locations $(z_{i})_{i=1}^{m}$ based on
environment-specific information, which can significantly perform the accuracy
of these methods for a given atom count $m$.

Figure 3: Wasserstein-1 approximation error and wallclock time for a variety
of distributional RL methods, discount factors, numbers of atoms, and numbers
of environment samples.

##  7 Related work and directions for future work

Here, we describe the most salient aspects of related work, and promising
directions for future work; an extended discussion is given in Appendix A.

Alternative distributional RL algorithms.  Our analysis has focused on
categorical distributional RL, and in particular has made use of linearity of
the corresponding CDF DP operator both as a function of input CDFs, and of
transition probabilities. Our experiments also compare against QDP, a tabular
dynamic programming version of a widely used approach to distributional RL
based on quantiles; these results suggest strong theoretical bounds may also
be possible for QDP, although the non-linearity of the operator may complicate
the analysis. Investigation of other approaches such as fitted likelihood
evaluation (Wu et al., 2023, FLE), in the generative model setting, is another
interesting direction.

Other statistical questions in distributional RL.  Zhang et al. (2023) make
several other contributions on the topic of statistical estimation in
distributional RL, including analysis of asymptotic fluctuations and limit
theorems. Wu et al. (2023) study the offline evaluation problem, in which
state-action pairs are sampled from the stationary distribution of the policy,
via FLE and focus on generalisation bounds. Böck & Heirzinger (2022) propose a
_model-free_ algorithm for distributional RL, _speedy categorical TD-learning_
, motivated by categorical TD learning (Rowland et al., 2018) and speedy
Q-learning (Azar et al., 2011), and prove a sample complexity bound of
$\widetilde{O}(\varepsilon^{-2}(1-\gamma)^{-3})$ for high probability
$\varepsilon$-accurate estimation in Cramér distance, implying a sample
complexity of $\widetilde{O}(\varepsilon^{-2}(1-\gamma)^{-4})$ in $w_{1}$, by
Lemma 5.2.

Risk-sensitive control. The theory developed in this paper has focused on
estimation of return distributions for individual policies. A natural
direction for future work is to analyse identification of near-optimal
policies for risk-sensitive decision criteria. Bastani et al. (2022); Wang et
al. (2023a) study efficient algorithms for CVaR optimisation, while Fei et al.
(2021a, b); Liang & Luo (2022) study entropic risk maximisation, and Du et al.
(2022); Lam et al. (2022) study iterated CVaR optimisation, all in the
_online_ setting.

##  8 Conclusion

We have introduced a new algorithm, DCFP, for directly computing the fixed
point of CDP, a widely used distributional reinforcement learning algorithm.
We then showed that this algorithm, with an appropriately chosen number of
categories $m$, achieves the minimax lower bound (up to logarithmic factors)
for sample complexity of return-distribution estimation in Wasserstein
distance obtained by Zhang et al. (2023). Thus, this paper closes an open
question raised by Zhang et al. (2023) by exhibiting an algorithm that obtains
this lower bound, and shows that estimation of return distributions via a
generative model is no harder statistically than the analogous task for mean
returns. Our analysis also requires the introduction of the novel stochastic
categorical CDF Bellman equation, which we expect to be of independent
interest, and our empirical study reveals several important take-aways for
practitioners, including findings that the optimal choice distributional RL
algorithm may depend on several particular aspects of the environment.

## Acknowledgements

We thank Dave Abel for detailed comments on a draft version of the paper. We
also thank Mohammad Gheshlaghi Azar for useful conversations, and Arthur
Gretton for advice on Bernstein-like inequalities in Hilbert space.

## References

  * Agarwal et al. (2020) Agarwal, A., Kakade, S., and Yang, L. F.  Model-based reinforcement learning with a generative model is minimax optimal.  In _Conference on Learning Theory_ , 2020.
  * Azar et al. (2011) Azar, M. G., Munos, R., Ghavamzadeh, M., and Kappen, H.  Speedy Q-learning.  In _Advances in Neural Information Processing Systems_ , 2011.
  * Azar et al. (2013) Azar, M. G., Munos, R., and Kappen, H. J.  Minimax PAC bounds on the sample complexity of reinforcement learning with a generative model.  _Machine learning_ , 91:325–349, 2013.
  * Bastani et al. (2022) Bastani, O., Ma, J. Y., Shen, E., and Xu, W.  Regret bounds for risk-sensitive reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Bellemare et al. (2017) Bellemare, M. G., Dabney, W., and Munos, R.  A distributional perspective on reinforcement learning.  In _Proceedings of the International Conference on Machine Learning_ , 2017.
  * Bellemare et al. (2020) Bellemare, M. G., Candido, S., Castro, P. S., Gong, J., Machado, M. C., Moitra, S., Ponda, S. S., and Wang, Z.  Autonomous navigation of stratospheric balloons using reinforcement learning.  _Nature_ , 588(7836):77–82, 2020.
  * Bellemare et al. (2023) Bellemare, M. G., Dabney, W., and Rowland, M.  _Distributional Reinforcement Learning_.  MIT Press, 2023.  http://www.distributional-rl.org.
  * Bernstein (1946) Bernstein, S.  _The Theory of Probabilities_.  Gastehizdat Publishing House, 1946.
  * Bhat & Prashanth (2019) Bhat, S. P. and Prashanth, L. A.  Concentration of risk measures: A Wasserstein distance approach.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Böck & Heirzinger (2022) Böck, M. and Heirzinger, C.  Speedy categorical distributional reinforcement learning and complexity analysis.  _SIAM Journal on Mathematics of Data Science_ , 4(2):675–693, 2022.
  * Böck et al. (2022) Böck, M., Malle, J., Pasterk, D., Kukina, H., Hasani, R., and Heitzinger, C.  Superhuman performance on sepsis MIMIC-III data by distributional reinforcement learning.  _PLoS One_ , 17(11):e0275358, 2022.
  * Bodnar et al. (2020) Bodnar, C., Li, A., Hausman, K., Pastor, P., and Kalakrishnan, M.  Quantile QT-Opt for risk-aware vision-based robotic grasping.  In _Robotics: Science and Systems_ , 2020.
  * Chandak et al. (2021) Chandak, Y., Shankar, S., and Thomas, P. S.  High-confidence off-policy (or counterfactual) variance estimation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2021.
  * Chatalic et al. (2022) Chatalic, A., Schreuder, N., Rosasco, L., and Rudi, A.  Nyström kernel mean embeddings.  In _Proceedings of the International Conference on Machine Learning_ , 2022.
  * Cramér (1928) Cramér, H.  On the composition of elementary errors.  _Scandinavian Actuarial Journal_ , 1928(1):13–74, 1928.
  * Dabney et al. (2018a) Dabney, W., Ostrovski, G., Silver, D., and Munos, R.  Implicit quantile networks for distributional reinforcement learning.  In _Proceedings of the International Conference on Machine Learning_ , 2018a.
  * Dabney et al. (2018b) Dabney, W., Rowland, M., Bellemare, M. G., and Munos, R.  Distributional reinforcement learning with quantile regression.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2018b.
  * Dabney et al. (2020) Dabney, W., Kurth-Nelson, Z., Uchida, N., Starkweather, C. K., Hassabis, D., Munos, R., and Botvinick, M.  A distributional code for value in dopamine-based reinforcement learning.  _Nature_ , 577(7792):671–675, 2020.
  * Doan et al. (2018) Doan, T., Mazoure, B., and Lyle, C.  GAN Q-learning.  _arXiv_ , 2018.
  * Du et al. (2022) Du, Y., Wang, S., and Huang, L.  Provably efficient risk-sensitive reinforcement learning: Iterated CVaR and worst path.  In _Proceedings of the International Conference on Learning Representations_ , 2022.
  * Dvoretzky et al. (1956) Dvoretzky, A., Kiefer, J., and Wolfowitz, J.  Asymptotic minimax character of the sample distribution function and of the classical multinomial estimator.  _The Annals of Mathematical Statistics_ , pp. 642–669, 1956.
  * Fawzi et al. (2022) Fawzi, A., Balog, M., Huang, A., Hubert, T., Romera-Paredes, B., Barekatain, M., Novikov, A., R Ruiz, F. J., Schrittwieser, J., Swirszcz, G., Silver, D., Hassabis, D., and Kohli, P.  Discovering faster matrix multiplication algorithms with reinforcement learning.  _Nature_ , 610(7930):47–53, 2022.
  * Fei et al. (2021a) Fei, Y., Yang, Z., Chen, Y., and Wang, Z.  Exponential bellman equation and improved regret bounds for risk-sensitive reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2021a.
  * Fei et al. (2021b) Fei, Y., Yang, Z., and Wang, Z.  Risk-sensitive reinforcement learning with function approximation: A debiasing approach.  In _Proceedings of the International Conference on Machine Learning_ , 2021b.
  * Freirich et al. (2019) Freirich, D., Shimkin, T., Meir, R., and Tamar, A.  Distributional multivariate policy evaluation and exploration with the Bellman GAN.  In _Proceedings of the International Conference on Machine Learning_ , 2019.
  * Gretton et al. (2012) Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A.  A kernel two-sample test.  _The Journal of Machine Learning Research_ , 13(1):723–773, 2012.
  * Harris et al. (2020) Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, J. F., Wiebe, M., Peterson, P., Gérard-Marchant, P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C., and Oliphant, T. E.  Array programming with NumPy.  _Nature_ , 585(7825):357–362, September 2020\.  doi: 10.1038/s41586-020-2649-2.
  * Hunter (2007) Hunter, J. D.  Matplotlib: A 2d graphics environment.  _Computing in Science & Engineering_, 9(3):90–95, 2007.
  * Kakade (2003) Kakade, S. M.  _On the sample complexity of reinforcement learning_.  PhD thesis, University College London, 2003.
  * Kearns et al. (2002) Kearns, M., Mansour, Y., and Ng, A. Y.  A sparse sampling algorithm for near-optimal planning in large Markov decision processes.  _Machine learning_ , 49:193–208, 2002.
  * Lam et al. (2022) Lam, T., Verma, A., Low, B. K. H., and Jaillet, P.  Risk-aware reinforcement learning with coherent risk measures and non-linear function approximation.  In _Proceedings of the International Conference on Learning Representations_ , 2022.
  * Li et al. (2020) Li, G., Wei, Y., Chi, Y., Gu, Y., and Chen, Y.  Breaking the sample size barrier in model-based reinforcement learning with a generative model.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Liang & Luo (2022) Liang, H. and Luo, Z.-Q.  Bridging distributional and risk-sensitive reinforcement learning with provable regret bounds.  _arXiv_ , 2022.
  * Massart (1990) Massart, P.  The tight constant in the Dvoretzky-Kiefer-Wolfowitz inequality.  _The Annals of Probability_ , pp. 1269–1283, 1990.
  * Morimura et al. (2010a) Morimura, T., Sugiyama, M., Kashima, H., Hachiya, H., and Tanaka, T.  Nonparametric return distribution approximation for reinforcement learning.  In _Proceedings of the International Conference on Machine Learning_ , 2010a.
  * Morimura et al. (2010b) Morimura, T., Sugiyama, M., Kashima, H., Hachiya, H., and Tanaka, T.  Parametric return density estimation for reinforcement learning.  In _Proceedings of the Conference on Uncertainty in Artificial Intelligence_ , 2010b.
  * Nguyen-Tang et al. (2021) Nguyen-Tang, T., Gupta, S., and Venkatesh, S.  Distributional reinforcement learning via moment matching.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2021.
  * Pananjady & Wainwright (2020) Pananjady, A. and Wainwright, M. J.  Instance-dependent $\ell_{\infty}$-bounds for policy evaluation in tabular reinforcement learning.  _IEEE Transactions on Information Theory_ , 67(1):566–585, 2020.
  * Peyré & Cuturi (2019) Peyré, G. and Cuturi, M.  _Computational optimal transport: With applications to data science_.  Foundations and Trends® in Machine Learning. Now Publishers, Inc., 2019.
  * Puterman (2014) Puterman, M. L.  _Markov decision processes: Discrete stochastic dynamic programming_.  John Wiley & Sons, 2014.
  * Rockafellar & Uryasev (2002) Rockafellar, R. T. and Uryasev, S.  Conditional value-at-risk for general loss distributions.  _Journal of banking & finance_, 26(7):1443–1471, 2002.
  * Rockafellar et al. (2000) Rockafellar, R. T., Uryasev, S., et al.  Optimization of conditional value-at-risk.  _Journal of risk_ , 2:21–42, 2000.
  * Rowland et al. (2018) Rowland, M., Bellemare, M. G., Dabney, W., Munos, R., and Teh, Y. W.  An analysis of categorical distributional reinforcement learning.  In _Proceedings of the International Conference on Artificial Intelligence and Statistics_ , 2018.
  * Rowland et al. (2023) Rowland, M., Munos, R., Azar, M. G., Tang, Y., Ostrovski, G., Harutyunyan, A., Tuyls, K., Bellemare, M. G., and Dabney, W.  An analysis of quantile temporal-difference learning.  _arXiv_ , 2023.
  * Sejdinovic et al. (2013) Sejdinovic, D., Sriperumbudur, B., Gretton, A., and Fukumizu, K.  Equivalence of distance-based and RKHS-based statistics in hypothesis testing.  _The Annals of Statistics_ , pp. 2263–2291, 2013.
  * Shahriari et al. (2022) Shahriari, B., Abdolmaleki, A., Byravan, A., Friesen, A., Liu, S., Springenberg, J. T., Heess, N., Hoffman, M., and Riedmiller, M.  Revisiting gaussian mixture critics in off-policy reinforcement learning: A sample-based approach.  _arXiv_ , 2022.
  * Sidford et al. (2018) Sidford, A., Wang, M., Wu, X., Yang, L., and Ye, Y.  Near-optimal time and sample complexities for solving Markov decision processes with a generative model.  _Advances in Neural Information Processing Systems_ , 2018.
  * Sobel (1982) Sobel, M. J.  The variance of discounted Markov decision processes.  _Journal of Applied Probability_ , 19(4):794–802, 1982.
  * Székely (2003) Székely, G. J.  E-statistics: The energy of statistical samples.  Technical Report 05, Bowling Green State University, Department of Mathematics and Statistics Technical Report, 2003.
  * Székely & Rizzo (2013) Székely, G. J. and Rizzo, M. L.  Energy statistics: A class of statistics based on distances.  _Journal of statistical planning and inference_ , 143(8):1249–1272, 2013.
  * Van Rossum & Drake (2009) Van Rossum, G. and Drake, F. L.  _Python 3 Reference Manual_.  CreateSpace, Scotts Valley, CA, 2009.
  * Villani (2009) Villani, C.  _Optimal transport: Old and new_ , volume 338.  Springer, 2009.
  * Virtanen et al. (2020) Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M., Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, J., van der Walt, S. J., Brett, M., Wilson, J., Millman, K. J., Mayorov, N., Nelson, A. R. J., Jones, E., Kern, R., Larson, E., Carey, C. J., Polat, İ., Feng, Y., Moore, E. W., VanderPlas, J., Laxalde, D., Perktold, J., Cimrman, R., Henriksen, I., Quintero, E. A., Harris, C. R., Archibald, A. M., Ribeiro, A. H., Pedregosa, F., van Mulbregt, P., and SciPy 1.0 Contributors.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ , 17:261–272, 2020.
  * Wang et al. (2023a) Wang, K., Kallus, N., and Sun, W.  Near-minimax-optimal risk-sensitive reinforcement learning with CVaR.  In _Proceedings of the International Conference on Machine Learning_ , 2023a.
  * Wang et al. (2023b) Wang, K., Zhou, K., Wu, R., Kallus, N., and Sun, W.  The benefits of being distributional: Small-loss bounds for reinforcement learning.  _arXiv_ , 2023b.
  * Wang et al. (2023c) Wang, Z., Gao, Y., Wang, S., Zavlanos, M. M., Abate, A., and Johansson, K. H.  Policy evaluation in distributional LQR (extended version).  _arXiv_ , 2023c.
  * Waskom (2021) Waskom, M. L.  seaborn: statistical data visualization.  _Journal of Open Source Software_ , 6(60):3021, 2021.
  * Wu et al. (2023) Wu, R., Uehara, M., and Sun, W.  Distributional offline policy evaluation with predictive error guarantees.  In _Proceedings of the International Conference on Machine Learning_ , 2023.
  * Wurman et al. (2022) Wurman, P. R., Barrett, S., Kawamoto, K., MacGlashan, J., Subramanian, K., Walsh, T. J., Capobianco, R., Devlic, A., Eckert, F., Fuchs, F., Gilpin, L., Khandelwal, P., Kompella, V., Lin, H., MacAlpine, P., Oller, D., Seno, T., Sherstan, C., Thomure, M. D., Aghabozorgi, H., Barrett, L., Douglas, R., Whitehead, D., Dürr, P., Stone, P., Spranger, M., and Kitano, H.  Outracing champion Gran Turismo drivers with deep reinforcement learning.  _Nature_ , 602(7896):223–228, 2022.
  * Yang et al. (2019) Yang, D., Zhao, L., Lin, Z., Qin, T., Bian, J., and Liu, T.-Y.  Fully parameterized quantile function for distributional reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Yue et al. (2020) Yue, Y., Wang, Z., and Zhou, M.  Implicit distributional reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Yurinsky (2006) Yurinsky, V.  _Sums and Gaussian vectors_.  Springer, 2006.
  * Zhang et al. (2023) Zhang, L., Peng, Y., Liang, J., Yang, W., and Zhang, Z.  Estimation and inference in distributional reinforcement learning.  _arXiv_ , 2023.

## APPENDICES:
Near-Minimax-Optimal Distributional Reinforcement Learning with a Generative
Model

For convenience, we provide a summary of the contents of the appendices below.

  * •

Appendix A expands on the discussion of related work given in Section 7 in the
main paper.

  * •

Appendix B provides additional convenient notation for the categorical CDF
operator $T_{P}$, in particular expressing it as the combination of a
scaling/shifting/projecting operation, and a mixing operation over the state
to be bootstrapped from. Several additional self-contained contraction results
are given that are used within the main proofs of the paper.

  * •

Appendix C provides proofs for the results in Section 4 concerning the
formation of the linear system that is solved by the DCFP algorithm, and
verification that the derived system has the unique desired solution.

  * •

Appendix D provides proofs relating to the stochastic categorical CDF Bellman
equation, in particular establishing the desired unique solution to this new
distributional Bellman equation, and proving a key bound on the local squared-
Cramér variation that is used within the proof of the main theorem of the
paper.

  * •

Appendix E provides the proof of the main result of the paper, Theorem 5.1,
giving full details for the steps traced through in the sketch proof in the
main paper.

  * •

Appendix F provides full details relating to the experiments in the main
paper, as well as additional empirical comparisons between methods for
distributional RL with a generative model.

##  Appendix A Related work

Other families of model-based distributional RL algorithms. There are many
approaches to distributional RL; important choices studied previously include
moments (Sobel, 1982), exponential families (Morimura et al., 2010b),
categorical distributions (Bellemare et al., 2017), collections of particles
(Morimura et al., 2010a; Nguyen-Tang et al., 2021), quantiles (Dabney et al.,
2018b), and generative models (Doan et al., 2018; Dabney et al., 2018a;
Freirich et al., 2019; Yang et al., 2019; Yue et al., 2020). Our choice of
categorical representations in this work is motivated by several
considerations: (i) the existence of principled dynamic programming methods
for these representations, with corresponding convergence theory (Rowland et
al., 2018); (ii) the flexibility of these representations to trade-off
computational complexity with accuracy (by varying $m$) (Rowland et al.,
2018); (iii) the mathematical structure of the dynamic programming operator
(linear), as described in this work; and (iv) the availability of an efficient
algorithm to exactly compute the DP fixed point (the DCFP algorithm proposed
in Section 4.1). An interesting and important direction, given our empirical
findings in Section 6, is whether analyses can also be carried out for other
approaches to distributional dynamic programming, such as quantile dynamic
programming (QDP; Dabney et al., 2018b; Rowland et al., 2023; Bellemare et
al., 2023), and fitted likelihood estimation (FLE; Wu et al., 2023). We expect
challenges in extending the analysis to result due to the fact that, for
example, the QDP operator is non-linear, and FLE operator applications
typically do not have closed forms. Nevertheless, it would be particularly
interesting to understand whether it is possible to obtain instance-dependent
bounds for QDP, particularly given its strong empirical performance.
Similarly, as described in Section 5, an interesting question for future work
is whether it is possible to improve on the computational complexity of model-
based DCFP for high-probability return-direction estimation.

Other statistical questions in distributional RL. Böck & Heirzinger (2022)
propose a _model-free_ algorithm for distributional RL, speedy categorical TD-
learning, motivated by categorical TD learning (Rowland et al., 2018) and
speedy Q-learning (Azar et al., 2011), and prove a sample complexity bound of
$\widetilde{O}(\varepsilon^{-2}(1-\gamma)^{-3})$ for high probability
$\varepsilon$-accurate estimation in Cramér distance (which implies a non-
minimax sample complexity of $\widetilde{O}(\varepsilon^{-2}(1-\gamma)^{-4})$
in Wasserstein-1 distance, per Lemma 5.2). Wu et al. (2023) study the offline
evaluation problem, in which state-action pairs are sampled from the
stationary distribution of the policy, via fitted likelihood estimation (FLE)
and focus on generalisation bounds, allowing for policy evaluation in
environments with uncountable state spaces. Wang et al. (2023c) also study
policy evaluation in this context, focusing on the LQR model. Aside from
studying minimax optimality, Zhang et al. (2023) also make several other
contributions on the topic of statistical efficiency of estimation for
distributional RL, including analysis of asymptotic fluctuations and limit
theorems, and analysis of approximations in more general metrics, including
Wasserstein-$p$ , Kolmogorov-Smirnov, and total variation metrics. Their
sample complexity results rely on careful analysis of the behaviour of the
_unprojected_ distributional Bellman operator $\mathcal{T}$ on certain
subspaces of probability/signed measures.

Sample complexity of mean-return estimation. The sample complexity of
estimating mean returns, and the related task of obtaining a near-optimal
policy, has been considered by Azar et al. (2013); Sidford et al. (2018);
Pananjady & Wainwright (2020); Agarwal et al. (2020); Li et al. (2020).
Interestingly, Agarwal et al. (2020), while treating the mean-return case,
make use of return-binning as a proof technique. Li et al. (2020) also obtain
bounds for mean-return sample complexity that apply with
$\varepsilon>(1-\gamma)^{-1/2}$ for modifications of certainty-equivalent
model-based algorithms; an interesting direction for future work would be to
check whether the restrictions on $\varepsilon$ in Theorem 5.1 can be lifted
by incorporating ideas from this mean-return analysis to the distributional
setting. Additionally, Chandak et al. (2021) consider the task of estimating
the variance of returns from off-policy data, and Wang et al. (2023b) study
regret minimisation properties (with respect to the expected return criterion)
of distributional RL algorithms in the online setting.

Further analysis. In this paper, we have resolved a conjecture of Zhang et al.
(2023), obtaining a near-minimax-optimal algorithm for estimation of return
distributions in Wasserstein-1 distance. Zhang et al. (2023) make
contributions to several other important statistical questions regarding
distributional RL, including approximation in stronger metrics such as
Kolmogorov-Smirnov and total variation metrics, as well as studying asymptotic
fluctuations of estimators; it will be interesting to see whether the analysis
presented here can be extended to these other settings as well.

##  Appendix B Additional CDF operator notation and contractivity results

In this section, we introduce finer-grained notation for the CDF operator
$T_{P}$ that allows us to straightforwardly refer to the operations that
correspond to shifting/scaling/projecting, and to mixing over transition
states, separately. We also establish several additional contraction lemmas
that will be useful in the sections that follow.

###  B.1 Categorical hat function definition

For convenience, we provide the full mathematical definition of the hat
functions $h_{i}:[0,(1-\gamma)^{-1}]\rightarrow[0,1]$ used in defining the
categorical projection described in the main paper, as given by Rowland et al.
(2018). For $i=2,\ldots,m-1$, we have

$\displaystyle h_{i}(z)=\begin{cases}\frac{z-z_{i-1}}{z_{i}-z_{i-1}}&\text{
for }z\in[z_{i-1},z_{i}]\\\ \frac{z_{i+1}-z}{z_{i+1}-z_{i}}&\text{ for
}z\in[z_{i},z_{i+1}]\\\ 0&\text{ otherwise. }\end{cases}$

For the edge case $h_{1}$, we have

$\displaystyle h_{1}(z)=\begin{cases}\frac{z_{2}-z}{z_{2}-z_{1}}&\text{ for
}z\in[z_{1},z_{2}]\\\ 0&\text{ otherwise, }\end{cases}$

and similarly for the edge case $h_{m}$, we have

$\displaystyle h_{m}(z)=\begin{cases}\frac{z-z_{m-1}}{z_{m}-z_{m-1}}&\text{
for }z\in[z_{m-1},z_{m}]\\\ 0&\text{ otherwise. }\end{cases}$

###  B.2 Finer-grained operator expression

Recall that the CDF operator $T_{P}:\mathbb{R}^{\mathcal{X}\times
m}\rightarrow\mathbb{R}^{\mathcal{X}\times m}$ is represented by a matrix with
elements given by

$\displaystyle T_{P}(x,i;j,y)=P(y|x)(H^{x}_{i,j}-H^{x}_{i,j+1})\,.$ (17)

We can therefore conceptualise the application $T_{P}F$ as standard matrix-
vector multiplication in the vector space $\mathbb{R}^{\mathcal{X}\times m}$.
However, the expression for matrix elements in Equation (17) has additional
structure that means we can express the application of $T_{P}$ in a different
manner, which will be convenient in several proofs below, particularly as it
separates out the influence of rewards (which remain fixed) and transition
dynamics (which are estimated via samples).

In particular, we will regard $F\in\mathbb{R}^{\mathcal{X}\times m}$ itself as
a matrix, with rows indexed by states in $\mathcal{X}$, and columns indexed by
indices $i=1,\ldots,m$. For each state $x\in\mathcal{X}$, we then introduce
the matrix $B_{x}\in\mathbb{R}^{m\times m}$, with $(i,j)$ element given by

$\displaystyle H^{x}_{i,j}-H^{x}_{i,j+1}\,.$

We then have that $(T_{P}F)(x)\in\mathbb{R}^{m}$ can alternatively be
expressed in matrix notation as

$\displaystyle P_{x}FB_{x}^{\top}\,,$

where $P_{x}$ is the row vector given by the row of $P$ corresponding to state
$x\in\mathcal{X}$.

###  B.3 Additional results

Below, we provide several additional results regarding contractivity
properties of $T_{P}$. To do so, it is useful to introduce the norm
$\|\cdot\|_{\ell_{2}}$ on $\mathbb{R}^{m}$, which we define by

$\displaystyle\|F\|_{\ell_{2}}=\bigg{[}\frac{1}{m(1-\gamma)}\sum_{i=1}^{m}F_{i}(x)^{2}\bigg{]}^{1/2}\,.$

The motivation for this definition is that if we have two distributions
$\nu,\nu^{\prime}\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})$ with corresponding
CDF vectors $F,F^{\prime}\in\mathbb{R}^{m}$, then
$\ell_{2}(\nu,\nu^{\prime})=\|F-F^{\prime}\|_{\ell_{2}}$, as
$\nu,\nu^{\prime}$ are supported on $[0,(1-\gamma)^{-1}]$. Thus, under the
abuse of notation $\ell_{2}(F,F^{\prime})$ introduced in the main paper, we
have $\ell_{2}(F,F^{\prime})=\|F-F^{\prime}\|_{\ell_{2}}$. We also introduce a
supremum version of this norm on the space $\mathbb{R}^{\mathcal{X}\times m}$,
which we denote by $\|\cdot\|_{\ell_{2},\infty}$, and define by

$\displaystyle\|F\|_{\ell_{2},\infty}=\max_{x\in\mathcal{X}}\|F(x)\|_{\ell_{2}}\,,$

for all $F\in\mathbb{R}^{\mathcal{X}\times m}$. This norm is defined so that
if we have RDFs
$\eta,\eta^{\prime}\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})^{\mathcal{X}}$,
and $F,F^{\prime}\in\mathbb{R}^{\mathcal{X}\times m}$ are the corresponding
CDF values, then

$\displaystyle\overline{\ell}_{2}(\eta,\eta^{\prime})=\|F-F^{\prime}\|_{\ell_{2},\infty}\,.$

With this norm defined, we can now state and prove our first result, which
essentially translates the contraction result in Proposition 2.2, expressed
purely in terms of probability distributions and the Cramér distance
$\ell_{2}$, into a slightly more general result expressed over
$\mathbb{R}^{\mathcal{X}\times m}$ and the norm $\|\cdot\|_{\ell_{2}}$.

######  Proposition B.1.

The operator $T_{P}:\mathbb{R}^{\mathcal{X}\times
m}\rightarrow\mathbb{R}^{\mathcal{X}\times m}$ is a contraction when
restricted to the subspace $\\{F\in\mathbb{R}^{\mathcal{X}\times
m}:F_{m}(x)=0\text{ for all }x\in\mathcal{X}\\}$ with respect to the norm
$\|\cdot\|_{\ell_{2},\infty}$, with contraction factor $\sqrt{\gamma}$.

###### Proof.

By Proposition 2.2, for any two RDF approximations
$\eta,\eta^{\prime}\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})^{\mathcal{X}}$
with corresponding CDF values $F,F^{\prime}\in\mathbb{R}^{\mathcal{X}}$, we
have

$\displaystyle\overline{\ell}_{2}({\Pi_{m}}\mathcal{T}\eta,{\Pi_{m}}\mathcal{T}\eta^{\prime})\leq\sqrt{\gamma}\overline{\ell}_{2}(\eta,\eta^{\prime})\,,$

and hence we also have

$\displaystyle\|T_{P}F-T_{P}F^{\prime}\|_{\ell_{2},\infty}\leq\sqrt{\gamma}\|F-F^{\prime}\|_{\ell_{2},\infty}\,.$

Hence, $T_{P}$ is a contraction map on the set

$\displaystyle\\{F-F^{\prime}:0\leq F_{1}(x)\leq\cdots\leq F_{m}(x)=1\,,\
0\leq F^{\prime}_{1}(x)\leq\cdots\leq F^{\prime}_{m}(x)=1\text{ for all
}x\in\mathcal{X}\\}\,.$

This set contains a basis for the subspace
$\\{F\in\mathbb{R}^{\mathcal{X}\times m}:F_{m}(x)=0\text{ for all
}x\in\mathcal{X}\\}$. Namely, the one-hot vector at coordinate $(x,j)$ (for
$j<m$) can be exhibited as lying in this subspace since it can be expressed as
the difference between the vectors
$F,F^{\prime}\in\mathbb{R}^{\mathcal{X}\times m}$ defined by
$F_{j}(y)=F^{\prime}_{j}(y)=1$ for all $y\not=x$, and all $j=1,\ldots,m$, and
$F_{i}(x)=1$ for $i\geq j$ (and 0 otherwise), and $F^{\prime}_{i}(x)=1$ for
$i\geq j+1$ (and 0 otherwise). Since $T_{P}$ is linear, it therefore holds
that

$\displaystyle\|T_{P}F\|_{\ell_{2},\infty}\leq\sqrt{\gamma}\|F\|_{\ell_{2},\infty}\,,$
(18)

for any $F$ in the subspace $\\{F\in\mathbb{R}^{\mathcal{X}\times
m}:F_{m}(x)=0\text{ for all }x\in\mathcal{X}\\}$, as required. ∎

######  Proposition B.2.

For each $x\in\mathcal{X}$, the matrix $B_{x}$ is a contraction mapping on the
space $\\{F\in\mathbb{R}^{m}:F_{m}=0\\}$ with respect to
$\|\cdot\|_{\ell_{2}}$, with contraction factor $\sqrt{\gamma}$.

###### Proof.

Consider a related one-state MRP, for which the reward at the single state is
$r(x)$, and the state transitions to itself with probability 1. The
categorical Bellman operator associated with this MRP and the support
$\\{z_{1},\ldots,z_{m}\\}$ is precisely $B_{x}$, and the statement of the
result therefore follows as a special case of Proposition B.1. ∎

The next result serves as a counterpoint to Proposition B.2; it shows that if
$m$ is sufficiently large, the map $B_{x}$ does not contract by too much in
$\|\cdot\|_{\ell_{2}}$.

######  Proposition B.3.

For any $F,F^{\prime}\in\mathscr{F}$, we have

$\displaystyle\|B_{x}F-B_{x}F^{\prime}\|^{2}_{\ell_{2}}\geq\gamma\|F-F^{\prime}\|^{2}_{\ell_{2}}-\frac{2}{m(1-\gamma)^{1/2}}-\frac{1}{m^{2}(1-\gamma)^{2}}\,.$

This is proven via the following lemma.

######  Lemma B.4.

For any $\nu\in\mathscr{P}([0,(1-\gamma)^{-1}])$ and the projection
${\Pi_{m}}:\mathscr{P}([0,(1-\gamma)^{-1}])\rightarrow\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})$,
we have

$\displaystyle\ell_{2}(\nu,{\Pi_{m}}\nu)\leq\frac{1}{2\sqrt{m(1-\gamma)}}\,,$

Further, for any $\nu,\nu^{\prime}\in\mathscr{P}([0,(1-\gamma)^{-1}])$, we
have

$\displaystyle\ell_{2}({\Pi_{m}}\nu,{\Pi_{m}}\nu^{\prime})\geq\ell_{2}(\nu,\nu^{\prime})-\frac{1}{m(1-\gamma)}\,,$

and

$\displaystyle\ell^{2}_{2}({\Pi_{m}}\nu,{\Pi_{m}}\nu^{\prime})\geq\ell^{2}_{2}(\nu,\nu^{\prime})-\frac{2}{m(1-\gamma)^{1/2}}-\frac{1}{m^{2}(1-\gamma)^{2}}\,.$

###### Proof.

By Rowland et al. (2018, Proposition 6), we have that the CDF values
$F_{{\Pi_{m}}\nu}(z_{i})$ for $i=1,\ldots,m-1$ are equal to the average of the
CDF values of $F_{\nu}$ on the interval $[z_{i},z_{i+1}]$. Therefore, in
computing the squared Cramér distance $\ell_{2}^{2}(\nu,{\Pi_{m}}\nu)$, the
worst-case contribution to the integral

$\displaystyle\ell^{2}_{2}(\nu,{\Pi_{m}}\nu)=\int_{0}^{(1-\gamma)^{-1}}(F_{\nu}(t)-F_{{\Pi_{m}}\nu}(t))^{2}\;\mathrm{d}t$

from the interval $[z_{i},z_{i+1}]$, holding $F_{\nu}(z_{i})$ and
$F_{\nu}(z_{i+1})$ constant, is

$\displaystyle(z_{i+1}-z_{i})\bigg{(}\frac{F(z_{i+1})-F(z_{i})}{2}\bigg{)}^{2}=\frac{1}{4m(1-\gamma)}(F(z_{i+1})-F(z_{i}))^{2}\,.$

Thus, the worst-case value for the entire integral is

$\displaystyle\frac{1}{4m(1-\gamma)}\sum_{i=1}^{m-1}(F(z_{i})-F(z_{i+1}))^{2}\,.$

The worst-case value for the sum is 1, by interpreting this as a sum of
squared probabilities, and we therefore deduce that

$\displaystyle\ell_{2}^{2}(\nu,{\Pi_{m}}\nu)\leq\frac{1}{4m(1-\gamma)}\,,$

leading to

$\displaystyle\ell_{2}(\nu,{\Pi_{m}}\nu)\leq\frac{1}{2\sqrt{m(1-\gamma)}}\,,$

as required for the first stated inequality. For the second and third
inequalities, we apply the triangle inequality twice to obtain

$\displaystyle\ell_{2}(\nu,\nu^{\prime})$
$\displaystyle\leq\ell_{2}(\nu,{\Pi_{m}}\nu)+\ell_{2}({\Pi_{m}}\nu,{\Pi_{m}}\nu^{\prime})+\ell_{2}({\Pi_{m}}\nu^{\prime},\nu^{\prime})\leq\ell_{2}({\Pi_{m}}\nu,{\Pi_{m}}\nu^{\prime})+\frac{1}{m(1-\gamma)}\,;$

rearrangement then gives the second statement. Squaring both sides of the
inequality above gives

$\displaystyle\ell_{2}^{2}(\nu,\nu^{\prime})\leq\ell_{2}^{2}({\Pi_{m}}\nu,{\Pi_{m}}\nu^{\prime})+\frac{2}{m(1-\gamma)}\ell_{2}({\Pi_{m}}\nu,{\Pi_{m}}\nu^{\prime})+\frac{1}{m^{2}(1-\gamma)^{2}}\,.$

Bounding the instance of $\ell_{2}({\Pi_{m}}\nu,{\Pi_{m}}\nu^{\prime})$ in the
cross-term on the right-hand side by $(1-\gamma)^{1/2}$ and rearranging then
yields the result. ∎

###### Proof of Proposition B.3.

Let $\nu,\nu^{\prime}\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})$ be the
distributions with CDF values $F,F^{\prime}$, respectively, and let
$G,G^{\prime}$ be random variables with CDFs $F,F^{\prime}$ respectively, and
$(x,X^{\prime})$ an independent random transition in the MRP beginning at
state $x$. Recall from the notation introduced earlier in this section that
$B_{x}F$ and $B_{x}F^{\prime}$ are the CDF values of the distributions of
$r(x)+\gamma G$ and $r(x)+\gamma G^{\prime}$ after projection onto the support
grid $\\{z_{1},\ldots,z_{m}\\}$ by the projection map ${\Pi_{m}}$. Following
common notation in distributional RL (Bellemare et al., 2023), we denote the
distributions of $r(x)+\gamma G$ and $r(x)+\gamma G^{\prime}$ by
$(b_{r(x),\gamma})_{\\#}\nu$ and $(b_{r(x),\gamma})_{\\#}\nu^{\prime}$,
respectively. Here, $b_{r(x),\gamma}:\mathbb{R}\rightarrow\mathbb{R}$ is the
bootstrap function $b_{r(x),\gamma}(z)=r(x)+\gamma z$, and
$(b_{r(x),\gamma})_{\\#}\nu$ is the push-forward distribution of $\nu$ through
$b_{r(x),\gamma}$ (intuitively, the distribution obtained by transforming the
support of $\nu$ according to $b_{r(x),\gamma}$). With this notation
introduced, we therefore have

$\displaystyle\|B_{x}F-B_{x}F^{\prime}\|^{2}_{\ell_{2}}$
$\displaystyle=\ell^{2}_{2}({\Pi_{m}}(\mathrm{b}_{r(x),\gamma})_{\\#}\nu,{\Pi_{m}}(\mathrm{b}_{r(x),\gamma})_{\\#}\nu^{\prime})$
$\displaystyle\overset{(a)}{\geq}\ell_{2}^{2}((\mathrm{b}_{r(x),\gamma})_{\\#}\nu,(\mathrm{b}_{r(x),\gamma})_{\\#}\nu^{\prime})-\frac{2}{m(1-\gamma)^{1/2}}-\frac{1}{m^{2}(1-\gamma)^{2}}$
$\displaystyle\overset{(b)}{=}\gamma\ell_{2}^{2}(\nu,\nu^{\prime})-\frac{2}{m(1-\gamma)^{1/2}}-\frac{1}{m^{2}(1-\gamma)^{2}}$
$\displaystyle=\gamma\|F-F^{\prime}\|^{2}_{\ell_{2}}-\frac{2}{m(1-\gamma)^{1/2}}-\frac{1}{m^{2}(1-\gamma)^{2}}\,,$

as required, where (a) follows from Lemma B.4, and (b) follows from
homogeneity of Cramér distance (see Rowland et al. (2018, Proof of Proposition
2)). ∎

##  Appendix C Proofs of results in Section 4

See 4.1

###### Proof.

Beginning by restating Equation (6), we have that if
$\eta\in\mathscr{P}(\\{z_{1},\ldots,z_{m}\\})^{\mathcal{X}}$ is an RDF with
corresponding probability mass values $p\in\mathbb{R}^{\mathcal{X}\times m}$,
then the updated RDF $\eta^{\prime}={\Pi_{m}}\mathcal{T}\eta$ has
corresponding probability mass values
$p^{\prime}\in\mathbb{R}^{\mathcal{X}\times m}$ given by

$\displaystyle
p^{\prime}_{l}(x)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m}P(y|x)h^{x}_{l,j}p_{j}(y)\,.$

First, for $i\in\\{1,\ldots,m\\}$, we sum $l$ from $1$ to $i$ to yield

$\displaystyle F^{\prime}_{i}(x)=\sum_{l\leq
i}p^{\prime}_{l}(x)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m}P(y|x)\sum_{l\leq
i}h^{x}_{l,j}p_{j}(y)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m}P(y|x)H^{x}_{i,j}p_{j}(y)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m}P(y|x)H^{x}_{i,j}(F_{j}(y)-F_{j-1}(y))\,,$

where by convention we take $F_{0}(y)\equiv 0$. By reorganising the terms on
the right-hand side, the claim now follows. ∎

See 4.2

###### Proof.

First, we consider a row of $F=T_{P}F$ that corresponds to the index $(x,i)$,
with $i\not=m$. Expanding under the definition of $T_{P}$, we have

$\displaystyle
F_{i}(x)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m-1}P(y|x)(H^{x}_{i,j}-H^{x}_{i,j+1})F_{j}(y)+\sum_{y\in\mathcal{X}}P(y|x)H^{x}_{i,m}F_{m}(y)\,.$

Since we assume the additional constraints $F_{m}(y)\equiv 1$ for all
$y\in\mathcal{X}$, the final term on the right-hand side can be simplified to
yield

$\displaystyle
F_{i}(x)=\sum_{y\in\mathcal{X}}\sum_{j=1}^{m-1}P(y|x)(H^{x}_{i,j}-H^{x}_{i,j+1})F_{j}(y)+H^{x}_{i,m}\,.$

This is precisely the row of
$\widetilde{F}=\widetilde{T}_{P}\widetilde{F}+\widetilde{H}$ corresponding to
index $(x,i)$.

Now, we consider a row of $F=T_{P}F$ that corresponds to the index $(x,m)$.
Again making the substitution $F_{m}(y)\equiv 1$ for all $y\in\mathcal{X}$, we
have

$\displaystyle 1$ $\displaystyle\equiv
F_{m}(x)=\sum_{y}\sum_{j=1}^{m-1}P(y|x)\overbrace{(H^{x}_{m,j}-H^{x}_{m,j+1})}^{=0}F_{j}(y)+\sum_{y}P(y|x)H^{x}_{m,m}F_{m}(y)=\sum_{y}P(y|x)H^{x}_{m,m}F_{m}(y)\equiv
1\,,$

which shows that the equation is redundant, and hence can be removed from the
system. The claim $H^{x}_{m,j}-H^{x}_{m,j+1}=0$ follows since in fact
$H^{x}_{m,j}=\sum_{i=1}^{m}h_{i}(r(x)+\gamma z_{j})$, and the sum over the hat
functions for any input argument is 1. In the final equality, we have used the
fact that $H^{x}_{m,m}=1$ similarly. Thus, we have deduced the claim of the
proposition. ∎

See 4.3

###### Proof.

By Proposition 4.2, we have that the CDF values $\widetilde{F}^{*}$ of the
categorical fixed-point solve Equation (11). Now, let us suppose that
$\widetilde{F}_{1},\widetilde{F}_{2}$ are distinct solutions to Equation (11),
aiming to obtain a contradiction. Proposition 4.2 also establishes that
Equation (11) is equivalent to Equation (10) with the additional conditions
that $F_{m}(x)=1$ for all $x\in\mathcal{X}$, so we will denote the
corresponding two solutions to Equation (10) built from
$\widetilde{F}_{1},\widetilde{F}_{2}$ (by setting the unspecified $(x,m)$
coordinates to 1) by $F_{1},F_{2}$, respectively.

The idea is now to use the contractivity of the projected operator
$\Pi_{m}\mathcal{T}$ in $\ell_{2}$, as established in Proposition 2.2, to
obtain a contradiction. Notationally, it is useful to phrase things in terms
of contractivity of the CDF operator $T_{P}$ itself. We use the norms
$\|\cdot\|_{\ell_{2}}$ and $\|\cdot\|_{\ell_{2},\infty}$, defined on
$\mathbb{R}^{m}$ and $\mathbb{R}^{\mathcal{X}\times m}$ in Appendix B, which
we recall here for convenience:

$\displaystyle\|F\|_{\ell_{2}}=\left[\frac{1}{m(1-\gamma)}\sum_{i=1}^{m}F_{i}^{2}\right]^{1/2}\text{
for }F\in\mathbb{R}^{m}\,,\ \ \text{ and }\ \
\|F\|_{\ell_{2},\infty}=\max_{x\in\mathcal{X}}\|F(x)\|_{\ell_{2}}\text{ for
}F\in\mathbb{R}^{\mathcal{X}\times m}$

We therefore have

$\displaystyle\|T_{P}(F_{1}-F_{2})\|_{\ell_{2},\infty}=\|T_{P}F_{1}-T_{P}F_{2}\|_{\ell_{2},\infty}=\|F_{1}-F_{2}\|_{\ell_{2},\infty}\,.$

However, Proposition B.1 establishes that $T_{P}$ is a contraction on
$\\{F\in\mathbb{R}^{\mathcal{X}\times m}:F_{m}(x)=0\text{ for all
}x\in\mathcal{X}\\}$ with respect to $\|\cdot\|_{\ell_{2},\infty}$, which
contradicts the statement above, as required. ∎

##  Appendix D Proofs relating to the stochastic CDF Bellman equation

Before giving the proofs, we present a version of the SC-CDF Bellman equation
in a purely distributional form, which will streamline the arguments. This
mirrors the development of the form of the distributional Bellman equation
given purely in terms of distributions (Rowland et al., 2018), rather than in
random variable form as in Bellemare et al. (2017). We define the stochastic
categorical CDF Bellman operator
$\mathcal{T}_{\mathrm{SCC}}:\mathscr{P}(\mathscr{F})^{\mathcal{X}}\rightarrow\mathscr{P}(\mathscr{F})^{\mathcal{X}}$
for each $\psi\in\mathscr{P}(\mathscr{F})^{\mathcal{X}}$ by

$\displaystyle(\mathcal{T}_{\mathrm{SCC}}\;\psi)(x)=\mathcal{D}((\hat{T}_{P}\Phi)(x))\,,$

where $\Phi(y)\sim\psi(y)$ independently of $\hat{T}_{P}$, and $\mathcal{D}$
extracts the distribution of the input random variable. The purely
distributional form of the SC-CDF Bellman equation is then written as a fixed
point condition on $\mathscr{P}(\mathscr{F})^{\mathcal{X}}$:

$\displaystyle\psi=\mathcal{T}_{\mathrm{SCC}}\;\psi\,.$ (19)

We also write

$\displaystyle\overline{w}_{\|\cdot\|_{\ell_{2}}}(\psi,\psi^{\prime})=\max_{x\in\mathcal{X}}w_{\|\cdot\|_{\ell_{2}}}(\psi(x),\psi(x^{\prime}))$

for the supremum-Wasserstein distance over
$\mathscr{P}(\mathscr{F})^{\mathcal{X}}$ with base metric
$\|\cdot\|_{\ell_{2}}$ on $\mathscr{F}$.

See 5.7

###### Proof.

We first show that the operator $\mathcal{T}_{\mathrm{SCC}}$ is a contraction
on $\mathscr{P}(\mathscr{F})^{\mathcal{X}}$ with respect to the metric
$\overline{w}_{\|\cdot\|_{\ell_{2}}}$. Suppose
$\psi,\psi^{\prime}\in\mathscr{P}(\mathscr{F})^{\mathcal{X}}$, and let
$(\Phi(y),\Phi^{\prime}(y))$ be an optimal coupling between $\psi(y)$ and
$\psi^{\prime}(y)$ with respect to $w_{\|\cdot\|_{\ell_{2}}}$ for each
$y\in\mathcal{X}$ (existence of such couplings is guaranteed by Villani (2009,
Theorem 4.1)). Then, letting $(x,X^{\prime})$ be a random transition from $x$,
independent of $\Phi,\Phi^{\prime}$, we have that
$(B_{x}\Phi(X^{\prime}),B_{x}\Phi^{\prime}(X^{\prime}))$ is a valid coupling
of $(\mathcal{T}_{\mathrm{SCC}}\;\psi)(x)$ and
$(\mathcal{T}_{\mathrm{SCC}}\;\psi^{\prime})(x)$. Then we have, using the
operator notation defined in Appendix B,

$\displaystyle
w_{\|\cdot\|_{\ell_{2}}}((\mathcal{T}_{\mathrm{SCC}}\;\psi)(x),(\mathcal{T}_{\mathrm{SCC}}\;\psi^{\prime})(x))$
$\displaystyle\overset{(a)}{\leq}\mathbb{E}\bigg{[}\|B_{x}\Phi(X^{\prime})-B_{x}\Phi^{\prime}(X^{\prime})\|_{\ell_{2}}\bigg{]}$
$\displaystyle\overset{(b)}{\leq}\sqrt{\gamma}\mathbb{E}\bigg{[}\|\Phi(X^{\prime})-\Phi^{\prime}(X^{\prime})\|_{\ell_{2}}\bigg{]}$
$\displaystyle=\sqrt{\gamma}\sum_{y\in\mathcal{X}}P(y|x)\mathbb{E}\bigg{[}\|\Phi(y)-\Phi^{\prime}(y)\|_{\ell_{2}}\bigg{]}$
$\displaystyle\overset{(c)}{=}\sqrt{\gamma}\sum_{y\in\mathcal{X}}P(y|x)w_{\|\cdot\|_{\ell_{2}}}(\psi(y),\psi^{\prime}(y))$
$\displaystyle\leq\sqrt{\gamma}\overline{w}_{\|\cdot\|_{\ell_{2}}}(\psi,\psi^{\prime})\,,$

where (a) follows since
$(B_{x}\Phi(X^{\prime}),B_{x}\Phi^{\prime}(X^{\prime}))$ is a valid coupling
of $(\mathcal{T}_{\mathrm{SCC}}\;\psi)(x)$ and
$(\mathcal{T}_{\mathrm{SCC}}\;\psi^{\prime})(x)$, (b) follows by contractivity
of $B_{x}$ with respect to $\|\cdot\|_{\ell_{2}}$, as shown in Proposition
B.2, and (c) follows since $(\Phi(y),\Phi^{\prime}(y))$ was chosen to be an
optimal coupling of $\psi(y)$ and $\psi(y^{\prime})$. Hence we have

$\displaystyle\overline{w}_{\|\cdot\|_{\ell_{2}}}(\mathcal{T}_{\mathrm{SCC}}\psi,\mathcal{T}_{\mathrm{SCC}}\psi^{\prime})\leq\sqrt{\gamma}\overline{w}_{\|\cdot\|_{\ell_{2}}}(\psi,\psi^{\prime})\,,$

proving contractivity.

The metric space
$(\mathscr{P}(\mathscr{F})^{\mathcal{X}},\overline{w}_{\|\cdot\|_{\ell_{2}}})$
is complete, since the base space $(\mathscr{F},\|\cdot\|_{\ell_{2}})$ is
separable and complete (Villani, 2009, Theorem 6.18). Hence, by Banach’s fixed
point theorem, we obtain that there is a unique fixed point
$\psi^{*}\in\mathscr{P}(\mathscr{F})^{\mathcal{X}}$ of
$\mathcal{T}_{\mathrm{SCC}}$. Thus, a collection of random CDFs
$(\Phi^{*}(x):x\in\mathcal{X})$ satisfy the SC-CDF Bellman equation in
Equation (16) if and only if we have $\Phi^{*}(x)\sim\psi^{*}(x)$ for all
$x\in\mathcal{X}$. ∎

See 5.8

###### Proof.

We take expectations on both sides of the random-variable stochastic
categorical CDF Bellman equation in Equation (16), yielding:

$\displaystyle\mathbb{E}[\Phi^{Q}(x)]=\mathbb{E}[(\hat{T}_{Q}\Phi^{Q})(x)]\,.$

Since $\hat{T}_{Q}$ is a random linear map, independent of $\Phi^{Q}$, we have

$\displaystyle\mathbb{E}[\Phi^{Q}(x)]$
$\displaystyle=(\mathbb{E}[\hat{T}_{Q}]\mathbb{E}[\Phi^{Q}])(x)$
$\displaystyle=T_{Q}\mathbb{E}[\Phi^{Q}](x)\,.$

This states that $\mathbb{E}[\Phi^{Q}]\in\mathbb{R}^{\mathcal{X}\times m}$
satisfies the standard categorical Bellman equation in Equation (10), and
hence $\mathbb{E}[\Phi^{Q}]=F^{Q}$, by Proposition 2.2, as required. ∎

See 5.11

###### Proof.

We calculate, writing $\Phi^{Q}$ as $\Phi$, $F^{Q}$ as $F$, and $\hat{T}_{Q}$
as $\hat{T}$ to lighten notation:

$\displaystyle\Sigma_{Q}(x)$
$\displaystyle=\mathbb{E}[\ell_{2}^{2}(\Phi(x),F(x))]$
$\displaystyle\overset{(a)}{=}\mathbb{E}[\ell_{2}^{2}((\hat{T}\Phi)(x),F(x))]$
$\displaystyle\overset{(b)}{=}\mathbb{E}[\ell_{2}^{2}((\hat{T}F)(x),F(x))]+\mathbb{E}[\ell_{2}^{2}((\hat{T}\Phi)(x),(\hat{T}F)(x))]\,.$
(20)

Here, (a) follows since $\Phi$ satisfies the random-variable version of the
stochastic categorical CDF Bellman equation, (b) is a result of a Pythagorean
identity for squared Cramér distance, which we derive below:

$\displaystyle\mathbb{E}[\ell_{2}^{2}((\hat{T}\Phi)(x),F(x))]$
$\displaystyle=\mathbb{E}\bigg{[}\int_{0}^{(1-\gamma)^{-1}}((\hat{T}\Phi)(x)(t)-F(x)(t))^{2}\;\mathrm{d}t\bigg{]}$
$\displaystyle=\int_{0}^{(1-\gamma)^{-1}}\mathbb{E}[((\hat{T}\Phi)(x)(t)-F(x)(t))^{2}]\;\mathrm{d}t\,,$

where the integral switch follows from Fubini’s theorem. Now, focusing on the
integrand above, it can be written as

$\displaystyle\mathbb{E}[(Y-\mathbb{E}[Y])^{2}]$

with $Y=(\hat{T}\Phi)(x)(t)$, since
$F(x)=\mathbb{E}[(\hat{T}\Phi)(x)]=\mathbb{E}[\Phi(x)]$, by Lemma 5.8. We have

$\displaystyle\mathbb{E}[(Y-\mathbb{E}[Y])^{2}]$
$\displaystyle=\mathbb{E}[(Y-\mathbb{E}[Y|\hat{T}]+\mathbb{E}[Y|\hat{T}]-\mathbb{E}[Y])^{2}]$
$\displaystyle=\mathbb{E}[(Y-\mathbb{E}[Y|\hat{T}])^{2}]+\mathbb{E}[(\mathbb{E}[Y|\hat{T}]-\mathbb{E}[Y])^{2}]\,.$

Now,
$\mathbb{E}[(\hat{T}\Phi)(x)(t)\mid\hat{T}]=(\hat{T}\mathbb{E}[\Phi])(x)(t)=(\hat{T}F)(x)(t)$,
by linearity, which concludes the validation of step (b) above. We recognise
the first term in Equation (D) as the local squared-Cramér variation, and
hence have

$\displaystyle\Sigma_{Q}(x)=\sigma_{Q}(x)+\mathbb{E}[\ell_{2}^{2}((\hat{T}\Phi)(x),(\hat{T}F)(x))]\,.$

We now explicitly write the evaluation of the application of $\hat{T}$ at
coordinate $x$ in terms of the random transition $(x,X^{\prime})$ used to
construct the single-sample random transition matrix $\hat{Q}$ described in
Definition 5.4, so that we obtain, with the operator notation of Appendix B,

$\displaystyle\mathbb{E}[\ell_{2}^{2}((\hat{T}\Phi)(x),(\hat{T}F)(x))]$
$\displaystyle=\mathbb{E}[\ell_{2}^{2}(B_{x}\Phi(X^{\prime}),B_{x}F(X^{\prime}))]$
$\displaystyle=\mathbb{E}[\mathbb{E}[\ell_{2}^{2}(B_{x}\Phi(X^{\prime}),B_{x}F(X^{\prime}))\mid\Phi]]$
$\displaystyle=\mathbb{E}\bigg{[}\sum_{y\in\mathcal{X}}Q(y|x)\ell_{2}^{2}(B_{x}\Phi(y),B_{x}F(y))\bigg{]}$
$\displaystyle\overset{(a)}{\geq}\sum_{y\in\mathcal{X}}Q(y|x)\mathbb{E}[\gamma\ell_{2}^{2}(\Phi(y),F(y))-\alpha]$
$\displaystyle=\gamma\sum_{y\in\mathcal{X}}Q(y|x)\mathbb{E}[\ell_{2}^{2}(\Phi(y),F(y))]-\alpha$
$\displaystyle=\gamma\sum_{y\in\mathcal{X}}Q(y|x)\mathbb{E}[\ell_{2}^{2}(\Phi(y),F(y))]-\alpha$
$\displaystyle=\gamma(Q\Sigma_{Q})(x)-\alpha\,,$

where (a) follows from Proposition B.3, with
$\alpha=\frac{2}{m(1-\gamma)^{1/2}}+\frac{1}{m^{2}(1-\gamma)^{2}}$, meaning
that we deduce

$\displaystyle\Sigma_{Q}(x)=\sigma_{Q}(x)+\gamma(Q\Sigma_{Q})(x)-\alpha\,,$

as required. ∎

See 5.12

###### Proof.

By Proposition 5.11 applied with $Q=\hat{P}$,

$\displaystyle\Sigma_{\hat{P}}\geq\sigma_{\hat{P}}+\gamma\hat{P}\Sigma_{\hat{P}}-\bigg{(}\frac{2}{m(1-\gamma)^{1/2}}+\frac{1}{m^{2}(1-\gamma)^{2}}\bigg{)}\mathbf{1}\,,$

where $\mathbf{1}\in\mathbb{R}^{\mathcal{X}}$ is the vector of ones. We first
note that from the condition $m\geq 4\varepsilon^{-2}(1-\gamma)^{-2}+1$ from
the statement of Theorem 5.3, we have

$\displaystyle\frac{2}{m(1-\gamma)^{1/2}}+\frac{1}{m^{2}(1-\gamma)^{2}}\leq\frac{2}{4\varepsilon^{-2}(1-\gamma)^{-2}(1-\gamma)^{1/2}}+\frac{1}{4\varepsilon^{-2}}\leq\frac{\varepsilon^{2}(1-\gamma)^{3/2}}{2}+\frac{\varepsilon^{2}}{4}<1\,,$

since $\varepsilon\in(0,1)$, from the statement of Theorem 5.3.

We therefore have

$\displaystyle(I-\gamma\hat{P})\Sigma_{\hat{P}}\geq\sigma_{\hat{P}}-\mathbf{1}\,.$

Now, $(I-\gamma\hat{P})^{-1}$ is a monotone operator (in the sense that
$v_{1}\geq v_{2}$ coordinatewise implies
$(I-\gamma\hat{P})^{-1}v_{1}\geq(I-\gamma\hat{P})^{-1}v_{2}$; this claim
follows as all elements of $(I-\gamma\hat{P})^{-1}$ are non-negative, since it
can also be written $\sum_{k\geq 0}\gamma^{k}\hat{P}^{k}$), we can apply it to
both sides of the inequality above to obtain

$\displaystyle\Sigma_{\hat{P}}+(1-\gamma)^{-1}\mathbf{1}\geq(I-\gamma\hat{P})^{-1}\sigma_{\hat{P}}\,.$

Finally, we note that since $\Sigma_{\hat{P}}(x)$ is an expected squared-
Cramér distance between two CDFs supported on an interval of length
$(1-\gamma)^{-1}$, we have

$\displaystyle\Sigma_{\hat{P}}(x)\leq(1-\gamma)^{-1}\,,$

from which the claim follows. ∎

##  Appendix E Proof of Theorem 5.1

We begin by restating the result we seek to prove.

See 5.1

We arrange the proof into a sequence of smaller results, in analogy with the
presentation of the sketch proof in the main paper.

###  E.1 Reduction to high-probability bounds in Cramér distance

Following the sketch provided in the main paper, we first restate and prove
Lemma 5.2.

See 5.2

###### Proof.

We begin by writing

$\displaystyle
w_{1}(\nu,\nu^{\prime})=\int_{0}^{(1-\gamma)^{-1}}|F_{\nu}(t)-F_{\nu^{\prime}}(t)|\;\mathrm{d}t=(1-\gamma)^{-1}\left[(1-\gamma)\int_{0}^{(1-\gamma)^{-1}}|F_{\nu}(t)-F_{\nu^{\prime}}(t)|\;\mathrm{d}t\right]\,,$

where $F_{\nu},F_{\nu^{\prime}}$ are the CDFs of $\nu,\nu^{\prime}$,
respectively. The quantity inside the squared brackets can be interpreted as
an expectation (with $t$ ranging over the values of a uniform variate on
$\mathbb{E}_{T\sim\text{Unif}([0,(1-\gamma)^{-1}])}[|F_{\nu}(T)-F_{\nu^{\prime}}(T)|]$,
and we can therefore apply Jensen’s inequality with the map $z\mapsto z^{2}$.
This then yields

$\displaystyle w_{1}(\nu,\nu^{\prime})\leq$
$\displaystyle(1-\gamma)^{-1}\left[(1-\gamma)\int_{0}^{(1-\gamma)^{-1}}(F_{\nu}(t)-F_{\nu^{\prime}}(t))^{2}\;\mathrm{d}t\right]^{1/2}$
$\displaystyle=$ $\displaystyle(1-\gamma)^{-1/2}\ell_{2}(\nu,\nu^{\prime})\,,$

as required. ∎

The first main step of the proof of Theorem 5.1 is a reduction to Theorem 5.3
via Lemma 5.2. To see this, we first restate Theorem 5.3 here for convenience.

See 5.3

For the proof of the reduction, suppose the statement of Theorem 5.3 holds.
Now, let us take $\varepsilon\in(0,(1-\gamma)^{-1/2})$, and $m\geq
4(1-\gamma)^{-2}\varepsilon^{-2}+1$, as in the assumptions of Theorem 5.1. We
then define $\widetilde{\varepsilon}=(1-\gamma)^{1/2}\varepsilon$; note that
from the assumption on $\varepsilon$, we have
$\widetilde{\varepsilon}\in(0,1)$. Applying the result of Theorem 5.3, we
therefore obtain that with
$N=\widetilde{\Omega}(\widetilde{\varepsilon}^{-2}(1-\gamma)^{-2}\text{polylog}(|\mathcal{X}|/\delta))$,
we have (with probability at least $1-\delta$)

$\displaystyle\max_{x\in\mathcal{X}}\ell_{2}(\eta^{*}(x),\hat{F}(x))\leq\widetilde{\varepsilon}\,.$

By Lemma 5.2, we therefore have (with probability at least $1-\delta$, for all
$x\in\mathcal{X}$)

$\displaystyle\overline{w}_{1}(\eta^{*}(x),\hat{F}(x))\leq(1-\gamma)^{-1/2}\ell_{2}(\eta^{*}(x),\hat{F}(x))\leq(1-\gamma)^{-1/2}(1-\gamma)^{1/2}\varepsilon=\varepsilon\,,$

which is the desired inequality in Theorem 5.1. Finally, we note that the
sample complexity term $\widetilde{\varepsilon}^{-2}(1-\gamma)^{-2}$ can be
rewritten as

$\displaystyle\widetilde{\varepsilon}^{-2}(1-\gamma)^{-2}=((1-\gamma)^{1/2}\varepsilon)^{-2}(1-\gamma)^{-2}=\varepsilon^{-2}(1-\gamma)^{-3}\,.$

Thus, we obtain the stated sample complexity in Theorem 5.1, and we have
established that to prove Theorem 5.1, it is sufficient to prove Theorem 5.3.

###  E.2 Reduction to categorical fixed-point error

The first step of the proof of Theorem 5.3 is to show that with $m$ taken
sufficiently large (as described in the statement of Theorem 5.3), the Cramér
distance between the true return distributions and the categorical fixed
points is small, and it is therefore sufficient to focus solely on the sample-
based error in estimating the categorical fixed point.

By applying the triangle inequality, we have

$\displaystyle\overline{\ell}_{2}(\eta^{*},\hat{F})$
$\displaystyle\leq\overline{\ell}_{2}(\eta^{*},F^{*})+\overline{\ell}_{2}(F^{*},\hat{F})$
$\displaystyle\overset{(a)}{\leq}\frac{1}{(1-\gamma)\sqrt{m-1}}+\overline{\ell}_{2}(F^{*},\hat{F})$
$\displaystyle\overset{(b)}{\leq}\frac{\varepsilon}{2}+\overline{\ell}_{2}(F^{*},\hat{F})\,,$

where (a) follows from the fixed-point approximation bound in Equation (5),
which itself is the result of Proposition 2 of Rowland et al. (2018), and (b)
follows from substituting the assumed inequality for $m$ in the statement of
Theorem 5.3. Thus, to establish that $\overline{\ell}_{2}(\eta^{*},\hat{F})$
is bounded by $\varepsilon$ with probability at least $1-\delta$, it suffices
to show that

$\displaystyle\overline{\ell}_{2}(F^{*},\hat{F})<\varepsilon/2\,,$

with probability at least $1-\delta$, as claimed.

###  E.3 Propagation of local errors

To begin analysing $\overline{\ell}_{2}(F^{*},\hat{F})$, we analyse the
difference of vectors $\hat{F}-F^{*}$ directly. We proceed in an analogous
manner to Azar et al. (2013) in the mean-return case, rearranging as follows:

$\displaystyle\hat{F}-F^{*}$
$\displaystyle\overset{(a)}{=}T_{\hat{P}}\hat{F}-T_{P}F^{*}$
$\displaystyle\overset{(b)}{=}T_{\hat{P}}\hat{F}-T_{\hat{P}}F^{*}+T_{\hat{P}}F^{*}-T_{P}F^{*}$
$\displaystyle\implies(I-T_{\hat{P}})(\hat{F}-F^{*})$
$\displaystyle=(T_{\hat{P}}-T_{P})F^{*}\,.$ (21)

Here, (a) follows since $\hat{F},F^{*}$ are fixed points of
$T_{\hat{P}},T_{P}$, respectively, (b) follows by adding and subtracting
$T_{\hat{P}}F^{*}$, and the implication follows from straightforward
rearrangement.

We would next like to rearrange Equation (E.3) to leave the term
$\hat{F}-F^{*}$ on its own. This requires some care, in checking that the
operator $(I-T_{\hat{P}})$ is invertible in an appropriate sense.

######  Lemma E.1.

The operator $I-T_{\hat{P}}:\mathbb{R}^{\mathcal{X}\times
m}\rightarrow\mathbb{R}^{\mathcal{X}\times m}$ is invertible on the subspace
$\\{F\in\mathbb{R}^{\mathcal{X}\times m}:F_{m}(x)=0\text{ for all
}x\in\mathcal{X}\\}$, with inverse $\sum_{k\geq 0}T_{\hat{P}}^{k}$.

###### Proof.

By Proposition B.1, $T_{\hat{P}}$ maps $\\{F\in\mathbb{R}^{\mathcal{X}\times
m}:F_{m}(x)=0\text{ for all }x\in\mathcal{X}\\}$ to itself, and is a
contraction on this subspace with respect to $\ell_{2}$, with contraction
factor $\sqrt{\gamma}$. It therefore follows that $I-T_{\hat{P}}$ maps this
subspace to itself, and is invertible on this subspace. Since $\sum_{k\geq
0}T_{\hat{P}}^{k}$ also maps this subspace to itself, it follows that on this
subspace, we have $(I-T_{\hat{P}})^{-1}=\sum_{k\geq 0}T_{\hat{P}}^{k}$, as
required. ∎

Now, $(T_{\hat{P}}-T_{P})F^{*}\in\\{F\in\mathbb{R}^{\mathcal{X}\times
m}:F_{m}(x)=0\text{ for all }x\in\mathcal{X}\\}$, and hence from Equation
(E.3) and Lemma E.1, it follows that

$\displaystyle\hat{F}-F^{*}=\sum_{k\geq
0}T_{\hat{P}}^{k}(T_{\hat{P}}-T_{P})F^{*}\,.$ (22)

The result is that we have expressed the difference in CDFs in terms of local
errors $(T_{\hat{P}}-T_{P})F^{*}$, which are then propagated via the operator
$\sum_{k\geq 0}T_{\hat{P}}^{k}$.

###  E.4 Bernstein concentration bounds

The next step in the proof is to establish a concentration bound for the local
errors $(T_{\hat{P}}-T_{P})F^{*}$.

One potential approach is to use a concentration inequality for each
individual coordinate of $(T_{\hat{P}}-T_{P})F^{*}$, indexed by $(x,i)$. Azar
et al. (2013) note that in the mean-return case, using a Hoeffding-style bound
is insufficient to obtain optimal dependence of the sample complexity on
$(1-\gamma)^{-1}$, and Zhang et al. (2023) also note this in their (non-
categorical) distributional analysis. Using a Bernstein concentration
inequality on each coordinate and then using a union bound can be made to
work, although this then incurs a $\log(m)$ factor in the sample complexity.
If such a dependence on $m$ were tight, this would suggest that the sample
complexity depends on $m$ (albeit only logarithmically), and hence there would
be a trade-off between picking $m$ sufficiently large so as to obtain a low
representation approximation error, as in Section E.2, and taking $m$ low so
as to not unduly increase the sample complexity. However, such a dependence on
$m$ can be avoided by working with concentration inequalities at the level of
CDFs themselves. The Dvorestsky-Kiefer-Wolfowitz (DKW) inequality (Dvoretzky
et al., 1956; Massart, 1990) behaves analogously to Hoeffding’s bound, and is
insufficient for obtaining a sample complexity bound with optimal
$(1-\gamma)^{-1}$ dependence. Instead, we make use of a Hilbert space
Bernstein-style inequality (Yurinsky, 2006; Chatalic et al., 2022), by
interpreting the Cramér distance $\ell_{2}$ as a maximum mean discrepancy
(Gretton et al., 2012), which measures distances between distributions via
embeddings in a (reproducing kernel) Hilbert space, allowing the Bernstein
result to be applied.

First, we precisely describe the connection between Cramér distance and
Hilbert space. Székely (2003) shows that for any distributions
$\nu,\nu^{\prime}\in\mathscr{P}([0,(1-\gamma)^{-1}])$, we have

$\displaystyle\ell^{2}_{2}(\nu,\nu^{\prime})=\mathbb{E}_{X\sim\nu,Y\sim\nu^{\prime}}[|X-Y|]-\frac{1}{2}\mathbb{E}_{X,X^{\prime}\sim\nu}[|X-X^{\prime}|]-\frac{1}{2}\mathbb{E}_{Y,Y^{\prime}\sim\nu^{\prime}}[|Y-Y^{\prime}|]\,.$
(23)

Sejdinovic et al. (2013) show that there exists an affine embedding
$\phi:\mathscr{P}([0,(1-\gamma)^{-1}])\rightarrow\mathcal{H}$ of distributions
into a Hilbert space $\mathcal{H}$, such that the right-hand side of the
equation can also be written as the squared norm
$\|\phi(\nu)-\phi(\nu^{\prime})\|_{\mathcal{H}}^{2}$; in other words, they
show that the Cramér distance is equal to a squared maximum mean discrepancy
(Gretton et al., 2012). We can then appeal to the following Bernstein-style
inequality for Hilbert-space-valued random variables, which is due to Chatalic
et al. (2022), and based largely on Yurinsky (2006, Theorem 3.3.4). We state
the result here in the usual Bernstein style of assuming almost-sure
boundedness of the the random variables concerned, while Chatalic et al.
(2022) use the more general formulation of assuming particular bounds on all
moments.

######  Lemma E.2.

(Chatalic et al., 2022, Lemma E.3) Let $Z_{1},\ldots,Z_{N}$ be i.i.d. mean-
zero random variables taking values in a Hilbert space
$(\mathcal{H},\|\cdot\|_{\mathcal{H}})$ such that $\|Z_{1}\|_{\mathcal{H}}\leq
M$ almost surely. Then for any $\delta\in(0,1)$, we have

$\displaystyle\bigg{\|}\frac{1}{N}\sum_{i=1}^{N}Z_{i}\bigg{\|}_{\mathcal{H}}\leq\frac{2M\log(2/\delta)}{N}+\sqrt{\frac{2\mathbb{E}[\|Z_{1}\|_{\mathcal{H}}^{2}]\log(2/\delta)}{N}}\,,$

with probability at least $1-\delta$.

To apply this result in our setting, first note that, using the notation
introduced in Appendix B, we have

$\displaystyle\bigg{\|}[(T_{\hat{P}}-T_{P})F^{*}](x)\bigg{\|}_{\ell_{2}}$
$\displaystyle=\ell_{2}((T_{\hat{P}}F^{*})(x),F^{*}(x))$
$\displaystyle=\|\phi((T_{\hat{P}}F^{*})(x))-\phi(F^{*}(x))\|_{\mathcal{H}}$
$\displaystyle\overset{(a)}{=}\bigg{\|}\phi\left(\frac{1}{N}\sum_{i=1}^{N}F^{*}(X^{x}_{i})B^{\top}_{x}\right)-\phi(F^{*}(x))\bigg{\|}_{\mathcal{H}}$
$\displaystyle\overset{(b)}{=}\bigg{\|}\frac{1}{N}\sum_{i=1}^{N}\Big{(}\phi\left(F^{*}(X^{x}_{i})B^{\top}_{x}\right)-\phi(F^{*}(x))\Big{)}\bigg{\|}_{\mathcal{H}}\,,$

where (a) follows from the expressions for $T_{\hat{P}}$ described in Appendix
B, and (b) follows from the affineness of the embedding $\phi$. Here again, we
make use of a slight abuse of notation by writing $\phi(F^{*}(x))$ for the
embedding of the distribution supported on $\\{z_{1},\ldots,z_{m}\\}$, with
CDF values $F^{*}(x)$, and similarly for embeddings of other CDF vectors.

We can now apply the result above, noting that

$\displaystyle\|\phi\left(F^{*}(X^{x}_{i})B^{\top}_{x}\right)-\phi(F^{*}(x))\|_{\mathcal{H}}=\ell_{2}(F^{*}(X^{x}_{i})B^{\top}_{x},F^{*}(x))\leq(1-\gamma)^{-1/2}$

almost surely, and

$\displaystyle\mathbb{E}\Big{[}\|\phi\left(F^{*}(X^{x}_{i})B^{\top}_{x}\right)-\phi(F^{*}(x))\|_{\mathcal{H}}^{2}\Big{]}=\mathbb{E}\Big{[}\ell_{2}^{2}(F^{*}(X^{x}_{i})B^{\top}_{x},F^{*}(x))\Big{]}=\sigma(x)\,,$

where we write $\sigma$ as shorthand for the local squared-Cramér variation at
$P$, introduced in Definition 5.5 with the notation $\sigma_{P}$. Hence, from
Lemma E.2 and a union bound over the state space $\mathcal{X}$, we obtain that
with probability $1-\delta/2$, we have for all $x\in\mathcal{X}$ that

$\displaystyle\bigg{\|}[(T_{\hat{P}}-T_{P})F^{*}](x)\bigg{\|}_{\ell_{2}}\leq
C_{\mathrm{B}1}\frac{\sqrt{\sigma(x)}}{\sqrt{N}}+C_{\mathrm{B}2}\frac{1}{(1-\gamma)^{1/2}N}\,,$
(24)

where

$\displaystyle C_{\mathrm{B}1}=\sqrt{2\log(4|\mathcal{X}|/\delta)}\,,\ \ \
\text{ and }C_{\mathrm{B}2}=2\log(4|\mathcal{X}|/\delta)\,.$

###  E.5 From population to empirical squared-Cramér variation

Next, we show that the population local squared-Cramér variation in Equation
(24) can be replaced with the corresponding _empirical_ quantity; that is, the
local-squared Cramér variation at $\hat{P}$, $\sigma_{\hat{P}}$, incurring
some additional terms in the bound. We start by calculating as follows:

$\displaystyle\sigma(x)$
$\displaystyle=\mathbb{E}\bigg{[}\|(\hat{T}_{P}F^{*})(x)-(T_{P}F^{*})(x)\|_{\ell_{2}}^{2}\bigg{]}$
$\displaystyle=\mathbb{E}\bigg{[}\|F^{*}(X^{\prime})B_{x}^{\top}\|_{\ell_{2}}^{2}\bigg{]}-\|P_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}$
$\displaystyle=P_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\|P_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}\,,$

where above we use the notation
$\|F^{*}B_{x}^{\top}\|^{2}_{\ell_{2}}\in\mathbb{R}^{\mathcal{X}}$ for the
vector indexed by state, so that the element corresponding to state
$y\in\mathcal{X}$ is $\|F^{*}(y)B_{x}^{\top}\|^{2}_{\ell_{2}}$. To relate this
quantity to $\hat{\sigma}(x)$, we add and subtract a term, motivated by the
derivation applied to standard return variance by Azar et al. (2013, Lemma 5),
and rearrange as follows:

$\displaystyle\phantom{=}\sigma(x)$
$\displaystyle=P_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\|P_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}$
$\displaystyle=P_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\|P_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\bigg{(}\hat{P}_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\|\hat{P}_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}\bigg{)}+\bigg{(}\hat{P}_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\|\hat{P}_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}\bigg{)}$
$\displaystyle=(P_{x}-\hat{P}_{x})\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}+\bigg{(}\|\hat{P}_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\|P_{x}F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}\bigg{)}+\mathbb{E}\bigg{[}\|(\hat{T}_{\hat{P}}F^{*})(x)-(T_{\hat{P}}F^{*})(x)\|_{\ell_{2}}^{2}\bigg{|}\hat{P}\bigg{]}$
$\displaystyle=(P_{x}-\hat{P}_{x})\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}+\langle\hat{P}_{x}F^{*}B_{x}^{\top}-P_{x}F^{*}B_{x}^{\top},\hat{P}_{x}F^{*}B_{x}^{\top}+P_{x}F^{*}B_{x}^{\top}\rangle_{\ell_{2}}+\mathbb{E}\bigg{[}\|(\hat{T}_{\hat{P}}F^{*})(x)-(T_{\hat{P}}F^{*})(x)\|_{\ell_{2}}^{2}\bigg{|}\hat{P}\bigg{]}$
$\displaystyle=(P_{x}-\hat{P}_{x})\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}+\langle(\hat{P}_{x}-P_{x})F^{*}B_{x}^{\top},\hat{P}_{x}F^{*}B_{x}^{\top}+P_{x}F^{*}B_{x}^{\top}\rangle_{\ell_{2}}+\mathbb{E}\bigg{[}\|(\hat{T}_{\hat{P}}F^{*})(x)-(T_{\hat{P}}F^{*})(x)\|_{\ell_{2}}^{2}\bigg{|}\hat{P}\bigg{]}\,,$
(25)

where $\langle\;,\;\rangle_{\ell_{2}}$ is the inner product on
$\mathbb{R}^{m}$ corresponding to $\|\cdot\|_{\ell_{2}}$, defined by

$\displaystyle\langle
F,F^{\prime}\rangle_{\ell_{2}}=\frac{1}{m(1-\gamma)}\sum_{i=1}^{m}F_{i}F^{\prime}_{i}\,.$

Focusing on the final term on the right-hand side of Equation (E.5), we have

$\displaystyle\mathbb{E}\left[\|(\hat{T}_{\hat{P}}F^{*})(x)-(T_{\hat{P}}F^{*})(x)\|_{\ell_{2}}^{2}\;\Big{|}\;\hat{P}\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}\left[\|(\hat{T}_{\hat{P}}F^{*})(x)-(\hat{T}_{\hat{P}}\hat{F})(x)+(\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}F^{*})(x)+(T_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}\hat{F})(x)\|_{\ell_{2}}^{2}\;\Big{|}\;\hat{P}\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}\left[\|(\hat{T}_{\hat{P}}F^{*}-\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}F^{*}-T_{\hat{P}}\hat{F})(x)+(\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}\;\Big{|}\;\hat{P}\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}\Big{[}\|(\hat{T}_{\hat{P}}F^{*}-\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}F^{*}-T_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}+2\langle(\hat{T}_{\hat{P}}F^{*}-\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}F^{*}-T_{\hat{P}}\hat{F})(x),(\hat{T}_{\hat{P}}\hat{F}-T_{\hat{P}}\hat{F})(x)\rangle_{\ell_{2}}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad+\|(\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}\;\Big{|}\;\hat{P}\Big{]}$
$\displaystyle\overset{(a)}{\leq}$
$\displaystyle\bigg{(}\mathbb{E}[\|(\hat{T}_{\hat{P}}F^{*}-\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}F^{*}-T_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}\;|\;\hat{P}]^{1/2}+\mathbb{E}[\|(\hat{T}_{\hat{P}}\hat{F})(X)-(T_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}\;|\;\hat{P}]^{1/2}\bigg{)}^{2}$
$\displaystyle=$
$\displaystyle\bigg{(}\mathbb{E}[\|(\hat{T}_{\hat{P}}F^{*}-\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}F^{*}-T_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}\;|\;\hat{P}]^{1/2}+\sqrt{\hat{\sigma}(x)}\bigg{)}^{2}\,.$

where (a) follows from the Cauchy-Schwarz inequality, and we write
$\hat{\sigma}$ as a shorthand for $\sigma_{\hat{P}}$. Finally, we note that

$\displaystyle\mathbb{E}[\|(\hat{T}_{\hat{P}}F^{*}-\hat{T}_{\hat{P}}\hat{F})(x)-(T_{\hat{P}}F^{*}-T_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}\;|\;\hat{P}]$
$\displaystyle\leq\mathbb{E}[\|(\hat{T}_{\hat{P}}F^{*}-\hat{T}_{\hat{P}}\hat{F})(x)\|^{2}_{\ell_{2}}\;|\;\hat{P}]$
$\displaystyle\overset{(a)}{\leq}\gamma\mathbb{E}_{X^{\prime}\sim\hat{P}_{x}}[\|F^{*}(X^{\prime})-\hat{F}(X^{\prime})\|^{2}_{\ell_{2}}\;|\;\hat{P}]$
$\displaystyle\leq\gamma\|F^{*}-\hat{F}\|^{2}_{\ell_{2},\infty}\,,$

with (a) following from contractivity of $B_{x}$ in $\|\cdot\|_{\ell_{2}}$, as
established in Proposition B.2. Bringing these bounds together with Equation
(E.5), we have

$\displaystyle\sigma(x)\leq(P_{x}-\hat{P}_{x})\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}+\langle(\hat{P}_{x}-P_{x})F^{*}B_{x}^{\top},\hat{P}_{x}F^{*}B_{x}^{\top}+P_{x}F^{*}B_{x}^{\top}\rangle_{\ell_{2}}+\bigg{(}\gamma\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+\sqrt{\hat{\sigma}(x)}\bigg{)}^{2}\,.$
(26)

We now apply concentration bounds to each of the first two terms on the right-
hand side of Equation (26).

######  Lemma E.3.

With probability at least $1-\delta/2$, we have for all $x\in\mathcal{X}$,

$\displaystyle\bigg{|}P_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\hat{P}_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}\bigg{|}\leq
C_{\text{H}}\frac{1}{(1-\gamma)\sqrt{N}}\,,$

where

$\displaystyle C_{\mathrm{H}}=\sqrt{\frac{\log(4|\mathcal{X}|/\delta)}{2}}\,.$

###### Proof.

The expression
$P_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\hat{P}_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}$
is equal to the negative of the average of $N$ i.i.d. copies of the random
variable $\|F^{*}(X^{\prime})B_{x}^{\top}\|_{\ell_{2}}^{2}$, where
$X^{\prime}\sim P_{x}$, minus its expectation. Since
$\|F^{*}(X^{\prime})B_{x}^{\top}\|_{\ell_{2}}^{2}$ is bounded in
$[0,(1-\gamma)^{-1}]$, we apply Hoeffding’s inequality (for a given
$x\in\mathcal{X}$) to obtain

$\displaystyle\bigg{|}P_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}-\hat{P}_{x}\|F^{*}B_{x}^{\top}\|_{\ell_{2}}^{2}\bigg{|}\leq\frac{1}{(1-\gamma)\sqrt{N}}\sqrt{\frac{\log(4|\mathcal{X}|/\delta)}{2}}\,,$

with probability at least $1-\delta/(2|\mathcal{X}|)$. Taking a union bound
over $x\in\mathcal{X}$ then yields the statement. ∎

For the second term, we may re-use the Bernstein concentration bound derived
above to deduce the following.

######  Lemma E.4.

Suppose Equation (24) holds. Then we have, for all $x\in\mathcal{X}$,

$\displaystyle\bigg{|}\langle(\hat{P}_{x}-P_{x})F^{*}B_{x}^{\top},(\hat{P}_{x}+P_{x})F^{*}B_{x}^{\top}\rangle_{\ell_{2}}\bigg{|}\leq
2C_{\mathrm{B}1}\frac{1}{(1-\gamma)\sqrt{N}}+2C_{\mathrm{B}2}\frac{1}{(1-\gamma)N}\,.$

###### Proof.

First, by the Cauchy-Schwarz inequality, we have

$\displaystyle\bigg{|}\langle(\hat{P}_{x}-P_{x})F^{*}B_{x}^{\top},(\hat{P}_{x}+P_{x})F^{*}B_{x}^{\top}\rangle_{\ell_{2}}\bigg{|}\leq\|(\hat{P}_{x}-P_{x})F^{*}B_{x}^{\top}\|_{\ell_{2}}\|(\hat{P}_{x}+P_{x})F^{*}B_{x}^{\top}\|_{\ell_{2}}\,.$

The first of the two factors in the product on the right-hand side is
precisely the term bounded on the left-hand side of Equation (24). On the
right-hand side, the vector inside the $\ell_{2}$-norm has components in
$[0,2]$, so the norm can be straightforwardly bounded by

$\displaystyle\left[\frac{1}{m(1-\gamma)}\sum_{i=1}^{m}2^{2}\right]^{1/2}=\frac{2}{(1-\gamma)^{1/2}}\,.$

Combining the inequalities for these two factors, and using the trivial bound
$\sqrt{\sigma}(x)\leq(1-\gamma)^{-1/2}$, we obtain the stated result. ∎

Combining all these bounds together in Equation (26), and taking a union
bound, we conclude that with probability at least $1-\delta/2$, for all
$x\in\mathcal{X}$ we have

$\displaystyle\sigma(x)\leq(2C_{\mathrm{B}1}+C_{\mathrm{H}})\frac{1}{(1-\gamma)\sqrt{N}}+2C_{\mathrm{B}2}\frac{1}{(1-\gamma)N}+\left(\gamma\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+\sqrt{\hat{\sigma}(x)}\right)^{2}\,,$

We now take square-roots of both sides, using the inequality
$\sqrt{a+b}\leq\sqrt{a}+\sqrt{b}$ on the right-hand side to obtain

$\displaystyle\sqrt{\sigma(x)}\leq\sqrt{\hat{\sigma}(x)}+\gamma\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+\sqrt{2C_{\mathrm{B}1}+C_{\mathrm{H}}}\frac{1}{(1-\gamma)^{1/2}N^{1/4}}+\sqrt{2C_{\mathrm{B}2}}\frac{1}{(1-\gamma)^{1/2}N^{1/2}}\,.$

Substituting this into Equation (24) and taking a union bound then yields that
with probability at least $1-\delta$, we have

$\displaystyle\bigg{\|}[(T_{\hat{P}}-T_{P})F^{*}](x)\bigg{\|}_{\ell_{2}}\leq
C_{\mathrm{B}1}\frac{1}{\sqrt{N}}\sqrt{\hat{\sigma}(x)}+C_{\mathrm{B}1}\frac{1}{\sqrt{N}}\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+C^{\prime}\frac{1}{(1-\gamma)^{1/2}N^{3/4}}\,,$
(27)

where

$\displaystyle
C^{\prime}=C_{\mathrm{B}1}(\sqrt{2C_{\mathrm{B}1}+C_{\mathrm{H}}}+\sqrt{2C_{\mathrm{B}2}})+C_{\mathrm{B}2}\,,$

so that our concentration inequality is expressed in terms of the empirical
local squared-Cramér variation $\hat{\sigma}$.

###  E.6 Converting local bounds to global bounds

To convert the local concentration results obtained above into a bound on
$\ell_{2}(\hat{F}(x),F^{*}(x))$, we first prove the following lemma.

######  Lemma E.5.

For $U\in\\{F\in\mathbb{R}^{\mathcal{X}\times m}:F_{m}(x)=0\text{ for all
}x\in\mathcal{X}\\}$, and any $k\geq 1$, we have

$\displaystyle\|(T_{\hat{P}}^{k}U)(x)\|_{\ell_{2}}\leq\gamma^{k/2}\sum_{x^{\prime}\in\mathcal{X}}\mathbb{P}_{\hat{P}}(X_{k}=x^{\prime}\mid
X_{0}=x)\|U(x^{\prime})\|_{\ell_{2}}\,,$

where $\mathbb{P}_{\hat{P}}$ denotes state transition probabilities for the
MRP with transition matrix $\hat{P}$.

###### Proof.

We proceed by induction. For the base case $k=1$, we have, using the operator
notation introduced in Appendix B,

$\displaystyle\|(T_{\hat{P}}U)(x)\|_{\ell_{2}}$
$\displaystyle=\bigg{\|}\sum_{y\in\mathcal{X}}\hat{P}(y|x)\sum_{j=1}^{m-1}B_{x}U(y)\bigg{\|}_{\ell_{2}}$
$\displaystyle\overset{(a)}{\leq}\sum_{y\in\mathcal{X}}\hat{P}(y|x)\bigg{\|}B_{x}U(y)\bigg{\|}_{\ell_{2}}$
$\displaystyle\overset{(b)}{\leq}\sum_{y\in\mathcal{X}}\hat{P}(y|x)\gamma^{1/2}\bigg{\|}U(y)\bigg{\|}_{\ell_{2}}\,,$

as required. where (a) follows from the triangle inequality, and (b) follows
from contractivity of $B_{x}$ in $\|\cdot\|_{\ell_{2}}$, as established by
Proposition B.2.

For the inductive step, we suppose that the result holds for some
$l\in\mathbb{N}$. Now, letting $k=l+1$, we have

$\displaystyle\|(T_{\hat{P}}^{l+1}U)(x)\|_{\ell_{2}}$
$\displaystyle=\|(T_{\hat{P}}T_{\hat{P}}^{l}U)(x)\|_{\ell_{2}}$
$\displaystyle=\bigg{\|}\sum_{y\in\mathcal{X}}\hat{P}(y|x)B_{x}(T_{\hat{P}}^{l}U)(y)\bigg{\|}_{\ell_{2}}$
$\displaystyle\leq\sum_{y\in\mathcal{X}}\hat{P}(y|x)\bigg{\|}B_{x}(T_{\hat{P}}^{l}U)(y)\bigg{\|}_{\ell_{2}}$
$\displaystyle\leq\sum_{y\in\mathcal{X}}\hat{P}(y|x)\gamma^{1/2}\bigg{\|}(T_{\hat{P}}^{l}U)(y)\bigg{\|}_{\ell_{2}}$
$\displaystyle\overset{(a)}{\leq}\sum_{y\in\mathcal{X}}\hat{P}(y|x)\gamma^{1/2}\gamma^{l/2}\sum_{x^{\prime}\in\mathcal{X}}\mathbb{P}_{\hat{P}}(X_{l}=x^{\prime}\mid
X_{0}=y)\|U(x^{\prime})\|_{\ell_{2}}$
$\displaystyle=\gamma^{(l+1)/2}\sum_{x^{\prime}\in\mathcal{X}}\bigg{[}\sum_{y\in\mathcal{X}}\mathbb{P}_{\hat{P}}(X_{l}=x^{\prime}\mid
X_{0}=y)\hat{P}(y|x)\bigg{]}\|U(x^{\prime})\|_{\ell_{2}}$
$\displaystyle=\gamma^{(l+1)/2}\sum_{x^{\prime}\in\mathcal{X}}\mathbb{P}_{\hat{P}}(X_{l+1}=x^{\prime}\mid
X_{0}=x)\|U(x^{\prime})\|_{\ell_{2}}\,,$

as required, where (a) follows by the inductive hypothesis. ∎

We therefore have, with probability at least $1-\delta$:

$\displaystyle\phantom{=}\ell_{2}(\hat{F}(x),F^{*}(x))$ (28)
$\displaystyle=\|\hat{F}(x)-F^{*}(x)\|_{\ell_{2}}$
$\displaystyle\overset{(a)}{=}\bigg{\|}\Big{[}\sum_{k\geq
0}T_{\hat{P}}^{k}(T_{\hat{P}}-T_{P})F^{*}\Big{]}(x)\bigg{\|}_{\ell_{2}}$
$\displaystyle\overset{(b)}{\leq}\sum_{k\geq
0}\bigg{\|}\big{[}T_{\hat{P}}^{k}(T_{\hat{P}}-T_{P})F^{*}\Big{]}(x)\bigg{\|}_{\ell_{2}}$
$\displaystyle\overset{(c)}{\leq}\sum_{k\geq
0}\sum_{x^{\prime}\in\mathcal{X}}\mathbb{P}_{\hat{P}}(X_{k}=x^{\prime}\mid
X_{0}=x)\gamma^{k/2}\bigg{\|}\Big{[}(T_{\hat{P}}-T_{P})F^{*}\Big{]}(x^{\prime})\bigg{\|}_{\ell_{2}}$
$\displaystyle\overset{(d)}{\leq}\sum_{k\geq
0}\sum_{x^{\prime}\in\mathcal{X}}\mathbb{P}_{\hat{P}}(X_{k}=x^{\prime}\mid
X_{0}=x)\gamma^{k/2}\bigg{[}C_{\mathrm{B}1}\frac{1}{\sqrt{N}}\sqrt{\hat{\sigma}(x^{\prime})}+C_{\mathrm{B}1}\frac{1}{\sqrt{N}}\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+C^{\prime}\frac{1}{(1-\gamma)^{1/2}N^{3/4}}\bigg{]}$
$\displaystyle\overset{(e)}{\leq}\frac{C_{\mathrm{B}1}}{\sqrt{N}}\|(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}\|_{\infty}+\frac{C_{\mathrm{B}1}}{(1-\sqrt{\gamma})\sqrt{N}}\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+\frac{C^{\prime}}{(1-\sqrt{\gamma})(1-\gamma)^{1/2}N^{3/4}}$
$\displaystyle\overset{(f)}{\leq}\frac{C_{\mathrm{B}1}}{\sqrt{N}}\|(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}\|_{\infty}+\frac{2C_{\mathrm{B}1}}{(1-\gamma)\sqrt{N}}\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+\frac{2C^{\prime}}{(1-\gamma)^{3/2}N^{3/4}}\,.$
(29)

Here, (a) follows from Equation (22), (b) follows from the triangle
inequality, and (c) follows from Lemma E.5. The step (d) is the one step in
the derivation that does not hold with probability 1, but rather with
probability $1-\delta$, and follows from the Bernstein-style bounds in
Equation (27), (e) follows from algebraic manipulation using the identity
$\sum_{k\geq 0}(\gamma^{1/2}\hat{P})^{k}=(I-\sqrt{\gamma}\hat{P})^{-1}$, as
$\hat{P}$ is a stochastic matrix, and bounding the elements of
$(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}$ by the $L^{\infty}$ norm of
the vector, and (f) follows from the inequality $(1-\sqrt{\gamma})^{-1}\leq
2(1-\gamma)^{-1}$ for $\gamma\in[0,1)$.

Focusing now on the coefficient
$\|(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}\|_{\infty}$, we note that
similar to the analysis of Agarwal et al. (2020, Lemma 4) in the mean-return
case, we can write

$\displaystyle(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}=(1-\sqrt{\gamma})^{-1}(1-\sqrt{\gamma})(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}\,,$

so that each row of

$\displaystyle(1-\sqrt{\gamma})(I-\sqrt{\gamma}\hat{P})^{-1}$

forms a probability distribution, and so we may apply Jensen’s inequality with
the map $z\mapsto\sqrt{z}$, to obtain

$\displaystyle(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}\leq(1-\sqrt{\gamma})^{-1}\sqrt{(1-\sqrt{\gamma})(I-\sqrt{\gamma}\hat{P})^{-1}\hat{\sigma}}=(1-\sqrt{\gamma})^{-1/2}\sqrt{(I-\sqrt{\gamma}\hat{P})^{-1}\hat{\sigma}}\,,$
(30)

where the inequality above holds coordinate-wise. We thus obtain the
inequality

$\displaystyle\|F^{*}-\hat{F}\|_{\ell_{2},\infty}\leq\frac{2C_{\mathrm{B}1}}{(1-\gamma)^{1/2}\sqrt{N}}\sqrt{\|(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}\|_{\infty}}+\frac{2C_{\mathrm{B}1}}{(1-\gamma)\sqrt{N}}\|F^{*}-\hat{F}\|_{\ell_{2},\infty}+\frac{2C^{\prime}}{(1-\gamma)^{3/2}N^{3/4}}\,.$

Note that the term $\|F^{*}-\hat{F}\|_{\ell_{2},\infty}$ appears on both sides
of the inequality above. By taking $N\geq
16C_{\mathrm{B}1}^{2}(1-\gamma)^{-2}$, we ensure that the coefficient in front
of $\|F^{*}-\hat{F}\|_{\ell_{2},\infty}$ on the right-hand side is made
smaller than $1/2$, similar to the approach taken in the proof of instance-
dependent sample complexity bounds for mean-return estimation by Pananjady &
Wainwright (2020, Theorem 1(a)). Under this assumption, rearrangement yields

$\displaystyle\|F^{*}-\hat{F}\|_{\ell_{2},\infty}\leq\frac{4C_{\mathrm{B}1}}{(1-\gamma)^{1/2}\sqrt{N}}\sqrt{\|(I-\sqrt{\gamma}\hat{P})^{-1}\sqrt{\hat{\sigma}}\|_{\infty}}+\frac{4C^{\prime}}{(1-\gamma)^{3/2}N^{3/4}}$
(31)

as described in Section 5.

###  E.7 The stochastic categorical CDF Bellman equation

We now aim to use our developments regarding the stochastic categorical CDF
Bellman equation (see Section 5.2) to bound the quantity appearing within the
square root. To be able to use the bound obtained in Corollary 5.12, we first
replace the factor $\sqrt{\gamma}$ inside the square-root on the right-hand
side of Equation (30) with $\gamma$, using the bound
$(I-\sqrt{\gamma}P)^{-1}\leq 2(I-\gamma P)^{-1}$, shown by Agarwal et al.
(2020, Lemma 4). This, together with Corollary 5.12, yields

$\displaystyle\sqrt{(I-\sqrt{\gamma}\hat{P})^{-1}\hat{\sigma}}\leq\sqrt{2}\sqrt{(I-\gamma\hat{P})^{-1}\hat{\sigma}}\leq\sqrt{2}\sqrt{\frac{2}{1-\gamma}}\mathbf{1}=\frac{2}{(1-\gamma)^{1/2}}\mathbf{1}\,.$

Thus, returning to Equation (31) and substituting this bound in, we obtain
(again, with probability at least $1-\delta$, for all $x\in\mathcal{X}$):

$\displaystyle\ell_{2}(\hat{F}(x),F^{*}(x))$
$\displaystyle\leq\frac{8C_{\mathrm{B}1}}{(1-\gamma)\sqrt{N}}+\frac{4C^{\prime}}{(1-\gamma)^{3/2}N^{3/4}}$
(32)

###  E.8 Final steps of the proof of Theorem 5.3

We now let $N\geq c_{0}\varepsilon^{-2}(1-\gamma)^{-2}$, with $c_{0}$ any
positive number satisfying

$\displaystyle c_{0}\geq 2^{10}C_{\mathrm{B}1}^{2}\text{ and
}c_{0}>16^{4/3}(C^{\prime})^{4/3}\,.$ (33)

Note that the first of these conditions implies the earlier assumption $N\geq
16C_{\mathrm{B}1}^{2}(1-\gamma)^{-2}$ made in order to arrive at Equation
(31). We then conclude from Equation (32) that (with probability at least
$1-\delta$, for all $x\in\mathcal{X}$):

$\displaystyle\ell_{2}(\hat{F}(x),F^{*}(x))$
$\displaystyle\leq\frac{8C_{\mathrm{B}1}}{\sqrt{c_{0}}}\varepsilon+\frac{4C^{\prime}}{c_{0}^{3/4}}\varepsilon^{3/2}$
$\displaystyle\leq\frac{8C_{\mathrm{B}1}}{\sqrt{c_{0}}}\varepsilon+\frac{4C^{\prime}}{c_{0}^{3/4}}\varepsilon$
$\displaystyle\leq\varepsilon/2\,,$

with the final inequality following due to the assumption on $c_{0}$ making
both coefficients of $\varepsilon$ bounded by $1/4$. This concludes the proof
of Theorem 5.3.

To derive a concrete sample complexity bound, including logarithmic terms, we
may bound the terms in Equation (33) as follows; we emphasise that we do not
aim to be as tight as possible in the following analysis, but merely to
provide an concrete, valid bound on $c_{0}$. First, we have that the

$\displaystyle 2^{10}C_{\mathrm{B}1}^{2}$
$\displaystyle=2^{10}(\sqrt{2\log(4|\mathcal{X}|/\delta)})^{2}$
$\displaystyle=2^{11}\log(4|\mathcal{X}|/\delta)\,.$

Next, we have (introducing the shorthand $L=\log(4|\mathcal{X}|/\delta)$):

$\displaystyle 16^{4/3}(C^{\prime})^{4/3}$ $\displaystyle\leq
2^{6}\bigg{(}C_{\mathrm{B}1}(\sqrt{2C_{\mathrm{B}1}+C_{\mathrm{H}}}+\sqrt{2C_{\mathrm{B}2}})+C_{\mathrm{B}2}\bigg{)}^{4/3}$
$\displaystyle=2^{6}\big{(}\sqrt{2L}\big{(}\sqrt{2\sqrt{2L}+\sqrt{L/2}}+\sqrt{4L}\big{)}+2L\big{)}^{4/3}$
$\displaystyle\leq
2^{6}\big{(}L\sqrt{2}\big{(}\sqrt{2\sqrt{2}+\sqrt{1/2}}+2\big{)}+2L\big{)}^{4/3}$
$\displaystyle=2^{6}L^{4/3}\big{(}\sqrt{2}\big{(}\sqrt{2\sqrt{2}+\sqrt{1/2}}+2\big{)}+2\big{)}^{4/3}$
$\displaystyle\leq 2^{6}L^{4/3}2^{4}$ $\displaystyle=2^{10}L^{4/3}\,.$

We therefore conclude, for example, that taking $c_{0}\geq
2^{11}\log(4|X|/\delta)^{4/3}$ is sufficient.

##  Appendix F Further experimental results and details

In this section, we give full details for the experiment presented in the main
paper, and additionally present extended results on several additional
environments. All experiments were run using the Python 3 language (Van Rossum
& Drake, 2009), and made use of NumPy (Harris et al., 2020), SciPy (Virtanen
et al., 2020), Matplotlib (Hunter, 2007), and Seaborn (Waskom, 2021)
libraries.

###  F.1 Environments

We report results on four MRPs defined as follows:

  1. 1.

Chain: 10 states arranged in a chain
$x_{1}\leftrightarrow{}x_{2}\leftrightarrow{}\cdots\leftrightarrow{}x_{10}$.
Each state transitions to its neighbours with equal probability. States 1 and
10 are terminal. Only state 10 has a reward of 1, and all other states have
reward 0.

  2. 2.

Low random: There are five states, and the transition probabilities from each
state to all five states are drawn independently from a Dirichlet distribution
with concentration parameter $(0.01,\ldots,0.01)$. The rewards for each state
is drawn i.i.d. from the uniform distribution over $[0,1]$. We draw these
transition probability and rewards using the same random seed to yield the
same MRP for all experiments.

  3. 3.

High random: Same as the low random environment, but with transitions drawn
from a Dirichlet distribution with concentration parameter $(10,\ldots,10)$.

  4. 4.

Two-state: A 2-state MRP modified from Example 6.5 of Rowland et al. (2023).
The transition matrix is

$\begin{pmatrix}0.6&0.4\\\ 0.8&0.2\\\ \end{pmatrix}$

where the $(i,j)$ entry is the probability of transitioning from state $i$ to
state $j$. State 1 has reward 0, and state 2 has reward 1.

We chose this set of environments to include classic tabular settings such as
the chain, two environments with very different levels of stochasticity (low
and high random), as well as the two-state example from Rowland et al. (2023),
which illustrates the nature of fixed-point approximation error incurred by
QDP in environments with tight bootstrapping loops. We vary the discount
factor $\gamma\in\\{0.8,0.9,0.95,0.99\\}$. For reference, we plot Monte Carlo
approximations of return distributions for each environment in Figure 4, with
$\gamma=0.9$. The return samples are computed using the first-visit Monte
Carlo algorithm. For each initial state, we run at least $T$ transitions such
that the maximum error after $T$ transitions
$\gamma^{T}r_{\text{max}}/(1-\gamma)<\epsilon$, where $r_{\text{max}}=1$ and
we set $\epsilon=10^{-4}$. This is repeated $10^{4}$ times for each state,
giving at least $10^{4}$ return samples per state.

(a) Chain

(b) Low random

(c) High random

(d) Two-state

Figure 4: Monte Carlo approximations of return distributions in each of the
four environments tested.

###  F.2 Algorithms

As described in the main paper, we compare DCFP, QDP, and CDP algorithms. For
DCFP and CDP, we make use of the sparse structure of the update matrix (see
Equation (9)), and implement these algorithms using SciPy’s sparse matrix-
vector multiplication and linear system solution methods (Virtanen et al.,
2020). The mathematical details regarding the sparsity of the update matrix is
described in Appendix F.3. For comparison, we also run implementations of DCFP
and CDP that do not exploit this sparsity, and instead use standard dense
NumPy implementations for matrix-vector multiplication and linear system
solution methods; these are denoted by d-DCFP and d-CDP, respectively.

By default the categorical methods (all variants of DCFP and CDP) use the atom
locations described in the main paper: $z_{i}=\tfrac{i-1}{m-1}(1-\gamma)^{-1}$
for $i=1,\ldots,m$. These atom locations are sufficient to establish the
theoretical results in the main paper, though we note that in practice,
particularly when true return distributions are localised within a small sub-
interval of $[0,(1-\gamma)^{-1}]$, this setting can be overly conservative.
However, in many cases, there are straightforward ways in which the choice of
support can be improved, essentially by replacing the uniformly valid return
range $[0,(1-\gamma)^{-1}]$ with an _a priori_ known environment-specific
reward range. As an example, if the known range of immediate rewards forms a
sub-interval $[r_{\text{min}},r_{\text{max}}]\subseteq[0,1]$, then the return
must fall into the environment-specific reward range
$[r_{\text{min}}(1-\gamma)^{-1},r_{\text{max}}(1-\gamma)^{-1}]$, and improved
atom locations
$z_{i}=r_{\text{min}}(1-\gamma)^{-1}+\tfrac{i-1}{m-1}(r_{\text{max}}-r_{\text{min}})(1-\gamma)^{-1}$