# Distributional Bellman Operators over Mean Embeddings

Li Kevin Wenliang1, Grégoire Déletang1, Matthew Aitchison1, Marcus Hutter1,
Anian Ruoss1, Arthur Gretton1,2, Mark Rowland1 1Google DeepMind; 2Gatsby Unit,
University College London  {kevinliw,gdelt,matthewai,mhutter,
<EMAIL_ADDRESS>
###### Abstract

We propose a novel algorithmic framework for distributional reinforcement
learning, based on learning finite-dimensional mean embeddings of return
distributions. We derive several new algorithms for dynamic programming and
temporal-difference learning based on this framework, provide asymptotic
convergence theory, and examine the empirical performance of the algorithms on
a suite of tabular tasks. Further, we show that this approach can be
straightforwardly combined with deep reinforcement learning, and obtain a new
deep RL agent that improves over baseline distributional approaches on the
Arcade Learning Environment.

##  1 Introduction

In distributional approaches to reinforcement learning (RL), the aim is to
learn the full probability distribution of future returns (Morimura et al.,
2010a; Bellemare et al., 2017; 2023), rather than just their expected value,
as is typically the case in value-based reinforcement learning (Sutton &
Barto, 2018). Distributional RL was proposed in the setting of deep
reinforcement learning by Bellemare et al. (2017), with a variety of precursor
work stretching back almost as far as Markov decision processes themselves
(Jaquette, 1973; Sobel, 1982; Chung & Sobel, 1987; Morimura et al., 2010a; b).
Beginning with the work in Bellemare et al. (2017), the distributional
approach to reinforcement learning has been central across a variety of
applications of deep RL in simulation and in the real world (Bodnar et al.,
2020; Bellemare et al., 2020; Wurman et al., 2022; Fawzi et al., 2022).

Typically, predictions of return distributions are represented directly as
approximate probability distributions, such as categorical distributions
(Bellemare et al., 2017). Rowland et al. (2019) proposed an alternative
framework where return distributions are represented via the values of
_statistical functionals_ , called a _sketch_ by Bellemare et al. (2023). This
provided a new space of distributional reinforcement learning algorithms,
leading to improvements in deep RL agents, and hypotheses regarding
distributional RL in the brain (Dabney et al., 2020; Lowet et al., 2020). On
the other hand, a potential drawback of this approach is that each
distributional Bellman update to the representation, these values must be
“decoded” back into an approximate distribution via an _imputation strategy_.
In practice, this can introduce significant computational overhead to Bellman
updates, and is unlikely to be biologically plausible for distributional
learning in the brain (Tano et al., 2020).

Here, we focus on a notable instance of the sketch called the _mean embedding
sketch_. In short, the mean embedding is the expectation of nonlinear
functions under the distribution represented (Smola et al., 2007;
Sriperumbudur et al., 2010; Berlinet & Thomas-Agnan, 2011), and is related to
_frames_ in signal processing (Mallat, 1999) and _distributed distributional
code_ in neuroscience (Sahani & Dayan, 2003; Vértes & Sahani, 2018). The core
contributions of this paper are to revisit the approach to distributional
reinforcement learning based on sketches (Rowland et al., 2019), and to
propose the _sketch Bellman operator_ that updates the implicit distributional
representation as a simple linear operation, obviating the need for the
expensive imputation strategies converting between sketches and distributions.
This provides a rich new space of distributional RL algorithms that operate
entirely in the space of sketches. We provide theoretical convergence analysis
to accompany the framework, investigate the practical behaviour of various
instantiations of the proposed algorithms in tabular domains, and demonstrate
the effectiveness of the sketch framework in deep reinforcement learning,
showing that our approach is robust enough to serve as the basis for a new
variety of deep distributional reinforcement learning algorithms.

##  2 Background

We consider a Markov decision process (MDP) with state space $\mathcal{X}$,
action space $\mathcal{A}$, transition probabilities
$P:\mathcal{X}\times\mathcal{A}\rightarrow\mathscr{P}(\mathcal{X})$, reward
distribution function
$P_{R}:\mathcal{X}\times\mathcal{A}\rightarrow\mathscr{P}(\mathbb{R})$, and
discount factor $\gamma\in[0,1)$. Given a policy
$\pi:\mathcal{X}\rightarrow\mathscr{P}(\mathcal{A})$ and initial state
$x\in\mathcal{X}$, a random trajectory $(X_{t},A_{t},R_{t})_{t\geq 0}$ is the
sequence of random states, actions, and rewards encountered when using the
policy $\pi$ to select actions in this MDP. More precisely, we have $X_{0}=x$,
$A_{t}\sim\pi(\cdot|X_{t})$, $R_{t}\sim P_{R}(X_{t},A_{t})$, $X_{t+1}\sim
P(\cdot|X_{t},A_{t})$ for all $t\geq 0$. We write $\mathbb{P}^{\pi}_{x}$ and
$\mathbb{E}^{\pi}_{x}$ for probabilities and expectations with respect to this
distribution, respectively. The performance along the trajectory is measured
by the discounted return, defined by

$\displaystyle\sum_{t=0}^{\infty}\gamma^{t}R_{t}\,.$ (1)

In typical value-based reinforcement learning, during policy evaluation, the
agent learns the expectation of the return for each possible initial state
$x\in\mathcal{X}$, which is encoded by the value function
$V^{\pi}:\mathcal{X}\rightarrow\mathbb{R}$, given by
$V^{\pi}(x)=\mathbb{E}^{\pi}_{x}[\sum_{t=0}^{\infty}\gamma^{t}R_{t}]$.

###  2.1 Distributional RL and the distributional Bellman equation

In distributional reinforcement learning, the problem of policy evaluation is
to learn the probability distribution of return in Equation (1) for each
possible initial state $x\in\mathcal{X}$. This is encoded by the return-
distribution function
$\eta^{\pi}:\mathcal{X}\rightarrow\mathscr{P}(\mathbb{R})$, which maps each
initial state $x\in\mathcal{X}$ to the corresponding distribution of the
random return, e.g. $\eta^{\pi}(x)$ is the return distribution of state $x$. A
central result in distributional reinforcement learning is the distributional
Bellman equation, which relates the distribution of the random return under
different combinations of initial states and actions.

To build the random variable formulation of the returns, we let
$(G^{\pi}(x):x\in\mathcal{X})$ be a collection of random variables with the
property that $G^{\pi}(x)$ is equal to Equation (1) in distribution,
conditioned on the initial state $X_{0}=x$. This formulation implies that the
random variable $G^{\pi}(x)$ is distributed as $\eta^{\pi}(x)$, introduced
above, for all $x\in\mathcal{X}$. Consider a random transition
$(x,R,X^{\prime})$ generated by $\pi$, independent of the $G^{\pi}$ random
variables. Then, the (random variable) distributional Bellman equation states
that for each initial state $x$,

$\displaystyle
G^{\pi}(x)\overset{\mathcal{D}}{=}R+\gamma{G}^{\pi}(X^{\prime})\quad|\,X=x\,.$

Here, we use the slight abuse of the conditioning bar to set the distribution
of $X$ in the random transition. It is also useful to introduce the
distributional Bellman operator
$\mathcal{T}^{\pi}:\mathscr{P}(\mathbb{R})^{\mathcal{X}}\rightarrow\mathscr{P}(\mathbb{R})^{\mathcal{X}}$
to describe the transformation that occurs on the right-hans side (Morimura et
al., 2010a; Bellemare et al., 2017). If
$\eta\in\mathscr{P}(\mathbb{R})^{\mathcal{X}}$ is a collection of probability
distributions, and $(G(x):x\in\mathcal{X})$ is a collection of random
variables such that $G(x)\sim\eta(x)$ for all $x$, and $(X,R,X^{\prime})$ is
random transition generated by $\pi$, independent of $(G(x):x\in\mathcal{X})$,
then $(\mathcal{T}^{\pi}\eta)(x)=\text{Dist}(R+\gamma G(X^{\prime})|X=x)$.

To implement algorithms of distributional RL, one needs to approximate the
infinite-dimensional return-distribution function $\eta^{\pi}$ with finite-
dimensional representations. This is typically done via direct approximations
in the space of distributions; see e.g. Bellemare et al. (Chapter 5; 2023).

###  2.2 Statistical functionals and sketches

Rather than using approximations in the space of distributions, Rowland et al.
(2019) proposed to represent return distributions indirectly via _functionals_
of the return distribution, called _sketches_ by Bellemare et al. (2023). In
this work we consider a specific class of sketches, defined below.

######  Definition 2.1 (Mean embedding sketches).

A _mean embedding sketch_ $\psi$ is specified by a function
$\phi:\mathbb{R}\rightarrow\mathbb{R}^{m}$, and defined by

$\displaystyle\psi(\nu)~{}\vcentcolon=~{}\mathbb{E}_{Z\sim\nu}[\phi(Z)]\,.$
(2)

For a given distribution $\nu$, the embedding $\psi(\nu)$ can therefore be
thought of as providing a _lossy_ summary of the distribution. The name here
is motivated by the kernel literature, in which Equation (2) can be viewed as
embedding the distribution $\nu$ into $\mathbb{R}^{m}$ based on the mean value
of $\phi$ under this distribution (Smola et al., 2007; Sriperumbudur et al.,
2010; Berlinet & Thomas-Agnan, 2011).

Figure 1: The statistical functional framework proposed by Rowland et al.
(2019) (top), and the framework proposed in this paper, directly updating
sketches, avoiding the imputation step (bottom).

Statistical functional dynamic programming and temporal-difference learning
(SFDP/SFTD; Rowland et al. (2019), see also Bellemare et al. (2023)) is an
approach to distributional RL in which sketch values, rather than approximate
distributions, are the primary object learned. Given a sketch $\psi$ and
estimated sketch values $U:\mathcal{X}\rightarrow\mathbb{R}^{m}$, these
approaches proceed by first defining an _imputation strategy_
$\iota:\mathbb{R}^{m}\rightarrow\mathscr{P}(\mathbb{R})$ mapping sketch values
back to distributions, with the aim that $\psi(\iota(U))\approx U$, so that
$\iota$ behaves as an approximate pseudo-inverse to $\psi$. The usual Bellman
backup is then applied to this _imputed_ distribution, and the sketch value
extracted from this updated distribution. Thus, if
$U:\mathcal{X}\rightarrow\mathbb{R}^{m}$ represents approximations to sketch
values, a typical update in SFDP takes the form
$U\leftarrow\psi((\mathcal{T}^{\pi}\iota(U))(x))$ (Bellemare et al., 2023);
see Figure 1.

This approach led to expectile-regression DQN, a deep RL agent that aims to
learn the sketch values associated with certain expectiles (Newey & Powell,
1987) of the return, and influenced a distributional model of dopamine
signalling in the brain (Dabney et al., 2020). An important consideration is
that computation of the imputation strategy is often costly in machine
learning applications, and considered biologically implausible in neuroscience
models (Tano et al., 2020).

##  3 The Bellman sketch framework

Our goal is to derive a framework for approximate computation of the sketch
$\psi$ (with corresponding feature function $\phi$) of the return
distributions corresponding to a policy $\pi$, without needing to design,
implement, and compute an imputation strategy as in the case of SFDP/TD; see
Figure 1 for a visual comparison of the two approaches. That is, we aim to
compute the object $U^{\pi}:\mathcal{X}\rightarrow\mathbb{R}^{m}$, given by

$\displaystyle\textstyle
U^{\pi}(x)~{}\vcentcolon=~{}\psi(\eta^{\pi}(x))~{}=~{}\mathbb{E}^{\pi}_{x}[\phi(\sum_{t=0}^{\infty}\gamma^{t}R_{t})]\,.$

We begin by considering environments with a finite set of possible rewards
$\mathcal{R}\subseteq\mathbb{R}$; we discuss generalisations later. To
motivate our method, we first consider a special case; suppose that for each
possible return $g\in\mathbb{R}$, and each possible immediate reward
$r\\!\in\\!\mathcal{R}$, there exists a matrix $B_{r}$ such that

$\displaystyle\phi(r+\gamma g)~{}=~{}B_{r}\phi(g)\,;$ (3)

note that $B_{r}$ does not depend on $g$, and $\gamma$ is a constant. In
words, this says that the feature function $\phi$ evaluated at the bootstrap
return $r+\gamma g$ is expressible as a linear transformation of the feature
function evaluated at $g$ itself. If such a relationship holds, then we have

$\displaystyle
U^{\pi}(x)~{}\overset{(a)}{=}~{}\mathbb{E}^{\pi}_{x}[\phi(R+\gamma
G^{\pi}(X^{\prime}))]~{}\overset{(b)}{=}~{}\mathbb{E}^{\pi}_{x}[B_{R}\phi(G^{\pi}(X^{\prime}))]~{}\overset{(c)}{=}~{}\mathbb{E}^{\pi}_{x}[B_{R}U^{\pi}(X^{\prime})]\,,$
(4)

where (a) follows from the distributional Bellman equation, (b) follows from
Equation (3), and (c) from exchanging the linear map $B_{r}$ and the
conditional expectation given $(R,X^{\prime})$, crucially relying on the
linearity of the approximation in Equation (3). Note that for example with
$\phi(g)=(1,g)^{\top}$ we have $B_{r}=({1~{}0\atop r~{}\gamma})$, and Equation
(4) reduces to the classical Bellman equation for $V^{\pi}$, with
$U^{\pi}(x)=(1,V^{\pi}(x))^{\top}$.

Thus, $U^{\pi}(x)$ satisfies its own linear Bellman equation, which motivates
algorithms that work directly in the space of sketches, without recourse to
imputation strategies. In particular, a natural dynamic programming algorithm
to consider is based on the recursion

$\displaystyle
U(x)~{}\leftarrow~{}\mathbb{E}^{\pi}_{x}[B_{R}U(X^{\prime})]\,.$ (Sketch-DP)

As this is an update applied directly to sketch values themselves, we
introduce the _sketch Bellman operator_
$\mathcal{T}^{\pi}_{\phi}:(\mathbb{R}^{m})^{\mathcal{X}}\rightarrow(\mathbb{R}^{m})^{\mathcal{X}}$,
with $(\mathcal{T}^{\pi}_{\phi}U)(x)$ defined according to the right-hand side
of Equation (Sketch-DP). Note that $\mathcal{T}^{\pi}_{\phi}$ is a _linear_
operator, in contrast to the standard expected-value Bellman operator, which
is affine. We recover the affine case by taking one component of $\phi$ to be
constant, e.g. $\phi_{1}(g)\equiv 1$, and enforcing $U_{1}(x)\equiv 1$.

The right-hand side of Equation (Sketch-DP) can be unbiasedly approximated
with a sample transition $(x,r,x^{\prime})$. Stochastic approximation theory
(Kushner & Yin, 1997; Bertsekas & Tsitsiklis, 1996) then naturally suggests
the following temporal-difference learning update, given a learning rate
$\alpha$:

$\displaystyle U(x)~{}\leftarrow~{}(1-\alpha)U(x)+\alpha
B_{r}V(x^{\prime})\,,$ (Sketch-TD)

Rowland et al. (2019) introduced the term _Bellman closed_ for sketches for
which an _exact_ dynamic programming algorithm is available, and provided a
characterisation of Bellman closed mean embedding sketches. The notion of
Bellman closedness is closely related to the relationship in Equation (3), and
from Rowland et al. (Theorem 4.3; 2019), we can deduce that the only mean
embedding sketches that satisfy Equation (3) are invertible linear
combinations of first-$m$ moments.

Thus, our discussion above serves as a way of re-deriving known algorithms for
computing moments of the return (Sobel, 1982; Lattimore & Hutter, 2014), but
is insufficient to yield algorithms for computing other sketches.
Additionally, since moments of the return distribution are naturally of widely
differing magnitudes, it is difficult to learn a high-dimensional mean
embedding based on moments; see Appendix D.3 for further details. To go
further, we must weaken the assumption made in Equation (3).

###  3.1 General sketches

To extend our framework to a much more general family of sketches, we relax
our assumption of the exact predictability of $\phi(r+\gamma g)$ from
$\phi(g)$ in Equation (3), by defining a matrix of _Bellman coefficients_
$B_{r}$ for each possible reward $r\in\mathcal{R}$ as the solution of the
linear regression problem:

$\displaystyle
B_{r}~{}\vcentcolon=~{}\operatorname*{arg\,min}_{B}\mathbb{E}_{G\sim\mu}\Big{[}\|\phi(r+\gamma
G)-B\phi(G)\|_{2}^{2}\Big{]}\,,$ (5)

so that, informally, we have $\phi(r+\gamma g)\approx B_{r}\phi(g)$ for each
$g$. Here, $\mu$ is a distribution to be specified that weights the returns
$G$. Using the same motivation as in the previous section, we therefore obtain

$\displaystyle
U^{\pi}(x)~{}\overset{(a)}{=}~{}\mathbb{E}^{\pi}_{x}[\phi(R+\gamma
G^{\pi}(X^{\prime}))]~{}\approx~{}\mathbb{E}^{\pi}_{x}[B_{R}\phi(G^{\pi}(X^{\prime}))]~{}\overset{(c)}{=}~{}\mathbb{E}^{\pi}_{x}[B_{R}U^{\pi}(X^{\prime})]\,,$
(6)

noting that informally we have _approximate_ equality in the middle of this
line. This still motivates the approaches expressed in Equations (Sketch-DP)
and (Sketch-TD), though we have lost the property that the exact sketch values
$U^{\pi}$ are a fixed point of the dynamic programming procedure.

Algorithm 1 Sketch-DP/Sketch-TD

# Precompute Bellman coefficients

Compute $C$ as in Equation (7)

for $r\in\mathcal{R}$ do

Compute $C_{r}$ as in Equation (7)

Set $B_{r}=C_{r}C^{-1}$

end for

Initialise $U:\mathcal{X}\rightarrow\mathbb{R}^{m}$

# Main loop

for $k=1,2,\ldots$ do

if DP then

$U(x)\leftarrow\\!\sum\limits_{r,x^{\prime},a}P(r,x^{\prime}|x,a)\pi(a|x)B_{r}U(x^{\prime})\
\ \forall x$

else if TD then

Observe transition $(x_{k},a_{k},r_{k},x^{\prime}_{k})$.

$U(x_{k})\leftarrow(1-\alpha_{k})U(x_{k})+\alpha_{k}B_{r_{k}}U(x^{\prime}_{k})$

end if

end for

Computing Bellman coefficients. Under mild conditions (invertibility of $C$ as
follows) the matrix of Bellman coefficients $B_{r}$ defined in Equation (5)
can be expressed as $B_{r}=C_{r}C^{-1}$, where $C,C_{r}\in\mathbb{R}^{m\times
m}$ are defined by

$\displaystyle C\vcentcolon=\mathbb{E}_{G\sim\mu}[\phi(G)\phi(G)^{\top}]\,,$
(7) $\displaystyle C_{r}\vcentcolon=\mathbb{E}_{G\sim\mu}[\phi(r+\gamma
G)\phi(G)^{\top}]\,.$

The elements of these matrices are expressible as integrals over the real
line, and hence several possibilities are available for (approximate)
computation: if $\mu$ is finitely-supported, direct summation is possible; in
certain cases the integrals may be analytically available, and otherwise
numerical integration can be performed. Additionally, for certain feature maps
$\phi$, the Bellman coefficients $B_{r}$ have particular structure that can be
exploited computationally; see Appendix B.3 for further discussion. Detailed
properties of $B_{r}$ are studied in Section B.5.

Algorithms. We summarise the two core algorithmic contributions, sketch
dynamic programming (Sketch-DP) and sketch temporal-difference learning
(Sketch-TD), that arise from our proposed framework in Algorithm 1. Pausing to
take stock, we have proposed an algorithm framework for computing
approximations of _lossy_ mean embeddings for a wide variety of feature
functions $\phi$. Further, these algorithms operate directly within the space
of sketch values.

Selecting feature maps. A natural question is what effects the choice of
feature map $\phi$ has on the performance of the algorithm. There are several
competing concerns. First, the richer the map $\phi$, the more information
about the return distribution can be captured by the corresponding mean
embedding. However, the computational costs (both in time and memory) of our
proposed algorithms scale in the worst case cubically with $m$, the
dimensionality of the mean embedding. In addition, the accuracy of the
algorithm in approximating the mean embeddings of the true return
distributions relies on having a low approximation error in Equation (6),
which in turn relies on a low regression error in Equation (5) (see
Proposition 4.1 below). Selecting an appropriate feature map is therefore
somewhat nuanced, and involves trading off a variety of computational and
approximation concerns.

A collection of feature maps we will use throughout the paper that offer the
potential for trade-offs along the dimensions identified above is given by the
translation family

$\displaystyle\phi_{i}(z)\vcentcolon=\kappa(s(z-z_{i})),~{}~{}\forall~{}i\in\\{1,\cdots,m\\}\,,$
(8)

where $\kappa:\mathbb{R}\to\mathbb{R}$ is a _base feature function_ ,
$s\in\mathbb{R}^{+}$ is the _slope_ , and the set
$\\{z_{1},\ldots,z_{m}\\}\subseteq\mathbb{R}$ is the _anchors_ of the feature
map. We will often take $\kappa$ to be commonly used bounded and smooth
nonlinear functions, such as the Gaussian or the sigmoid functions, and spread
the anchor points over the return range. We emphasise that in principle there
are no restrictions on the feature maps that can be considered in the
framework; see Section B.2 for other possible choices.

######  Remark 3.1 (Invariance).

Given the $m$-dimensional function space obtained from the span of the
coordinate functions $\phi_{1},\ldots,\phi_{m}$, the algorithms proposed above
are essentially independent of the choice of basis for this space. For any
invertible matrix $M\in\mathbb{R}^{m\times m}$, replacing $\phi$ by
$M^{-1}\phi$, and _also_ $||\cdot||_{2}$ by $||\cdot||_{M^{\top}M}$ in
Equation (5) gives an equivalent algorithm.

######  Remark 3.2 (The need for _linear_ regression).

It is tempting to try and obtain a more general framework by allowing _non-
linear_ regression of $\phi(r+\gamma g)$ on $\phi(g)$ in Equation (5), to
obtain a more accurate fit, for example fitting a function
$H:\mathbb{R}\times\mathbb{R}^{m}\rightarrow\mathbb{R}^{m}$ so that
$\phi(r+\gamma g)\approx H(r,\phi(g))$. The issue is that if $H$ is not linear
in the second argument, then generally
$\mathbb{E}[H(r,\phi(G(X^{\prime})))]\not=H(r,\mathbb{E}[\phi(G(X^{\prime}))])$,
and so step (c) in Equation (6) is not valid. However, there may be settings
where it is desirable to _learn_ such a function $H$, to avoid online
computation of Bellman coefficients every time a new reward is encountered in
TD learning.

###  3.2 Sketch-DP at work

To provide more intuition for the Bellman sketch framework, we provide a walk-
through of using Algorithm 1 to estimate the return distributions for the
environment in Figure 2A; full details for replication are given in Appendix
C. We take a feature map $\phi$ of the form given in Equation (8), taking
$\kappa$ to be the sigmoid function, and $m=13$ anchors evenly spaced between
$-4.5$ and $4.5$ (Figure 2B). The Bellman regression problem in Equation (5)
is set with $\mu=\text{Uniform}([-4,4])$, based on the typical returns
observed in the environment. The anchors and the choice of $\mu$ for
regression are important but can be set following simple heuristics; see
Sections B.2 and B.3. We then run the Sketch-DP algorithm with the initial
estimates $U(x)$ set to $\phi(0)$ for all $x\in\mathcal{X}$.

We compare the estimates produced by Sketch-DP against ground-truth by
estimating the true mean embeddings from a large number of Monte Carlo samples
of the returns from each state. Figure 2C (top) shows illustrates the ground-
truth embeddings for each state, and Figure 2D (top) compares these ground-
truth embeddings with those computed by Sketch-DP as the algorithm progresses;
by 30 iterations, the mean embeddings are very close to the ground-truth.

To aid interpretation of these results, we also include a comparison in which
we “decode” the mean embeddings back into probability distributions (via an
imputation strategy (Rowland et al., 2019)), and compare with the ground-truth
return distributions, projected onto the anchor locations of the features
(Rowland et al., 2018). Full details of the imputation strategy are in Section
B.1. These results are shown in the bottom panels of Figure 2C & D. Initially,
the imputed distributions of the Sketch-DP mean embedding estimates reflect
the initialisation to the mean embedding of $\delta_{0}$, though as more
iterations of Sketch-DP are applied, the imputed distributions become close to
the ground-truth. This indicates that, in this example, not only does Sketch-
DP compute accurate mean embeddings of the return, but that this embedding is
rich enough to recover a lot of information regarding the return distributions
themselves.

Concluding the introduction of the Sketch-DP algorithmic framework, there are
several natural questions that arise. Can we quantify how accurately Sketch-DP
algorithms can approximate mean embeddings of return distributions? What
effects do choices such as the feature map $\phi$ have on the algorithms in
practice? The next sections are devoted to answering these questions in turn.

A

B

C D

Figure 2:  A: State transitions and rewards in the environment. B: The feature
functions $\phi$ for the sketch-DP. Dotted lines indicate anchors. The
regression Equation (5) is performed under a densely spaced grid over the
light region $[-4,4]$. C: The ground-truth mean embeddings under the sigmoid
features in B, and the categorical projection of the ground-truth distribution
onto the anchors of $\phi$ in B. D: The evolution of the estimated mean
embeddings (bright blue lines) and imputed distributions (bright red lines)
during Sketch-DP. The stems are the respective ground-truth from panel C.

##  4 Convergence analysis

${\bar{\eta}}$${\bar{\eta}^{\prime}}$${\eta}$${\eta^{\prime}}$${\bar{U}}$${\bar{U}^{\prime}}$${U}$${\tilde{U}}$${U^{\prime}}$$\scriptstyle{\mathcal{T}^{\pi}}$$\scriptstyle{\Phi}$$\scriptstyle{\Phi}$$\scriptstyle{\mathcal{T}^{\pi}}$$\scriptstyle{\Phi}$$\scriptstyle{\definecolor[named]{.}{rgb}{.5,.5,.5}\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\delta+\varepsilon_{\text{R}}}$$\scriptstyle{\definecolor[named]{.}{rgb}{.5,.5,.5}\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\gamma^{c}(\delta+\varepsilon_{\text{R}})}$$\scriptstyle{\Phi}$$\scriptstyle{\mathcal{T}^{\pi}_{\phi}}$$\scriptstyle{\definecolor[named]{.}{rgb}{.5,.5,.5}\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\delta}$$\scriptstyle{\definecolor[named]{.}{rgb}{.5,.5,.5}\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\varepsilon_{\text{B}}}$$\scriptstyle{\definecolor[named]{.}{rgb}{.5,.5,.5}\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\gamma^{c}(\delta+\varepsilon_{\text{R}})+\varepsilon_{\text{E}}}$
Figure 3: The objects and structure used to analyse the Sketch-DP algorithm.

We analyse the Sketch-DP procedure described in Algorithm 1, which can be
mathematically described in the following succinct manner. We let
$U_{0}:\mathcal{X}\rightarrow\mathbb{R}^{m}$ denote the initial sketch value
estimates, and then note from Algorithm 1 that the collection of estimates
after each DP update form a sequence $(U_{k})_{k=0}^{\infty}$, with
$U_{k+1}=\mathcal{T}^{\pi}_{\phi}U_{k}$. Our convergence analysis therefore
focuses on the asymptotic behaviour of this sequence. We introduce the
notation $\Phi:\mathscr{P}(\mathbb{R})\rightarrow\mathbb{R}^{m}$ for the
sketch associated with the feature function $\phi$, so that
$\Phi\mu=\mathbb{E}_{Z\sim\mu}[\phi(Z)]$, and define $\Phi$ for return-
distribution functions by specifying for
$\eta\in\mathscr{P}(\mathbb{R})^{\mathcal{X}}$ that
$(\Phi\eta)(x)=\Phi(\eta(x))$. Ideally, we would like these iterates to
approach $U^{\pi}:\mathcal{X}\rightarrow\mathbb{R}^{m}$, the sketch values of
the true return distributions, given by
$U^{\pi}(x)=\mathbb{E}^{\pi}_{x}[\phi(\sum_{t=0}^{\infty}\gamma^{t}R_{t})]$.
As already described, typically this is not possible when the sketch $\Phi$ is
not Bellman closed, and so we can only expect to approximate $U^{\pi}$.
Mathematically, this is because in general we have
$\Phi\mathcal{T}^{\pi}\not=\mathcal{T}^{\pi}_{\phi}\Phi$ when $\phi$ is not
Bellman closed.

The first step is to bound the error incurred in a single step of dynamic
programming due to using $\mathcal{T}^{\pi}_{\phi}$ directly on the sketch
values, rather taking sketch values after applying the true distributional
Bellman operator to the underlying distributions; this corresponds to the
foreground of Figure 3.

######  Proposition 4.1.

(Regression error to Bellman approximation.) Let $\|\cdot\|$ be a norm on
$\mathbb{R}^{m}$. Then for any return-distribution function
$\eta\in\mathscr{P}([G_{\text{min}},G_{\text{max}}])^{\mathcal{X}}$, we have

$\displaystyle\max_{x\in\mathcal{X}}\|\Phi(\mathcal{T}^{\pi}\eta)(x)-(\mathcal{T}^{\pi}_{\phi}\Phi\eta)(x)\|\leq\sup_{g\in[G_{\text{min}},G_{\text{max}}]}\max_{r\in\mathcal{R}}\|\phi(r+\gamma
g)-B_{r}\phi(g)\|\,.$ (9)

The second step of the analysis is to chain together the errors that are
incurred at each step of dynamic programming, so as to obtain a bound on the
asymptotic distance of the sequence $(U_{k})_{k=0}^{\infty}$ from $U^{\pi}$,
motivated by error propagation analysis in the case of function approximation
(Bertsekas & Tsitsiklis (1996); Munos (2003); see also Wu et al. (2023) in the
distributional setting). The next proposition provides the technical tools
required for this; the notation is chosen to match the illustration in Figure
3.

######  Proposition 4.2.

(Error propagation.) Consider a norm $\|\cdot\|$ on $\mathbb{R}^{m}$, and let
$\|\cdot\|_{\infty}$ be the norm on $(\mathbb{R}^{m})^{\mathcal{X}}$ defined
by $\|U\|_{\infty}=\max_{x\in\mathcal{X}}\|U(x)\|$. Let $d$ be a metric on
return-distribution functions (RDFs) such that $\mathcal{T}^{\pi}$ is a
$\gamma^{c}$-contraction with respect to $d$. Suppose the following bounds
hold.

  * •

(Bellman approximation bound.) For any
$\eta\in\mathscr{P}([G_{\text{min}},G_{\text{max}}])^{\mathcal{X}}$,

$\displaystyle\max_{x\in\mathcal{X}}\|\Phi(\mathcal{T}^{\pi}\eta)(x)-(\mathcal{T}^{\pi}_{\phi}\Phi\eta)(x)\|\leq\varepsilon_{\text{B}}\,.$

  * •

(Reconstruction error bound.) For any
$\eta,\bar{\eta}\in\mathscr{P}([G_{\text{min}},G_{\text{max}}])^{\mathcal{X}}$
with sketches $U,\bar{U}$, we have
$d(\eta,\bar{\eta})\leq\|U-\bar{U}\|_{\infty}+\varepsilon_{\text{R}}$.

  * •

(Embedding error bound.) For any
$\eta^{\prime},\bar{\eta}^{\prime}\in\mathscr{P}([G_{\text{min}},G_{\text{max}}])^{\mathcal{X}}$
with sketches $U^{\prime},\bar{U}^{\prime}$, we have
$\|U^{\prime}-\bar{U}^{\prime}\|_{\infty}\leq
d(\eta^{\prime},\bar{\eta}^{\prime})+\varepsilon_{\text{E}}$.

Then for any two return-distribution functions
$\eta,\bar{\eta}\in\mathscr{P}([G_{\text{min}},G_{\text{max}}])^{\mathcal{X}}$
with sketches $U,\bar{U}$ satisfying $\|U-\bar{U}\|\leq\delta$, we have

$\displaystyle\|\Phi\mathcal{T}^{\pi}\eta-\mathcal{T}^{\pi}_{\phi}\bar{U}\|_{\infty}\leq\gamma^{c}(\delta+\varepsilon_{\text{R}})+\varepsilon_{\text{R}}+\varepsilon_{\text{E}}\,.$

A formal proof is given in Appendix A; Figure 3 (bottom) shows the intuition,
propagating bounds through different intermediate stages of the analysis of
the update. We now state the main error bound result, which combines the two
earlier results.

######  Proposition 4.3.

Suppose the assumptions of Proposition 4.2 hold, that $\mathcal{T}^{\pi}$ maps
$\mathscr{P}([G_{\text{min}},G_{\text{max}}])^{\mathcal{X}}$ to itself, and
suppose $\mathcal{T}^{\pi}_{\phi}$ maps
$\\{\Phi\nu:\nu\in\mathscr{P}([G_{\text{min}},G_{\text{max}}])^{\mathcal{X}}\\}$
to itself. Then for a sequence of sketches $(U_{k})_{k=0}^{\infty}$ defined
iteratively via $U_{k+1}=\mathcal{T}^{\pi}_{\phi}U_{k}$, we have

$\displaystyle\limsup_{k\rightarrow\infty}\|U_{k}-U^{\pi}\|\leq\frac{1}{1-\gamma^{c}}(\gamma^{c}\varepsilon_{\text{R}}+\varepsilon_{\text{B}}+\varepsilon_{\text{E}})\,.$

###### Proof.

For each $U_{k}$, let $\eta_{k}$ be an RDF with the property
$\Phi\eta_{k}=U_{k}$. Applying Proposition 4.2 to sketches $U^{\pi}$ and
$U_{k}$, we obtain
$\|U_{k+1}-U^{\pi}\|_{\infty}\leq\gamma^{c}\|U_{k}-U^{\pi}\|_{\infty}+\gamma^{c}\varepsilon_{\text{R}}+\varepsilon_{\text{B}}+\varepsilon_{\text{E}}\,.$
Taking a limsup on both sides over $k$ and rearranging yields the result. ∎

###  4.1 Concrete example

The analysis presented above is abstract; it provides a generic template for
conducting error propagation analysis to show that Sketch-DP converges to a
neighbourhood of the true values, and moreover illustrates the dependence of
this error on the “richness” of the sketch, and accuracy of the Bellman
coefficients. To apply this abstract result to a concrete algorithm, we are
required to establish the three error bounds that appear in the statement of
Proposition 4.2. The result below shows how this can lead to a concrete result
for a novel class of sketches; in particular, proving that computed mean
embeddings under these features become arbitrarily accurate as the number of
features increases.

######  Proposition 4.4.

Consider a sketch $\phi$ whose coordinates are feature functions of the form
$\phi_{i}(z)=\mathbbm{1}\\{z_{1}\leq z<z_{i+1}\\}$ ($i=1,\ldots,m-1$), and
$\phi_{m}(z)=\mathbbm{1}\\{z_{1}\leq z\leq z_{m+1}\\}$, where
$z_{1},\ldots,z_{m+1}$ is an equally-spaced grid over
$[G_{\text{min}},G_{\text{max}}]$, with
$G_{\text{min}}=\min\mathcal{R}/(1-\gamma)$,
$G_{\text{max}}=\max\mathcal{R}/(1-\gamma)$. Let $\mathcal{T}^{\pi}_{\phi}$ be
the corresponding Sketch-DP operator given by solving Equation (5) with
$\mu=\text{Unif}([G_{\text{min}},G_{\text{max}}])$, and define a sequence
$(U_{k})_{k=0}^{\infty}$ by taking $U_{0}(x)$ to be the sketch of some initial
distribution in $\mathscr{P}([G_{\text{min}},G_{\text{max}}])$, and
$U_{k+1}=\mathcal{T}^{\pi}_{\phi}U_{k}$ for all $k\geq 0$. Let
$U^{\pi}\in(\mathbb{R}^{m})^{\mathcal{X}}$ be the mean embeddings of the true
return distributions. Finally, let $\|\cdot\|$ be the norm on $\mathbb{R}^{m}$
defined by
$\|u\|=\frac{G_{\text{max}}-G_{\text{min}}}{m}\sum_{i=1}^{m}|u_{i}|\,.$ Then
we have

$\displaystyle\limsup_{k\rightarrow\infty}\|U_{k}-U^{\pi}\|_{\infty}\leq\frac{(G_{\text{max}}-G_{\text{min}})(3+2\gamma)}{(1-\gamma)m}\,.$

##  5 Experiments

We first conduct a broad empirical investigation into the effects of three key
factors in Equation (8): the base feature $\kappa$, the number of features
$m$, and the slope $s$, using three tabular MRPs (details in Section C.1,
extended results in Section D.1). As in the example in Section 3.2, we compare
the mean embeddings estimated by Sketch-DP with ground-truth mean embeddings,
reporting their squared $L^{2}$ distance (mean embedding squared error), and
also compare the Cramér distance
$\max_{x\in\mathcal{X}}\ell_{2}^{2}(\hat{\eta}(x),\eta^{\pi}(x))$ (see e.g.
Rowland et al. (2018)) between the distribution $\hat{\eta}(x)$ imputed from
the Sketch-DP estimate, and the ground-truth return distribution
$\eta^{\pi}(x)$. To aid interpretation of the Cramér distance results, we also
report the Cramér distance between the ground truth $\eta^{\pi}(x)$ and two
baselines. First, the Dirac delta $\delta_{V^{\pi}(X)}$ at the mean return; we
expect Sketch-DP to outperform this naïve baseline by better capturing
properties of the return distribution beyond the mean. Second, the return
distribution estimate computed by categorical DP (Rowland et al., 2018;
Bellemare et al., 2023), a well-understood approach to distrbutional RL based
on categorical distributions.

The results for sweeps over feature count $m$ and slope $s$ are shown in
Figure 4. By sweeping over $m$, we see that the estimated mean embedding goes
towards the ground-truth as we use more features. Further, the Cramér distance
also decreases as $m$ increases, suggesting that the distribution represented
also approaches the ground-truth. To highlight differences between various
Sketch-DP algorithms, we also compute the excess Cramér: the Cramér distance
$\max_{x\in\mathcal{X}}\ell_{2}^{2}(\hat{\eta}(x),\eta^{\pi}(x))$ as above,
minus the corresponding distance between the categorical projection of
$\eta^{\pi}$ (c.f. the red stems in Figure 2) and $\eta^{\pi}$ itself. All
distributional methods perform well on these tasks, and significantly
outperform the Dirac estimator in stochastic environments; we note that all
methods have tunable hyperparameters (bin locations for CDRL, feature
parameters for Sketch-DP), which should inform the interpretation of these
results, and in particular direct comparison between methods. The results of
the sweep on the slope parameter $s$ show different trends depending on the
metric. For smoother $\phi$, generally we can obtain smaller error on the mean
embeddings, but the Cramér distances are only small for intermediate range of
slope values. This result is expected: when the features are too smooth or too
sharp, there exists regions within the return range where the feature values
do not vary meaningfully. This results in a more lossy encoding of the return
distribution, indicating the importance of tuning the slope parameter of the
translation family (Equation (8)).

Figure 4: Results of running Algorithm 1 on tabular environments.

###  5.1 Deep reinforcement learning

We also verify that the Bellman sketch framework is robust enough to apply in
combination with deep reinforcement learning. To do so, we aim to learn
neural-network predictions $U_{\theta}(x,a)$ of sketch values for each state-
action pair $(x,a)$ in the environment. To be able to define greedy policy
improvements based on estimated sketch values, we precompute _value-readout
coefficients_ $\beta\in\mathbb{R}^{m}$ by solving
$\operatorname*{arg\,min}_{\beta}\mathbb{E}_{G\sim\mu}[(G-\langle\beta,\phi(G)\rangle)^{2}]\,,$
so that we can predict expected returns from the sketch value as
$\langle\beta,U_{\theta}(x,a)\rangle$. This allows us to define a greedy
policy, and therefore a Q-learning-style update rule, which given an observed
transition $(x,a,r,x^{\prime})$, first computes
$a^{\prime}=\operatorname*{arg\,max}_{\tilde{a}}\langle\beta,U_{\bar{\theta}}(x^{\prime},\tilde{a})\rangle$,
and then the gradient:
$\nabla_{\theta}\|U_{\theta}(x,a)-B_{r}U_{\bar{\theta}}(x^{\prime},a^{\prime})\|^{2}_{2}\,,$
where $\bar{\theta}$ are the target network parameters. In our experiments, we
parametrise $U_{\theta}$ according to the architecture of QR-DQN (Dabney et
al., 2018b), so that the $m$ outputs of the network predict the values of the
$m$ coordinates of the corresponding sketch value. We use the sigmoid function
as the base feature $\kappa$. Full experimental details for replication are in
Appendix C.2; further results are in Section D.2.

Figure 5: Median (left) and mean (right) human-normalised scores on the Atari
57 suite.

Figure 5 shows the mean and median human-normalised performance on the Atari
suite of environments (Bellemare et al., 2013) across 200M training frames,
and includes comparisons against DQN (Mnih et al., 2015), as well as the
distributional agents C51 (Bellemare et al., 2017), QR-DQN (Dabney et al.,
2018b), and IQN (Dabney et al., 2018a). Sketch-DQN attains higher performance
on both metrics relative to the comparator agents C51 and QR-DQN, and
approaches the performance of IQN, which uses a more complex prediction
network to make non-parametric predictions of the quantile function of the
return. These results indicate that the sketch framework can be reliably
applied to deep RL, and we believe further investigation of the combination of
this framework and deep RL agents is a promising direction for future work.

##  6 Related work

Typical approaches to distributional RL focus on learning approximate
distributions directly (see, e.g., Bellemare et al. (2017); Dabney et al.
(2018b); Yang et al. (2019); Nguyen-Tang et al. (2021); Wu et al. (2023)).
Much prior work has considered statistical functionals of the random return,
at varying levels of generality with regard to the underlying Markov decision
process model. See for example Mandl (1971); Farahmand (2019) for work on
characteristic functions, Chung & Sobel (1987) for the Laplace transform,
Tamar et al. (2013; 2016) for variance, and Sobel (1982) for higher moments.
Our use of finite-dimensional mean embeddings is inspired by distributed
distributional codes (DDCs) from theoretical neuroscience (Sahani & Dayan,
2003; Vértes & Sahani, 2018; Wenliang & Sahani, 2019), which can be regarded
as neural activities encoding return distributions. DDCs were previously used
to model transition dynamics and successor features in partially observable
MDPs (Vértes & Sahani, 2019). Tano et al. (2020) consider applying non-
linearities to rewards themselves, rather than the return, and learning with a
variety of discount factors, to encode the distribution of rewards at each
timestep. The sketches in this paper are in fact mean embeddings into finite-
dimensional reproducing kernel Hilbert spaces (RKHSs; the kernel corresponding
to the feature function $\phi$ is
$K(z,z^{\prime})=\langle\phi(z),\phi(z^{\prime})\rangle$). Kernel mean
embeddings have previously been used in RL for representing state-transition
distributions (Grünewälder et al., 2012; Boots et al., 2013; Lever et al.,
2016; Chowdhury & Oliveira, 2023), and maximum mean discrepancies in RKHSs
(Gretton et al., 2012) have been used to define losses in distributional RL by
Nguyen-Tang et al. (2021).

##  7 Conclusion

We have proposed a framework for distributional reinforcement learning based
on Bellman updates that take place entirely within the sketch domain. This has
yielded new dynamic programming and temporal-difference learning algorithms as
well as novel error propagation analysis, and we have provided further
empirical analysis in the context of a suite of tabular MRPs, as well as
demonstrating that the approach can be successfully applied at scale as a
variant of the DQN architecture. We expect that there will be benefits from
further exploration of algorithmic possibilities opened up by this framework,
as well as potential consequences for value representations in the nervous
system.

#### Acknowledgments

We thank Tim Genewein, Jiaxin Shi, Yayi Zou, Maneesh Sahani and Csaba
Szepesvari for helpful discussions on the contractivity and generalisations of
the sketch Bellman framework.

## References

  * Bellemare et al. (2013) Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 2013.
  * Bellemare et al. (2017) Marc G. Bellemare, Will Dabney, and Rémi Munos.  A distributional perspective on reinforcement learning.  In _Proceedings of the International Conference on Machine Learning_ , 2017.
  * Bellemare et al. (2020) Marc G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera S. Ponda, and Ziyu Wang.  Autonomous navigation of stratospheric balloons using reinforcement learning.  _Nature_ , 588(7836):77–82, 2020.
  * Bellemare et al. (2023) Marc G. Bellemare, Will Dabney, and Mark Rowland.  _Distributional Reinforcement Learning_.  MIT Press, 2023.  http://www.distributional-rl.org.
  * Berlinet & Thomas-Agnan (2011) Alain Berlinet and Christine Thomas-Agnan.  _Reproducing kernel Hilbert spaces in probability and statistics_.  Springer Science & Business Media, 2011.
  * Bertsekas & Tsitsiklis (1996) Dimitri Bertsekas and John N. Tsitsiklis.  _Neuro-dynamic programming_.  Athena Scientific, 1996.
  * Bodnar et al. (2020) Cristian Bodnar, Adrian Li, Karol Hausman, Peter Pastor, and Mrinal Kalakrishnan.  Quantile QT-Opt for risk-aware vision-based robotic grasping.  In _Robotics: Science and Systems_ , 2020.
  * Bondanelli & Ostojic (2020) Giulio Bondanelli and Srdjan Ostojic.  Coding with transient trajectories in recurrent neural networks.  _PLoS computational biology_ , 16(2):e1007655, 2020.
  * Boots et al. (2013) Byron Boots, Arthur Gretton, and Geoffry J. Gordon.  Hilbert space embeddings of predictive state representations.  In _Proceedings of the Conference on Uncertainty in Artificial Intelligence_ , 2013.
  * Chowdhury & Oliveira (2023) Sayak Ray Chowdhury and Rafael Oliveira.  Value function approximations via kernel embeddings for no-regret reinforcement learning.  In _Proceedings of The Asian Conference on Machine Learning_ , 2023\.
  * Chung & Sobel (1987) Kun-Jen Chung and Matthew J. Sobel.  Discounted MDPs: Distribution functions and exponential utility maximization.  _SIAM Journal on Control and Optimization_ , 25(1):49–62, 1987.
  * Dabney et al. (2018a) Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos.  Implicit quantile networks for distributional reinforcement learning.  In _Proceedings of the International Conference on Machine Learning_ , 2018a.
  * Dabney et al. (2018b) Will Dabney, Mark Rowland, Marc G. Bellemare, and Rémi Munos.  Distributional reinforcement learning with quantile regression.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2018b.
  * Dabney et al. (2020) Will Dabney, Zeb Kurth-Nelson, Naoshige Uchida, Clara Kwon Starkweather, Demis Hassabis, Rémi Munos, and Matthew Botvinick.  A distributional code for value in dopamine-based reinforcement learning.  _Nature_ , 577(7792):671–675, 2020.
  * Doan et al. (2018) Thang Doan, Bogdan Mazoure, and Clare Lyle.  GAN Q-learning.  _arXiv preprint arXiv:1805.04874_ , 2018.
  * Farahmand (2019) Amir-massoud Farahmand.  Value function in frequency domain and the characteristic value iteration algorithm.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Fawzi et al. (2022) Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco J. R. Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, David Silver, Demis Hassabis, and Pushmeet Kohli.  Discovering faster matrix multiplication algorithms with reinforcement learning.  _Nature_ , 610(7930):47–53, 2022.
  * Freirich et al. (2019) Dror Freirich, Tzahi Shimkin, Ron Meir, and Aviv Tamar.  Distributional multivariate policy evaluation and exploration with the Bellman GAN.  In _Proceedings of the International Conference on Machine Learning_ , 2019.
  * Gretton et al. (2012) Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola.  A kernel two-sample test.  _The Journal of Machine Learning Research_ , 13(1):723–773, 2012.
  * Grünewälder et al. (2012) Steffen Grünewälder, Guy Lever, Luca Baldassarre, Massi Pontil, and Arthur Gretton.  Modelling transition dynamics in MDPs with RKHS embeddings.  In _Proceedings of the International Conference on Machine Learning_ , 2012.
  * Hennequin et al. (2012) Guillaume Hennequin, Tim P Vogels, and Wulfram Gerstner.  Non-normal amplification in random balanced neuronal networks.  _Physical Review E_ , 86(1):011909, 2012.
  * Jaquette (1973) Stratton C. Jaquette.  Markov decision processes with a new optimality criterion: Discrete time.  _The Annals of Statistics_ , 1(3):496–505, 1973\.
  * Kushner & Yin (1997) Harold J. Kushner and George Yin.  _Stochastic approximation and recursive algorithm and applications_.  Springer, 1997.
  * Lattimore & Hutter (2014) Tor Lattimore and Marcus Hutter.  Near-optimal PAC bounds for discounted MDPs.  _Theoretical Computer Science_ , 558:125–143, 2014.
  * Lever et al. (2016) Guy Lever, John Shawe-Taylor, Ronnie Stafford, and Csaba Szepesvari.  Compressed conditional mean embeddings for model-based reinforcement learning.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2016.
  * Lowet et al. (2020) Adam S. Lowet, Qiao Zheng, Sara Matias, Jan Drugowitsch, and Naoshige Uchida.  Distributional reinforcement learning in the brain.  _Trends in neurosciences_ , 43(12):980–997, 2020\.
  * Mallat (1999) Stéphane Mallat.  _A wavelet tour of signal processing_.  Elsevier, 1999.
  * Mandl (1971) Petr Mandl.  On the variance in controlled Markov chains.  _Kybernetika_ , 7(1):1–12, 1971.
  * Marthe et al. (2023) Alexandre Marthe, Aurélien Garivier, and Claire Vernade.  Beyond average return in markov decision processes.  In _Advances in Neural Information Processing Systems_ , 2023.
  * Mnih et al. (2015) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis.  Human-level control through deep reinforcement learning.  _Nature_ , 2015.
  * Morimura et al. (2010a) Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka.  Nonparametric return distribution approximation for reinforcement learning.  In _Proceedings of the International Conference on Machine Learning_ , 2010a.
  * Morimura et al. (2010b) Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka.  Parametric return density estimation for reinforcement learning.  In _Proceedings of the Conference on Uncertainty in Artificial Intelligence_ , 2010b.
  * Munos (2003) Rémi Munos.  Error bounds for approximate policy iteration.  In _Proceedings of the International Conference on Machine Learning_ , 2003.
  * Newey & Powell (1987) Whitney K. Newey and James L. Powell.  Asymmetric least squares estimation and testing.  _Econometrica: Journal of the Econometric Society_ , pp. 819–847, 1987.
  * Nguyen-Tang et al. (2021) Thanh Nguyen-Tang, Sunil Gupta, and Svetha Venkatesh.  Distributional reinforcement learning via moment matching.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2021.
  * Rowland et al. (2018) Mark Rowland, Marc Bellemare, Will Dabney, Rémi Munos, and Yee Whye Teh.  An analysis of categorical distributional reinforcement learning.  In _Proceedings of the International Conference on Artificial Intelligence and Statistics_ , 2018.
  * Rowland et al. (2019) Mark Rowland, Robert Dadashi, Saurabh Kumar, Rémi Munos, Marc G. Bellemare, and Will Dabney.  Statistics and samples in distributional reinforcement learning.  In _Proceedings of the International Conference on Machine Learning_ , 2019.
  * Rowland et al. (2023) Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Marc G Bellemare, and Will Dabney.  An analysis of quantile temporal-difference learning.  _arXiv preprint arXiv:2301.04462_ , 2023.
  * Sahani & Dayan (2003) Maneesh Sahani and Peter Dayan.  Doubly distributional population codes: Simultaneous representation of uncertainty and multiplicity.  _Neural Computation_ , 2003.
  * Smola et al. (2007) Alex Smola, Arthur Gretton, Le Song, and Bernhard Schölkopf.  A Hilbert space embedding for distributions.  In _Proceedings of the International Conference on Algorithmic Learning Theory_ , 2007.
  * Sobel (1982) Matthew J. Sobel.  The variance of discounted Markov decision processes.  _Journal of Applied Probability_ , 19(4):794–802, 1982.
  * Song et al. (2008) Le Song, Xinhua Zhang, Alex Smola, Arthur Gretton, and Bernhard Schölkopf.  Tailoring density estimation via reproducing kernel moment matching.  In _Proceedings of the International Conference on Machine Learning_ , 2008.
  * Sriperumbudur et al. (2010) Bharath K. Sriperumbudur, Arthur Gretton, Kenji Fukumizu, Bernhard Schölkopf, and Gert R. G. Lanckriet.  Hilbert space embeddings and metrics on probability measures.  _Journal of Machine Learning Research_ , 11:1517–1561, 2010\.
  * Sun et al. (2022) Ke Sun, Yingnan Zhao, Yi Liu, Wulong Liu, Bei Jiang, and Linglong Kong.  Distributional reinforcement learning via sinkhorn iterations.  _arXiv preprint arXiv:2202.00769_ , 2022.
  * Sutton & Barto (2018) Richard S. Sutton and Andrew G. Barto.  _Reinforcement learning: An introduction_.  MIT Press, 2nd edition, 2018.
  * Tamar et al. (2013) Aviv Tamar, Dotan Di Castro, and Shie Mannor.  Temporal difference methods for the variance of the reward to go.  In _Proceedings of the International Conference on Machine Learning_ , 2013.
  * Tamar et al. (2016) Aviv Tamar, Dotan Di Castro, and Shie Mannor.  Learning the variance of the reward-to-go.  _The Journal of Machine Learning Research_ , 17(1):361–396, 2016.
  * Tano et al. (2020) Pablo Tano, Peter Dayan, and Alexandre Pouget.  A local temporal difference code for distributional reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Vértes & Sahani (2018) Eszter Vértes and Maneesh Sahani.  Flexible and accurate inference and learning for deep generative models.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Vértes & Sahani (2019) Eszter Vértes and Maneesh Sahani.  A neurally plausible model learns successor representations in partially observable environments.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Virtanen et al. (2020) Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, İlhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ , 17:261–272, 2020.
  * Wenliang & Sahani (2019) Li Kevin Wenliang and Maneesh Sahani.  A neurally plausible model for online recognition and postdiction in a dynamical environment.  _Advances in Neural Information Processing Systems_ , 2019.
  * Wu et al. (2023) Runzhe Wu, Masatoshi Uehara, and Wen Sun.  Distributional offline policy evaluation with predictive error guarantees.  In _Proceedings of the International Conference on Machine Learning_ , 2023.
  * Wurman et al. (2022) Peter R. Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J. Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, Leilani Gilpin, Piyush Khandelwal, Varun Kompella, HaoChih Lin, Patrick MacAlpine, Declan Oller, Takuma Seno, Craig Sherstan, Michael D. Thomure, Houmehr Aghabozorgi, Leon Barrett, Rory Douglas, Dion Whitehead, Peter Dürr, Peter Stone, Michael Spranger, and Hiroaki Kitano.  Outracing champion Gran Turismo drivers with deep reinforcement learning.  _Nature_ , 602(7896):223–228, 2022.
  * Yang et al. (2019) Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie-Yan Liu.  Fully parameterized quantile function for distributional reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Zhang et al. (2021) Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Tie-Yan Liu.  Distributional reinforcement learning for multi-dimensional reward functions.  _Advances in Neural Information Processing Systems_ , 2021.

Distributional Bellman Operators over Mean Embeddings: Supplementary Material

##  Appendix A Proofs

See 4.1

###### Proof.

Let $(G(x):x\in\mathcal{X})$ be an instantiation of $\eta$ (Bellemare et al.,
2023); that is, a collection of random variables such that for each
$x\in\mathcal{X}$, we have $G(x)\sim\eta(x)$. First, note that the
distribution $(\mathcal{T}^{\pi}\eta)(x)$ is exactly the distribution of
$R+\gamma G(X^{\prime})$ (when the transition begins at $x$ and is generated
by $\pi$). So we have

$\displaystyle\Phi(\mathcal{T}^{\pi}\eta)(x)=\mathbb{E}_{Z\sim(\mathcal{T}^{\pi}\eta)(x)}[\phi(Z)]=\mathbb{E}^{\pi}_{x}[\phi(R+\gamma
G(X^{\prime}))]\,.$

It then follows that:

$\displaystyle\max_{x\in\mathcal{X}}\|\Phi(\mathcal{T}^{\pi}\eta)(x)-(\mathcal{T}^{\pi}_{\phi}\Phi\eta)(x)\|$
$\displaystyle=\max_{x\in\mathcal{X}}\Big{\|}\mathbb{E}^{\pi}_{x}\Big{[}\phi(R+\gamma
G(X^{\prime}))\Big{]}-\mathbb{E}^{\pi}_{x}\Big{[}B_{R}\mathbb{E}[\phi(G(X^{\prime}))|X^{\prime}]\Big{]}\Big{\|}$
$\displaystyle=\max_{x\in\mathcal{X}}\Big{\|}\mathbb{E}^{\pi}_{x}\Big{[}\phi(R+\gamma
G(X^{\prime}))-B_{R}\phi(G(X^{\prime}))\Big{]}\Big{\|}$
$\displaystyle\leq\max_{x\in\mathcal{X}}\mathbb{E}^{\pi}_{x}\Big{[}\big{\|}\phi(R+\gamma
G(X^{\prime}))-B_{R}\phi(G(X^{\prime}))\big{\|}\Big{]}$
$\displaystyle\leq\max_{x\in\mathcal{X}}\max_{g\in[G_{\text{min}},G_{\text{max}}]}\max_{r\in\mathcal{R}}\|\phi(r+\gamma
g)-B_{r}\phi(g)\|\,,$

as required. ∎

See 4.2

###### Proof.

We follow the illustration laid out in Figure 3:

$\displaystyle\|\mathcal{T}^{\pi}_{\phi}U-\Phi\mathcal{T}^{\pi}\bar{\eta}\|_{\infty}$
$\displaystyle\overset{(a)}{\leq}\|\mathcal{T}^{\pi}_{\phi}U-\Phi\mathcal{T}^{\pi}\eta\|_{\infty}+\|\Phi\mathcal{T}^{\pi}\eta-\Phi\mathcal{T}^{\pi}\bar{\eta}\|_{\infty}$
$\displaystyle\overset{(b)}{\leq}\varepsilon_{\text{B}}+\|\Phi\mathcal{T}^{\pi}\eta-\Phi\mathcal{T}^{\pi}\bar{\eta}\|_{\infty}$
$\displaystyle\overset{(c)}{\leq}\varepsilon_{\text{B}}+d(\mathcal{T}^{\pi}\eta,\mathcal{T}^{\pi}\bar{\eta})+\varepsilon_{\text{E}}$
$\displaystyle\overset{(d)}{\leq}\varepsilon_{\text{B}}+\gamma^{c}d(\eta,\bar{\eta})+\varepsilon_{\text{E}}$
$\displaystyle\overset{(e)}{\leq}\varepsilon_{\text{B}}+\gamma^{c}(\delta+\varepsilon_{\text{R}})+\varepsilon_{\text{E}}\,,$

as required, where (a) follows from the triangle inequality, (b) follows from
the Bellman approximation bound, (c) follows from the embedding error bound,
(d) follows from $\gamma^{c}$-contractivity of $\mathcal{T}^{\pi}$ with
respect to $d$, and (e) follows from the reconstruction error bound. ∎

See 4.4

###### Proof.

We begin by obtaining reconstruction and embedding error bounds for this
sketch. We introduce the shorthand $\Delta=(G_{\text{max}}-G_{\text{min}})/m$.
To obtain a reconstruction error bound, for any distribution
$\nu\in\mathscr{P}([z_{1},z_{m+1}])$, define $\Pi\nu$ to be the distribution
obtained by mapping each point of support $z$ of $\nu$ to the greatest $z_{i}$
less than or equal to $z$. Mathematically, if we define
$f(z)=\max\\{z_{i}:z_{i}\leq z\\}$, then $\Pi\nu=f_{\\#}\nu$, i.e. $\Pi\nu$ is
the pushforward of $\nu$ through $f$. We then have
$w_{1}(\nu,\Pi\nu)\leq\Delta$ for all $\nu$ supported on $[z_{1},z_{m}]$,
where $w_{1}$ is the 1-Wasserstein distance, since $f$ transports mass by at
most $\Delta$. Introducing another distribution $\nu^{\prime}$ and the
projection $\Pi\nu^{\prime}$, we note that
$w_{1}(\Pi\nu,\Pi\nu^{\prime})=\|\Phi\nu-\Phi\nu^{\prime}\|$. Combining these
observations with the triangle inequality yields

$\displaystyle w_{1}(\nu,\nu^{\prime})\leq
w_{1}(\nu,\Pi\nu)+\|\Phi\nu-\Phi\nu^{\prime}\|+w_{1}(\nu^{\prime},\Pi\nu^{\prime})\leq\|\Phi\nu-\Phi\nu^{\prime}\|+2\Delta\,,$

which gives the required form of reconstruction bound, with
$\varepsilon_{\text{R}}=2\Delta$, for the supremum-Wasserstein distance
$\overline{w}_{1}(\eta,\eta^{\prime})=\max_{x\in\mathcal{X}}w_{1}(\eta(x),\eta^{\prime}(x))$
defined over RDFs
$\eta,\eta^{\prime}\in\mathscr{P}(\mathbb{R})^{\mathcal{X}}$. We can also
essentially reverse the argument to get

$\displaystyle\|\Phi\nu-\Phi\nu^{\prime}\|=w_{1}(\Pi\nu,\Pi\nu^{\prime})\leq
w_{1}(\Pi\nu,\nu)+w_{1}(\nu,\nu^{\prime})+w_{1}(\nu^{\prime},\Pi\nu^{\prime})\leq
w_{1}(\nu,\nu^{\prime})+2\Delta$

which gives the required form of the embedding error bound, with
$\varepsilon_{\text{E}}=2\Delta$.

Additionally, we can analyse the worst-case regression error $\|\phi(r+\gamma
g)-B_{r}\phi(g)\|$ to get a bound on the Bellman approximation
$\varepsilon_{\text{B}}$, by Proposition 4.1. Observe that $\phi(g)$ is
constant for $g\in[z_{i},z_{i+1})$, and equal to

$\displaystyle(\underbrace{1,\ldots,1}_{i\text{ times}},0,\ldots,0)^{\top}\,.$

The minimum regression error in

$\displaystyle\mathbb{E}_{G\sim\text{Unif}([z_{1},z_{m}]}[\|\phi(r+\gamma
G)-B_{r}\phi(G)\|]$ (10)

is therefore obtained by setting the $i$th column of $B_{r}$ so that

$\displaystyle
B_{r}\phi(z_{i})=\mathbb{E}_{G\sim\text{Unif}([z_{i},z_{i+1}))}[\phi(r+\gamma
G)]\,;$

note the support of the distribution in the line above. Since $r+\gamma G$ in
this expectation varies over an interval of width $\gamma\Delta$, the
integrand $\phi(r+\gamma G)$ takes on at most two distinct values. It then
follows that we can bound the minimum regression error in Equation (10) by
$\Delta$, and hence we can take $\varepsilon_{\text{B}}=\Delta$.

Finally, we observe that $\mathcal{T}^{\pi}$ maps
$\mathscr{P}([G_{\text{min}},G_{\text{max}}])$ to itself, since for any
$g\in[G_{\text{min}},G_{\text{max}}]$ and any $r\in\mathcal{R}$, we have by
construction of $G_{\text{min}}$, $G_{\text{max}}$ that $r+\gamma
g\in[G_{\text{min}},G_{\text{max}}]$. In addition, we have
$\\{\Phi\nu:\nu\in\mathscr{P}([G_{\text{min}},G_{\text{max}}]_{\\}}=\\{u\in\mathbb{R}^{m}:0\leq
u_{1}\leq\cdots\leq u_{m-1}\leq u_{m}=1\\}$, and by the inspection of the
columns of $B_{r}$ above, it follows that $\mathcal{T}^{\pi}_{\phi}$ maps
$\\{\Phi\nu:\nu\in\mathscr{P}([G_{\text{min}},G_{\text{max}}]_{\\}}^{\mathcal{X}}$
to itself. Therefore the conclusion of Proposition 4.3 holds, and we obtain

$\displaystyle\limsup_{k\rightarrow\infty}\|U_{k}-U^{\pi}\|_{\infty}$
$\displaystyle\leq\frac{1}{1-\gamma}(\gamma\varepsilon_{\text{R}}+\varepsilon_{\text{B}}+\varepsilon_{\text{E}})$
$\displaystyle\leq\frac{1}{1-\gamma}(\gamma 2\Delta+\Delta+2\Delta)$
$\displaystyle=\frac{\Delta(3+2\gamma)}{1-\gamma}$
$\displaystyle=\frac{(G_{\text{max}}-G_{\text{min}})(3+2\gamma)}{(1-\gamma)m}$

as required. ∎

##  Appendix B Further details and extensions

In this section, we collect further details on a number of topics raised in
the main paper.

###  B.1 Categorical imputation

In the tabular experiments in Sections 3.2 and 5, we include comparisons of
distributions imputed from the learned mean embeddings, to provide an
interpretable comparison between the different Sketch-DP methods studied.
Here, we provide a detailed description of the imputation method.

For a given feature map $\phi$, and a learned sketch value $u$, the goal is to
define an imputation strategy
$\iota:\mathbb{R}^{m}\rightarrow\mathscr{P}(\mathbb{R})$ (Rowland et al.,
2019; Bellemare et al., 2023); that is, a function with the property
$\mathbb{E}_{Z\sim\iota(u)}[\phi(Z)]\approx u$, so that $\iota$ serves as an
approximate pseudo-inverse to the mean embedding. Here, we follow the approach
of Song et al. (2008), and impute probability distributions supported on a
finite support set $\\{z_{1},\ldots,z_{n}\\}$. We define $\iota(s)$ implicitly
through the following (convex) quadratic program

$\displaystyle\operatorname*{arg\,min}_{p\in\Delta_{n}}\Big{\|}\sum_{i=1}^{n}p_{i}\phi(z_{i})-s\Big{\|}_{2}^{2}\,.$

Note that the left-hand term inside is the expectation of $\phi(Z)$ with
$Z\sim\sum_{i=1}^{n}p_{i}\delta_{z_{i}}$, and so the objective is simply
aiming to minimise the squared error between the learned sketch value and the
sketch value from this discrete distribution. Since this quadratic program is
convex, it is solvable efficiently; in our implementations, we use SciPy’s
minimize algorithm (Virtanen et al., 2020).

###  B.2 Choices of feature function

In the main paper, we note that any set of features spanning the degree-$m$
polynomials is Bellman closed, as described by Rowland et al. (2019), and
hence exact dynamic programming is possible with this feature set, as shown by
Sobel (1982). However, in preliminary experiments we found these features
difficult to learn with temporal-difference methods beyond small values of
$m$, due to the widely varying magnitude of moments as $m$ grows, making
learning rate selection problematic in stochastic environments; further
details are provided in Appendix D.3.

In this paper, we tested the following functions as the base feature $\kappa$
in the translation family Equation (8):

  * •

Sigmoid: $\kappa(x)=\frac{1}{1+\exp(-x)}$;

  * •

Gaussian: $\kappa(x)=\exp(-x^{2}/2)$;

  * •

Hyperbolic tangent: $\kappa(x)=\tanh(x)$

  * •

Morbid: $\kappa(x)=\exp(-x^{2}/2)\cos(x)$

  * •

Triangular: $\kappa(x)=1-|x|$ for $|x|\leq 1$, zero otherwise.

  * •

Parabolic: $\kappa(x)=1-x^{2}$ for $|x|\leq 1$, zero otherwise.

  * •

Quartic: $\kappa(x)=(1-x^{2})^{2}$ for $|x|\leq 1$, zero otherwise.

In Appendix D.1, we provide further experimental results for a wide variety of
feature maps from this family. In addition, we also consider the indicator
feature used in Proposition 4.4.

We found that the anchor points must be chosen carefully so that the features
produce variations within the return range. This can be done by choosing the
range of the anchors to be slightly wider than and estimated return range, and
setting the slope so that there does not exist a region of the return range
that produce no change in the feature functions. In the tabular experiments,
we set the extremum anchor points to be $\hat{G}_{\text{min}}-a\hat{L}$ and
$\hat{G}_{\text{max}}+a\hat{L}$, where
$\hat{L}=\hat{G}_{\text{min}}-\hat{G}_{\text{min}}$ is the estimated return
range, and $a$ is a small positive value around 0.4 casually chosen and not
tuned. The return limits $\hat{G}_{\text{min}}$ and $\hat{G}_{\text{max}}$ are
estimated by the sample minimum and maximum from samples collected by first-
visit Monte Carlo; see Section C.1.

The slope parameter should depend on the feature and the return range, and we
applied the following intuition. Most of the base feature $\kappa$ have an
“non-trivial support” that produces the most variations in the function value.
For example, the “non-trivial” support for the sigmoidal and Gaussian $\kappa$
can be chosen as $[-2,2]$; and for base features that are nonzero only in
$[-1,1]$, this support is $[-1,1]$. We define the width $w$ of a base function
as the length of the non-trivial support. Crudely, the feature with slope $s$
has width $w/s$, as sharper features tend to have shorter non-trivial support.
In addition, for the set of features to cover return range uniformly, we set
each adjacent feature functions to overlap by 50%. Finally, we want the union
of the non-trivial supports of 10 (arbitrary chosen) such overlapping features
to equal the return range. We must then have $0.5\times
10w/s=\hat{G}_{\text{min}}-\hat{G}_{\text{max}}$. This is the _default slope_
for each feature and each environment with known return range. As such, the
sigmoidal and Gaussian base features have default slope equal to
$s=20/(\hat{G}_{\text{min}}-\hat{G}_{\text{max}})$.

###  B.3 Choices of regression distribution $\mu$

In the main paper, we note that the one-dimensional integrals defining the
matrices $C$ and $C_{r}$ which in turn define the Bellman coefficients $B_{r}$
can be computed in a variety of ways, depending on the choice of $\mu$ and
feature map $\phi$. In our experiments, we take $\nu$ to be a finitely-
supported grid in the range
$[\hat{G}_{\text{min}}-b\hat{L},\hat{G}_{\text{max}}+b\hat{L}]$, where $b$ is
casually chosen to be around 0.2. The support of $\nu$ is thus slightly wider
than the estimated return range and slightly narrower than the anchor range
described in Section B.2. The intuition for using a wider anchor range is that
we need the features to cover the return distribution (and $\nu$) with the
non-trivial support of the features. We validate this intuition in an
additional experiment in Section D.1. With this $\nu$, Equation (5) is a
standard regression problem, and $C$ and $C_{r}$ can be computed with standard
linear-algebraic operations.

Another possibility, particularly if one wishes to use $\mu$ which is not
finitely supported, is to use numerical integration to compute these
integrals. Additionally, in certain settings the integrals may be computed
analytically. For example, with Gaussian
$\phi_{i}(x)=\exp(-s^{2}(x-z_{i})^{2}/2)$ and Gaussian $\mu$, or $\mu$ as
Lebesgue measure (in which case, technically, we modify the expectation in
Equation (5) into an integral against an unnormalised measure), $C$ and
$C_{r}$ can be computed analytically. In the case of $\mu$ as Lebesgue
measure, we have

$\displaystyle
C_{ij}~{}=~{}\sqrt{\frac{\pi}{2s}}~{}\exp\left(-\frac{s}{2}(z_{i}-z_{j})^{2}\right)~{}\text{and}~{}(C_{r})_{ij}~{}=~{}\sqrt{\frac{\pi}{s(1+\gamma^{2})}}~{}\exp\left(-\frac{s(r+\gamma
z_{i}-\gamma z_{j})^{2})}{1+\gamma^{2}}\right)\,.$

###  B.4 Knowledge of rewards

In distributional approaches to dynamic programming, it is necessary to know
all aspects of the environment’s transition structure and reward structure in
advance, including the set $\mathcal{R}$ required for precomputing the Bellman
coefficients. However, in temporal-difference learning, this is a non-trivial
assumption. In many environments, this information is available in advance (in
the Atari suite with standard reward clipping post-processing (Mnih et al.,
2015), rewards are known to lie in $\\{-1,0,1\\}$, for example). When this
information is not available, one may modify Algorithm 1 to instead compute
Bellman coefficients for observed rewards _just-in-time_ ; that is, when these
rewards are encountered in a transition. This makes the algorithm more broadly
applicable, but clearly incurs a significant cost of computing Bellman
coefficients for rewards for which these coefficients are not already cached.
As Remark 3.2, one possibility in this setting is to learn an approximator
$H:\mathbb{R}\rightarrow\mathbb{R}^{m\times m}$ that maps from rewards to
Bellman coefficients, and use the predictions of the approximator as proxies
for the true Bellman coefficients to reduce the need to solve for the Bellman
coefficients every time a new reward is encountered.

###  B.5 Mathematical properties of Bellman coefficients

The Bellman coefficients $B_{r}$ play a crucial role in our Bellman sketch
framework. Here, we present various properties of $B_{r}$ in a worked example,
derived from both a sigmoid and a Gaussian base feature in Figure 6. For each
base feature, we choose 20 evenly spaced anchors in $[-8,8]$, and find $B_{r}$
for $r=1$ and $\gamma=0.8$, and $\mu$ uniformly supported on a dense grid of
10,000 evenly spaced points in $[-5,5]$. We apply a small $L^{2}$ regulariser
with weight $10^{-6}$ in the regression problem.

Figure 6: In-depth analysis of Bellman coefficients in the setting described
in Section B.5. In the third column, solid curves are the most significant
input/right singular vectors, and dashed lines with matching colours are the
corresponding output/left singular vectors.

First, we assess how accurate approximation in Equation (6) is when $B_{r}$ is
found via the regression problem in Equation (5). In the left two columns of
Figure 6, we show the feature functions $\phi(z)$ and $\phi(r+\gamma z)$ in
the first two columns. In the second column, we also show $B_{r}\phi(z)$ in
dashed lines evaluated on $[-8,8]$, wider than the grid over which we
minimised the error. The error is tiny and virtually invisible within the
interval $[-5,5]$, but is larger outside. Quantitatively, the maximum absolute
difference between $\phi(r+\gamma z)$ and $B_{r}\phi(z)$ over the dense grid
in $[-5,5]$ is less than 0.002 for both base features. By Proposition 4.1, we
expect a small error in a single step of dynamical programming.

Had $B_{r}$ been a contraction, we would have been able to prove contraction
for Algorithm 1. However, we show empirically that $B_{r}$ is not in general a
contraction in $L^{2}$ norm, but the dynamics from repeated multiplication of
$B_{r}$ may converge to a stable fixed point. First, we performed a singular
value decomposition of the $B_{r}$ for the two base features. We see that the
singular vectors in Figure 6 (third column) are similar to harmonic functions.
Importantly, in the largest singular values (operator norms) in Figure 6
(fourth column) are greater than 1, suggesting that a single application of
$B_{r}$ may expand the input. Further, we show the eigenvalues of $B_{r}$ in
the fifth column of Figure 6. Interestingly, all eigenvalues have real parts
less or very close to 1.0 suggesting that there exists fixed points in the
dynamics induced by $B_{r}$. As such, the Bellman coefficients $B_{r}$ exhibit
transient dynamics (typical for non-normal matrices) but is stable after
repeated applications to an initial vector. Further studies into these
dynamical properties are important for future work. Given the important role
of non-normal dynamics hypothesised to be present in the nervous system
(Hennequin et al., 2012; Bondanelli & Ostojic, 2020), these observations
allude to the possibility that the Bellman sketch framework could contribute
to a biological implementation of distributional RL.

###  B.6 Computational properties of Bellman coefficients

Under many choices of feature maps $\phi$, the matrix $B_{r}$ has structure
that may be exploited computationally. We provide sketches of several cases of
interest. For “binning features”, even for overlapping bins, $B_{r}$ is a very
narrow band matrix, and hence is sparse, leading to linear-time matrix-vector
product computation. This remains approximately true for other forms of
localised features, such as low-bandwidth Gaussians and related bump-like
functions, and in particular applying truncation to near-zero coefficient in
the Bellman coefficients in such cases will also lead to sparse matrices.

Bellman coefficients as least-squares coefficients. The closed-form solution
for the Bellman coefficients $B_{r}$ in Equation (7) can be derived by viewing
the optimisation problem in Equation (5) as a vector-valued linear regression
problem, and using the usual expression for the optimal prediction
coefficients. The derivation is the same in content to the usual derivation of
least-squares coefficients, which we provide below for completeness, to
illustrate how it is obtained in our case. We begin by differentiating the
(quadratic) objective in Equation (7) with respect to $B$, and setting the
resulting expression equal to the zero vector, to obtain

$\displaystyle-2\mathbb{E}_{G\sim\mu}[\phi(r+\gamma
G)\phi(G)^{\top}]+2\mathbb{E}_{G\ \sim\mu}[B_{r}\phi(G)\phi(G)^{\top}]=0\,.$

Rearranging, we obtain

$\displaystyle B_{r}\mathbb{E}_{G\
\sim\mu}[\phi(G)\phi(G)^{\top}]=\mathbb{E}_{G\sim\mu}[\phi(r+\gamma
G)\phi(G)^{\top}]\,.$

Finally, under the assumption of invertibility of $\mathbb{E}_{G\
\sim\mu}[\phi(G)\phi(G)^{\top}]$, we obtain the expression for the Bellman
coefficients in Equation (7):

$\displaystyle B_{r}=\mathbb{E}_{G\sim\mu}[\phi(r+\gamma
G)\phi(G)^{\top}]\mathbb{E}_{G\ \sim\mu}[\phi(G)\phi(G)^{\top}]^{-1}\,.$

Online computation of Bellman coefficients in the case of unknown rewards. In
settings where the set of possible rewards $\mathcal{R}$ is not known in
advance, is infinite, or is too large to cache Bellman coefficients for all
possible rewards $r\in\mathcal{R}$, we may exploit the structure of the
Bellman coefficients described above to speed up the computation of the
coefficients online. Rather than solving the regression problem from scratch,
an alternative is to cache the matrix $\mathbb{E}_{G\
\sim\mu}[\phi(G)\phi(G)^{\top}]^{-1}$ above, and construct the matrix
$\mathbb{E}_{G\sim\mu}[\phi(r+\gamma G)\phi(G)^{\top}]$ as required, upon
observing a new reward $r$. This reduces the marginal cost of computing the
Bellman coefficients $B_{r}$ to a matrix-matrix product.

###  B.7 Comparison with other approaches to distributional RL

In this section, we provide additional comparisons against existing approaches
to distributional RL, including their theoretical analysis. As distributional
RL is a quickly evolving field, we focus our comparison on a few main classes
of algorithms related to our work, which illustrate some key axes of variation
within the field: (i) categorical approaches (Bellemare et al., 2017); (ii)
quantile approaches (Dabney et al., 2018a; b; Yang et al., 2019); (iii)
approaches related to maximum mean discrepancy (MMD; Gretton et al., 2012),
such as Nguyen-Tang et al. (2021); Zhang et al. (2021); Sun et al. (2022); and
(iv) sketch-based approaches (Sobel, 1982; Rowland et al., 2019).

Distribution representation. Categorical, quantile, and MMD approaches are
typically presented as learning approximate return distributions directly. In
categorical approaches, the approximate distribution is parametrised as

$\displaystyle\sum_{i=1}^{m}p_{i}\delta_{z_{i}}\,,$

with fixed particle locations $(z_{i})_{i=1}^{m}$, and learnable probabilities
$(p_{i})_{i=1}^{m}$ for each state-action pair at which the return
distribution is to be approximated. In contrast, quantile and MMD approaches
learn fixed-weight particle approximations, of the form

$\displaystyle\sum_{i=1}^{m}\frac{1}{m}\delta_{z_{i}}\,,$

in which the particle locations $(z_{i})_{i=1}^{m}$ are learnable. Work on
sketches has instead focused on learning the values of particular statistical
functionals of the return, rather than explicitly approximating return
distributions. Rowland et al. (2019) also shows that standard categorical- and
quantile-based algorithms can also be viewed through the lens of sketch-based
distributional RL. The approach proposed in this paper sits firmly in the camp
of sketch-based approaches, without ever representing approximated
distributions directly. We highlight generative models of distributions as
another prominent class of (non-parametric) representation (see, e.g., Doan et
al. 2018; Freirich et al. 2019; Dabney et al. 2018a; Yang et al. 2019; Wu et
al. 2023).

Algorithm types. Most prior algorithmic contributions to distributional
reinforcement learning have focused on sample-based temporal-difference
approaches, in which prediction parameters are iteratively and incrementally
updated based on the gradient of a sampled approximation to a loss function.
These approaches include the original C51 (Bellemare et al., 2017), QR-DQN
(Dabney et al., 2018b), MMDRL (Nguyen-Tang et al., 2021), and EDRL (Rowland et
al., 2019) algorithms. Dynamic programming algorithms, in which parameters are
not updated incrementally via loss gradients, but instead according to the
application of an implementable operator, have also been considered (see
Rowland et al. (2018; 2023); Wu et al. (2023) for categorical dynamic
programming, quantile dynamic programming, and fitted likelihood estimation,
respectively). Marthe et al. (2023) discussed learning of certain risk-
sensitive functionals in finite-horizon MDPs. In this paper, our algorithmic
contributions include both DP and TD methods.

Losses and projections. One of the core axes of variation across
distributional RL approaches is the loss used to define updates in incremental
algorithms, and to define projections in dynamic programming. Categorical
approaches use a projection in Cramér metric (Rowland et al., 2018) to define
a target distribution for both dynamic programming and incremental versions of
the algorithm; the incremental algorithm updates predictions via the gradient
of a KL loss between the current and target distributions. Quantile-based
approaches use either the quantile regression loss in incremental settings, or
a Wasserstein-1 projection in dynamic programming (Dabney et al., 2018b). The
MMDRL approach has been proposed only in the incremental setting, and
considers sample-based gradients through an MMD loss, specifically taking the
form

$\displaystyle\text{MMD}_{K}^{2}\Bigg{(}\sum_{i=1}^{m}\frac{1}{m}\delta_{z_{i}(x,a)},\sum_{i=1}^{m}\frac{1}{m}\delta_{r+\gamma
z_{i}(x^{\prime},a^{\prime})}\Bigg{)}\,,$ (11)

for some choice of kernel $K$. In contrast, the Sketch-DP and Sketch-TD
algorithms introduced in this paper define updates via matrix-vector products
with the Bellman coefficients derived in Equation (5). With the connection to
RKHS noted in Section 6, one can view the squared norm appearing in Equation
(5) as the squared norm in the RKHS generated by the feature functions $\phi$.
From this perspective, it is clear that both the Bellman coefficients that
feature in Sketch-DP and Sketch-TD, and the gradients of the loss in Equation
(11), both relate to the mathematical structure of RKHS.

Contrasting against the Sketch-DP/Sketch-TD approaches described above,
earlier approaches to sketched-based distributional RL, both in dynamic
programming and incremental forms, have defined losses via imputation
strategies, which compute updates by converting sketches into approximate
distributions (Rowland et al., 2019; Bellemare et al., 2023). Foreshadowing
the remarks on theoretical analysis below, we remark that neither MMDRL nor
the earlier sketch-based approach described above have been analysed for
convergence, while Section 4 in this paper deals with convergence analysis of
the approach proposed in this paper.

Theoretical analysis. The analysis of both dynamic programming and sample-
based incremental algorithms requires analysis of the interaction between
distributional updates, and approximations made to the distributions due to
the choice of representation mentioned above.

Convergence analysis of dynamic programming algorithms has been obtained for
several classes of distributional algorithms; see Rowland et al. (2018) for
the case of categorical dynamic programming, Dabney et al. (2018b) for the
case of a quantile dynamic programming algorithm, and Bellemare et al. (2023);
Rowland et al. (2023) for later generalisations of this work. This analysis
typically centres around (i) proving contractivity of the distributional
Bellman operator $\mathcal{T}^{\pi}$ with respect to some metric d, and
proving non-contractivity of the specific distributional projection used by
the dynamic programming algorithm under this same metric. Notably, the metric
d used in the _analysis_ need not be the same as any metrics used in
_defining_ the algorithm; this is the case for quantile dynamic programming,
for which Wasserstein-1 distance is used to define the algorithm, while
Wasserstein-$\infty$ distance is used to analyse the algorithm (Dabney et al.,
2018b; Bellemare et al., 2023; Rowland et al., 2023). Our proof technique, in
particular, makes use of contraction of the distributional Bellman operator in
Wasserstein distances, though such distances do not feature in the definition
of the Sketch-DP/TD algorithms.

In general, there has been less work on the convergence analysis of sample-
based incremental algorithms. Rowland et al. (2023) recently showed
convergence of quantile temporal-difference learning, though the question of
convergence for many other sample-based incremental distributional
reinforcement learning algorithms is currently open. In particular, we note as
a point of comparison to Nguyen-Tang et al. (2021), that although this work
provides a contraction analysis of $\mathcal{T}^{\pi}$ (Theorem 2), and MMD
approximation bounds for fixed target distributions (Theorem 3 and Proposition
2), these do not constitute a proof of convergence of the incremental
algorithm described therein. The analysis of incremental algorithms is
generally more mathematically involved than in the dynamic programming case,
principally owing to the fact that rather than analysing the iterated
application of a fixed operator, one needs to analyse the continuous dynamical
system associated with incremental updates.

##  Appendix C Experimental details

In this section, we provide additional details on the experimental results
reported in the main paper.

###  C.1 Tabular environments

We describe the setup in the main paper. In Section D.1, we show extended
results of more features and more environments.

Environments.  In the main paper, we reported results on the on the following
environments,

  * •

Random chain: Ten states $\\{x_{1},x_{2},\ldots,x_{10}\\}$ are arranged in a
chain. There is equal probability of transitioning to either neighbour at each
state, and state $x_{10}$ has a deterministic reward of +1;

$\mathrm{terminal}\leftarrow x_{1}\longleftrightarrow x_{2}\longleftrightarrow
x_{3}\longleftrightarrow\cdots\longleftrightarrow
x_{10}\rightarrow\mathrm{terminal}\,.$

  * •

Directed chain (DC): Five states are arranged in a directed chain, but the
agent can only move along the arrow deterministically until termination. A
deterministic reward of $+1$ is given at state $x_{5}$;

$x_{1}\longrightarrow x_{2}\longrightarrow
x_{3}\longrightarrow\cdots\longrightarrow
x_{5}\rightarrow\mathrm{terminal}\,.$

  * •

DC with Gaussian reward: A variant of the directed chain above, with the only
difference that $x_{5}$ has a Gaussian reward with mean 1 and unit variance.

The discount factor is $\gamma=0.9$. These environments cover stochastic and
deterministic rewards and state transitions, giving a range of different types
of return distributions.

Feature functions. We use features of translation family Equation (8), with
$\kappa$ chosen from a subset of base features described in equation (B.2). We
also include the indicator features used in Proposition 4.4. For the sweep
over slope, we set the slope $s$ to be the default slope (described in Section
B.2) multiplied by a scaling factor, and sweep over this factor. from 0.001 to
10.0. This is done primary because the return range varies a lot across
different environments, and the default slope is adjusted to the return range.
The results serve as justification for the heuristics on choosing the default
slope.

Ground-truth distribution. We approximate the ground-truth mean embeddings and
the ground-truth return distributions by collecting a large number of return
samples from the MRPs. To do so, we use first-visit Monte Carlo with a
sufficiently long horizon (after the first visit to each state) to ensure that
the samples are unbiased and has bounded error caused by truncating the
rollout to a finite horizon. For environments with deterministic rewards,
truncating the horizon at $L$ steps induces maximum truncation error
$|r|_{\mathrm{max}}\gamma^{L}/(1-\gamma)$, where $|r|_{\mathrm{max}}$ is the
maximum reward magnitude. We bound this error at $10^{-4}$, giving $L>110$, so
we set the horizon after the first visit to 110. For environments with
Gaussian rewards, we set the horizon to 200. We initialise the rollout at each
state in the environment, and for initial each state this is repeated $10^{5}$
times. This gives us at least $10^{5}$ samples each state.

Sketch DP under conditional independence. Many RL environments, including the
tabular environments tested in this paper, have the property that
$R\perp\\!\\!\\!\\!\perp X^{\prime}|X$ for the trajectory $X,A,R,X^{\prime}$,
so the Sketch-DP update Equation (Sketch-DP) simplifies to

$U(x)\leftarrow\mathbb{E}_{x}^{\pi}[B_{R}]\mathbb{E}_{x}^{\pi}[U(X^{\prime})]=\mathbb{E}_{x}^{\pi}[B_{R}]\sum_{x^{\prime}\in\mathcal{X}}P(x^{\prime}|x)U(x^{\prime}).$

This means we need to evaluate the expected Bellman coefficient
$\mathbb{E}_{x}^{\pi}[B_{R}]$. This is trivial for deterministic rewards. For
stochastic rewards with known distributions, we approximate the expectation
via numerical integration. We run all DP methods for 200 iterations.

Jittered imputation support. The support on which we impute the distribution
are the anchors of the features. Some tabular environments have states with
deterministic returns that directly align with the feature anchors, which
interferes in unintuitive ways with the finite support on which we impute
distributions, producing non-monotonic trends in the results. To avoid this
unnecessary complication, we jitter the support before imputing the
distribution: for points in the support, we add noise uniformly distributed
over $[-\Delta/2,\Delta/2]$, where $\Delta$ is the distance between
consecutive support points. Likewise, we project ground truth distribution
using the same jittered support to. In Figure 4, we report the average of the
metrics computed from 100 independent jitters. Note that since the imputation
(from mean embedding) and projection (from ground-truth) share the same
support for each of the 100 jitters, the average Cramér distance between the
projected and the ground-truth still lower-bounds the average Cramér distance
between the imputed distribution and the ground-truth.

###  C.2 Deep reinforcement learning implementation details

In this section, we provide further details on the deep reinforcement learning
experiments described in the main paper, in particular describing
hyperparameters and relevant sweeps.

Environment. We used the exact same Atari suite environment for benchmarking
QR-DQN (Bellemare et al., 2013; Dabney et al., 2018b). In all experiments, we
run three random seeds per environment.

Feature map $\phi$. The results in Figure 5 uses the sigmoid base feature
$\kappa(x)=1/(e^{-x}+1)$ with slope $s=5$, and the anchors to be 401 (tuned
from 101, 201 and 401) evenly spaced points between $-12$ and $12$. These
values are loosely motivated by the range used in C51 (Bellemare et al.,
2017). We did not use the heuristic in Section B.2 to choose the slope
parameter, instead we set this to 10 by tuning from $\\{1,2,\ldots,12\\}$.
Larger slope values typically resulted in Bellman coefficients with a worst-
case regression error
$\max_{r\in\mathcal{R}}\max_{g\in\text{supp}(\mu)}\|\phi(r+\gamma
g)-B_{r}\phi(g)\|$ greater than 0.01. In these cases, we regard the regression
error as too large, and did not perform agent training with these
hyperparameter settings. In addition, $\phi$ is appended with a constant
feature of ones, which we found to be very crucial for a good performance; as
noted in the main paper, this ensures that the sketch operator is truly
affine, not linear. We also tried a several other feature functions, including
the Gaussian, the hyperbolic tangent and the (Gaussian) error function, and
found that the sigmoid reliably performed the best.

Solving the regression problems. To compute the Bellman coefficients as well
as the value readout weights $\beta$ described in Section 5.1, we solve the
corresponding regression problem with $\mu$ set to be 100,000 points evenly
spaced between $-10$ and $10$. We also add a $L^{2}$ regulariser with strength
set to $10^{-9}$ to avoid numerical issues, which is tuned from
$\\{10^{-15},10^{-12},10^{-9},10^{-6},10^{-3}\\}$.

Neural network. We implement Sketch-DQN based on the QR-DQN architecture,
using exactly the same convolutional torso and fully-connected layers to
estimate the mean embeddings. The differences are:

  * •

We add a sigmoid or tanh nonlinearity, depending on the base feature output
range to the final layer. This helps bound the predicted mean embedding and
improved the results.

  * •

The network only predicts the non-constant dimensions of the mean embedding,
and the constant feature is appended as a hard-coded value.

  * •

We use the pre-computed mean readout coefficients $\beta$ to predict state-
action values for state-action pairs in the Q-learning objective, and at
current states to determine the greedy policy.

Training. We use the exact same training procedure as QR-DQN (Dabney et al.,
2018b). Notably, the learning rate, exploration schedule, buffer design are
all the same. We tried a small hyperparameter sweep on the learning rate, and
found the default learning rate $0.00005$ to be optimal for performance taken
at 200 million frames.

Evaluation. The returns are normalised against random and human performance,
as reported by Mnih et al. (2015). We use the mean and median over all games
and three random seeds for each game.

Baseline methods. We also tuned the number of atoms of the approximating
distributions in the baseline methods. In particular, we found that C-51
performed the best compared to using more atoms; IQN did best when using 51
quantiles; and QR-DQN did best using 201 quantiles. Increasing the number of
atoms in these methods lead to worse performance. We report the results of
these best variants of the corresponding baseline methods in Figure 5.

##  Appendix D Further experiments

In this section, we collect further experimental results to complement those
reported in the main paper.

###  D.1 Extended tabular results

We tested Sketch-DP using the following additional MRPs. If unspecified, the
default reward distribution for each state is a Dirac delta at 0, and the
transition probabilities from a state to its child states are equal.

  * •

Tree: State 1 transitions to states 2 and 3; state 3 transitions to states 4
and 5. State 2 has mean reward 5; state 4 has mean reward -10, and state 5 has
mean reward 10. All leaf states are terminal.

  * •

Loopy tree: Same as Tree, but with a connection from state 2 back to state 1.

  * •

Cycle: Five states arranged into a cycle, with only a single state having mean
reward 1.

  * •

Rowland ’23: The environment in Example 6.5 of Rowland et al. (2023).

  * •

S&B ’18: The environment in Example 6.4 of Sutton & Barto (2018).

  * •

Loopy fork: The environment shown in Figure 2(A).

All environments have discount factor $\gamma=0.9$. For each environment,
except Rowland ’23 and S&B’18, the reward distributions for the non-zero-
reward states are either Dirac deltas at the specified mean, or Gaussian with
specified mean and unit standard deviation.

The results, extending those in Figure 4, are illustrated in Figure 7. The
results are in general consistent with the main Figure 4. In particular, the
Cramér distances decay as the number of features increases, and can be closer
to the corresponding projected ground-truths than the CDRL baseline.

In Section B.2, we suggested that the range of the anchors should be wider
than the range of the uniform grid $\mu$ on which we measure the regression
loss. We validate this intuition by performing another experiment, sweeping
the ratio of the width of the anchor range relative to the width of the grid
$\mu$, fixing the mid points between these two ranges the same. Here, we use
$m=50$ features for each base feature function and apply the default slope
described in Section B.2. The results in Figure 8 shows that a slightly wider
anchor range produces reliably small Cramér distances for almost all base
features. When the anchor range decreases from 1 to 0, there is a much sharper
increase in the Cramér distance, because the support on which we impute the
distribution is too narrow and can miss substantial probability mass outside
the support. On the other hand, when the anchor range increases from 1, the
Cramér distance also increases because the support points get further away
from each other, lowering the resolution of the imputed distribution. Since
the grid is chosen to be slightly wider than the true return range, we see
that the smallest Cramér distances can be attained at anchor range ratio
slightly less than 1, but this does not hold universally (see, e.g., random
chain and cycle results). Choosing this ratio to be slightly greater than 1,
as suggested in Sections B.2 and B.3, is more reliable at the cost of a small
increase in distributional mismatch.

(a) Environments with deterministic rewards, sweeping over feature count $m$.

(b) Environments with deterministic rewards, sweeping over slope $s$.

(c) Environments with Gaussian rewards, sweeping over feature count $m$.

(d) Environments with Gaussian rewards, sweeping over slope $s$.

Figure 7: Additional tabular results extending Figure 4.

(a) Environments with deterministic rewards.

(b) Environments with Gaussian rewards.

Figure 8: Additional tabular results from sweeping over the range of the
anchors relative to the width of the uniform grid given by the support of
$\mu$.

###  D.2 Extended results on Atari suite

We show in Figure 9 the advantage of the Sketch-DQN method to other baselines.
On one hand, we see that Sketch-DQN surpasses DQN on almost all games.
Compared to IQN and QR-DQN, Sketch-DQN is consistently better on crazy
climber, space invaders, river raid, road runners and video pinball, while
worse on assault, asterix, double dunk, krull, phoenix, and star gunner.

Figure 9: Advantage of Sketch-DQN, measured as Sketch-DQN’s normalised return
minus the returns of other baseline methods. Positive values means Sketch-DQN
is better. We show the mean and standard error over 3 seeds for each game.

To show how sensitive are the results depending on the feature parameters, we
ran the full Atari suite using a few sigmoidal and Gaussian features, and show
the results in Figure 10. Overall, we see that the choice on feature
parameters is important; in particular, the sigmoidal feature outperforms
Gaussian features. For the sigmoidal feature, the performance improved from
using 101 to 201 features. On the contrary, for Gaussian features, increasing
the feature count does not produce much change.

(a) Mean normalised return.

(b) Median normalised return.

Figure 10: Results on Atari suite for different feature parameters.

###  D.3 Temporal-difference learning with polynomial features

As noted in the main text, the sketch corresponding to the polynomial feature
function $\phi(g)=(1,g,g^{2},\ldots g^{m})$ is Bellman closed, and the Bellman
coefficients $B_{r}$ obtain zero regression error in Equation (5). In
addition, there has been much prior work on dynamic programming (Sobel, 1982)
for moments of the return, and temporal-difference learning specifically in
the case of the first two moments (Tamar et al., 2013; 2016). However, such
polynomial feature functions are difficult to use as the basis of learning
high-dimensional feature embeddings. This stems from several factors,
including that the typical scales of the coordinates of the feature function
often vary over many orders of magnitude, making tuning of learning rates
difficult, as well as the fact that polynomial features are _non-local_ ,
making it more difficult to decode distributional information via an
imputation strategy.

To quantify these informal ideas, we ran an experiment comparing Sketch-TD
updates for a 50-dimensional mean embedding based on the translation family
(Equation (8)) with a sigmoid base feature $\kappa$, as well as polynomial
features with $m=5$, and $m=50$. The sigmoid features are chosen according to
the intuitions in Section B.2. We ran 100,000 synchronous TD updates on mean
embedding estimates initialised at $\phi(0)$ for all state. Each run uses a
fixed learning rate chosen from $10^{-6}$ to $1$.

In Figure 11, we plot the Cramér distance of imputed distributions from
ground-truth after running TD for each of these three methods, on a variety of
the environments described in Section D.1. In all environments, there is a
similar pattern. For the Sketch-TD algorithm based on sigmoid non-linearities,
there is a reasonably wide basin of good learning rates, with performance
degrading as the learning rate becomes too small or too large. On several
environments this pattern is reflected also in the performance of the degree-5
polynomial embedding, though the minimal Cramér error is generally
significantly worse than that of the sigmoid embedding. This supports our
earlier observations; this mean embedding captures relatively coarse
information about the return distribution, and in addition different feature
components have different magnitudes, meaning that a constant learning rate
cannot perform well. The mean embedding with degree-50 polynomials generally
performs very badly on all environments, as the components of the embedding
are at such different magnitudes that no appropriate learning rate exists.

Figure 11: Results comparing sketch-TD with sigmoidal and polynomial features.

###  D.4 Comparison with statistical functional dynamic programming

The Sketch-DP methods developed in this paper were motivated in Section 2 with
the aim of having distributional dynamic programming algorithms that operate
directly on sketch values, without the need for the computationally intensive
imputation strategies associated with SFDP algorithms. In this section, we
empirically compare Sketch-DP and SFDP methods, to quantitatively measure the
extent to which this has been achieved. We give details below of the Sketch-DP
and SFDP algorithms we compare, and provide comparisons of per-update
wallclock time to assess computational efficiency, and distribution
reconstruction error to assess accuracy.

Sketch-DP. We consider the Sketch-DP algorithm based on sigmoid features as
described in Section 3.1, and implemented as described in Section 5 and
Appendix C.1.

SFDP. We consider the SFDP algorithm for learning expectile values as
described by Bellemare et al. (2023, Section 8.6). We use SciPy’s default
minimize implementation (Virtanen et al., 2020) to solve the imputation
strategy optimisation problem given in Equation (8.15) in Bellemare et al.
(2023). For a given sketch dimension $m$, we use expectiles at linearly spaced
levels $\tau_{i}=(2i-1)/(2m)$ for $i=1,\ldots,m$.

Results. These two algorithms, for varying numbers of feature/expectiles $m$,
were run on a selection of deterministic-reward environments as described in
Section D.1. In Figure 12, we plot the Cramér distance and the excess Cramér
distance of reconstructed distributions to ground truth, as described in
Section 5 and plotted in Figure 4. In addition, we also plot two wallclock
times in each case: the average time it takes to run one iteration of the
dynamic programming procedure, and the time it takes to setup the Bellman
operator, which includes solving for $B_{r}$ for Sketch-DP. As predicted, the
run time is significantly higher for the SFDP algorithm, due to its use of
imputation strategies. The approximation errors measured by Craḿer distances
are also smaller for Sketch-DP, particularly as the number of
features/expectiles is increased. Considering the per-update wallclock times
in the third row of the figure, there is consistently a speed up of at least
100x associated with the Sketch-DP algorithm relative to SFDP. This is due to
the fact that the Sketch-DP update consists of simple linear-algebraic
operations, while the SFDP update includes calls to an imputation strategy,
which must solve an optimisation problem. The one-off computation of the
Bellman coefficients takes around 0.1–4 seconds, depending on the number of
features $m$, which is only at most a couple of SFDP iterations, and hence a
small fraction of the total run time of the SFDP algorithm.

Figure 12: Results comparing Cramér distances as in Figure 4 (first two rows),
and wallclock runtimes for each DP iteration (third row) and for setting the
corresponding Bellman operator (bottom row), for Sketch-DP and SFDP
algorithms, varying the numbers of features/expectiles $m$.