Near-Minimax-Optimal Distributional Reinforcement Learning with a Generative Model
Authors: Mark Rowland
Li Kevin Wenliang
Remi Munos
Clare Lyle
Yunhao Tang
Will Dabney
Venue: NeurIPS 2024 poster (2024)
Keywords: Reinforcement learning
Distributional reinforcement learning
dynamic programming
TD learning
sample complexity
theory
Primary area: reinforcement_learning
TL;DR: Roughly speaking, we show that distributional reinforcement learning is of comparable statistical difficulty to learning a value function, in the generative model setting.

Abstract: We propose a new algorithm for model-based distributional reinforcement learning (RL), and prove that it is minimax-optimal for approximating return distributions in the generative model regime (up to logarithmic factors), the first result of this kind for any distributional RL algorithm. Our analysis also provides new theoretical perspectives on categorical approaches to distributional RL, as well as introducing a new distributional Bellman equation, the stochastic categorical CDF Bellman equation, which we expect to be of independent interest. Finally, we provide an experimental study comparing a variety of model-based distributional RL algorithms, with several key takeaways for practitioners.

Near-Minimax-Optimal Distributional Reinforcement Learning with a Generative Model
Mark Rowland Google DeepMind
Li Kevin Wenliang Google DeepMind
Rémi Munos FAIR, Metal
Clare Lyle Google DeepMind
Yunhao Tang Google DeepMind
Will Dabney Google DeepMind
Abstract
We propose new algorithm for model- based distributional reinforcement learning (RL), and prove that it is minimax-optimat for approximating return distributions in the generative model regime (up to logarithmic factors), the first result of this kind for any distributional RI. algorithm. Our analysis also provides new theoretical perspectives on categorical approaches to distributional RI., as well introducing new distributional Bellman equation, the stochastic categorical CDF Bellman equation, which we expect to be of independent interest. Finally, we provide an experimental study comparing variety of model-based distributional RL algorithms, with several key takeaways for practitioners.
Introduction
In distributional reinforcement learning, the aim is to predict the full probability distribution of possible returns at each state, rather than just the mean return (Morimura et al. 2010a| Bellemare et al., 2017, 2023). Applications of distributional reinforcement learning range from dopamine response modelling neuroscience (Dabney et 2020), to driving risk-sensitive decision-making and exploration in domains such as robotics (Bodnar et al. 2020), healthcare (Böck et al. 2022). and algorithm discovery (Fawzi et al., 2022), as well as forming a core component of many deep
reinforcement learning architectures (Bellemare et al.] [017]] Dabney et al. 2018h a Yang et al.
2019 Bellemare et al. 2020 Shahriari et al. 2022 Wurman et al.) 2022).
The full distribution of returns is a much richer signal than the expectation to predict. A foundational, as-yet-unanswered problem is how many sampled transitions are required accurately estimate return distributions, and in particular, whether this task is statistically harder than estimating just the value function. We study these questions in the setting where sampled transitions are given by generative model (Kearns et al. 2002 Kakade 2003 Azar et al. 2013).
We provide new distributional RI. algorithm, the direct categorical fixed-point algorithm (DCFP), and prove that the number of samples required by this algorithm for accurate return distribution estimation matches the lower bound established by Zhang et al. (2023), up to logarithmic factors. This resolves the foundational question above, and, perhaps surprisingly. shows that in this setting, distributional RL essentially no harder, statistically speaking, than learning a value function.
In addition to this central result, our analysis provides new perspectives on categorical approaches to distributional RL (Bellemare al.) (017), including new distributional Bellman equation, the stochastic categorical CDF Bellman equation (see Section (.2), which we expect to broad in future work categorical distributional RL. We also provide empirical study, comparing
Correspondence markrrovland@goolle.coom
Work done while Google DeepMind.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the newly-prov proposed DCFP algorithm to existing approaches to distributional RL such as quantile dynamic programming (QDP; [abnnyy al., 2018b] Rowland al., 2024). and identify several
key findings for practitioners, including the importance of levels of environment stochasticity and discount factor for the relative performance of these algorithms.
2 Background
Throughout the puper, we consider the problem of evaluation in an infinite-horizon Markov reward process (MRP), with finite state space X' transition probabilities P E RXX* reward function :: + [0. 1]. and discount factor y E [0. 1); this encompasses the problem of policy evaluation in Markov decision processes (Sutton and Barto 2018). A random trajectory (X,, Re)e20 is generated from an initial state Xo I according to the conditional distributions Xx (Xp..... Xt-1) ~ P(-[X,1)), and R r(X,). The return associated with the trajectory is given by the quantity Ei20? 2 R4. In RL, a central task is to estimate the value function V* X - R, defined by
(*()) = B|Er20
Xo=xi x]
given some form of observations from the MRP. The value function defines the expected return. conditional on each possible starting state in the MRP. The value function satisfies the Bellman equation V' = TVV, where T : R - R is defined by
(2) where (a, RR X') is a random transition in the environment, distributed as described above. When the transition probabilities of the MRP are known, the right-hand side can be evaluated as an affine transformation of V. MRP theory (see, c.g., Puterman) 2014 for an overview) then shows how (an approximation to) V can be obtained. For example, a dynamic programming approach takes an initial approximation Vo E [0. (1 - --) 10)11pp 11 and the sequence (V&) is computed via the update V = TV;; it is guaranteed that [|/ V* I S -- 7)_ Alternatively, the linear system V TV can be solved directly with linear- angar-algeeraic methods obtain V* its unique solution.
(IV) = E
+ +V()))),
2.1 Reinforcement learning with generative model
In many settings the transition probabilities of the MRP are unknown, and the value function must be estimated based data comprising sampled transitions, introducing statistical element to the problem. A commonly used model for this data is a generative model (Kearns et al., 2002 Kakade 2003 In this setting, for each state z E X, we observe N i.i.d. samples (X;)) from P(- and this collection of N [ samples may then be used by an algorithm to estimate the value function. Azar et al. (203)) showed that least N l2s 21 -)- -3 log((X()/5) samples are required to obtain g-accurate estimates of the value function with high probability (measured in I norm), and also showed that this bound is attained (up logarithmic factors) by certainty equivalence algorithm, which treats the empirically observed transition frequencies as the true ones, and solves for the value function of the corresponding MRP
2.2 Distributional reinforcement learning
Distributional RI. aims to capture the full probability distribution of the random return each state, not just its mean. Mathematically, the object of interest is the return-distribution function (RDF) 17 : X 5 (R), defined by
(xxx)= - D
X. x),
(3)
where DD extracts the probability distribution of the input random variable. The distributional perspective reinforcement learning has proved practically useful in wide variety of applications, including healthcare (Böck et al. (2022), navigation (Bellemare et al. (202)), and algorithm discovery
(Fawzi et al. (0022). The central equation behind dynamic programming approaches to approximating
the return distribution function is the distributional Bellman equation (Sobell 1982. Morimura
al. 2010a, Bellemare et al.) 2017). given by 17" = Tn"", where T distributional Bellman operator, defined by
is the
(Tn)(z) = D(R 7G(X') x
where independent from the random transition (X, R, X'), we have G(z) 17xx) for each E X.

---

2.3 Categorical dynamic programming
Given an initial RDF approximation 17 ((O,, (1-)))1))-1 it also holds that the update 1 +- Tn converges to '7" in an appropriate sense (e.g., in Wasserstein distance; see Bellemare et al. 2017), in analogy with dynamic programming algorithms for the value function, as described above. However, generally it is not possible to tractably implement repeated computation of the update < Tn as a means of computing approximations to return distributions; probability distributions are infinite- dimensional objects, and as such computational costs quickly become prohibitive. Instead, the use of some kind of approximate, tractable representation of probability distributions is typically required.
Representations. In this paper, we focus on the categorical approach to distributional reinforcement learning (Bellemare al. 2017), in which estimates of return distributions are represented as categorical distributions over fimite number of outcomes 21 < < 2m We will take {z1..... 8m} 2-1
to be an equally spaced grid over the range of possible returns [0, (1 - =) so that 2; m
7) 1 for i = 1,..., Approximations of the RDF are then represented in the form
nTx) - pi(2))..
(4)
Here, &g is the Dirac distribution at the outcome 2, and p ((p, E X) are adjustable probability mass parameters; see Figure Ia). The number of categories m can be interpreted as controlling the expressivity of the representation, and should be carefully chosen in practice to trade off between increased accuracy (larger mi) and computational tractability (smaller m.).
Dynamic programming. The iteration 17 < Tn cannot be used to update the parameters p in Equation (4) directly, since the distributions (Tn) no longer supported {21,.... 2m ). and so cannot be expressed in the form given in Equation 4 see Figure [ellemare et al. (2017) circumvent this issue by projecting the resulting distributions back onto the support set (z1, 3m
via a map Tmm ([0, (1 - -1) + (((2), Zm) })). Intuitively, Ih can be thought of as allocating each outcome E [2,, 1 its neighbouring gridpoints 2; and 2441, in proportion to their proximity, so that the projection of the Dirac distribution 8,, is defined by
2++1 2
2, 82 5--1-
m=
62,
2++1 2g Zi+1
In this paper, we work with the equivalent definition of the projection 11m given by Rowland et al. (2018) Proposition 6), in which the probability mass assigned to z; by Im is given by the expectation Ez~p (h, [Z)]. where h2 [0, (1->) 101---111--- -1 [0, 1] is the "hat function" at 2,, which linearly interpolates between a value of at 2g. and 0 at neighbouring gridpoints 2,-1, 21+1. and is outside this range.
Figure I c) illustrates h4 and hm see Appendix B for restatement of the full definition given by
Rowland et al. (2018).
The projected update 7 < IIm T7 is thus guaranteed to keep in the space of approximations of the form given in Equation 44. and can be viewed as a tractable alternative to the update rj t 'Tn described above. Rowlland et al. (2018) show that despite the introduction of the additional projection map Im,, repeated computation of the update 7) + II T 1j, referred to as categoricul dynamic programminn (CDP), is guaranteed to convergence to a categorical fixed point, and further, the categorical fixed point can be made arbitrarily close to the true RDF 17" by increasing m, as measured by Cramér distance (Cramér 1928 Szčkely 2003 Szzkelyy and Rizzo 2013).
(a)
(b)
(c)
21
2m
1 Zi
2i41
Zm
Figure 1: (a) The density of distribution V' (grey), and its categorical projection Im e ((zz).... 2m.)) (blue). (b) categorical distribution (blue); its update after being scaled by y and shifted by r by the distributional Bellman operator T, moving its support off the grid (=1,..., 2m } (pink): the resulting realigned distribution supported on the grid {2i,... 2m after projection via IT, (green). (c) Hat functions h; (solid) and hm (dashed).
3

---

Definition 2.1. The Cramér distance l2 (P(R) x (R) - R is defined by
(F)(t))Fx((())))))) (2(1,0)= [ (F (t) dt
1/2
where Fu, For are the CDFs of v, respectively. The supremum- Cramér distance &g on (())* is defined by
1/) max (2))) xEX
The central convergence results concerning CDP are summarised below.
Proposition 2.2. (Rowland et al. 2018). The operator II ( (1 - 1) ( (1
-)> a contraction mapping with respect to Ez. with contraction factor VT, and has a unique fixed point, ne E ({21,-... (/(21,-.,,zp) 2m
1)²² As a result, for any no E { (1 -)) with qh+1 =
Im Tijs we have (2(T Tc) S 7) -2gh. Further, the distance between tic and the true RDF ni*
can be bounded as
(5)
This establishes CDP as a principled approach to approximating return distributions, and also quantifies the accuracy achievable with CDP using categories, which will be central in informing our choice of m to obtain a sample-efficient, accurate algorithm below.
3 Distributional reinforcement learning with generative model
The central problem we study in this paper is how to do sample-efficient distributional RL with generative model. That is, given the samples ((x))) : x E X²) described in Section 2.1 how accurate of an approximation to the return disttriittibttion function in Equation (3) can one compute?
This question was raised by Zhang et al. (023). who proposed to perform distributional dynamic programming updates 1) < Tn as described in Section 2.2 using the empirical distributional Bellman operator T derived from the empirical transition probabilities P, defined by P(y (r) = N- -i11 [[x,- = y). producing an estimate r E { (1 (1-2)-1) of the true RDF ny such that for any > and å E 1), we have wa (fi(z), (z)) S with probability least 1-6 for all æ E X. whenever N = Ô( 2(1--(11))1 7) "polylog (1/)). Here, wi denotes the Wasserstein-1 distance between probability distributions, defined for any v, E (R) with CDFs F,, F, by wi (v, fg (t) (t) dt. We focus on the Wasserstein-1 1 distance as the main metric of interest in this paper it is particularly compatible with existing methods for analysing categorical approaches to distributional RL, and it provides upper bounds for differences many statistical functionals of interest, such as expectations of Lipschitz functions (Villani 2009 Peyré and Cuturi
2019); and conditional-valuee -vala-t---- (Rockafellar and Uryasev 2000 2002 Bhat and Prashanth) 2019
3) samples required to obtain such an accurate prediction with high probability, which follows from a reduction to the mean-return case (Azar et al. 2013)..
CVaR). Zhang et al. (202)) also prove a lower bound of N = Q( 2(1 >)-
There are two natural questions that the analysis of Zhang et al. (023) leaves open. Firstly, can the gap between the lower bound and upper bound as a function of (1
described above be closed? Zhang et al. (023) conjecture that their analysis is loose, and that this gap can indeed be closed.
Second, also note that the distributional dynamic programming procedure 17 < T1 proposed by Zhang et al. (2023), without incorporating any restrictions on the representations of distributions, runs into severe space and memory issues, and not practical to run (and indeed Zhang al. (2023) introduce approximations to the algorithm when running empirically for these reasons). A remaining question then whether there are tractable algorithms that can achieve the lower bound on sample complexity described above. Our contributions below provide a new, tractable distributional RL algorithm that attains (up to logarithmic factors) the lower bound on sample complexity provided by Zhang et al. (2023), resolving these questions.

---

Direct categorical fixed-point computation
Our approach to obtaining sample-efficient algorithm that is near-minimax-opiiaal in the sense de- scribed above begins with the categorical approach to distributional dynamic programming described in Section 2.3 We begin first by introducing a new algorithm for computing the categorical fixed point referred to in Proposition| 2.2 directly that avoids computing an approximate solution via dynamic programming iterations < mm Ty. We expect this algorithm to be of independent interest within the field of distributional reinforcement learning.
4.1 Direct categorical fixed-point computation
Our first contribution to develop a new computational perspective on the projected categorical Bellman operator II T, which results in a new algorithm for computing the fixed point ne exactly, without requiring the iterative CDP algorithm described in Section 2.3
CDF operator and fixed-point equation. We first formulate the application of the projected categorical Bellman operator IT 7 as linear map, and give explicit expression for the matrix representing this linear map when the input RDFs are represented with cumulative distribution functions (CDFs). We consider the effect of applying I to an RDF dj ((=1.-.... Zm ) with nzz) = >E p.(r)),, By Rowland et al. 20118 Proposition 6), the updated probabilities for (IIm (z)(r) -i ( can be expressed as
m P(ll))) yEN
p((x)) = Σ
where hr - h,((()) + yzz). We convert this into expression for cumulative probabilities, rather than individual probability masses, to obtain a simpler analysis below. To do so, we introduce the encoding of 17 (f=1......... ..., 2m))) 1)2²
into corresponding CDF values F X where F2 = (z) ([21, 2/]) denotes the cumulative mass at state x over the set {21....,
Proposition 4.1. If E P(12),.... Zm 1m))x* })* is an RDF with corresponding CDF values F E Xm
then the corresponding CDF values F E Xmm for Im TT satisfy
m
P((() = Σ
P())))(H) 11)+))
VEN 1=1
where
H = Eisi h((((x)) + 722)
(8)
for j = 1,...,m, m. and by convention we take Hi = O.
Under this notation, we can rewrite Equation
7 simply as matrix-vector multiplication in IR [m],
F' TTF,
where Tp is the (X' [m]) x (X' x [m]) square matrix, with entries given by
Tp(z, i. = P( Hf
and F, F" E Xmm above are interpreted in vectorised form. We drop dependence on from the notation Tp for conciseness. Thus, CDP can implemented via simple matrix multiplication on CDF values, and the CDF values F* for the categorical fixed point nje solve the equation
F = TpF, or cquivalently (I - FFTP = 0.
(10)
Obtaining system with unique solution. Equation (00) suggests that we can directly solve linear system to obtain the exact categorical fixed point, rather than performing the iterative CDP algorithm to obtain an approximation. Note, however, that F* is not the unique solution of Equation (10): for example, F = 0 is also a solution. This arises because in Equation (10). the distribution masses at each state (that is, (x)),, are unconstrained. By contrast, Proposition 2.2 establishes that yjc the unique solution of 77 11m T in the space P ([2,, 2m.)) where each element 17(x:) is constrained
5

---

to be probability distribution a priori. Thus, Equation (TO) requires some modification to obtain linear system with a unique solution. This is obtained by removing F. (z) variable from the system (for each resulting linear system, as the following proposition describes; the "axis-aligned" nature of these constraints is the benefit of working with CDF values.
X)),, replacing it by the constant 1, and removing redundant rows from the
Proposition 4.2. The linear system in Equation (TO), with the additional linear constraints Fm [m) = 1 for all x E X, is equivalent to the following linear system in F E Rxx//m-11;
(-TTp)F H.
(11)
where the (x,1: coordinate of (for S i,j S m 1)
Tp(x, 1; = P(yl P(yl)))(H H1 and for each x X, liisiss m 1, we have H(x, i) = H..m
Having moved to the inhomogeneous system over IR [m-11 Equation 11 we can deduce the following via the contraction theory in Proposition 2.2
Proposition 4.3. The linear system in Equation CDF values ((P) (=)) E X) of the categorical fixed point.
has unique solution, which is precisely the
The direct categorical fixed-point algorithm (DCFP) consists of solving the linear system in Equa- tion (II) to obtain the exact categorical fixed point; see Algorithm I for a summary, and Appendix G.5 for more details on implementations.
Algorithm 1: The direct categorical fxed-point algorithm (DCFP).
Calculate matrices (H" x E X)) via Equation
2 Calculate matrix Tp and vector H via Equation
Call linear system solver Equation 11
4 Obtain resulting solution F E (m--)
Return F", obtained by appending the values F% (x) 1 to the solution F"
Complexity and implementation details. Representing Tp dense matrix requires O(Ix(m²) space, and solving the corresponding linear system in Equation (11) with standard linear solver requires (mx))mm)) time. However, in many problems there are cases where the DCFP algorithm
can be implemented more efficiently. Crucially, Tp often has sparse structure, and so sparse linear solvers may afford opportunity make substantial improvements in computational efficiency. We explore this point further empirically in Section 0 and give a theoretical perspective in Appendix G
4.2 DCFP with a generative model
We now return to the setting where the Markov reward process in which are performing evaluation is unknown, and instead we have access to the random next-state samples ((X:) x X). as described in Section 3 Our model-based DCFP algorithm proceeds by first constructing the corresponding empirical transition probabilities P, that P(yla) N- (( y). and then calls the DCFP procedure outlined in Algorithm treating P as the true transition probabilities of the MRP when constructing the matrix in Line which we denote here by Tp. to reflect the fact that it is built from P. This produces the output CDF values É. from which estimated return distributions can be decoded (with the convention Fo(x) = 0) as
m
(xx) = Σ (f((x) - (x)
(12)
5 Sample complexity analysis
Our goal now to analyse the sample complexity model-based DCFP, as described in the previous section. We first introduce notational shorthand that will be of extensive use in the statement and

---

proof of this result: when writing distances between distributions, such as wn w1(n" identify CDF vectors F(z) E R with the distributions they represent in Equation (2) The core theoretical result of the paper is as follows.
(z), (1(7"(x),
(x))), we
Theorem 5.1. Let ε E (0, (1 - 1/21) and E (0,1), and suppose the number of categories
satisfies m 2 4(1 - -2)-2--)222-- +
Then the oulput F of model-based DCFP with N = =s(s 21 samples satisfies, with probability at least 1 8,
polypooy
max w1 (17 (x), SE. xEX
This result establishes that model-based DCFP attains the minimax lower-bound (up logarithmic factors) for proh-probability return distribution estimation in Wasserstein distance, resolving an open question raised by Zhang et al. (2023); in a certain sense, estimation of return distributions is no more statistically difficult than that of mean returns with a generative model. Note there is no direct dependence of N on m, so there is no statistical penalty to using large number of categories m.
Extensions. We note that this core result can be straightforwardyy extended in several directions. In particular, similar bounds apply in the case of predicting returns for learnt near-optimal policies in MDPs (see Section F2). for the iterative categorical DP algorithm in place of DCFP (when using sufficiently many DP updates; see Section F. 1. and in the case of stochastic rewards (see Section F3)
5.1 Structure of the proof of Theorem 5.1
The remainder of this section provides a sketch proof of Theorem 5.1 a complete proof is provided in the appendix. The proof is broadly motivated by the approaches of Azar et al. (2013), AAgaraa et al. (2020), and[Panajjady and Wainwright (2020), who analyse the mean-return case, and we highlight where key ideas and new mathematical objects are required in this distributional setting. In particular, we highlight the use the stochastic categorical CDF Bellman equation, new distributional Bellman equation that plays key role in our analysis, which we expect to be of independent interest.
Reduction to Cramér distance. The first step of the analysis is to reduce Theorem 5.1 to a statement about approximation in Cramér distance, which is much better suited to the analysis of DCFP, owing to the results described in Section 2.3] This can be done by upper-bounding Wasserstein distance by Cramér distance using the following result, which is proven in the appendix via Jensen's 's inequality.
Lemma 5.2. For any two distributions U, // E 9 (10, (1 - (1-7)--11) I1), we have
W1 (v. v') S - (1-)))11/22(1,, v').
Theorem 5.1 is now reducible to the following, stated in terms of the Cramér distance.
Theorem 5.3. Let E (0,1) 1) and ê 1), and suppose the number of categories satisfies m 2 4(1 - 22-2 + 1. Then the output F. of model-based DCFP with N 21
7) - (0=)poylog((xx118) samples satisfies, with probability at least 1 - 5,
max Ea("2 (a))nn"()),,))))))) (())) SE. zEX
(13)
Reduction categorical fixed-point error. Our first step in proving Theorem 5.3 is use the triangle inequality to split the Cramér distance on the left-hand side of Equation into representation approximation error, and sample-basse error:
Ee F) < 22(n" +2)+ F (1- 11-r)/m-1 + (2(FF,F) F),
1
with the second inequality following from the fixed-point quality bound in Equation (5) With m as specified in the theorem statement, the first term in the right-hand side above is bounded by e/2. Thus, it suffices to focus on the second term on the right-hand side, which quantifies the sample-based error in estimating the categorical fixed point F*
Concentration. Through a combination of the use a version of Bernstein's 's inequality in Hilbert space (Chatalic et al. (202)) and propagation of this inequality across time steps in the MRP, we next arrive at the following inequality with probability at least 5:
()(F 2
II( (111-ン))-1 op + (1 (1- (1-2)3/2N3/4)
1
(14)
VN(1-
7

---

Here, dp. E R is an instance of new class of distributional object. the local squared- Cramér variation. The general definition is given below, for the case of a general transition matrix Q E R-*'x to avoid conflation with the specific transition matrices P and P that Theorem 5.3 is concerned with.
Definition 5.4. For given transition matrix Q, the single-sample operator To R + R. is the random operator given by: (i) constructing a random transition matrix Q by, for each x E X, sampling X' ~ Q[[[z), and setting [x) = 1; (ii) setting To = =To-
Definition 5.5. For given transition matrix with corresponding CDP fixed point F e R.m the local squared-Cramér variation at Q, oq IRx is defined by
oo(x) (() (x),
Intuitively, the local squared-Cramee variation oQ encodes the variability of the fixed point FQ after a sample-based, rather than exact, dynamic programming update. From this point of view, it is natural quantity arise in Equation (14), and plays similar role the variance in the classical Bernstein inequality (Bernstein 1946
In Corollary 5.12 below, will deduce that under the conditions of Theorem 5.3 we have
2
-P olloo S
(15)
--7
Substituting this into Equation 14 gives
(2(F(c), (x) s o
1 (1 VVN (1 7)3/2/33/4
+
with probability at least 1 6, for all æ E X. Now, taking N = Õ ?) 22 yields that this
expression is O()), which completes the sketch proof of Theorem 5.3 What remains be described
is how to arrive at the bound in Equation (15); the section below provides high-level overview of the technical details involved in obtaining it.
5.2 The stochastic categorical CDF Bellman equation
The central idea is to relate the local squared-Cran -Crméééé variation to a corresponding global notion of variation, in analogy with the variance Bellman equation (Sobel) 1982) used by Azar et al. (2013) in the mean-return case. To define this corresponding global notion, we begin by defining a new type of distributional Bellman equation, which can be intuitively thought of encoding the result of repeatedly applying a sequence of independent single-sample operators. Again, we work with a general transition matrix Q.
Definition 5.6. For general transition matrix Q, the stochastic categorical CDF (SC-CDF) Bellman equation is given by
(xx) 2 (z),
(16)
D where denotes equality in distribution. Here, To is a single-sample operator with respect to Q. in Definition 5.4 Each ((x) is random variable taking values in the space of valid CDF values
= {F E RT :0" < Fi S S Fm-1 I Fm = 1} for distributions in (1 2m H), and is taken to be independent of the random operator To-
The intuition is that "unravelling" Equation (6) should lead to a solution of the form
+ lim (ik) -00 F,
where (I(()) y are independent Ainle-sample operators, so that ((x)) encodes the fluctuations due repeated CDP updates with randomly-sample transitions. To make this intuition precise, we first. verify that the SC-CDF Bellman equation has a unique solution.
Proposition 5.7. The SC-CDF Bellman equation in Equation (6) has a unique solution, in the sense that there is unique distribution for each ((x) such that Equation (16) holds for each E X.
8

---

0.75
0.25
(0))
T
I 1
0
4
6
10
Figure 2: Left: Example MRP with r(z0) 1, rx1) 0, 7 0.9. Right: Categorical fixed point F* (xo) with m 15, and 5 independent samples from the random CDF 4 (zo).
We write Φ for the collection of random CDFs that satisfy the SC-CDF Bellman equation in Equation (16). whose existence is guaranteed by Proposition 5.7 We next relate & to the standard categorical fixed point F
Proposition 5.8. For all x E X, we have
Proposition 5.8 shows that & can indeed thought of as encoding random variation around the usual categorical fixed point F see Figure 2 This motivates the following.
Definition 5.9. For given transition matrix with corresponding CDP fixed point FQ E RR the global squared- Sruared-Cramér Cramér variation at Q Σq E R is defined by
[x)x)= =
Remark 5.10. Note that (x() a doubly distributional object. represents a probability distribu- tion centred around the object F (x), which itself already provides summary of the distribution of the return. This reveals a dual perspective on distributional RL itself. The distributional predictions made serve several purposes: (i) modelling the aleatoric uncertainty in the return, as is the in the work of Bellemare et al.) (2017) and much subsequent algorithmic work, and/or (ii) serving to model specific types of epistemic uncertainty in the estimation of a non-random object from random data, as used in the analysis Azar et al. ( (2013) and much subsequent work on the sample complexity of reinforcement learning. The object 44 is motivated by both of these concerns simultaneously.
The following Bellman-like inequality draws a relationship between local and global squared-Crante variation, allowing make progress from Equation 15
Proposition 5.11. We have
Eq 2 aQ + - -(mV/-/ m²(1-7)²2
where 1 x is a vector of ones, and the inequality above is interpreted component-wise.
Rearrangement and bounding of the quantities in the inequality of Proposition 5.11 the specific case Q = P. yields the required inequality in Equation (15) that completes the proof of Theorem 5.3
Corollary 5.12. We can bound the term [|(Z -Po-oo appearing in Equation 14 under the
assumptions of Theorem 5.3 as follows:
< 2 (11 (1- - oplloo & 1E + 1- 1-7
Empirical evaluation
To complement our theoretical analysis, which focuses on worst-case sample complexity bounds, we report empirical findings for implementations of several model-based distributional RL algorithms in the generative model setting. We compare the new DCFP algorithm introduced in Section 4.1 with
quantile dynamic programming (Dabney et al. 201 8b Bellemare et al.) 2023 QDP) distinct ap- proach to distributional RL which return distributions are approximated via dynamic programming with finite number of quantiles.

---

V=0.9
y=
y=0.9
V=0.99
Wassersteln distance
Method
Wasserstein distance
100
DCFP
10A
Atoms
10
100
10-1
300
1000
10 10-1
DCFP CDP
10
10-2
QDP
102
10
10 102
104
10
10
3
10
10
10
10 10a
# Samples / state, N
Wallclock time
/s
Figure 3: Approximation error/wallclock time for a variety distributional RL methods, discount factors, numbers of atoms, and numbers of environment samples.
In Figure 3 (left), we report results of running DCFP and QDP on a 5-state environment with transition matrix randomly sampled from Dirichlet distributions, and entries of the immediate reward function r E RR* randomly sampled from Unif([0, [0. 1)), with varying numbers of atoms m, environment samples per state N, and discount y. We report the maximum Wi -ererror against true return distributions (estimated via Monte Carlo sampling). All runs are repeated 30 times, and error bars are 95% bootstrapped confidence intervals. Sufficient DP iterations ensure approximate convergence to their fixed points. In Figure 3 (right), we plot estimation error against wallelock time for 100, N 10 including results for CDP (which approximates the solution of DCFP via dynamic programming: Section 2.3); line plots indicate the estimation error/wallclock time trade-off as we increase the number of DP iterations. Both DCFP and CDP methods benefit from setting the atom support based on maximal/minima. values of r, as described Appendix G
For low discount factors and atom counts QDP generally outperforms DCFP in terms of asymptotic estimation error, due QDP's ability modify its atom support to regions of the interval [0, (1 y)= where mass is concentrated. However, we note that DCFP is generally faster than QDP. Further. DCFP generally outperforms QDP, in terms of both speed and estimation error, under larger discounts and/or larger atom count. We also note that particularly at high discounts, DCFP outperforms CDP terms of wallelock time, due to DP methods requiring many iterations to converge in these cases (Rowland et al. 2018). Full results, including given in Appendix G key findings include that QDP works particularly well near-deterministic environments, and DCFP works particularly well in settings where there are short high-probability paths from a state to itself.
several further environments, are
7 Conclusion
We have introduced a new algorithm, DCFP, for directly computing the fixed point of CDP, widely used distributional reinforcement learning algorithm. We then showed that this algorithm, with an appropriately chosen number of categories m, achieves the minimax lower bound (up to logarithmic factors) for sample complexity of return-distribution estimation in Wasserstein distance with generative model. Thus, this paper closes an open question raised by Zhang et al. (2023) by exhibiting algorithm that obtains this lower bound, and shows that estimation of return distributions via generative model essentially no harder statistically than the task of estimating a value function.
Our analysis also casts new light on categorical approaches to distributional reinforcement learning in general. The newly introduced stochastic categorical CDF Bellman equation serves to encode information about statistical fluctuations of categorical approaches to distributional RL, and we expect to be further use in theoretical work distributional RL generally. Our experimental results also highlight salient differences in performance for distributional RI. algorithms making use of distinct representations, depending on levels of environment stochasticity and discount factor in particular. We believe further investigation of these phenomena is an interesting direction for future work.
10

---

Acknowledgments and Disclosure of Funding
We thank Dave Abel for detailed comments on draft version of the paper. We also thank Mohammad Gheshlaghi Azar for useful conversations, and Arthur Gretton for advice on Bernstein-like inequalities in Hilbert space.
References
Alekh Agarwal, Sham Kakade, and Lin F. Yang. Model-based reinforcement learning with generative model is minimax optimal. In Proceedings of the Conference on Learning Theory, 2020.
Mohammad Gheshlaghi Azar, Rémi Munos, Mohammad Ghavamzadeh, and Hilbert Kappen. Speedy Q-learning. In Advances in Neural Information Processing Systems, 2011.
Mohammad Gheshlaghi Azar, Rémi Munos, and Hilbert J. Kappen. Minimax PAC bounds on the sample complexity of reinforcement learning with a generative model. Machine learning, 91: 325-349, 2013.
Osbert Bastani, Jason Yecheng Ma, Estelle Shen, and Wanqiao Xu. Regret bounds for resk-sensitive reinforcement learning. In Advances in Neural Information Processing Systems, 2022.
Mare Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforcement learning. In Proceedings of the International Conference on Machine Learning, 2017.
Marc G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera Ponda, and Ziyu Wang. Autonomous navigation of stratospheric balloons using reinforcement learning. Nature, 588(7836):77-82, 2020.
Mare G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning, MIT Press, 2023. http://www. www.diiiuiiiionrrri.....
Sergei Bernstein. The Theory of Probabilities. Gastehizdat Publishing House, 1946.
Sanjay P. Bhat and I.. A. Prashanth. Concentration of risk measures: A Wasserstein distance approach. In Advances in Neural Information Processing Systems, 2019.
Markus Böck and Clemens Heirzinger. Speedy categorical distributional reinforcement learning and complexity analysis. SIAM Journal on Mathematics of Data Science, 4(2):675-693, 2022.
Markus Böck, Julien Malle, Daniel Pasterk, Hrvoje Kukina, Ramin Hasani, and Clemens Heitzinger. Superhuman performance sepsis MIMIC-III data by distributional reinforcement learning. PLoS One, 17(1 17(11))0075335) 2022.
Cristian Bodnar, Adrian Li, Karol Hausman, Peter Pastor, and Mrinal Kalakrishnan. Quantile QT- for risk-aware vision-bact based robotic grasping. In Robotics: Science and Systems, 2020.
Yash Chandak, Shiv Shankar, and Philip Thomas. High-confidence off-policy (or counterfactual) variance estimation. In Proceedings of the AAAI Conference Artificial Intelligence, 2021.
Antoine Chatalic, Nicolas Schreuder, Lorenzo Rosasco, and Alessandro Rudi. Nyström kernel mean embeddings. In Proceedings of the International Conference on Machine Learning, 2022.
Harald Cramér. On the composition of elementary errors. Scandinavian Actuarial Journal, 1928(1): 13-74, 1928.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In Proceedings of the International Conference Machine Learning, 10.
Will Dabney, Mark Rowland, Mare G. Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018b.
11

---

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- bbaed reinforcement learning. Nature, 57777722:)))667166 2020.
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju. Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Clau- dio Fantacei, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojevié, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL
http://github.ct htt.:ggituuuu.oooooggggeeo www.:////oooooeeeeeiiiiooo
Thang Doan, Bogdan Mazoure, and Clare Lyle. GAN Q-learning- arXiv, 2018.
Yihan Du, Siwei Wang, and Longbo Huang. Provably efficient risk-sensitive reinforcement learning: Iterated CVaR and worst path. In Proceedings of the International Conference Learning Representations, 2022.
Aryeh Dvoretzky, Jack Kiefer, and Jacob Wolfowitz. Asymptotic minimax character of the sample distribution function and of the classical multinomial estimator. The Annals of Mathematical Statistics, pages 642-669, 1956.
Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Moham- madamin Barekatain, Alexander Novikov, Francisco J R Ruiz, Julian Schrittwieser, Grzegorz Swirszez, David Silver, Demis Hassabis, and Pushmeet Kohli. Discovering faster matrix multipli- cation algorithms with reinforcement learning. Nature, 610(7930):47-53, 2022.
Yingjie Fei, Zhuoran Yang, Yudong Chen, and Zhaoran Wang. Exponential Bellman equation and improved regret bounds for risk-sensitive reinforcement learning. In Advances in Neural Information Processing Systems, 2021a.
Yingjie Fei, Zhuoran Yang, and Zhaoran Wang. Risk-sensitive reinforcement learning with function approximation: A debiasing approach. Proceedings of the International Conference on Machine Learning, 2021b.
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.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(1):723-773, 2012.
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825))3577-32 September 2020.
J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, 9()):
90-99 95, 2007.
Sham Machandranath Kakade. On the sample complexity of reinforcement learning. PhD thesis, University College London, 2003.
Michael Kearns, Yishay Mansour, and Andrew Y. Ng. sparse sampling algorithm for near-optimal planning in large Markov decision processes. Machine learning, 49: 193-208, 2002.
Thanh Lam, Arun Verma, Bryan Kian Hsiang Low, and Patrick Jaillet. Risk-aware reinforcement learning with coherent risk measures and non-linear function approximation. In Proceedings of the International Conference Learning Representations, 2022.
12

---

Gen Li, Yuting Wei, Yuejie Chi, Yuantao Gu, and Yuxin Chen. Breaking the sample size barrier in model-based reinforcement learning with a generative model. In Advances in Neural Information Processing Systems, 2020.
Hao Liang and Zhi QuaiQQaa Luo. Bridging distributional and risk sensitive reinforcement learning with provable regret bounds. Journal of Machine Learning Research, 2022.
Pascal Massart. The tight constant in the Dvoretaky-Kieelr-Woofle Wolfowitt inequality. The Annals of Probability, pages 1269-1283, 1990.
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.
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.
Thanh Aguyen-Tang. Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning via moment matching. In Proceedings of the AAAI Conference on Artificial Intelligence 2021.
Ashwin Pananjady and Martin Wainwright. Instance-dependent - for policy evaluation in tabular reinforcement learning, IEEE Transactions Information Theory, 67(1))566-55) 2020.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning. Now Publishers, Inc., 2019.
Martin L. Puterman. Markov decision processes: Discrete stochastic dynamic programming. John Wiley & Sons, 2014.
R. Tyrrell Rockafellar and Stanislav Uryasev. Optimization of conditional value-at-risk. Journal of Risk, 2:21-42, 2000.
R. Tyrrell Rockafellar and Stanislav Uryasev. Conditional value-at-risk for general loss distributions. Journal of banking & finance,
2002.
Mark Rowland, Marc G. 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.
Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Mare G. Bellemare, and Will Dabney. An analysis of quantile temporal- difference learning. Journal of Machine Learning Research, 2024.
Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equivalence of distance- based and RKHS-based statistics in hypothesis testing. The Annals of Statistics, pages 2263-2291, 2013.
Bobak Shahriari, Abbas Abdolmaleki, Arunkumar Byravan, Abe Friesen, Siqi Liu, Jost Tobias Springenberg, Nicolas Heess, Matt Hoffman, and Martin Riedmiller. Revisiting Gaussian mixture critics in off-policy reinforcement learning: A sample-based approach. arXiv, 2022.
Aaron Sidford, Mengdi Wang, Xian Wu, Lin Yang, and Yinyu Ye. Near-optimal time and sample complexities for solving Markov decision processes with a generative model. Advances Neural Information Processing Systems, 2018.
Matthew J. Sobel. The variance of discounted Markov decision processes. Journal of Applied Probability, 19(4):79-800, 1982.
Richard Sutton and Andrew Barto. Reinforcement learning: An introduction. MIT Press, 2nd edition, 2018.
Gábor Székely. E-statistics: The energy of statistical samples. Technical Report 05, Bowling Green State University, Department of Mathematics and Statistics Technical Report, 2003.
13

---

Gábor Székely and Maria L Rizzo. Energy statistics: A class of statistics based on distances. Journal of Statistical Planning and Inference, 143(8):124--1272, 2013.
Guido Van Rossum and Fred L. Drake. Python Reference Manual. CreateSpace, Scotts Valley, CA, 2009.
Cédric Villani. Optimal transport: Old and new, volume 338. Springer, 2009.
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. Nelson, Eric Jones, Robert Kern, Eric Larson, C. J Carey, ilhan 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.
Kaiwen Wang. Nathan Kallus, and Wen Sun. Near-minimax-optinall risk-sensitive reinforcement learning with CVaR. In Proceedings of the International Conference on Machine Learning, 2023a.
Kaiwen Wang, Kevin Zhou, Runzhe Wu, Nathan Kallus, and Wen Sun. The benefits of being distributional: Small-loss bounds for reinforcement learning. In Advances in Neural Information Processing Systems, 2023b.
Zifan Wang, Yulong Gao, Siyi Wang, Michael M Zavlanos, Alessandro Abate, and Karl H Johansson. Policy evaluation in distributional LQR (extended version). arXiv, 2023c.
Michael L. Waskom, seaborn: statistical data visualization. Journal of Open Source Software, 6(60):
3021, 2021.
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.
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. Outrac- ing champion Gran Turismo drivers with deep reinforcement learning. Nature, 602(7896):223-228. 2022.
Derek Yang. Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie- Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. Advances Neural Information Processing Systems, 2019.
Yuguang Yue, Zhendong Wang, and Mingyuan Zhou. Implicit distributional reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Vadim Yurinsky. Sums and Gaussian vectors. Springer, 2006.
Liangyu Zhang, Yang Peng, Jiadong Liang, Wenhao Yang, and Zhihua Zhang. Estimation and inference in distributional reinforcement learning. arXiv, 2023.
14

---

APPENDICES:
Near-Minimax-ppiine Distributional Reinforcement Learning with a Generative Model
For convenience, we provide a summary of the contents of the appendices below.
Appendix A provides detailed discussion of related work.
Appendix provides additional convenient notation for the categorical CDF operator Tr, in particular expressing it as the combination of a acaling/shitinn/prooeciin operation, and a mixing operation over the state to be bootstrapped from. Several additional self-contained contraction
results are given that are used within the main proof's of the paper.
Appendix C provides proofs for the results in Section 4 concerning the formation of the linear system that is solved by the DCFP algorithm, and verilication that the derived system has the unique desired solution.
Appendix D provides proofs relating to the stochastic categorical CDF Bellman equation, in particular establishing the desired unique solution to this new distributional Bellman equation, and proving a key bound on the local squared-Cramér variation that is used within the proof of the main theorem of the paper.
Appendix provides the proof of the main result of the paper, Theorem 5.1 giving full details for the steps traced through in the sketch proof in the main paper. Appendix F
provides discussion of several straightforward extensions of the main result, Theo-
r.m
Appendix G provides full details relating to the experiments in the main paper, well as additional empirical comparisons between methods for distributional RL with a generative model.
A Related work
Other families of model-based distributional RL algorithms. There are many approaches to
distributional RL; important choices studied previously include moments (Sobel 1982) . exponential
families (Morimura et al.] 2010b), categorical distributions (Bellemare et al. 2017). ) collections
of particles (Morimura al.J 2010aa) Nguyen- Tang et al., (021). quantiles (Dabney et al. 20186),
and generative models (Doan et al.) 2018 Dabney et al. 2018a Freirich et al., 2019] Yang et al.
2019: Yue et al. 2020 Our choice of categorical representations in this work is motivated by several considerations: (i) the existence of principled dynamic programming methods for these representations, with corresponding convergence theory (Rowland et al., 2018); (ii) the flexibility of these representations to trade-off computational complexity with accuracy (by varying m. (Rowland) ct al. 2018); (iii) the mathematical structure of the dynamic programming operator (linear), as described in this work; and (iv) the availability of an efficient algorithm to exactly compute the DP fixed point (the DCFP algorithm proposed in Section 4.1
An interesting and important direction, given our empirical findings in Section o is whether analyses can also be carried out for other approaches to distributional dynamic programming. such as quantile dynamic programming (QDP; [abneyy et al., 2018b Rowland et al. 2024 Bellemare et al. 2023), and fitted likelihood estimation (FLE; Wu et al., (023). We expect chal challenges in extending the analysis to result due the fact that, for example, the QDP operator is non-linear, and FLE operator applications typically do not have closed forms. Nevertheless, it would be particularly interesting to understand whether it is possible to obtain instance-dependent bounds for QDP, particularly given its strong empirical performance. Similarly, as described in Section [ an interesting question for future work is whether is possible improve the computational complexity of model-based DCFP for high-probability return-direction estimation.
propose model- Other statistical questions in distributional RL. Böck and Heirzinger 2022)
free algorithm for distributional RL, speedy categorical TD-Iearning, motivated by categorical TD learning (Rowland et al. 2018) and speedy Q-learning (Azar et al. 2011), and prove sample complexity bound of Ō( (1 -)- for high probability e-accurate estimation in Cramér distance (which implies non-minimax sample complexity (1 -4) in Wasserstein-1 distance, per I Lemma 5.2) Wu et al. (2023) study the offline evaluation problem, in which state-action pairs are sampled from the stationary distribution of the policy, via fitted likelihood estimation (FLE) and focus on generalisation bounds, allowing for policy evaluation in environments with uncountable state
15

---

spaces, Wang et al. (2023c) also study policy evaluation in this context, focusing on the LQR model. Aside from studying minimax optimality, Zhang et al.] (202)) also make several other contributions on the topic of statistical efficiency of estimation for distributional RI., including analysis of asymptotic fluctuations and limit theorems, and analysis of approximations in more general metrics, including Wasserstein-p. Aalmogorov-Smmroon and total variation metrics. Their sample complexity results rely on careful analysis the behaviour of the unprojected distributional Bellman operator 7 on certain subspaces of probability/signed measures.
Sample complexity of mean-return estimation. The sample complexity of estimating mean returns, and the related task of obtaining near-optimal policy, has been considered by Azar et al. (2013);
Sidford al. (2018); Pananjady jady and Wainwright (2020); [Agarwal al. (2020): Li et (2020).
Interestingly, Agarwal et al. (2020), while treating the mean-return case, make use of return binning as proof technique, Li et al. 2020) also obtain bounds for mean-return sample complexity that apply model-based algorithms; an interesting
with (1-?) 1/2 for modifications of certainty-equival direction for future work would be to check whether the restrictions on in Theorem [ can be lifted by incorporating ideas from this mean-return analysis to the distributional setting. Additionally, Chandak et al. (2021) consider the task of estimating the variance of returns from off- off-policy data, and anng et al. (2023b) study regret minimisation properties (with respect to the expected return criterion) of distributional RI algorithms in the online setting.
Risk-sensitive control. The theory developed in this paper has focused estimation of return distributions for individual policies. A natural direction for future work is to analyse identification of near-optimal policies for risk-sensitive decision criteria. Bastani al. (2022); Wang al. (2023a)
study efficient algorithms for CVaR optimisation, while Fei et al. (202 la b Liang and Luo (0222
study entropic risk maximisation, and Du et al. (2022) Lam optimisation, all in the online setting.
(2022) study iterated CVaR
Further analysis. In this paper, have resolved conjecture of Zhang al. 2023). obtaining
a near- minummm mainimxx-pptimal optimal algorithm for estimation of return distributions in Wasserstein distance. Zhang et al.] (2023) make contributions to several other important statistical questions regarding distributional RI, including approximation in stronger metrics such as Kolmogorov-Smimon and total variation metrics, as well as studying asymptotic fluctuations of estimators; it will be interesting to see whether the analysis presented here can be extended to these other settings as well.
B Additional CDF operator notation and contractivity results
In this section, we introduce finer-grained notation for the CDF operator Tp that allows us to straightforwarddy refer to the operations that correspond to phiftng//saalin/proooeet and to mixing over transition states, separately. We also establish several additional contraction lemmas that will be useful in the sections that follow.
B.1 Categorical hat function definition
For convenience, provide the full mathematical definition of the hat functions h; [0, (1- [0, 1] used in defining the categorical projection described in the main paper, as given by Rowland al. 2018)). For 2,..., 1, we have
for [ =i]
hx(2)
for E [ze, #141] otherwise.
For the edge case h1, we have
=2-2
for 2 [21, 22] otherwise,
h1(2)
22-21
and similarly for the edge case hms we have
-1
for E [=m-11 2m]]
hmm (m(2)) (2) =
Zm - Am-1
0
otherwise.
16

---

B.2 Finer-grained operator expression
Recall that the CDF operator Tp: R IR xm is represented by a matrix with elements given by
Tp(z, (: = ((y H()+1))
(17)
We can therefore conceptualise the application T'F' as standard matrix-vector multiplication in the vector space x However, the expression for matrix elements in Equation (17) has additional structure that means we can express the application of Tp in a different manner, which will be convenient in several proofs below, particularly it separates out the influence of rewards (which remain fixed) and transition dynamics (which are estimated via samples).
In particular, will regard F xx itself matrix, with rows indexed by states in X, and columns indexed by indices = .,....m. m. For each state x E X, we then introduce the matrix B, Rmxm with (4, j) element given by
H
H
We then have that (T) R can alternatively expressed matrix notation as
PaFB
where Pa is the row vector given by the row of P corresponding to state æ E X.
B.3 Additional results
Below, we provide several additional results regarding contractivity properties of Tp. To do so, is useful to introduce the norm || lle, on I,",, which we define by
m
P2(1)2 1/2
Fl
m(1
The motivation for this definition is that if we have two distributions v.v/ ({ ......... 2m }) with corresponding CDF vectors F' E IR",, then ( |F as v, are supported on [0, (1-1) Thus, under the abuse of notation &2(F, F*) introduced in the main paper, have (2 F- F" leg- We also introduce a supremum version of this norm on the space IR X
=
which we denote olt and define by
Fl = max [())))
sEX
for all F E RXxm This norm is defined that if have RDFs n.1/ / ((s),....., =m)) and F, F E Rm are the corresponding CDF values, then
2 (7. v/) = | F 122.00
With this norm defined, we can now state and prove our first result, which essentially translates the contraction result in Proposition 2.2, expressed purely in terms probability distributions and Cramér distance lg, into slightly more general result expressed over x and the norm || -llez-
Proposition B.1. The operator Tp Xx - xx is a contraction when restricted to the subspace (F Xmm (m) for all X} with respect to the norm 1leg,00. with contraction factoor V7.
Proof. By Proposition 2.2 for any two RDF approximations n' E P(22) .m 2m)) with corresponding CDF values F, F X. we have
2 (II Tn, I TT/) (i2().)))
and hence we also have
TpF TPF V F' F'lle,.0
Hence, Tp is a contraction map on the set
[F- - F':0 < (z( S Fm (x) = 0 F) < ... (x))) for all X).
17

---

This set contains a basis for the subspace {F E RX (=) = 0 for all x E X}. Namely, the one-hot vector at coordinate (z. (for j < m) can be exhibited as lying this subspace since can
be expressed as the difference between the vectors F, F' E RR defined by F)(() = Ff(()) = 1
for all y F and all j = l..... m. and (x = 1 for i 2 j (and O otherwise), and (x) = 1 for i 122 j + 1 (and 0 otherwise). Since Tp is linear, therefore holds that
|Tp $ /
(18)
for any F in the subspace { {F E TR.xxmm
(m) = for all x E x), as required.
Proposition B.2. For each x X, the matrix By is contraction mapping on the space {F IR"
Fm = 0} with respect to llea. with contraction factor V7.
Proof. Consider related one-state MRP, for which the reward the single state is r(x:), and the state transitions to itself with probability 1. The categorical Bellman operator associated with this MRP and the support (Z1,..., Zm precisely Bg, and the statement of the result therefore follows
as special case of Proposition B.1
The next result serves as counterpoint Proposition B.2 it shows that if is sufficiently large, the map Bg does not contract by too much in || lle.. Proposition B.3. For any F,F's E F, we have
2
1
BBFF-BBFT 2 2 F'l
m(1 7/1/2
m²(1 7)2
This is proven via the following lemma.
Lemma B.4. For any E ((O,, (1 - -7)) and the projection IT... P((0, (1 -7)-11)
((2), 2m }), we have
12(v,1mm/ S
2 2/m(1 -7)
Further, for any v, V' E ( (1 - 7) we have
(2 Imme 2 (2(1,()) - m(1 77)²
and
(1mm,,nmm 2 (v,,v') m(1 - 2 m² 7)
Proof. By Rowland et al.
2018 Proposition 6), have that the CDF values Fi for i
1...., m are equal to the average of the CDF values of F, on the interval [2,, 21+1]. Therefore, in computing the squared Cramér distance (3(1, I), the worst-case contribution to the integral
EBV, Im =
(F, - Fuum (t))² dt
from the interval [#e, 41)1), holding F, (z; and F (2(41)) constant, is
F F( P
1
(2411 24)
4m 4m)117) (F(2+11)
Thus, the worst-case value for the entire integral
Σ
4m -
1
The worst-case value for the sum is 1, by interpreting this sum of squared probabilities, and therefore deduce that
1
1m))) < 4mi1 - 77)
18

---

leading to
(v, 1m 2
2 2/m(1= 7)
as required for the first stated inequality. For the second and third inequalities, we apply the triangle inequality twice to obtain
1
(2) I (2, =m + (I Mm + (2 M 1/) S l 1. mm + m(1 7)
rearrangement then gives the second statement. Squaring both sides of the inequality above gives
2
(m Im m² 7)² Bounding the instance of (2 Um in the cross-term on the right-hand side by (1 - -/1 and rearranging then yields the result.
(z v/) < (T (I Mm
m(1 7)
Proof of Proposition B.3 Let E ({2),..., 2m be the distributions with CDF values F, F', respectively, and let G. be random variables with CDFs F, F" respectively, and (x, X') an independent random transition in the MRP beginning at state æ. Recall from the notation introduced earlier in this section that F and F² the CDF values of the distributions of r(x) + +GG and r(z) + 7G' after projection onto the support grid {z1..... 2m by the projection map Im Following common notation distributional RL (Bellemare al. 2023). we denote the distributions of r(x) + yG and r(z) + GG by (br(z),7)44 and (br(x),)) respectively. Here, br(x)., R - R is the bootstrap function (z(=),((2) = r(x) yz, and (6)(z)) 4 is the push-forward distribution of through by()), (intuitively, the distribution obtained by transforming the support of v according to br(x),)). With this notation introduced, we therefore have
| - F, (br(s))))))))))))) # 1Im (b)(s))))
#M, (br(x),)) #/) (a) (b((br(e))) #M
2
1
m 7/1
m²(1 7)2
( (V
2
m² - 7)2 m(1 7)1/2
2
1
m 7)/
m²(1 >
as required, where (a) follows from Lemma B.4 and (b) follows from homogeneity of Cramér distance (see Rowland et al. (0018 Proof of Proposition 2)).
C Proofs of results in Section
4
Proposition 4.1. In P(2,,,... 2m3) ))x is an RDF with corresponding CDF values F Xx then the corresponding CDF values F" E Xmx for II TT satisfy
F(x)= = Σ
Ply (og)))
MEX 1=1
where
= = -igs hu((()) (ighh()t(()())+++ + +
(8)
for j 1,...,,m and by convention we take H = 0.
Proof. Beginning by restating Equation (6), we have that if nj ({ Em 1m is an RDF with corresponding probability mass values p E Xxm, then the updated RDF T/ II,, Tn has corresponding probability mass values p' E xmm given by
pT(x)) ΣΣ m
P()l))M i=
19

---

First, for i E m.) we sum from to i to yield
n
m
F((x) Σ M((x) = Σ
P(yla)
Σ
Plylt)H(,pp(0)))
15
MEN ==1
yEX 1=1
Σ
(F F)-1(1)),
yEX j=1
where by convention take Fo(y) = 0. By reorganising the terms the right- Light-handaeg side, the claim now follows.
Proposition 4.2. The linear system in Equation (0) with the additional linear constraints F. (x)
1 for all x E X, is equivalent to the following linear system in F Rx mmmm)))
(T (T-Tr) (-Tr)F=1i
(11)
where the (I,1;y,5) coordinate of (for 1 E i,j S m - 1) is
Tp( (1),) = P(ula)(HH, H11+)), for each E X. Kxlsism - 1, we have H(x, i) H
Proof. First, we consider a row of F = TpF that corresponds to the index (x, with * m. Expanding under the definition of Tp, we have
m-1
Fx(x)
Σ
P(, HP)-1)P(()) Σ P BEA
yEX
Since we assume the additional constraints Fm = for all y E X. the final term on the right-hand side can be simplified yield
m-1
Px(x)=
Σ Σ
P( - H.
VEA j=1
This is precisely the row of F F H corresponding to index (xr i).
Now, we consider a row of F = -TpF that corresponds to the index (I, m). Again making the substitution Fm (y) = 1 for all y E X, we have
m
Fm (m)())
Σ
Plyla) (H H H Fg()) + Σ P( m.m Fm
y
1=1
y
Fmn = Σ H(ylz)HH Fnl(y) 11, y
which shows that the equation is redundant, and hence can be removed from the system. The claim H H H 0 follows since fact H Σm hr(r(z) + 72y). ). and the sum over the hat functions for any input argument is 1. In the final equality, we have used the fact that H m.m 1 similarly. Thus, have deduced the claim the proposition.
Proposition 4.3. The linear system in Equation 1I has unique solution, which is precisely the CDF values (=))
m- x E X)) of the categorical fixed point.
Proof. By Proposition
4.2 we have that the CDF values F of the categorical fixed-point solve
Equation (TT). Now, let suppose that Fi, F. distinct solutions Equation (11) aiming obtain
a contradiction. Proposition 4.2 also establishes that Equation (1I) equivalent to Equation (TO)
with the additional conditions that (x) 1 for all X, will denote the corresponding
two solutions to Equation (TOp built from F1, F2 (by setting the unspecified (ær, m) coordinates to 1)
by Fi, F2, respectively.
20

---

The idea is now to use the contractivity of the projected operator IIm T in l2, established in Proposition 2.2 to obtain contradiction. Notationally, it is useful phrase things in terms of contractivity of the CDF operator Tp itself. We use the norms lle, and || 12e defined on R and R-xxm in Appendix which we recall here for convenience:
1/2
Σ
F?
for F E R" and || max [[()l,,, for F E Rxm aex
m(1
We therefore have
[Tp(Fi - F2)lle = - |FI -22lez,00
-
However, Proposition B.1 establishes that Tp is contraction on {F Rxx=; Fm 0 for all E ) with respect to || -lle,,00 which contradicts the statement above, as required. O
D Proofs relating to the stochastic CDF Bellman equation
Before giving the proofs, we present version of the SC-CDF Bellman equation in purely distribu- tional form, which will streamline the arguments. This mirrors the development of the form of the distributional Bellman equation given purely in terms of distributions (Rowland et al. 2018), rather than in random variable form as in Bellemare et al. (2017). We define the stochastic categorical CDF Bellman operator Tsco: P
P )dt for each th E (FF) ) by
(Tsce (x) = (((T)))(z))
(yy ~ (y)) independently of py, and extracts the distribution of the input random variable. The purely distributional form of the SC-CDF Bellman equation is then written fixed point condition on P (
where
+ = Tscc 15.
(19)
We also write
1 (V) v = max lllle (v)(2)), (xx))
zEA
for the supremum- -aasserstenss distance over ( with base metric || lle on g.
Proposition 5.7. The SC-CDF Bellman equation in Equation (T6) has a unique solution, in the sense that there is unique distribution for each (xx) such that Equation (16) holds for each X.
Proof. We first show that the operator Tsco a contraction on ( with respect to the metric Wi-lle. Suppose vi, v E ((F)* and let (o()), F'(1)) be an optimal coupling between (h(y) and (l(y)) with respect to Wli-lle, for each y X' (existence of such couplings is guaranteed by Villani (2009 Theorem 4.1)). Then, letting (2r, X') be random transition from æ; independent of $ Φ have that (B, (X'), B. (X"))) is valid coupling of (Tsco (x) and (Tscc v/ (x)). Then we have, using the operator notation defined in Appendix
(a)
"I -I,,,, ((TscC (x), (Tsco v (x))
<
E
BBA((') B,
(b) VTB
[0((X))
('(x')lle
vr
Σ P(u [Q(()) क (v)))
Plylz) (F())
(2), v')
21

---

where (a) follows since (B,)(X), (,'(())) is valid coupling of (Tsce (x) and (Tsec (z), (b) follows by contractivity of B, with respect to || -lle. as shown in Proposition B. and (c) follows since (v)), Φ (v)) was chosen to be an optimal coupling of (b()) and (vy"). Hence we have
W llee (Tscot, Tscot') V (v) b²),
proving contractivity.
The metric space complete [Villani 2009 Theorem 6818). Hence, by Banach's fixed point theorem, we obtain that there is a unique fixed point 3* E ( of Tscc. Thus, a collection of random CDFs (Φ*(x) E X)
w Ille is complete, since the base space (.F. (Ie.) separable and
satisfy the SC-CDF Bellman equation Equation (16) if and only we have d (Ir) o (x)x)) for all x E X.
Proposition 5.8. For all x E X, we have
(x)] = = (x))
Proof. We take expectations on both sides the Lundo-vvrrabee stochastic categorical CDF Bellman equation in Equation (16). yielding:
E B B (x)]
Since To is random linear map, independent of Φ have
El [B[Tg]E[
= To B
This states that E Q R satisfies the standard categorical Bellman equation in Equation and hence E 2 F2, by Proposition 2.2 as required.
Proposition 5.11. We have
2
Σq σQ QQ0q
(m/I-7 m²(1 7)1 7/)))1
where 1 X is vector of ones, and the inequality above is interpreted component-wise.
Proof. We caleulate, writing Φ as Φ F as F. and TQq f to lighten notation:
Σạ(x) (l2( F(x))
@+) (2))
((G) (())) ((G((T)))))
(20)
Here, (a) follows since & satisfies the variale version of the stochastic categorical CDF Bellman equation, (b) is result of Pythagorean identity for squared Cramér distance, which we derive below:
(TT))) F(x)) E
((T))))(() F(x)())² -(()))) at
-()))))))))((( dt,
where the integral switch follows from Fubini's theorem. Now, focusing on the integrand above, can be written as
Ei(Y - E
with = (x),,, since F(x) = [x( -( (x)), by Lemma 5.8 We have
E - [E(T)) = B[(Y - - ][
= =E[(Y + - B
22

---

Now, (z)() |
7 (z) (z)), by linearity, which concludes the
validation of step (b) above. We recognise the first term Equation (20) as the local squared-Cramee variation, and hence have
EΦ(e) = ([(())))[(((())))))))) (TF)
We now explicitly write the evaluation of the application of 7 at coordinate æ terms of the random transition (x, X') used to construct the single-sample random transition matrix Q described in
Definition 5.4 so that we obtain, with the operator notation of Appendix B
(TF)) El (B, BxF(()))
R.
B. (X))), B, $]
Σ
(B,))) (,F())
VEN
(0)
Σ
( F()) a)
VEN
Σ (Ul)))
VEN
Σ (1)) - a
NEN
>(E a
where (a) follows from Proposition B.3
with a = m(1 m(1-7)'/2 + m= meaning that we deduce
Eq(x) - + (xx aa.
required.
Corollary 5.12. We can bound the term ((T-~p)) -P) - appearing in Equation under the
assumptions of Theorem 5.3 as follows:
1 1- S 2
(1(T -oplo 1EEploo +
1-2
Proof. By Proposition 5.11 applied with Q ==
Dp 2 op 7PEp
1,
m((1 - 2/22
m² 7)3
where 1 E R is the vector of ones. We first note that from the condition m 2 45
+1
from the statement of Theorem 5.3 we have
2
1
e²
7//2
3/22
g2
S
<1,
ml1 - 1/2 m² 7)2
4s 7) -2 7)1/2
4e-2
2
since (0, 1). from the statement Theorem 5.3
We therefore have
(I (ー>ン)EE200ー1
Now, (I -P) is monotone operator (in the sense that 11 02 coordinatewise implies U2;
(T P 2 (T - p) -12z; this claim follows as all elements of (7 - P) are non-negative, since it can also be written Σx20 n AD)), we can apply it to both sides of the inequality above to obtain
Ep (1-7)-11 71 (I PP)-opp.
Finally, note that since Σ [x(x) is expected squared- squared-Cramér distance between two CDFs supported interval of length (1 -1 have
Ep(x) (1 -1
from which the claim follows.
23

---

E Proof of Theorem 5.1
We begin by restating the result seek to prove.
Theorem 5.1. Let E E (0, (1)(1-)) 1 and E (0, 1), and suppose the number of categories
satisfies m 4(1 - +
1. Then the output F of model-based DCFP with N = -0(s 21
- polyyllog((111)) samples satisfies, with probability at least
max w1 (17 (I), F())) E
=Ex
We arrange the proof into a sequence of smaller results, in analogy with the presentation of the sketch proof in the main paper.
E.1 Reduction to high probbbbbbability bounds in Cramér distance
Following the sketch provided the main paper, we first restate and prove Lemma 5.2
Lemma 5.2. For any two distributions v. E (10, (1 -7) "1), we have
w1(v,v/) S (1 (1-7)-117) (2--7)
Proof. We begin by writing
wi wi(u,,/) 1-
[F((1) - (t) di (1-7) (1 - 7) (1--)-
(F((1) - (2)
where Fo, Fur are the CDFs of v,v/ respectively. The quantity inside the squared brackets can be interpreted as an expectation (with ranging over the values of a uniform variate on Er-Umir(O..C (1-))(001111)))))) [I For (T)I), and can therefore apply Jensen's inequality with the map + 22. This then yields
(1-1)- (FL(1) - (L) (L)) de
1/2
wi (1, 5(1-7) (1 (1-
-(1 (1-7)) -1/2)((112)
as required.
The first main step of the proof of Theorem 5. is a reduction Theorem 5.3 via Lemma [.. To see this, we first restate Theorem here for convenience.
Theorem 5.3. Let E (0, 1) and å (0, and suppose the number of categories satisfies m > 4(1 - -2 + Then the output F. of model-based DCFP with N (( 7) pooyyoog((2118) samples satisfies, with probability at least 1 - 5,
max E2(1" (x) F(x))) Se. ZEA
(13)
For the proof of the reduction, suppose the statement of Theorem 5.3 holds. Now, let take
as in the assumptions of Theorem 5.1 We then define 5 = (1 //2; note that from the assumption E, we have 3 E (0, 1). Applying the result of Theorem 5.3 therefore obtain that with N i 2(1 -11-)) 7) Poolylog((X1]//1 (8)), we have (with
E (0, (1 (1-7)-1/3) -1/2), and m 4(1 ---
probability at least 1 - ō)
max (2 (x), F()) E zEX
By Lemma 5.2 therefore have (with probability at least 1 8, for all X) - 1/2(1
wh (x), F(x)) S (1 (1- (x), F(x)) < (1 - 7)
//2
which is the desired inequality in Theorem 5.1 Finally, note that the sample complexity term
E 21 - can be rewritten as
2 (1-)) ((1
-
Thus, we obtain the stated sample complexity in Theorem 5.1 and we have established that to prove Thcorem 5.1 it is sufficient to prove Theorem 5.3
24

---

E.2 Reduction to categorical fixed-point error
The first step of the proof Theorem 5.3 is to show that with m taken sufficiently large (as described in the statement of (a.oremm5.3) the Cramér distance between the true return distributions and the categorical fixed points is small, and is therefore sufficient to focus solely on the sample-base based error in estimating the categorical fixed point.
By applying the triangle inequality, we have
22(n" F) - 22(1" FF)) F2)+ F))
(0)
(1 (1-2) 1--m--11 f2(F",
( 2 + (2F"F, F),
where (a) follows from the fixed-point approximation bound in Equation (5). which itself is the result of Proposition 2 of Rowland et al. (2018), and (b) follows from substituting the assumed inequality for
m in the statement of Theorem 5.3 Thus, to establish that (2 F') is bounded by with probability
at least 1 - 6, it suffices to show that
1F(F" F) < e/2.
with probability at least 8, claimed.
E.3 Propagation of local errors
To begin analysing (2 F), we analyse the difference of vectors F' F directly. We proceed in an analogous manner to Azar (2013) in the mean-return case, rearranging as follows:
F F (a) TT TpF
(b) Tpp -T + +TØFF -TpF"
= (1-T))(F (F)= (Tp - F+
(21)
Here, (a) follows since F, F fixed points of Tp. Tp, respectively, (b) follows by adding and subtracting TrF.,, and the implication follows from straightforward rearrangement.
We would next like to rearrange Equation 21 to leave the term F. - 1* on its own. This requires some care, in checking that the operator (I Tp) invertible in an appropriate sense.
Lemma E.1. The operator I - Tp R X is invertible the subspace {F E Rxm (m) for all X), with inverse EB20
Proof. By Proposition B.1 Tp maps {F E TR
F..()) for all X) itself, and is
a contraction on this subspace with respect to l2, with contraction factor v7. It therefore follows that / - Tp maps this subspace to itself, and is invertible on this subspace. Since Σx>0 T% also maps this subspace itself, it follows that this subspace, we have (1 -T- Σx20 T as required.
Now, (Tp (Tp-Tpp) FT E {F E RRxxm (x) for all X), and hence from Equation 21 and Lemma E. it follows that
F-F
Σ (T F*
(22)
K20
The result is that we have expressed the difference in CDFs in terms of local errors (T, F* which are then propagated via the operator Σ120 T
F*
25

---

E,4 Bernstein concentration bounds
The next step in the proof is establish concentration bound for the local errors (Tp -p)p
One potential approach is to use a concentration inequality for each individual coordinate of (T), FT indexed by (z, i). Azar et al. (2013) note that in the mean-return case, using a Hoeffiding style bound is insufficient to obtain optimal dependence of the sample complexity on (1 101-7 7)- and /Zhang al. (2023) also note this in their (non-categorical) distributional analysis. Using a Bernstein concentration inequality on each coordinate and then using a union bound can be made to work, although this then incurs a log(m) factor in the sample complexity. If such dependence on m were tight, this would suggest that the sample complexity depends on m (albeit only logarithmically). and hence there would be a trade-off between picking m sufficiently large so as to obtain low representation approximation eтrоr, as in Section E.2 and taking m low so as to not unduly increase the sample complexity. However, such dependence on m can be avoided by working with concentration inequalities at the level of CDFs themselves. The Dvorestsky- Kiefer-----e----- Wolfowitz (DKW) inequality (Dvoretzky et al. 1956, Massart 1990) behaves analogously to Hoeffding's bound, and is insufficient for obtaining a sample complexity bound with optimal (1 -)- dependence. Instead, we make use of Hilbert space Bernstein-sty -tyle inequality (Yurinsky) 2006; Chatalic et al..
2022), by interpreting the Cramér distance l as a maximum mean discrepancy (Gretton et 2012). which measures distances between distributions via embeddings in a (reproducing kernel) Hilbert space, allowing the Bernstein result to be applied.
First. we precisely describe the connection between Cramér distance and Hilbert space. (2003) shows that for any distributions V E (10, (1 - 7) 'I), we have
Székely
1
B(v, = Ex Ne ||- Yi] 2 =x x'I ErY Y'I
(23)
2
Sejdinovic et al. (2013) show that there exists an affine embedding o: P (10, (1 7) 'I) + H of distributions into Hilbert space H. such that the right-hand side of the equation
also be written as the squared norm (lő(v) <iv')l ? in other words, they show that the Cramér distance is equal to squared maximum mean discrepancy (Gretton et al., 2012). We can then appeal to the following Bernstein-style inequality for Hilbert-space-valued random variables, which is due to Chatalic et al. (022)), and based largely on Yurinsky (2006) Theorem 3.3.4). We state the result here in the usual Bernstein style of assuming almost-sure boundedness of the the random variables concerned, while Chatalic et al. 2022) use the more general formulation of assuming particular bounds on all moments.
Lemma E.2. Chatalic et al. 2022 Lemma E.3) Let Zi..... Zn be i.i.d. mean-zero random variables taking values in a Hilbert space (H, || 11ze) such that Zz 1/ 5 E (0, 1), we have
S M almost surely. Then for any
N
2M log(2/6)
2E Z1 [ log(2/8)
Z
+
N
N
H
with probability at least 1
8.
To apply this result in our setting, first note that, using the notation introduced in Appendix B have
[TT (2x)
=(()(TFF)()))FF(()
= =I1(((TFF))))) (x)) - (F)
1
F FX - 6(F -O(F -
N
(F(F(()))B)) -(F o(F
where (a) follows from the expressions for Tp described Appendix B and (b) follows from the
affineness of the embedding é. Here again, we make use of a slight abuse of notation by writing ((F"(x) for the embedding of the distribution supported on ...,...; 2m), with CDF values F* (x).
and similarly for embeddings of other CDF vectors.
26

---

We can now apply the result above, noting that
|| (F* - = (*())) S ( (1-7)-17
almost surely, and
E - -(F -
(x)) = o
where we write as shorthand for the local squared-Craméé variation at P. introduced in Defini- tion [.55 with the notation op. Hence, from Lemma E.2 and a union bound over the state space X, we obtain that with probability 1 - 8/22, we have for all x E X that
[CTp T())
[o)z
1
S CBI
CB2
(24)
VN
(1 (1-2)//2
where
CB1 = 22 and CB2 = 2 2log(41221/)
E,5
From population empirical Cramér variation
Next, we show that the population local squared-Craméé variation in Equation (24) can be replaced with the corresponding empirical quantity; that is, the local-squared Cramér variation at P. op incurring some additional terms the bound. We start by calculating as follows:
(z) E [(T=F)))) - (TpF (x))2
E
||FT (X B 12 - P B l1?
= B 112, -IPP 12,
where above we use the notation ||=BB] 11, E = for the vector indexed by state, so that the element corresponding state y E X' is (1) B3 1 To relate this quantity to (zz), we add and subtract term, motivated by the derivation applied to standard return variance by Azar et al. (2013) Lemma 5), and rearrange as follows:
(x)
Pa
B
F
B 112, - PPF BT -
+ (2all BT - - BT_ BTll2.
Fa) B -
+ (1 12, |PF
=
+ ((TT)F))) (z) - (+) p)
(Pa P3) B. 12 (..FB. - P.F BT P.F.B + +PF
+E
|| |[T (z) (TPF))()))) P
(P. Pa) F. B. 11, + ((. - (P)- BJJ P.F.B. P.F.B..
+ E [(TTPP)())
(TpF (TpF)(=)
(25)
where < )e, is the inner product on TM corresponding to || lle,, defined by
(F, (F,F),, = m(1 7) 2mm
Σ F.F
27

---

right-hand Focusing on the final term on the right- hand side of Equation 25 we have
E |[TT (z) (T)F)) P
=B [I(TFF"))) (z) - -1 -TppF)()) P]
-E - 1110::00000000--0000-0000-110000
p
=B [[(TFF - (x) - (TpF)
+ +2((T - (p)
- -TpF)tr) -Tpp-TTD(a)).
- ((pp))((11 Pl
+
E[ll - - - -T)F) 211/2
+ - - (TT,F) 1
(Tp (T
/| + vā(a))
where (a) follows from the Cauchy-Schwarz inequality, and we write as shorthand for og. Finally, note that
B (x))(x) (TpF" -TpF)xx)
EII(T (2)))2))1 P
(0)
S 7B x (X)) f(x') P)
1| F Fl1e,
with (a) following from contractivity of Bg in || Ile, as established in Proposition B.2 Bringing these bounds together with Equation 25 we have
(x) $ (Pa - Pa) 12, + ((Pa - P.).BT,, + B.F
(26)
(11 oo + vo(x))
We now apply concentration bounds to each of the first two terms the right-hand side of Equa- tion (26). Lemma E.3, With probability at least - 8/22, we have for all x E X,
P2I BJ 12, -P, PF 12,
< Ch (1-7)VN*
1
where
log (1)
CH
Proof. The expression P=||/*B] II,,, PF B lle, is equal to the negative of the average of N i..d. copies of the random variable |[** (X') 12,, where X' Pz. minus its expectation. Since [F"' (X') _ I1?, is bounded in [0, (1 - 7) "1,, we apply Hoeffding's inequality (for given r E X') to obtain
log log(4 (4)2)100)
Pa Pa B. 11,,
P. ||
(1 VVN
with probability at least 1-6/(2/x1). Taking union bound over X' then yields the statement.
28

---

For the second term, we may re-use the Bernstein concentration bound derived above to deduce the following.
Lemma E.4. Suppose Equation 24 holds.
Then we have, for all a E X,
((P: - P.)F P.a S 2BBB
+ 20n (1 (1-7) (1-
1
(1 (1-
Proof. First, by the Cauchy-Schwarz inequality, we have
((P: ((P-- (P2 P.)F D. s I1((2=
P.)F" -
1
The first of the two factors in the product on the right-hand side is precisely the term bounded on left-hand side of Equation (2).. On the right-hand side, the vector inside the fo-norm has components in [0, 2). so the norm can be straightforwardly bounded by
1/2
2
2
m(1
(1-7)1/2
Combining the inequalities for these two factors, and using the trivial bound vō(z) S S (1-7)) 11/1 we obtain the stated result.
1/2
Combining all these bounds together in Equation (6) and taking a union bound, we conclude that with probability at least - 8/2, for all x E X' we have
+CBB2 N) + (i11F* Flle, + / vā(z))
(x) s (2CB1 + Ch)
(1 - 7)-
We now take square-roots of both sides, using the inequality Va va+bs vā+ + vb on the right-hand side to obtain
Vo(x) S + - Fl 2CBI CH (1 (1-2) A 2CB2 (122)1/2N1/2
Substituting this into Equation 24 and taking union bound then yields that with probability least 1 - 8, we have
[CT-- Tp)
S CBI
1
vo(x)
CBI
- 1F -Fle + C"
VN
VN
(1 1/23311
(27)
where
C' = + CH + 22Cma) + CB2
so that our concentration inequality is expressed in terms of the empirical local gquared-Crance variation ō.
E.6
Converting local bounds to global bounds
To convert the local concentration results obtained above into a bound on (2(F(z), F*(x)), we first prove the following lemma.
Lemma E.S. For U E (F E RXxm Fm(xx) = -Oforr all X}, and any 2 1, we have
((TTU) S h/2 E = x' Xo = x))(U
where Pp denotes state transition probabilities for the MRP with transition matrix P.
29

---

Proof. We proceed by induction. For the base case h = 1, we have, using the operator notation introduced in Appendix B
m-1
I(T;U)())]l+
Σ P(ylx)
x))
Σ
BaU(m))
yEA
j=l
(a)
P()|z))
vE
Σ P(y)o)))1/
1/2 [
yEA
82
as required. where (a) follows from the triangle inequality, and (b) follows from contractivity of B. in || lley, established by Proposition B.2
For the inductive step. we suppose that the result holds for some I N. Now, letting k 1+ we have
||T ()))lle =
Σ
B
E
x Σ P(vlz)
Ba
vEX
( (T) 1
P(yla))7 1/2
VEAt
112
Σ P P = x' X. v))l(()))) Ile,
xeeX
-2112)/2
Σ
P (X = | = y) P(11a)) |1(23')le.
x'EN
VEX
= -2(1+1//2 Σ Pp((1+1 = x 1 = x)|(U<>)11e,,
x'EX'
as required, where (a) follows by the inductive hypothesis.
30

---

We therefore have, with probability at least - 8:
((((()),F'z))
(28)
= [F(()) - (x)lle
Σ TT FT) (2)
k20
(6)
>
[T - -T)) -TD))())))
R20
Px = 1 =))///1
[T, Tp) (1')
#'0
PA(X+= = o' 10 = x)> CBI 1 / Vā(z") + +BI
|F Fllez,0o
VN
R20
1
+
(1
(1-7)1/2N3/4
(e) CBI VN
CBI
0
11( (-P)-1
Volloc
F -Illez.o
(1 V
(1- -
CBI
2CBB
20"
|[[ -
(1 (1-7)VN || -ll
(29)
(1 (1-7)3/2N3/4
Here, (a) follows from Equation (2)). (b) follows from the triangle inequality, and (c) follows from
Lemma E.3 The step (d) is the one step in the derivation that does not hold with probability 1, but
6, and follows from the Bernstein-style bounds in Equation (27 (e) follows from algebraic manipulation using the identity Ex>o(>1/pp) = -I- (FP) as P is
rather with probability
stochastic matrix, and bounding the elements of (I () 'Vo by the L norm of the vector,
and (f) follows from the inequality (1 - -1) S 2(1 - -) for y E [0. 1).
Focusing now on the coefficient et al. (2020) Lemma in the mean-return case, can write
1
Volloo we note that similar to the analysis of Agarwal
(1 - (-VYP))-o - (1 (7) (1 -18)
so that each row of
(1 (T- (アP)-
forms probability distribution, and so we may apply Jensen's incquality with the map 2 + V2, to obtain
(T ( (1- <-
(1-)))-11----------------------- 11 v
(30)
where the inequality above holds coordinate-wise. We thus obtain the inequality
||F - Fole2.00 S
2CHl
2CBB | - Fleesox + (1- 3/22)3/2
2C"
|I(I -VP)- + (1 (1-7)VN
(1 1/N
Note that the term | Flles,oc appears on both sides of the inequality above. By taking N 2 16C2 (1 - -2 made smaller than 1/2, similar to the approach taken in the proof of mustance-dependeer sample complexity bounds for mean-return estimation by Pananjady and Wainwright 2020 Theorem ((a)). Under this assumption, rearrangement yields
ensure that the coefficient in front of || Flle,. on the right-hand side
4CBI
4C"
F - Plles < (1 1/2 |1(
Vo + (1 (1-2)8/2N3/4
(31)
as described in Section 5
31

---

E.7 The stochastic categorical CDF Bellman equation
We now aim to use our developments regarding the stochastic categorical CDF Bellman equation (see Section 5.2) to bound the quantity appearing within the square root. To be able to use the bound obtained in Corollary [5.12 we first replace the factor v inside the square-root on the right-hand side of Equation (30) with " using the bound (7 P)P) -1 2(7 y pP) shown by Agarwal
et al. (2020) Lemma 4). This, together with Corollary 5.12 yields
(I - (U-v (1- v2
2
1
(1 (1-7)1/2 1.
2
Thus, returning to Equation 31 and substituting this bound in, we obtain (again, with probability at least 1 8, for all X):
8CB1
40"
((F()) S
+
(32)
(1 (1-7)VN (1 (1-2)8/2N3/1
E.8
Final steps of the proof of Theorem 5.3
We now let N 2 CUE {1 7) ?. with co any positive number satisfying
co 2 2 and > (C
(33)
Note that the first of these conditions implies the earlier assumption N 160 (1 7) made in order to arrive Equation (3T). We then conclude from Equation (32 that (with probability least
1 - 8, for all X):
8CB1
4C"
(((()),FF (x))
Vco
+
3/4
3/2
8CB1
4C"
+ 3/4
40
with the final inequality following due to the assumption on co making both coefficients of E bounded by 1/4. This concludes the proof of Theorem 5.3
To derive a concrete sample complexity bound, including logarithmic terms, we may bound the terms in Equation (33) as follows; we emphasise that we do not aim to be tight possible in the following analysis, but merely to provide an concrete, valid bound on co- First, we have that the
2 Ch 2 2 (1/
2 log( 4 (1)))
Next, we have (introducing the shorthand L -= log(41 11x1)/5)):
4/3
164/3( < 20 (CBI(V 20BI + CH + 20B2) + CB2
= -V 2021 + VL/2+ v4L) + +2L)
4/3
220(LV2( 202+ 11/2+ 2L) 4/3
= (V2 1 + +
S 1//324
= 200
We therefore conclude, for example, that taking co
21l log /og sufficient.
F
Extensions
Below, we describe several straightforward extensions of Theorem 5.1 that allow us to give guarantees for related algorithms and problems concerning sample-efficient distrbuttooal reinforcement learning.
32

---

F1 Categorical dynamic programming
The core theoretical result, Theorem 5.1 can be used to provide guarantees for the iterative CDP algorithm. Let write F& [Rx for the result of applying the empirical CDP operator mI total & times to an initial estimate Fo e [x(mmm By the existing categorical theory described in Proposition 2.2 (Rowland et al. 2018), we can bound the distance between the CDP estimate F& and the DCFP solution
as
2/12
(Fk P) - (2
(1 (1- 7)/2
Thus, taking k > 2lox(1/9)++3oo((([11010 this error is smaller than e e(1-7) 1//2, Thus, by Lemma 5.2 have
lng(1/y)
(FR, F) S
By the triangle incquality, we therefore have that 1 (F, 17) S then 11 (Fk. 17*) )²==)= assumption k, and we therefore conclude that under the assumptions of Theorem 5.1 the CDP algorithm run with is [lo((1/))+ (og(1/(17)) also has sample complexity f?(= (1 7) 33).
< 2e under this
log(1/7)
F.2 Policy optimisation
Theorem 5.1 can also be straightforwardly adapted obtain sample complexity results for the policy optimisation problem in Markov decision processes, in which the optimal policy must also be learnt from data, prior estimating its return distribution. We describe one way in which this can be made concrete. Suppose we have a finite-state, finite-action MDP with reward function of state, with a unique optimal policy x² with an action gap E E [0, (1 -) meaning that at any state 2, if a E AA are the optimal action and sub-optimal action respectively, then Q" (x)b) Q" (z, 100--2 Then, for example, Theorem of Agarwal et al. 2020 yields that for E (0,1), with N Ó(E-2(1 - 7) log(1 (1/5)) sampled transitions per state- action pair we can correctly identify the optimal policy π² with probability at least 1 - 5. We may then form the MRP corresponding to executing the policy π², and may resample from the transitions used in identifying the optimal policy to compute approximate return distributions for the optimal policy; Theorem [.] and a union bound then guarantee e-accurate approximations Wasserstein distance with high probability.
To understand the role of the unique optimal policy assumption in our discussion above, recall that when there are multiple optimal policies, each generally has distinct collection of return distributions, and so there not a single object that one should hope to approximate. This is not the case with value functions, which are necessarily identical for all optimal policies. As described in earlier work on policy optimisation with distributional RL (Bellemare et al., 2017 2023), this lack of uniqueness of optimal return distributions also leads to interesting theoretical issues when analysing algorithms such as value iteration. However, we expect a result which one obtains an approximation to the return distributions of some optimal policy can also be obtained, even in the non-unique-optimallpoiii setting.
F.3 Stochastic rewards
In our main results, we have assumed that the immediate reward is a deterministic function of state. Results in mean-return RL are also commonly given under this assumption (Azar et al.) 2003), since this simplifies notation and the extension to stochastic rewards often straightforward to obtain (Pananjady and Wainwright 2020). In this section, describe how modify the argument presented in the deterministic-rewaar case to allow for stochastic rewards.
Assumptions. We now consider working instead with reward distribution function R X -> ((0, 1]), so that R(x) is the distribution of immediate rewards received at state æ. The i..d. transitions obtained from state are given by (z, R X where as before X² independently RT iid. R(x), for 1,
iid P(-[z), and
Algorithms. The categorical operator Tp defined in the stochastic reward case in direct analogy with the definition presented in Section 4 in the deterministic reward case. The generalisation is that
33

---

the quantities hi given in Equation (6, which represent the contribution to mass allocated to 2; at due to backing up from atom 2j. are now defined
he (R-
+
for each X and ز, 1,..., m. The quantities H and Tp defined exactly in the deterministic case, from this generalised definition the j. The model-based algorithms, based on the N sampled transitions at each state described above, now also approximate the h from the empirical reward distributions ((x)) = N E
1 Snz, well as P. In particular, we have
N (R, + 772)),
N
from which we define the corresponding estimators fit H.y. We now write the empirical categorical operator as pR to emphasise the dependence on the empirical reward distributions (R(x) : x E x) too, which is then defined as
(2, = ((y) - (1)+
Model-based DCFP in the stochastic reward setting then simply corresponds to solving the linear system built from T p appearing in Equation [TT) based on this empirical categorical operator.
Sample complexity. The masel-based DCFP algorithm also satisfies the sample complexity bound given in Theorem 5.3 (and hence in T5e in the case of stochastic rewards. Only few additional details are required extend the proof in the deterministic reward case the stochastic case, which we sketch below. As before, we write F for the CDF values of the true categorical fixed point and F for the estimated fixed point obtained by running DCFP with the operator I P.R. We also introduce the notation T P: for the categorical operator using estimated transition probabilities P but true reward distributions (R(x)= E X),, and write F for the corresponding categorical fixed point.
Proof modifications. First, when performing the initial reduction to categorical fixed-point error, we invoke the triangle inequality an additional time, to obtain
E2(n", F) S &2 &,
(34)
First term. As the proof in the deterministic case, the first term on the right-hand side O(e) due to the assumption on in the theorem statement; in particular, Proposition 2.2 proven by Row land ct al. (2018) holds in the case stochastic rewards too.
Second term. The second term on the right-hand hand side of Equation 34 is bounded by following the same argument as given in Appendices E.3 E.8 The only part of the argument in these sections that specifically relies on the immediate reward being deterministic, rather than relying on the contractivity of the operator (which also holds in the siochasticrrwwate case (Rowland et al., (018)). is Proposition B.3 which is used in establishing Proposition 5. and Corollary 5.12 We state and prove a version of Proposition B.3 that holds under the weaker assumption of stochastic rewards too.
Proposition F.1. the general case of stochastic immediate rewards described above, we have that for any F, F E E.
1
| B, 2 F F'lE,
4.
m(1 - 22/22
m² - 7)2
Proof. Following the proof of Proposition B.3 we obtain
2
1
BBF-BBFFIH 2 E(R+
m(1 7/122 m² 7)2
where R~ R(X'), and independently, G, are random variables taking values in { pm } with CDF values given by F", respectively, and we write E3 + + R +G') as shorthand for the Cramér distance between the two distributions of these random variables, ER~R(z) vv], ER~R()) ((b ( )/1/ respectively. We then use the characterisation of the Mnhared-Cramér distance
34

---

in Equation 23 to obtain, writing R1, R2 ~ R(x), G1, G2 u. Gú, G22 ~ v/, all independent of
one another.
G(R+,G,
=E|RR + (R2 + G22
2 +
(R2 + 7G2)
2 Bi + + (R2 + +G22)
(a)
2 Bl - G - 2 _ Eih - G + 2 Blnh - -GG2 + +
2
- 4= F -
F' 11 - 4,
required, where (a) follows from boundedness of the rewards, and the triangle inequality.
With Proposition F.I we obtain a weaker version of Corollary 5.12 in the stochastic reward case, yielding || PP oplc < 6(1 -) 7)= which is sufficient to conclude the claimed bound on the second term as described Appendices! E.7 & E.8 specifically yielding that this second is O(E) with probability at least 1 6.
Third term. Finally, the third term the right-hand side of Equation (34 specifically the stochastic-reward case, corresponding to errors in the fixed-point solely due to mis-estimated immediate reward distributions. First, by analogous argument that in Section E.3 we may write
term that emerges
Σ T p (Tp. Tp.R)
k=0
We have
[T p2 (5)
(a) S (R(()))
(8)
log(2 (1/))
2N
with probability at least 5. Here, (a) follows by homogeneity of the Cramér distance (Rowland
et al. 2018, Proof of Proposition 2), and (b) follows from the DKW inequality (Dvoretzky et al. 1956
Massart 1990) and union bound. Then, using contractivity of P. and the triangle inequality,
obtain
log 2
| - 0l.0000 S
1/2
= =ō((-)) (0(117)--N' -(11-)))--111
2N
k=
with probability at least - 5, which Oee) under the choice of N in the theorem statement. In summary, taking union over the events corresponding to all concentration bounds used, obtain that (2("", F) is Oss) with probability least 26, required.
G Further experimental results and details
In this section, we give full details for the experiment presented in the main paper, and additionally present extended results several additional environments. All experiments were run using the Python language (Van Rossum and Drake 2009), and made use of NumPy (Harris et al. 2020). SciPy (Virtanen al., (220)), Matplotlib (Hunter (2077). and Seaborn (Waskom) (022)) libraries. As all experiments are tabular, each run uses single CPU, and timings are reported within the experimental results.
G.1 Environments
We report results on four MRPs defined as follows:
1. Chain: 10 states arranged in chain
< 210- Each state transitions to its neighbours with equal probability. States 1 and 10 are terminal. Only state has reward of 1, and all other states have reward O.
35

---

State
State
State 3
State
State
1,0
0.5
0.0
State 6
State
State 8
State
State 10
1.0
0.5
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0 0.0
0.5
1.0
0.0
0.5
1.0
(a) Chain
State
State
State 3
State 4
State 5
0.25
(b) Low random
State
State
State
State
State
0.05
0.00
(c) High random
State
State
0.05
0.00
2
3
6
2
3
(d) Two-state
Figure 4: Monte Carlo approximations of return distributions in each of the four environments tested.
2. Low random: There are five states, and the transition probabilities from each state to all five states drawn independently from Dirichlet distribution with concentration parame- ter (0.01,.... 0.01). The rewards for each state is drawn i.i.d. from the uniform distribution over [0, 1]. We draw these transition prohability and rewards using the same random seed to yield the same MRP for all experiments.
3. High random: Same the low random environment, but with transitions drawn from Dirichlet distribution with concentration parameter (10,..., 10).
4. Two-state: 2-state MRP modified from Example 6.5 of Rowland et al. (2024). The transition matrix is
0.6 0.4
0.8 0.2)
where the ji entry is the probability of transitioning from state i to state j. State has reward O, and state 2 has reward
We chose this set of environments to include classic tabular settings such as the chain, two envi- ronments with very different levels of stochasticity (low and high random), well the two-state example from Rowland et al. (2024), which illustrates the nature of fixed-poing approximation er- ror incurred by QDP environments with tight bootstrapping loops. We vary the discount factor {0.8. 0.9, 0.95, 0.99]. For reference, we plot Monte Carlo approximations of return distributions for each environment Figure 4 with 7 = 0.9. The return samples are computed using the first-visit Monte Carlo algorithm. For each initial state, we run at least 7' transitions such that the maximum error after T transitions (TTruxx(( 7) < E, where Tmax and we set g = 10-4, This is repeated 10d times for each state, giving least 10 return samples per state.
36

---

G.2 Algorithms
As described in the main paper, we compare DCFP, QDP, and CDP algorithms. For DCFP and CDP, we make use of the sparse structure of the update matrix (see Equation (9), and implement these algorithms using SciPy's sparse matrix-vector multiplication and linear system solution methods
(Virtanen et al., 2020). The mathematical details regarding the sparsity of the update matrix is
described in Appendix G.3 For comparison, we also run implementations of DCFP and CDP that do not exploit this sparsity, and instead use standard dense NumPy implementations for matrix- vector multiplication and linear system solution methods; these are denoted by d-DCFP and d-CDP, respectively.
By default the categorical methods (all variants of DCFP and CDP) use the atom locations described in the main paper: 2; = - (1 - - for = 1.... m. These atom locations are sufficient to establish the theoretical results in the main paper, though we note that in practice, particularly when true return distributions are localised within a small sub-interval of (1- ---->>___ "J, this setting can be overly conservative. However, in many cases, there are straightforward ways in which the choice of support can be improved, essentially by replacing the uniformly valid return range [0, (1 >] with an priori known environment-speciffc reward range. As an example, if the known range of immediate rewards forms a sub-interval [rmin: mmax] S [0. 1), then the return must fall into the environment-specifne
reward range [T'in (1 7) 7)- .mmax (1 - 7)_ "1 and improved atom .mmx
m--1 (Fmax Tmmia - 7) (for i = I...,, m) can be used, whilst maintaining the guarantee that the support for the true return distributions lie within this interval. We thus additionally report results for versions of DCFP and CDP that
locations 2 ''iin (1 - 7)
11 + 1-1
environment-specine atom locations, to investigate what practical improvements can be gained with such additional information. Specifically, for the high random and low random environments, we use the tighter bounds on support given by known Ymin and T'ax, as described above. In the chain environment, we consider using the advanced knowledge that only the transition into the terminal state is rewarding, yielding sub-interval for possible returns of [o, 1]. In the two-state case, the range of possible returns is the full interval [0, (1 - "), and so we do not investigate an wwtrrommentsspecffiie set of atoms in this case. However, it is clear from the Monte Carlo approximations to the return distributions in this environment in Figure 4 that these distributions do have the vast majority of their probability mass in a much smaller interval than the worst case interval [0, (1 -))
We ran all DP methods with 30,000 iterations. For all categorical DP methods, we verify that the Wasserstein distance between the last iterate to the Monte Carlo ground truth is almost identical to the Wasserstein distance between the categorical fixed-point to the ground truth.
G.3 Efficient implementation of CDP and DCFP
straightforward implementation of CDP and DCFP is to vectorise F, treating it as a vector indexed by state-index pairs i)) and correspondingly treat Tpp as matrix whose rows and columns are indexed similarly. Iterations of CDP can then be performed by simple matrix-vector multiplications with this representation of T p. and the solution of the linear system appearing in Equation (TT) that forms the core of DCFP can be implemented with a call to a standard linear system solver, such as numpy. .in.lg.......................... solve (Harris et al. 2020)).
However, the operator Túp typically has some specific structure that can be cxploited in implementa- tions. In particular, we highlight here the sparse structure of p. allowing for potential speed-ups in implementation by making use of sparse linear solvers, such as (Virtanen et al., 2020). Recall that, viewing Tр matrix (c.f. Equation (9), the element corre- sponding to the (x, 2) row and (v,j) column is given by
linalg.spsee...... spsolve lve
P(y)))HH H1,+1).
(35)
In MRP with sparse transition matrix P, the empirical estimate f inherits this sparsity, inducing sparsity in Tp.. In addition, there is also sparsity induced by the atom index components of the row and column indices, as we now describe. Recall from Equation (8) that HE, Eis (u(((x) + +7z); see Figure 5 for an illustration the function 2 Σis hz(2).
Now, suppose that for some E X and 1 S ú,j m 1, Hy H is non-zero. This says that the function 2 + Eis, h2(2) takes different values at r(z) + yzy and r(æ) + 72/+1, and since
37

---

r(x) + 723
r(x) + Y
Zi
Zi+1
2++2
Figure 5: The function = + Eisi (z() (grey), and a possible configuration for r(x) + YRj.
r(z) + 72j+1 in the event of a non-zero H H H11 term.
the distance between these arguments is y times the grid width (1 - -7)
mmm -
it follows that at least
one of these two arguments must lie in the interval [2s, 2441]; see Figure 5
From this observation, we deduce two forms of sparsity for the elements given in Equation (35) First, since one of r(x) +22 and r(z) + +22++ must lie in [2,, 24- 1]. we deduce that neither of these points can lie in any interval [zy, 22+1] with i' <i- or > i + 1, and hence H H = 0 for such i'. From this reasoning, it follows that ranging over i in Equation (35). there are at most 2 non-zero elements. For large m. this means that Tp is very sparse.
Similarly, we can deduce row sparsity by noting that at most [2/7] indices j can have the property that rxx) + +22 or ((x) + +22++1 can lie in [2 2,-1], and it is only for these indices j that we can have H Hf,5+1 non-zero. So for y ze 1, this also implies sparsity of the elements in Equation (35) as we range over j.
G.4 Results
For each setting, we repeat the experiment 30 times with different sampled transitions. We display Suuremun-Wassersteein
trade-off plots (supremum- Wasserstein- error and wallclock time) for each of the four environments described above, in Figures 6 8 and 9 for the cases m E {30, 100, 300, 1000} atoms and using N 10th sample transitions from each state to estimate transition matrix. These curves are obtained by averaging across the 30 repetitions. Some central themes emerge from the results.
In the case when the categorical methods use the environment enviromment-indeeendenn standard atom locations 1-1
2, m- (1 (1-7)-1 for i 1,..., m, we find that for given atom count m, QDP often achieves the lowest asymptotic Wasserstein-1 I error. A notable exception is the two-state environment; Rowland et al. (2024) observed that such environments, in which there is a short, nigh-probabiliig path from state to itself, can cause high upproximation error for QDP. which we believe is the cause of the inaccuracy observed here, particularly at high discounts. However, the categorical approaches often deliver better performance as judged by wallclock time. This owing to the efficient implementations of the linear operator, and solution methods for the linear system, that are associated with these algorithms. By contrast, the QDP operator non- and requires a call sorting method. We tend to observe the greatest benefits of the sparse implementation for large atom counts, as expected, and also observe the greatest benefits DCFP over the iterative CDP algorithm settings with high discount factors. This is also to be expected, since the discount factor controls the rate of convergence of the DP algorithms. Note additionally that the use of environment--pecifie atom locations generally improves the results obtained with categorical approaches; the improvements in the chain environment are particularly strong, since the narrow return range allows for a denser packing of atom locations, by a factor of (1- for given atom count m.
We also compare how the supremum- Wasserstein distance changes while the number of samples used to estimate P increases from 10a to 105, Since the categorical methods all converge to the same fixed point, we only show the result of DCFP and QDP, using E {30, 100, 300, 1000) atoms. Figure a0 a) shows that the distance decreases as the number of samples increase. One exception the QDP with 30 atoms applied to high random at high 7, where N = 100 produced smallest distance on average. Figure b) shows the results when the environment-specific return range is given. This substantially reduced the supremum- Saaremum-Wassersted Wasserstein distance, especially when using a small number of
atoms.
38

---

y=0.8
V=
y=
y=0.99
109
10-2
109
m-100
10
10
Method
d-DCFF
DCFP
109
m-300
CDP
10.1
QDP
10
Wassersteln
109
10-1
10-2
109
Wallclock time / sec
(a) Global return range
y=0.8
y=0.9
y=0,95
y=0.99
m- 30
10-1
109
m-100
1
10-2
Methad
d-DCFF
DCFP
109
CDP
m-300
QDP
10-2
10a
m- Wassersteln
10-2
1da
1
100
10a
10-3
10a
Wallclock time / sec
(b) Environment--pecifcc return range
Figure 6: Distance vs. run time results for the chain environment, for P estimated from N 106 sample transitions.
39

---

y=0.8
y=0.9
y=
y=0.99
104
m-30
10-1
101
m-100
10-4
Method
d-DCFF
DCFP
10
CDP
m-300
QDP
10-1
Wassersteln
101
10
10a
10a
10-3
Wallclock time / sec
(a) Global return range
y=0.8
0=0.9
y=0,95
y=0.99
10
m- 30
10
10a
m-100
10-1
Method
d-DCFF
DCFP
10a
CDP
m-300
QDP
10-4
Wassersteln
10
m-
1-
10a
10
10-3
Wallclock time / sec
(b) Enviromment-specffc return range
Figure 7: Distance vs. run time results for the low random environment, for P estimated from N 10 sample transitions.
40

---

y=0.8
y=0.9
y=
y=0.99
104
m-30
10-1
101
m-100
10-4
Method
d-DCFF
DCFP
104
CDP
m-300
QDP
Wassersteln
101
10
Wallclock time / sec
(a) Global return range
y=0.8
0=0.9
y=0,95
y=0.9
104
m- 30
10-1
101
m-100
10-4
Methad
d-DCFF
DCFP
10
CDP
m-300
QDP
m- Wassersteln
101
10a
10a
10
Wallclock time / sec
(b) Enviromment-specffc return range
Figure 8: Distance vs. run time results for the high random environment, for P estimated from N 10 sample transitions.
41

---

y=0.8
y=0.9
y=
y=
10
10a
10-1
10a
m-100
109
10
d-DCFP
DCFP
d-CDP
10.
m-300
100
QDP
10-1
Wasserstein
10t
1000
10a
10a
10
1000
10
10a
107
1.0
10
10'
Wallclock time
Figure 9: Distance vs. run time results for the two-state environment, for P estimated from N 106 sample transitions. The return range of this environment coincides with the global return range.
42

---

0=0.8
y-
Y
10a
Chain
10-4
DCFP
random
10a
QDP
10-3
10-2
100
1000
High random
10c
10-4
10
10
10
10a 10?
10
Global return range
y=0.8
y=0.9
y=0.95
y= 0.99
10a
10-1
10-9.
Methed
random
10a
10-3
Atoms
10-2
100
300
High random
10a
10-2
10a
10a
10
10a
10a
10a
102
10
10a
10a
10
10c
(b) Envinomment-specfee return range
Figure 10: Supremum- Wasserstein distance convergence. Error envelope indicates 95% confidence interval by bootstrapping. The return range of the two-state environment coincides with the global return range.
43

---

G.5 Example implementations
We provide example implementations for the CDP and DCFP algorithms described the main paper. We work with updates over cumulative distribution functions, in Equation ( matching the form of the operator analysed in the paper. Our intention is to provide straightforward implementation of these key algorithms, and in particular we have not included optimisations such as exploiting sparsity in the linear solver, as described in Appendix G.3 As shown in our experimental results, there are often significant performance gains that can be obtained from such optimisations, and we encourage practitioners make use of such optimisations when running DCFP beyond the very smallest instances.
Listing D gives an implementation of the categorical Bellman operator itself, Listing 2 gives imple- mentation of a single CDP update, Listing 3 gives implementation of the DCFP algorithm in terms of CDFs, and Listing 4 provides a function for converting computed CDF values corresponding probability mass functions. The code snippets in this paper have the licence below.
Copyright 2025 Google Ltc. sPoI tiox-ticenae- Ideatifier: Apache-24
import zuzpy
up. ndarray, support: ap.ndarray) np. ndarray::::
"*"Coxpuee the values of the catogorical hat functions described in Section 2.3.
values: Values to projested, arbitrary shape.
support: Distribation support to project onto, of shape (n.atons,)..
Returns: Probabilities support projected values, shape shape, n_atoxs).
support_dit diff - support(1] support
prej_left (support values(.... Sene]) support__iis
prej_left ap. concateaatee(ppoo.... support [1) support,
like (values) (..., Nene)), axie=-1) (iseeueees
proj.rigat (values(..., None]
proj_right concateaaae ((ap like (values) [..., Sone], proj_zight), axisn-1) proj.right),
retura up maxinux www..iiiinuxpoooooo..
def construct_ ctt.gorcargrrraa.......o.. wet.ccaibelllaaalee operator)
trancition.probb. up ap.rcarryy, revards: ap.ndarray, discount: flest, support: adarray
adardyy **Constructe catogorical
operator.
Azgu:
traasitios.probs Transition probabilities, rewarde: Revard vector, of chape (a_etater,), diacouat: Discount factor.
a_statee).
support: Vector of aton locations for categorical discriburiozs, shape (a_atons, ,).
Returns:
The categorical Bellnan operator, as array with shape (n.statez, B_slosa, a_states, a_stoss).
*..
Indices:
bootstrap_returne savards (:,, Kene] discount aupport [Noss, :]
* Indices: (x,
functionn (wwotstrrap.rtuurrree.
support)
Indicea: (x. .J)
b up. tra.raaapeeeet (h, (0, 2, 1))
B cp..cusuun axia-1)
B_diff -ip. appead=O., exis-2)
return up. np.eiasux( ,xy xij- transition_probs. H_dist)
Listing 1: Construction of the categorical Bellman operator over CDFs.
def seplly_asegarieee bellmas, Sacimas__pppe opwratorrtttooooooooooo (operator: up ndarray, F: up..p-p ndarray) >> ap. ndarray: return wissus ( ("xiy) operator, F)
Listing 2: A single step of CDP implemented by application of the categorical Bellman operator.
44

---

tranaition.proba. ug- ndanday revarda: op.ndarray, flaat, support: ap.ndarray
-> sp.adarray:
h_states a cransitton. ptw.o....r.aaa.............. shape
spgpet..ahape(00
aperator construci, wtt.rreccee...aaaeriiee.... belleea pollaan_ogerater (transition, sevards, discount, support)
t Constructing the linear syssen is Equation (11)
H_tilde wgg.aunnogaaraatorotoo (operrattor :-1, -1], axix--1)
flti [_tilde.flatteur)
T_tilde aperatorii, :-l, :. :-1]
a_din n_atates (n_ators 1)
1_tilde_flat app.reeaape (7_tiié, (a_din, s_din))
Solving sp..lineg..oolvwdn (np.aool.e (up.aye(a_dia)
aystes Rquation (11)
f_ti F_rilde ap.reshape (Fiiappe (a.states, a_atons 1))
-lat (,tilde_
# batacck(F..tldd (p.onnx((n_atatar 1))))
Listing 3: example implementation of DCFP.
def odf_to_ppt cedf_ralues: up ndarray) mp. ndarray:
return di.diff (adf_valuse asiae-1, prepand-O.)
Listing 4: Conversion of computed CDF values probability masses.
It is also possible to implement the operator take probability rather than CDF values, as in Equation (6); this essentially amounts to a change basis, and be achieved by using rather than H_ didif in the np einsum call in the return line of the construct, _cffcctc___..ccc http:_ccccca_aooeeaeeee lct_tecrttccaaa saraaabelllaaae bellman, bel.lmm_opppratorr function; see Listing For implementations of CDP, this implementation of the operator may be preferable for its interpretability, more closely aligning with implementations of categorical temporal-difference learning (Bellemare et al.) (2117) that are implemented terms of probability mass functions, such in the RLax library (DeepMind et al. (020). It is also possible to implement DCFP with the probability mass function version of the operator. In comparison with the implementation Listing the normalisation constraints to be added (as discussed in Section 4.1 are no longer constraints on individual variables of the form F (m) I, but rather take on the form Σ Pz( =
function values as input
construct_paff ptt.epaaa_caaeirieeeeeeee bellsan belsssu.
operator(
tranxition.prohx: ndarray,,,, revards: .darray discount: flast. support: ap.ndarray
>> sp.ndarray:
Indices:
hootatrap_raturne rauArda(:, Kane] diacount. support [Nome,
# Indices: (s. 1)
fasctonctions (bortstraprreturns support)
("sy,8j1->23)) sjj1-21yy'' traxiicio_.praba
Listing 5: An implementation the categorical Bellman operator over probability mass functions.
45

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction describe the main contributions of the paper, and contextualise these against previous contributions in the literature.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the
results can be expected to generalize to other setrings. It fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: The paper discusses detail the assumptions required for the theoretical results to hold, such as using tabular representations and generative model for sampled transitions. The experimental results also make clear the relative performance of the analysed DCFP algorithm. and other approaches to model- based distributional RL, such as quantile dynamic programming, which is not analysed in the paper. Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed the paper.
authors are encouraged create a separate "Limitations" section in their paper. pa
The paper should point out any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-ten system might not be used reliably to
provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how
they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address
problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't I acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes]
Justification: All theoretical results are stated with precise assumptions required, and full proofs are provided in the appendix. We additionally provide a proof sketch for the main result in the paper, to aid readers with intuition. Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced. All assumptions should be clearly stated or referenced in the statement of any theorems.
46

---

The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided appendix or supplemental material. Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide precise algorithm descriptions, as well as all hyperparameters used. In addition, we discuss specific details for obtaining efficient implementations, such as exploiting sparsity, in the appendices.
Guidelines:
The answer NA means that the paper does not include experiments,
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g.. with an open-source dataset instructions for how construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed- source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material? Answer: [No]
Justification: We have described our tabular algorithms and hyperparameters so that the experi- ments can be reproduced solely within a numerical programming framework such as NumPy; we also provide example NumPy implementations of the key functions required to implement the proposed algorithms. Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nipp ceciips.cc/publli public/
cc/
guides/
guides/CodeSuumin Ctt////oeeuubeeec ssionPolicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to repro- duce the results. See the NeurIPS code and data submission guidelines (https://nips... ec/
for more details.
47

---

The authors should provide instructions on data access and preparation, including how to access
the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which
ones are omitted from the script and why. At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Detaile
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Our experiments are tabular, and we describe all implementation details and hyperpa- rameters in the appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix,
supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report Wasserstein distances along with 95% bootstrapped confidence intervals in experimental results. As our experiments are tabular, we are able to run many repetitions of the results with distinct random seeds, meaning that these confidence intervals are made extremely narrow.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main
claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given
experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to
library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors). It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma eтror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experi- ments?
Answer: [Yes]
Justification: We describe in the appendices that since our experiments are tabular, each run uses a single CPU, and run times are reported within the experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimen- tal runs well as estimate the total compute.
48

---

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:// httperriippppppppppepeeemm httppb/c///hcccccceeeeeeeeeeeeee
Answer: [Yes]
Justification: We comply with code of ethics from the conference. Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics. If the authors answer No, they should explain the special circumstances that require a deviation
from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The contributions are primarily theoretical and foundational, and do not lead to immediate societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed. If the authors answer NA No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. The conference expects that many papers will be foundational research and not tied par- ticular applications, let alone deployments. However, there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI...
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The paper concerns solely tabular methods. Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary
safeguards to allow for controlled use of the model, for example by requiring that users adhere
to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should
describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this. but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA]
49

---

Justification: Our experiments are solely tabular, and we cite the programming frameworks used
implement these.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if
include a URL.
The name of the license (c.g.. CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service of
that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should
be provided. For popular datasets, paperswi thcode. com/ hom/dataeets) has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. For existing datasets that are re-packaged, both the original license and the license the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged reach out to the asset's
creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets introduced. Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/code/model as part of their sub-
missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects Question: For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human
subjects.
Including this information in the supplemental material is fine, but the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Sub- jects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country which research conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
50


Reviews:
[{"review_id": "oIfAzKIQaN", "reviewer": "Reviewer_1p2t", "summary": "This paper proposes a min-max optimal algorithm for model-based distributional RL, which is used to approximate the return distributions under the assumption of having access to generative model. New theoretical analysis is provided for categorical approaches in distributional RL, with the introduction of a new distributional Bellman equation and the stochastic categorical CDF Bellman equation. Empirical studies are performed for benchmark comparison with other model-based distributional RL algorithms.", "questions": "1. In the studied setting, how does \"policy\" and \"actions\" come into play when evaluating return / value functions? Do you consider fixed policy? \n\n2. In Section 2, it seems that the randomness being considered comes from transition, reward function and initial state distribution. How about the randomness in policy?", "limitations": "This is a theoretical work, no potential negative social impact", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The main contribution of the paper is the theoretical results for distributional RL by assuming the access of a generative model, where the goal is to estimate the full distribution of returns at each state instead of just the expected returns. In particular, this work exhibits several promising results:\n\n1. Authors establish theoretical guarantees for the proposed distributional RL algorithm, i.e. the direct categorical fixed-point algorithm (DCFP), which directly computes the fixed point of CDP (an existing distributional RL algorithm). It achieves a sample complexity that matches with the lower bound up to logarithmic factors when estimating return distribution in Wasserstein distance.\n\n2. Theoretically, authors introduces a new distributional Bellman equation and the stochastic categorical CDF Bellman equation for categorical distributional RL.", "weaknesses": "While this paper provides insights into the distributional RL, there are still potential improvements that authors can further consider:\n\n1. It is unclear about the motivation of considering categorical approaches to distributional RL and how the proposed algorithm excels existing distributional RL approaches throughout the text. While it makes sense to choose a tractable representation of probability distributions for approximation, the representation considered as in line 81 - 88 appears to be simple and may have limited expressiveness in practice. \n\n2. The convergence result provided in Proposition 2 does not categorize the convergence rate, which can be essential for computational efficiency.\n\n3. While the theoretical results are interesting, the practical applicability is not quite clear. Numerical results are presented on a 5-state simple tabular environment. It will be beneficial to include a practical example to better illustrate how the proposed algorithm can be particularly useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback on our work. We believe we address all concerns raised in the review, and would welcome any further questions.\n\n**Weakness 1: Categorical approximations.** These approximations are well established in distributional reinforcement learning. They have been applied successfully in a variety of practical settings, such as in deep reinforcement learning in simulation (Bellemare et al., 2017), environments based on clinical decision-making (Böck & Heirzinger, 2022), and robotics (Haarnoja et al., 2024). Further, Proposition 2.2 (due to Rowland et al., 2018), and specifically the bound in Equation (5), establishes theoretically that these parametrizations can learn arbitrarily accurate approximations to the true return distributions, with error decaying as $O(1/\\sqrt{m})$ as measured by Cramér distance, so there is both empirical and theoretical support that the categorical approximations used are sufficiently rich.\n\n**Weakness 2: Convergence result.** We don't know what the reviewer is referring to by \"Proposition 2\". If Proposition 2.2 is meant, then please see the comment above regarding the fact that this bound *does* characterize the convergence rate of the approximate distribution to the true return distribution as a function of m. Note that Theorem 5.1, our main theoretical result, also characterizes the convergence rate as a function of number of samples.\n\n**Weakness 3: Applicability.** We have included a small suite of experiments on environments in the appendix, varying qualities of the environment such as levels of stochasticity, to exhibit the method in a range of settings. Having said this, this is essentially a theoretical paper, and our core contribution is to show the (perhaps surprising) result that, in the sense described in the paper, no more samples are needed to estimate full return distributions accurately compared to just their means. We expect many of the tools developed for our analyses (such as the linear structure of the categorical Bellman operator, the stochastic CDF Bellman equation, the availability of a closed-form solution to the categorical Bellman equation) will be useful to the RL community in general in developing implementations and analyses of categorical-based algorithms, and we have already seen our techniques applied in analyzing temporal-difference learning versions of categorical algorithms.\n\n**Question 1: Actions.** The reviewer is right, we consider a fixed policy, and this allows us to consider a Markov reward process (rather than a Markov decision process) as described in Section 2, and avoid including actions in our notation, which leads to more concise expressions throughout the paper. This is a common approach when describing algorithms for evaluation of a fixed policy in reinforcement learning, see for example Pananjady & Wainwright (2020) for a recent example, and Section 2.3 of Szepesvári (2010) for discussion in the context of a textbook. \n\nIn more detail, given an MDP with transition probabilities $p : \\mathcal{X} \\times \\mathcal{A} \\rightarrow \\mathscr{P}(\\mathcal{X})$ and a fixed policy $\\pi : \\mathcal{X} \\rightarrow \\mathscr{P}(\\mathcal{A})$, we can define a corresponding Markov reward process where the transition probabilities $P : \\mathcal{X} \\rightarrow \\mathscr{P}(\\mathcal{X})$ are defined by $P(x'|x) = \\sum_a \\pi(a|x) p(x'|x,a)$, so that the randomness owing to action selection is \"folded into\" the transition probabilities of the MRP. Reward distributions for the MRP are constructed similarly.\n\nWe would be happy to add discussion along these lines to the final version of the paper, please let us know if you have any further questions on this point.\n\n**Question 2: Randomness.** As described in the answer to the point above, the randomness in the policy is folded into randomness in the next-state and reward distributions in defining the Markov reward process.\n\n**Additional references**\n\nHaaronja et al. (2024). *Learning Agile Soccer Skills for a Bipedal Robot with Deep Reinforcement Learning.*\n\nSzepesvári (2010). *Algorithms for Reinforcement Learning.*", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Dnz6TLhWV8", "reviewer": "Reviewer_HAEX", "summary": "The authors propose a new algorithm for distributional reinforcement learning under the generative model setting with categorical representation. New upper bound for sample complexity is presented for the proposed algorithm. Some empirical results comparing the method with other alternatives are presented, showcasing its benefits over the alternatives.", "questions": "How about the all-gaussian setting where every reward distribution is a gaussian?", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main algorithm as well as the sample complexity results are solid contribution to the reinforcement learning community. Although I did not check the proofs, the results seem believable.", "weaknesses": "The generative model setting is rather limited. The results do not directly contribute towards learning a better policy, but I assume some of the tools presented will find their uses in many scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and positive assessment of our work.\n\nOn the significance of the work, we believe many of the tools we have developed in the paper (such as the interpretation of the categorical operator as a particular linear map, the stochastic categorical CDF Bellman equation, the availability of a closed-form solution to the categorical Bellman equation) will play a key role in implementations and analysis of categorical distributional RL in general. We have already seen our work used in this way to develop analyses for temporal-difference algorithms for distributional RL based on categorical parameterisations.\n\n**Gaussian setting.** This is a great question. The core algorithms described in the paper can be applied in this setting too. The algorithms, and the ways in which the analysis can be applied in this more general setting, are described in Section F.3 (referred to just before Section 5.1 in the main text). The central idea is that the variables $h^x_{ij}$ are now defined as expectations under these reward distributions (see Line 1054), which can be accurately approximated via standard numerical integration libraries such as scipy.integrate; this applies to many different classes of distributions, not only Gaussians.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TCYOCzNNH1", "reviewer": "Reviewer_Wuoe", "summary": "The paper presents a novel algorithm for model-based distributional RL and establishes that it achieves near-minimax-optimal performance for approximating return distributions in the generative model regime. This is a significant contribution, showing that the distributional RL is sample-efficient with a generative model, and being the first work providing practical implementation guidance to achieve this.", "questions": "Q1: The algorithm operates under the generative model assumption, where it can sample from any state-action pair. Since I am not particularly familiar with this field, could you elaborate on the justifiability of this assumption?  In my limited knowledge, this assumption seems quite far away from reality, and discussing its limitations and applicability would be beneficial.\n\nQ2 The proposed DCFP algorithm aims to align theoretical analysis with practical implementation and can be seen as a modified category-based distributional RL algorithm. However, However, is the sacrifice in computational efficiency deserved?  It would be better to discuss the scalability of the proposed DCFP algorithm in real-world, large-scale environments.\n\nQ3: The current algorithm is model-based and solving this algorithm relies on constructing the empirical transition probabilities. Does this suggest that DCFP is not easily applicable in the model-free setting? In other words,  does the idea of directly learning the cumulative probability mass rather than discrete probability mass face additional challenges in a model-free setting?  Furthermore, if sample-efficient distributional RL can only be achieved under a generative model assumption?", "limitations": "Please refer to the question and weakness part.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper introduces the new algorithm DCFP for distributional RL, which directly approximates the cumulative probability.  The development of the DCFP algorithm is a key contribution. The DCFP algorithm effectively leverages a generative model to achieve efficient performance. \n\nThe authors show that DCFP achieves minimax lower bounds for sample complexity in Wasserstein distance, thus addressing a question raised by Zhang et al. (2023).  Besides, the paper's theoretical contributions, such as the introduction of the stochastic categorical CDF Bellman equation and the detailed derivations of the minimax-optimal sample complexity, are significant but may be challenging for readers without a strong background. I believe this work provides a deeper understanding of the field of distributional RL.", "weaknesses": "I have some concerns about the experiments. The authors claim that DCFP is generally faster than QDP. I feel this conclusion is primarily due to computationally costly sort operations in QDP. I am also curious about whether both DCFP and QDP are generally lower than CDP, as DCFP also involves costly matrix inverse operations. It would be better to show some results of CDP in the main paper. Additionally, the authors mentioned  $T_P$ has some sparse structure allowing for potential speed-ups and making use of sparse linear solvers from the scipy package in implementation. If solving this linear system using these solvers is stable and whether the results remain reliable as the state space scales up.  It would be helpful to discuss these.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and positive assessment of our work. We are glad to hear you believe this work provides a deeper understanding of distributional RL. \n\n**Experiments.**\n * *Comparison with QDP.* The reviewer is correct that sorting in the QDP algorithm contributes to its higher computational complexity (and higher wall-clock time in experiments). One of the takeaways from the paper is to highlight the nice linear structure of categorical updates in contrast to the structure of the QDP update. We discuss this in Section G.4, and would be happy to add further discussion if there are any further points the reviewer would like to see highlighted.\n * *CDP results.* Our paper is primarily theoretical (focusing on sample complexity, and key tools for better understanding categorical distributional RL), and the main paper experiments are intended to provide a brief illustration of implementations of the algorithms described in the main paper. However, to complement the main theoretical contributions of the paper, we do include detailed comparisons against CDP in the appendix.\n * *Sparse structure and implementation details.* We discuss these aspects in more detail in Section G.3, and speaking generally have found the implementations to be very stable in practice, as they rely on straightforward linear algebra. Please let us know if you have any further questions on this aspect.\n\n**Question 1.**\n * The generative model setting is a standard assumption in statistical analysis of reinforcement learning algorithms. Some references included in the paper that study RL algorithms under this assumption include (Kearns et al., 2002; Kakade, 2003; Azar et al., 2013; Sidford et al., 2018; Agarwal et al., 2020; Pananjady & Wainwright, 2020; Li et al., 2020).\n * For further background, the generative model setting can be thought of as removing the exploration aspect of the RL problem, and allows a clear framing of the purely statistical question as to how many samples a given algorithm needs to compute accurate estimators. This is the approach taken in many analyses of RL algorithms (see references above), and insights in the generative model setting often reveal techniques in tackling other settings too. In our case, our analysis contributes key insights (such as the linear structure of the categorical operator, and the stochastic categorical CDF Bellman equation) which will be useful in improving analyses and algorithms in distributional RL. In fact, we have already seen these ideas being picked up and applied in analyzing temporal-difference learning distributional RL algorithms.\n\nWe would be happy to expand the discussion in the final version of the paper, please let us know if there are any further questions you have on this topic.\n\n**Question 2.**\n\nDCFP has cubic complexity in the state space, in common with approaches in non-distributional RL that aim to directly compute value functions by solving linear systems of equations. However, we emphasize that our general analysis approach is also applicable to standard categorical DP (with a sufficient number of updates), and so a sacrifice in computational efficiency is not necessarily required. We discuss extensions of the analysis to CDP in Appendix F.1, and provide an empirical comparison in Appendix G. In general, the relative computational merits of DCFP and CDP will depend on several factors such as state space size, discount factor, and any exploitable structure in the operator such as sparsity. Our further experimental results in Appendix G give a sense of these trade-offs.\n\n**Question 3.**\n * DCFP is intrinsically model-based. However, this is orthogonal to the idea of learning cumulative probabilities and discrete probabilities, which can be translated between as described in Section 4.1. We find the cumulative probability formulation neater in obtaining a linear system with a unique solution, which is why we settle on this presentation in Section 4.1. We hope this clears up the comment contrasting cumulative mass and discrete mass, but let us know if anything remains unclear.\n * To the best of our knowledge, there is indeed a gap in the sample complexity obtainable under the generative model assumption, and under e.g. online interaction. However, we emphasize that the technical tools we contribute here (linear structure of categorical operator, stochastic CDF Bellman equation) are already useful in analyzing distributional RL algorithms under other modeling assumptions, such as online interaction.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nsJ6FbvT54", "reviewer": "Reviewer_WbyR", "summary": "This paper proposes a novel model-based algorithm for policy evaluation in distributional RL and shows its near optimality within the generative model framework.  Additionally, they conduct an experimental study comparing their algorithm with quantile DP in a synthetic environment.", "questions": "Is the cubic dependence on the state space for computational efficiency good?", "limitations": "The authors discussed the limitations, and I found no potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Discretizing the distribution's support and then rewriting the update as a linear mapping is an important algorithmic contribution of this paper and is particularly interesting. It reveals a connection between categorical algorithms and linear systems---and, as they did, we can reduce the former to the latter. The paper is well-written. The algorithm and its main idea are clearly presented. The theory looks sound.", "weaknesses": "Why do you assume N iid samples for *each state* (line 56)? The common wisdom in policy evaluation under generative model is that more samples are needed for states frequently visited by the policy and fewer samples are needed for states that are less visited. As an extreme case, consider a state that is entirely unreachable---in this case no samples would be needed for that state, and the theoretical analysis should still hold. Thus, I wonder if the results in this paper can be achieved under that assumption.\n\nAccording to the experimental results, the proposed algorithm does not outperform quantile DP in terms of sample efficiency. However, it has significantly better running time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your summary of our work, your feedback and your questions. We are glad to hear you found the paper well-written, the algorithm and main idea clear, and the theory sound.\n\nWe are glad also that you found the connection between categorical DP and linear mappings interesting. This is one of the key observations underpinning our theoretical results in this paper, and we expect this observation alone will be important for both implementation and further analysis of categorical-based algorithms in the future.\n\n$N$  **iid samples at each state.**\n \n* The task we focus on is to estimate return distributions conditioned on each starting state for the environment (this is the typical case in the literature on estimation with generative models: see for example Azar et al., 2013; Pananjady and Wainwright, 2020), so there is no notion of an unreachable state in this problem.\n * We also note that in the standard problem set up, we cannot know in advance whether any particular state is unreachable from any other; we would have to infer this from samples we observe.\n * The reason that $N$ iid samples are used at each state is that this is sufficient to obtain a sample complexity bound that essentially matches the lower bound; this observation goes back to Azar et al. (2013) in the case of expected returns. This shows that more complex strategies for determining the number of samples to be taken at each state (for example, using adaptive strategies based on the estimated level of stochasticity in transition distributions based on samples seen so far) cannot improve the functional form of the sample complexity.\n\nHowever, we believe what the reviewer suggests may be an interesting direction for further research questions, such as:\n * For the distinct problem where some information about the environment is revealed in advance, before any samples are observed (e.g. states that are reachable with low, or zero, probability, from some fixed initial state), should the sampling strategy be adapted?\n * Even when no information is revealed in advance, can we adapt our sampling process on the basis of the samples observed? For example, if a particular state has never been transitioned to in our dataset, and we are not interested in estimating the return distribution conditioned on this state as the starting state, can we safely avoid sampling transitions from this state? As described above, such approaches cannot improve the functional form of the sample complexity, but may lead to empirical improvements in some settings. \nAnswering this question would likely constitute a full new paper, requiring substantial analysis of adaptive confidence bounds for simplex-valued random variables and their interaction with the algorithms described in our paper.\n\n**Cubic dependence on state space.**\n * This cubic dependence ultimately stems from the fact that even in the non-distributional case, practical solution of the linear system $V^\\pi = r^\\pi + \\gamma P^\\pi V^\\pi$ has cubic complexity in the state space size.\n * The focus of our paper is on sample complexity (in particular, in the sense described in the paper, showing that learning return distributions is not more statistically complex than learning mean returns, when using our categorical algorithm). However, we do describe possible computational improvements at various points in the paper, as we expected this to be a topic of interest to a variety of readers. See for example:\n   * Appendix G.3 for discussion of efficient implementation considerations, such as exploiting sparsity in the operators concerned.\n   * Appendix F.1 for discussion of categorical dynamic programming approaches which do not compute the exact fixed point, but for which theoretical guarantees are still obtainable (referred to just prior to Section 5.1 in the main text).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a min-max optimal algorithm for model-based distributional RL, which is used to approximate the return distributions under the assumption of having access to generative model. New theoretical analysis is provided for categorical approaches in distributional RL, with the introduction of a new distributional Bellman equation and the stochastic categorical CDF Bellman equation. Empirical studies are performed for benchmark comparison with other model-based distributional RL algorithms.

All reviewers believe this paper is above the bar of acceptance. The AC agrees and recommends acceptance.
Decision: Accept (poster)