Foundations of Multivariate Distributional Reinforcement Learning
Authors: Harley Wiltzer
Jesse Farebrother
Arthur Gretton
Mark Rowland
Venue: NeurIPS 2024 poster (2024)
Keywords: distributional reinforcement learning
rl theory
dynamic programming
temporal difference learning
successor features
successor representation
Primary area: reinforcement_learning
TL;DR: We introduce the first provably convergent oracle-free algorithms for distributional reinforcement learning with multivariate reward functions.

Abstract: In reinforcement learning (RL), the consideration of multivariate reward signals has led to fundamental advancements in multi-objective decision-making, transfer learning, and representation learning. This work introduces the first oracle-free and computationally-tractable algorithms for provably convergent multivariate *distributional* dynamic programming and temporal difference learning. Our convergence rates match the familiar rates in the scalar reward setting, and additionally provide new insights into the fidelity of approximate return distribution representations as a function of the reward dimension. Surprisingly, when the reward dimension is larger than $1$, we show that standard analysis of categorical TD learning fails, which we resolve with a novel projection onto the space of mass-$1$ signed measures. Finally, with the aid of our technical results and simulations, we identify tradeoffs between distribution representations that influence the performance of multivariate distributional RL in practice.

Foundations of Multivariate Distributional Reinforcement Learning
Harley Wiltzer
Jesse Farebrother
Mila
Québec AI Institute
Mila
Québec AI Institute
McGill University
McGill Unversity
harley wiltzer@mail wtt.zeermaalmmmmlll....
j wwwprrocssccggilll..
Arthur Gretton
Mark Rowland
Google DeepMindd
Google DeepMind
Gatsby Unit, University College London gretton@google. com
markrovllandgoogle..co com
Abstract
In reinforcement learning (RL), the consideration of multivariate reward signals has led to fundamental advancements in multi-objective decision-making. transfer learning, and representation learning. This work introduces the first oracle-free and computationally -ttactablee algorithms for provably convergent multivariate distribu- tional dynamic programming and temporal difference learning. Our convergence rates match the familiar rates in the scalar reward setting, and additionally provide new insights into the fidelity of approximate return distribution representations as function of the reward dimension. Surprisingly, when the reward dimension is larger than 1, we show that standard analysis of categorical TD learning fails, which we resolve with novel projection onto the space of mass-1 signed measures. Finally, with the aid of our technical results and simulations, we identify tradeoff's between distribution representations that influence the performance of multivariate distributional RL in practice.
Introduction
Distributional reinforcement learning [DRL; MSK* 10, BDM17b, BDR23] focuses the idea of learning probability distributions of an agent's random return, rather than the classical approach of learning only its mean. This has been highly effective in combination with deep reinforcement learning [ 19, BCC+20, WBK+ 22]. and DRL. has found applications in risk-sensitive decision making [LM22, KEF23], neuroscience [DKNU 20], and multi-agent settings [ROH+21, SLL21J.
In general, research in distributional reinforcement learning has focused on the classical setting of a scalar reward function. However, prior non-distributional approaches to multi-objective RL [RVWD13, HRB +22] and transfer learning [BDM* 17a, BHB + 20] model value functions mul- tivariate cumulants, rather than a sealar reward. Having learnt such multivariate value function, is then possible to perform zero-shot evaluation and policy improvement for any scalar reward signal contained in the span of the coordinates of the multivariate cumulants, opening up variety of applications in transfer learning, and multi-objective and constrained RI.
Multivariate distributional RI. combines these two ideas, and aims to learn the full probability distribution of returns given multivariate cumulant function. Successfully learning the multivariate
Cumulantt refer to accumulated quantities RL (e.g.. rewards or multivariate rewards) not to be confused with statistical cumulants.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

reward distribution opens up variety of unique possibilities, such as zero-shot return distribution estimation [WFG+ 24] and risk-sensitive policy improvement [CZZ* 24].
Pioneering works have already proposed algorithms for multivariate distributional RI.. While these works all demonstrate benefits from the proposed algorithmic approaches, each suffers from separate drawbacks, such as not modelling the full joint distribution [GBSL21], lacking theoretical guarantees [FSMT19, ZCZ + 211,, or requiring maximum-likelihood optimisation oracle for implementation [WUS23]. Concurrently, the work of [LK24] analyzed algorithms for DRL with Banach-space- valued rewards, and provided convergence guarantees for dynamic programming with non- paraammmtttic (intractable) distribution models.
Our central contribution in this paper is to propose algorithms for dynamic programming and temporal- difference learning in multivariate DRI which are computationally tractable and theoretically justified, with convergence guarantees. We show that reward dimensions strictly larger than introduce new computational and statistical challenges. To resolve these challenges, we introduce multiple novel algorithmic techniques, including randomized dynamic programming operator for efficiently approximating projected updates with high probability, and a novel TD-learning algorithm operating on mass-1 signed measures. These new techniques recover existing bounds even in the scalar reward case, and provide new insights into the behavior of distributional RL algorithms as a function of the reward dimension.
2 Background
We consider a Markov decision process with Polish state space X, action space A. transition kernel P : X' A + (X),, and discount factor 7 E [0. 1). Unlike the standard RL setting, we consider vector-valued reward function r X [0, Rmmx]] )d in the literature successor features [BDM+] + 1aa]. Given a policy π X' - P((A), we write the policy-conditioned transition kernel ( x) f P x, (dd x:).
Multi-variate return distributions. We write (Xe)r20 for a trajectory generated by setting Xo- = x. and for each t > 0. X,-1 (-)x)... The return obtained along this trajectory is then defined by (z(x)= = Ei rxx))) and the prettrr distribution function is ny" (x) = Law (x)).
Zero-shot evaluation. An intriguing prospect of estimating multivariate return distributions is the ability to predict (scalar) return distributions for any reward function in the span of the cumulants. Indeed, 1ZCZ++21, CZZ+ + 24] show that for any reward function f -> (r(x:), w) for some w E R.. (G*(r), w) Flaw Σizo (Xx, for Xp z. Likewise, one might consider r(z) 8g, in which case GH (1) corresponds the per-trajectory discounted state visitation measure, and [WFG+24] demonstrated methods for learning the distribution of G" to infer the return distribution for any bounded deterministic reward function.
Multivariate distributional Bellman equation. It was shown in 22] that multi-return distributions obey distributional Bellman equation, similar to the scalar case [MSK 10, BDM17b), and defines the multivariate distributional Bellman operator T P(R P(R by
(T" = x'xPF(-1x) [(b'))))))(X')),
(1)
where
= = 7z and Suu = 11 o f-1 is the pushforward of a measure μ through measurable function f. In particular, [ZC7+21] showed that 17" satisfies the malti-variate distributional Bellman equation T=1²" = 17= and that T is a 7-contraction in W p. where W, (T)().))) = supzex W,(n (z), n2())) and W, is the p-Wasserstein metric [Vi109]. This sug- gests convergent scheme for approximating 17" in W, p by distributional dynamic programmingg that is, computing the iterates 77++1 = TKT following Banach's fixed point thcorem.
Approximating multivariate return distributions. In practice, however, the iterates Mk+1 = T" cannot be computed efficiently, because the size of the support of πß may increase exponentially with it. A variety of alternative approaches that aim to circumvent this computational difficulty have been considered [FSMT19, 21, WUS23J. Many of these approaches have proven effective in combination with deep reinforcement learning. though as tabular algorithms, either lack theoretical guarantees, or rely on oracles for solving possibly intractable optimisation problems. A more complete account of multivariate DRL is given in Appendix A. A central motivation of our work

---

is the development of computationally-traaee algorithms for multivariate distributional RL with theoretically guarantees.
Maximum mean discrepancies. A core tool in the development of our proposed algorithms, as well as some prior work INTGV20, ZCZ++1 ], is the notion of distance probability distributions given by maximum mean discrepancies [GBR* + 12, MMDJ. A maximum mean discrepancy MMD,, R+ assigns a notion of distance to pairs of probability distributions, and is parametrised via a choice of kernel K y x y + R, defined by
(J)) X (V)
MM.,.(p q) = (V.. (Yi. Y2)] 2 (Y,7) [<:(Y, Z)] + (22z Z2)1
useful alternative perspective MMD is that the choice of kernel induces reproducing kernel Hilbert space (RKHS) of functions H. namely the closure of the span of functions of the form + k(y, 2) for each y E V, with respect to the norm || |/ induced by the inner product (r((11,), k(y,. -)) = x(y1, 32). With this interpretation, ,MD,, SIID,((2 (p. q) is equal to | 1 Mggl Ht. where My = Iy k(, (dy) H. the mean embedding of p (similarly for Mt)). When p + Pp injective, the kernel A called characteristic, and MMD,, is then a proper metric on P()) [GBR + 12]. In the remainder of this work, we will assume that all spaces of measures will be compact sets V: thus with continuous kernels, we are ensured that distances between probability measures are bounded. When comparing return distributions, this is achieved by asserting that rewards are bounded.
We conclude this section by recalling a particular family of kernels, introduced [SSGF13], that will be particularly useful for our analysis. These are the kernels induced by semimetrics.
Let y be a nonempty set, and consider a function p y x y + R... Then p is called a semimetric it is symmetric and p(y1: 12)
Definition 1.
- y1 V2- Additionally, p is said to have strong negative type if J' d(pp ql x (p- q)) < O whenever p, q E P(Y) with p & q.
Notably, certain semimetrics naturally induce characteristic kernels and probability metrics.
Theorem ([SSGF13, Proposition 291). Let be a semimetric on space y have strong negative type, in the sense that J pd q) Ip q) < O whenever p # are probability measures
compact set y. Moreover, let k: y x y + R denote the kernel induced by that is
k(y1, V2) 2 (p(x), vo) + p(y2, wo) -((y1 v(y),12))
for some Yo Y. Then K is characteristic, so MMD, a metric.
Remark 1. An important class of semimetries are the functions Pa Rd x RR given by (3( V2) = 11y1 y2ll2 for E (0, 2). It is known that these semimetrics have strong negative type, and thus the kernels Ka induced by Pr. are characteristic [SR13, SSGF13J. The resulling metric MMD c. is known the energy distance.
3 Multivariate Distributional Dynamic Programming
To warm up, we begin by demonstrating that indeed the (multivariate) distributional Bellman operator is contractive in supremal form MMD, of MMD, given by MMD,(m, (72) suPzex MMD, (7 (z), mh (x)), for a particular class of kernels к. Our first theorem generalizes the analogous results of [NTGV20] in the scalar case to multivariate cumulants. The proof of Theorem 2, as well as proofs of all remaining results, are deferred to Appendix B.
Theorem 2 (Convergent MMD dynamic programming for the multi-return distribution function). Let be kernel induced by semimetric p on [0, (1->)- "maax] ]d with strong negative type, satisfying
1. Shift-invariance For any 2 E Rd, p(z + v1, + U2) p(y1.82)
2. Homogeneity. For any 7 E [0, 1), there exists > 0 0.ror which p( Ty2) >
plyi, y2).
Consider the sequence [T 1oc given by m = Th Then mhe - 17" at a geometric rate of 7/22 in MMDx. as long as MMD, ")") S C DC.
Notably, the energy distance kernels K., satisfy the conditions of Theorem 2, and pa (aM1,,770) S 7 plyn, V/2 by the homogeneity of the Fuclidean norm, so T a conttracniii in the energy distances. This generalizes the analogous result of [NTGV20] in the one-dimensional case.
3

---

While Theorem 2 illustrates method for approximating 77" in MMD, it leaves lot to be desired. Firstly, even in tabular MDPs, just as in the case of scalar distributional RL, return distribution functions have infinitely many degrees of freedom, precluding tractable exact representation. As such, it will necessary to study approximate, finite parameterizations of the return distribution functions, requiring more careful convergence analysis. Moreover, in RL it is generally assumed that the transition kernel and reward function are not known analytically- we only have access to sampled state transitions and cumulants. Thus, T cannot be represented or computed exactly, and instead we must study algorithms for approximating 1y" from samples. We provide algorithms for resolving both of these concerns- the former in Section 5 and the latter in Section 6-where illustrate the difficulties that arise once the cumulant dimension exceeds unity.
Particle-Based Multivariate Distributional Dynamic Programming
Our first algorithmic contribution is inspired by the empirically successful equally-weighted particle (EWP) representations of multivariate return distributions employed by
Temporal-difference learning with EWP representations. EWP representations, expressed by the class 6EWP, m. are defined by
6WwP,m m = (6EWT, m )*
EWT,
50, 0, E R
(2)
For simplicity, we consider the case here where at each state x, the multi-return distribution is approximated by (²)) = m atoms. We can represent 7) E 6EWP.m by 77()) = E 50,(a) for 0: + R The work of [ZCZ+21] introduced TD-learning algorithm for learning CEwP, representation of ry",, computing iterates of the particles (0(*)) )m according to
g g( (x) 0 (z) Vo MMD?
8 (()) (t)
(
for step sizes (Ax )220 and sampled next states X' ~ PT (- I), where 0 = stop- tooeggradiinttttee is copy of g(k) that does not propagate gradients. Despite the empirical success of this method in combination with deep learning, no convergence analysis has been established, owing to the nonconvexity of the MMD objective with respect to the particle locations. In this section we aim to understand to what extent analysis is possible for dynamic programming and temporal-difference learning algorithms based on the EWP representations in Equation (2).
Dynamic programming with EWP representations. As is often the case in approximate distri- butional dynamic programming [RBD* 18, RMA + 24], we have T"WEP,.. & 'Ewp... in words, the distributional Bellman operator does not map EWP representations as long as there exists state x at which the support of P" (- x)) is not singleton, (T" will consist of more than m atoms even when 1) E 6gwr,,, mi and secondly, P(- x) is not uniform, (z) will not consist of equally-weighted particles.
themselves. Concretely,
Consequently, obtain DP algorithm over EWP representations, we must consider projected operator of the form TEwr""", for projection Tewr P(R lewr,m. A natural choice for this projection is the operator that minimizes the MMD of each multi-return distribution in 6EWP,..
(TEWP,.")) E argmin MM.
(4)
PE
Unfortunately, even in the sealar-reward (d atic; (Ie (x) is not uniquely defined, and IEWP, is not non-expansion in MMD, [LB22, RMA+24]. These pathologies present significant complications when analyzing even the convergence of dynamic programming routines for learning an EWP representation of the multi-return distribution in particular, it is not even clear that HEww, T has fixed point (let alone a unique one). Another complication arises due to the computational difficulty of computing the projection (4): even in the case where y)x)) has finite support for each state x, the projection (Tlma Ewr, "i) (2r) is very similar to clustering, which can intractable to compute exactly for large m [She21] Thus, the argmin projection in Equation (4) cannot be used directly to obtain a tractable DP algorithm.
case, the operator ITm E... is problem-

---

Randomised dynamic programming. Towards this end, we introduce a tractable randomized dy- namic programming algorithm for the EWP representation, by using randomized proxy BootProj for m T", that produces accurate return distribution estimates with high probability. The method produces the following iterates,
Mh+1 (z) = BBooPProj mo
(x) == Σ
(()))+2 E=
iid Z,~ (X,), X; ~ P (=(-12)) x)
(5)
similar algorithm for categorical representations was discussed in concurrent work [LK24] without convergence analysis.
The intuition that, particularly for large m, the Monte Carlo emror associated with the sample-based approximation to (7"))( (x) is small, and we can therefore expect the DP process, though randomised, be accurate with high probability. This is summarised by key theoretical result of this section; our proof of this result in the appendix provides general approach to proving convergence for algorithms using arbitrary accurate approximations (4) that we expect to be useful in future work.
Theorem 3. Consider kernel induced by the semimetric p(x,, y) || y122 with E (0, 2), and suppose rewards are bounded in each dimension within [0, Rmax]. For any To such that MMD, (7o. 7²") & D < oc, and any > 0, for the sequence (m=)220 defined in Equation (5), with probability at least 1 5 we have
d0 R2
(216--1
(MR,(KR θō
(2)0/)(((111 7) vm
max
log
logy
where M6+1 = BootProj and K =
log logy- a
and where Ô omits logarithmic factors m.
This shows that our novel randomised DP algorithm with EWP representations can tractably compute the true multivariate return distributions, with only polynomial depen- dence on the dimension d. Appendix C illustrates explicitly how this procedure is more memory efficient than unprojected EWP dynamic programming. However, the guarantees associated with this algorithm hold only in high probability, and are weaker than the pointwise convergence guarantees of mne-dimnniional distributional DP algorithms [RBD+ 18, RMA 24, BDR23]. Consequently, these guarantees do not provide a clear understanding of the EWP-TD WWP-TD method described at the beginning of this section. However, the sequel, we introduce DP and TD algorithms based on categorical representations, for which we derive dynamic programming and TD-learning convergence bounds.
accurate approximations
The proof of Theorem 3 is hinges the following proposition, which demonstrates that convergence of projected EWP dynamic programming is controlled by how far return distributions are transported under the projection map.
Proposition (Convergence of EWP Dynamic Programming). Consider kernel satisfying the hypotheses of Theorem 2, suppose rewards are globally bounded in each dimension in [0, Rmax), and let (TT(k) he2o be a sequence of maps II: P([O, (1 - - - CEwP,m satisfying
MMD
k.m ((T
())x)))()))))))))))))))))))))))) S f(d, m) < 00
v 0.
Then the iterates k20 given by Mc+1 ( m T nhe with MMo,((Mo 77" D converge B(n", (1 - -(( m)) in MMD,, where B denotes the closed ball in
a NEWP, MMD
B(n R) 2 {pr P(Rd)* (M.)
As an immediate corollary of Proposition and Theorem 3, we can derive an error rate for projected dynamic programming with We well.
Corollary For any kernel satisfying the hypotheses of Theorem 3, for any to0 CEwP for which MMD, (T0: r7 < <, co, the iterates Th+11 II TT converge set NEWP, C CEEP ma where
// Rp R
To D C 17
(1-00/2)(1 /
5

---

Categorical Multivariate Distributional Dynamic Programming
Our next contribution the introduction of convergent multiariatee distributional dynamic program- ming algorithm based on a categorical representation of return distribution functions, generalizing the algorithms and analysis of [RBD*] 18] to the multivariate setting.
Categorical representations. As outlined above, to model the multi-return distribution function in practice, it is necessary restrict each multi-retun distribution to miitely-parametenee class. In this work, we take inspiration from successful distributional RL algorithms [BDM17b, RBD* 1BDD11 18] and employ categorical representation. The work of [WUS23] proposed categorical representation for multivariate DRI., but their categorical projection was not justified theoretically, and it required a particular choice of fixed support. We propose novel categorical representation with finite (possibly
N(a) C R that models the multi-return distribution function y such that 77(x)) RR(.) for each x X. The notation {(x)); simply refers to the ith support point at state specified by R. and Aa denotes the probability simplex on the finite set A. We refer the mapping R. as the support map' and we denote the class of multi-return distribution
state-dependent) support R(x) = (E(z),
functions under the corresponding categorical representation as CR. A IILex AR(Z).
Categorical projection. Once again, the distributional Bellman operator is not generally closed over ER. that is, T-E z CR As such, we cannot actually employ the procedure described in Theorem 2 - rather, we need to project applications of 7 back onto CR. Roughly, we need an operator 11: )( + R for which 11/ &0 = id. Given such an operator, as in the literature on categorical distributional RL [BDM17b, RBD 18], we will study the convergence of iterates 1/ 117"T),. "//-
Projection operators used the scalar categorical distributional RL literature are specific to dis- tributions over R. must introduce new projection. We propose a projection similar to (4),
(TT (x) arginf (M,
pεAR(a)
We will now verify that II is well-defined, and that it satisfies the aforementioned properties.
Lemma 1. Let i. be a kernel induced by a semimetric p on [0, (1 - --)- Rangx] with strong negative type (ef. Theorem 1). Then IT is well-defined, Ran(TT C CR. and IT C, 16. id.
is worth noting that beyond simply ensuring the well well-posednnss: the projection I Lemma 1 also certifies an efficient algorithm for computing the projection namely, by solving the appropriate quadratic program (QP), as observed by [SZS+08), We demonstrate pseudocode for computing the projected Bellman operator IT& T with QP solver QPSolve in Algorithm
Algorithm 1 Projected Categorical Dynamie Programming
Require: Support map R, kernel k, transition kernel ps reward function r, discount y Require: Return distribution function E '6
for x X do
(T" + Evex EGER(H)) Pa (m
KF, k(E;,E,) for each (Ex, ) R(x)²
o3 + Esesupp(T (7*,), (T (F)))(()s(( (E) for each & R(x))
p - QPSolve (min (EIRRE) (pT K 2p v] subject to 2 0, Σ, P:
(11 1) EgTER()) prde.
end for return T.AT"
Lemma 2. Under the conditions of Lemma J. I.B.B..
k. is a nonexpansion in MMD That is, for any 77, E P([0, (1 (1-7) 11 dy we have MMD,, ( m, II m1) S MMD (m,, m2).
Categorical multivariate distributional dynamic programming, As an immediate consequence of Lemma 2, follows that projected dynamic programming under the projection I is convergent;
In many applications, the most natural support map is constant across the state space.

---

this is because T is a contraction in MMD,, x and I is a nonexpansion in MMD,, so the projected operator I T a contraction MMD,; a standard invokation of the Banach fixed point theorem appealing to the completenes of MMD,, certifies that repeated application of 11 T will result in convergence unique fixed point.
Corollary 2. Let K. be a kernel satisfying the conditions of Theorem 2.
Then for any no E CR. the iterates {mk k=1 given by T7t+1 ,,,,, TFTtk converge geometrically to a unique fixed point.
Beyond the result of Theorem 2, Corollary illustrates an algorithm for estimating 17" in MMD,, provided knowledge of the transition kernel and the reward function, which is computationally tractable in tabular MDPs. Indeed, the iterates (77))20 all lie in ER. having finitely-many degrees of freedom. Algorithm outlines computationally tractable procedure for learning nč,, in this setting.
We note that the work of [WUS23] provided an alternative multivariate categorical algorithm, which was not analyzed theoretically. Moreover, our method provides the additional ability to support state-dependent arbitrary support maps, while theirs requires support maps to be uniform grids.
Accurate approximations. We now provide bounds showing that the fixed point n from Corollary 2 can be made arbitrarily accurate by increasing the number of atoms.
To derive a bound on the quality of the fixed point, we provide a reduction via partitioning the space of returns to the covering number of this space. Proceeding, we define class of partitions P R(s). where each P E R(a))
PR(,)) satisfies
1. |P| = (xx);
2. For any 01, 02 E P. either 81 n n0= - or 01 =2;
3. U = P(0, (1 7) Rmax)] 1();
4. Each element 0; E P contains exactly one element 2; E R(x)..
For any partition P, we define a notion of mesh size relative to kernel x induced by a semimetric p.
mesh (P; k) max sup p(y1,
(8)
0EP Vi.vEE0
The following result confirms that nē recovers 17" as the mesh decreases. Theorem Let k. be kernel induced by a c-homogeneous and shift-invariant semimetric p conforming to the conditions of Theorem 2. Then the fixed point ne of 1f118 T satisfies
MMDx 7²) 7012
1
sup inf Vmesh(P:k).
(9)
DE PEPP(:
Thus, for any sequence of supports (R(#) Jm21 for which the maximal space (in p) between any two points in R(I) tends to 0 as m + oc, the fixed point nē approximates 17" to arbitrary precision for large enough m. The next corollary illustrates this in a familiar setting.
Corollary 3. Let R(x) =
Um where Um is a set of m mmiformly-space support points on [0, (1
7) Rmax]. For induced by the semimetric p(z, 3) = || y 12 for a E (0, 2).
d R
MMD 7²) (1 (1- (a-ya/3)(1 7)a/2 (m 1/d 2)a/2"
With and d = 1, the MMD in Corollary 3 is equivalent to the Cramér metric [SR13], the setting in which categorical (scalar) distributional dynamic programming is well understood. Our rate matches the known (m 1122) rate shown by [RBD* 18] this setting. Thus, our results offer a new perspective categorical DRL, and naturally generalizes the theory to the multivariate setting.
Theorem relies on the following lemma about the approximation quality of the categorical MMD projection, which may be of independent interest.
Lei x be kernel satisfying the conditions of Lemma 1, and for any finite R C Rd define II Lemma 3.
P(RR) - AR via TTp arginf ye MMD, q). Then MMD? (Tp.p) p) & inf inf P. mesh( P; k).
At this stage, we have shown definitively that categorical dynamic programming converges in the multivariate case. In the sequel, we build on these results to provide convergent multivariate categorical TD-leaming algorithm.
7

---

Monte Carlo
Monke Carlo Categorical-MMM
Monte Carlo EWP-MMD MMD
Monte Cario EWP-MMD
Categarical MMD
Monte Carlo Celegorical-XMD
Monte EWP MMD
0
Figure 1: Distributional SMs and associated predicted return distributions with the categorical (left) and EWP (right) representations. Simplex plots denote the distributional SM. Histograms denote the associated return distributions, predicted from a pair of held-out reward functions.
5.1 Simulation: The Distributional Successor Measure
As preliminary example, we consider 3-state MDPs with the cumulants r(i) (1-))
for e; the ith basis vector. In this setting, 1y" encodes the distribution over trajectory-wise discounted state occupancies, which was discussed in the recent work of [WFG+244 24] and called the distributional successor measure (DSM). Particularly, [WFG* 24] showed that + Law (G for G2 ~ 17" (x) the return distribution function for any scalar reward function F. Figure shows that the projected categorical dynamic programming algorithm accurately approximates the distribution over discounted state occupancies as well as distributions over returns on held-out reward functions.
Multivariate Distributional TD-Learning
Next, we devise an algorithm for approximating the multi-return distribution function when the transition kernel and reward function are not known, and are observed only through samples. Indeed, this is strong motivation for TD-learning algorithms [Sut88], wherein state transition data alone is used to solve the Bellman equation. In this section, we devise TD-learning algorithm for multivariate DRL, leveraging our results on categorical dynamic programming in MMD,-
setting, the categorical projection presented above is known to be affine IRBD+18), making scalar categorical TD-learning amenable to common techniques in stochastic approximation theory. However, the projection is not affine when d > 2; give explicit example in Appendix D. Thus, relax the categorical representation to include signed measures, which will provide us with an affine projection [BRCM19]- this crucial for proving main result, Theorem 6. We denote by M (2)) the set of all signed measures μ over y with (x)))= = 1. We begin by noting that the MMD endows M with metric structure.
Relaxation to signed measures. the
Lemma 4. Let K y x + R be a characteristic kernel over some space y. Then MMD ()) (D) R+ given by (p, q) F 1|p Mallx defines a metric M! (V), where Mp denotes the usual mean embedding of p and H is the RKHS with kernel K.
We define the relaxed projection
IS MM (10,
Rmax |d) 10(x Ilxx M (R(x)) =: FR.
(IT SC m) (x) E arginf (p. (R(M)
(10)
Note that (10) is very similar to the definition of the categorical MMD projection in (7)-the only difference that the optimization occurs over the larger class of signed mass-l measures. also worth noting that the distributional Bellman operator can be applied directly to signed measures, which yields the following convenient result.
Lemma 5. Under the conditions of Corollary 2, the projected operator TT SC T :RR - IR is
affine, is contractive with contraction factor e//2 and has unique fixed point 1SC.
While we have "relaxed" the projection, the fixed point VgC,.. is a good approximation 17 Theorem 5. Under the conditions of Lemma 5, we have that
8

---

1. MMD, x (TED,,(TCCww 11") - 1-g-/3 sup inf RMo
Ri0 v /mesh(P P:K); and
2. MMD,, TSC., 77) $ (1+ 1-7072 suprex inf r R(.) v meshh(P:k)
Notably, the second statement of Theorem 5 states that projecting Nšc., back onto the space of multi-return distribution functions yields approximately the same error ',2, when y is near 1.
In the remainder of the section, assume access stream of MDP transitions (Ti)
X Rd xX' consisting of elements (Xx, Ar, Re, X;) with the following structure,
X P F(-1)
A~
(- X²) R2 = r(x,)
X ~ P(: Xx, A,)
(1I)
where P is some probability measure and {Fo is the canonical filtration F, = =i)). Based on these transitions, we can define stochastic distributional Bellman backups by
Tnn()) = (n(x) ((Rx.)) == X₁
(12)
otherwise
which notably can be computed exactly without knowledge of P.r. Due to stronger convergence guarantees shown for projected multivariate distributional dynamic programming, we introduce an asynchronous incremental algorithm leveraging the categorical representation, which produces iterates {fde according
M++1 (1-aa)) au T.
(13)
for No E CR. where [ is any sequence of (possibly) random step sizes adapted to the filtration 0c The iterates of (13) closely resemble those of classic stochastic approximation algorithms [RM51] and particularly asynchronous TD learning algorithms [JJS93, Tsi94, BT96], but with iterates taking values in the space of state- indexed signed measures. Indeed, our next result draws on the techniques from these works establish convergence of TD-leaming on PR representations.
[Fr)
Theorem 6. For kernel induced by semimetric p of strong negative type, the sequence (fc) 1e given by (11)-(13) converges Tisc.x with probability 1.
6.1 Simulations: Distributional Successor Features
To illustrate the behavior of our categorical TD algorithm, employ it to learn the multi-return distributions for several tabular MDPs with random cumulants. We focus on the case 2- and 3-dimensional cumulants, which is the setting studied in recent works regarding multivariate distribu- tional RL [ZCZ*2] 21, WUS23J. Interpreting the multi-return distributions as joint distributions over successor features [BDM* 17,, SFs], additionally evaluate the return distributions for random reward functions in the span the cumulants. We compare our categorical TD approach with a tabular implementation of the EWP TD algorithm of [ZCZ+21]] for which convergence bounds are known.
0.50
Method
Method EWP-TD
0.45
0.6
EWP-TD
Cramer Distance
0.40
Cramer Distance
0.35
CatTD
0.5
Signed-Cat-TD
0.30
Signed-Cat-TD
0.4
0.25
0.20
0.3
0.15
0.10
0.2
100
200
300
400
500
100
200
300
400
500 600 700 800 900 1000 Atoms
Atoms
(a) 2
(b) d 3
Figure 2: Error of zero-shot return distribution predictions over random MDPs, measured by Cramér distance, and showing 95% confidence intervals.
Figure 2a compares TD learning approaches based on their ability to accurately infer (scalar) return distributions on held out reward functions, averaged over 100 random MDPs, with transitions drawn

---

MC
NC
Signodi Cac.D
EMP-TD
Cumutonts
0.5
Gumulont
0.5
farsum
MC
-0.5
-0.
-1.0
Signed-Cat-TD
-10
EWP-TD
Cumulant
Oumulant
EMPITD
MC
Figure 3: Distributional SFs and predicted return distributions with m = 400 atoms, in random MDP with known rectangular bound on cumulants. Left: Categorical TD. Right: TD.
from Dirichlet priors and 2-dimensional cumulants drawn from uniform priors. The performance of the categorical algorithms sharply increases as the number of atoms increases. On the other hand, the EWP TD algorithm performs well with few atoms, but does not improve very much with higher-resolution representations. We posit this is due to the algorithm getting stuck in local minima, given the non-convexity of the EWP MMD objective. This hypothesis is supported as well by Figure 3, which depicts the learned distributional SFs and return distribution predictions.
Particularly, we observe that the learned particle locations the EWP TD approach tend to cluster in some areas or get stuck in low-density regimes, which suggests the presence of a local optimum. On the other hand, our mrovably-converye the true multi-return distributions.
categorical TD method learns high fidelity quantization of
Naturally, however, the benefits of the poly(d) bounds for EWP suggested by Theorem become more present as we increase the cumulant dimension. Figure 2b repeats the experiment of Figure 2a with = 3, using randomized support points for the categorical algorithm to avoid cubic growth in the cardinality of the supports. Notably, our method is the first capable of supporting such unstructured supports. While the categorical TD approach can still outperform EWP, much larger number of atoms is required. This is unsurprising in light of our theoretical results.
7 Conclusion
We have provided the first provably convergent and computationally tractable algorithms for learning multivariate return distributions in tabular MDPs. Our theoretical results include convergence guarantees that indicate how accuracy scales with the number of particles used in distribution representations, and interestingly motivate the use signed measures to develop provably convergent TD algorithms.
While it difficult to scale categorical representations to high-dimensional cumulants, our algorithm is highly performant in the low d setting, which has been the focus of recent work in multivariate distributional RL. Notably, even the
setting has important applications- applications-indeed indeed, efforts in safe RL depend on distinguishing cost signal from a reward signal (see, e.g., (YSTS23]), which can be modeled by bivariate distributional RL. In this setting, our method can easily be scaled to large state spaces by approximating the categorical signed measures with neural networks; an illustrated example is given in Appendix F.
On the other hand, the prospect of learning multi-return distributions for high- high-dimensional cumulants also has many important applications, such as modeling close approximations distributional successor measures [WFG 24] for zero-shot risk-sensitive policy evaluation. In this setting, we believe based multivariate DRL will be highly impactful. Our results concerning EWP dynamic programming provide promising evidence that the accuracy of EWP representations scales gracefully with d for a fixed number of atoms. Thus, we believe that understanding convergence of EWP TD- Telaannn algorithms is very interesting and important open problem for future investigation.
10

---

Acknowledgements
The authors would like to thank Yunhao Tang, Tyler Kastner, Arna Jain, Yash Jhaveri, Will Dabney, David Meger, and Mare Bellemare for their helpful feedback, well as insightful suggestions from anonymous reviewers. This work was supported by the Fonds de Recherche du Québec, the National Sciences and Engineering Research Council of Canada, and the compute resources provided by Mila (mila. quebec).
References
[BB96] Steven J. Bradtke and Andrew G. Barto. Linear least- algorithms for temporal difference learning. Machine Learming, 22(1):33-57, 1996.
[BBC+21] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López. Fabian Pedregosa, and Jean- San-Philipppe Vert. Efficient and modular im- plicit differentiation. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[BCC*20] Marc G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera S. Ponda, and Ziyu Wang. Autonomous naviga- tion of stratospheric balloons using reinforcement learning. Nature, 588(7836):77-82, December 2020.
[BDM+17a] André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, Hado P. van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[BDM17b] Mare G. Bellemare, Will Dabney, and Rémi Munos. A Distributional Perspective on Reinforcement Learning. In International Conference on Machine Learning (ICML), 2017.
[BDR23] Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. The MIT Press, 2023.
[BEKS17] Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral B. Shah. Julia: A fresh approach numerical computing. SIAM Review, (9())-5-9 2017.
[BFH+ 18] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy pro- grams, 2018.
[BHB+20 22] André Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. Proceedings of the National Academy of Sciences (PNAS), 1 117448)33079--0088 2020.
[BRCM19] Marc G Bellemare, Nicolas Le Roux, Pablo Samuel Castro, and Subhodeep Moitra. Distributional reinforcement learning with linear function approximation. In Interna- tional Conference Artificial Intelligence and Statistics (AISTATS), 2019.
[BT96] Dimitri P. Bertsekas and John N. Tsitsiklis. Neuro-dynamic programming. Athena Scientific, 1996.
[CGH +9] Robert M. Corless, Gaston H. Gonnet, David E.G. Hare, David J. Jeffrey, and Donald E. Knuth. On the Lambert W function. Advances in Computational Mathematics, 5:329- 359, 1996.
[CZZ+24] Xin-Qiang Cai, Pushi Zhang, Li Zhao, Jiang Bian, Masashi Sugiyama, and Ashley Llorens. Distributional Pareto-optimal multi-objective reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2024.

---

[DKNU 20] Will Dabney, Zeb Kurth-Nelson, Naoshige Uchida, Clara Kwon Starkweather, Demis Hassabis, Rémi Munos, and Matthew Botvinick. A distributional code for value in dopamine-based reinforcement learning. Nature, 577(7792):61--675 2020.
[DRBM18] Will Dabney, Mark Rowland, Mare G. Bellemare, and Rémi Munos. Distributional Reinforcement Learning with Quantile Regression. In AAAI Conference Artificial Intelligence, 2018.
[FSMT19] Dror Freirich, Tzahi Shimkin, Ron Meir, and Aviv Tamar. Distributional Multivariate Policy Evaluation and Exploration with the Bellman GAN, International Conference Machine Learning (ICML), 2019.
[GBR+12] 12] Arthur Gretton, Karsten M. Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A Kernel Two-Sample Test. Journal of Machine Learning Research (JMLR), 13(25):723-77, 2012.
[GBSL21] Michael Gimelfarb, André Barreto, Scott Sanner, and Chi-Guhn Lee. Risk-Aware Transfer in Reinforcement Learning using Successor Features. In Advances Neural Information Processing Systems (NeurIPS), 2021.
22] Conor F. Hayes, Roxana Radulescu, Eugenio Bargiacchi, Johan Källström, Matthew Macfarlane, Mathieu Reymond, Timothy Verstraeten, Luisa M. Zintgraf, Richard Dazeley, Fredrik Heintz, Enda Howley, Athirai A. Irissappane, Patrick Mannion, Ann Nowé, Gabriel de Oliveira Ramos, Marcello Restelli, Peter Vamplew, and Diederik M. Roijers. A practical guide to multi-objective reinforcement learning and planning. In Proceedings of the International Conference on Autonomous Agents and Mulliagent Systems (AAMAS), 2022.
[JJS93] Tommi Jaakkola, Michael I. Jordan, and Satinder P. Singh. On the Convergence of Stochastic Iterative Dynamic Programming Algorithms. In Advances in Neural Information Processing Systems (NeurIPS), 1993.
[KEF23] Tyler Kastner, Murat A. Erdogdu, and Amir- massoud Farahmand. Distributional Model Equivalence for Risk-Sensitive Reinforcement Learning. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[Lax02] Peter D. Lax. Functional analysis. John Wiley & Sons, 2002.
[LB22] Alix Lhéritier and Nicolas Bondoux. A Cramér Distance perspective on Quantile Regression based Distributional Reinforcement Learning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
[LK24] Dong Neuck Lee and Michael R. Kosorok. Off-policy reinforcement learning with high dimensional reward. CoRR, ans/2408.07660, 2024.
[LM22] Shiau Hong Lim and Hyas Malik. Distributional Reinforcement Learning for Risk- Sensitive Policies. In Advances in Nearal Information Processing Systems (NeurIPS), 2022.
[MSK+ + 10] Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and 'oshiiyki Tanaka. Nonparametric return distribution approximation for reinforcement learning. In International Conference on Machine Learning (ICML), 2010.
[NTGV20] Thanh Nguyen-Tang Sunil Gupta, and Svetha Venkatesh. Distributional Reinforcement Learning via Moment Matching. In AAAI Conference on Artificial Intelligence, 2020.
[RBD+18] 18] Mark Rowland, Mare G. Bellemare, Will Dabney, Remi Munos, and Yee Whye Teh. An Analysis of Categorical Distributional Reinforcement Learning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2018.
[RM51] Herbert Robbins and Sutton Monro. Stochastic Approximation Method. The Annals of Mathematical Statistics, 2203):400-407, September 1951.
12

---

[RMA* 22] Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Mare Bellemare, and Will Dabney. An Analysis of Quantile Temporal-Difference Learning. Journal of Machine Learning Research (JMLR), 25(163):1 2024.
[ROH+21] Mark Rowland, Shayegan Omidshafiei, Daniel Hennes, Will Dabney, Andrew Jaegle, Paul Muller, Julien Pérolat, and Karl Tuyls. Temporal difference and return optimism in cooperative multi-agent reinforcement learning. In AAMAS ALA Workshop, 2021.
[RVWD13] Diederik M. Roijers, Peter Vamplew, Shimon Whiteson, and Richard Dazeley. A survey of multi-objective sequential decision-making Journal of Artificial Intelligence Research (JAIR), 48:67-113, 2013.
[Sch00] Bernhard Schölkopf. The Kernll Trick for Distances, In Advances in Neural Informa- lion Processing Systems (NeurIPS), 2000.
[SFS24] Eiki Shimizu, Kenji Fukumizu, and Dino Sejdinovic. Neural-kernel conditional mean embeddings. In International Conference on Machine Learning (ICML), 2024.
[She21] Vladimir Shenmaier. On the Complexity of the Geometric Median Problem with Outliers. CoRR, ahs/2111.005111 2021.
[SLL221] Wei-Fang Sun, Theng--Kang Lee, and Chun-Yi Lee. DFAC framework: Factorizing the value function via quantile mixture for multi-agent distributional Q-learning. In International Conference on Machine Learning (ICML), 2021.
[SR13] Gábor J. Székely and Maria L.. Rizzo. Energy statisties: A class of statistics based on distances. Journal of Statistical Planning and Inference, 143(8):1249-1277 August 2013.
[SSGF13] Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equiv- alence of distance-based and RKHS-based statisties in hypothesis testing. The Annals of Statistics, 41(5), October 2013.
[Sut88] Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 3:9-44, 1988.
08) Le Song, Xinhua Zhang, Alex Smola, Arthur Gretton, and Bernhard Schölkopf. Tai- loring density estimation via reproducing kernel moment matching. In International Conference on Machine Learning (ICML), 2008.
[Tsi94] John N. Tsitsiklis. Asynchronous stochastic approximation and Q-learning. Machine learning, 16:185-202, 1994.
[TSM17] Ilya Tolstikhin, Bharath K. Sriperumbudur, and Krikamol Mu. Minimax estimation of kernel mean embeddings. Journal of Machine Learning Research (JMLR), 18(86):1-47, 2017.
[TVR97] J.N. Tsitsiklis and B. Van Roy. An analysis of temporal-difference learning with function approximation. IEEE Transactions on Automatic Control, 42(5):674-690, 1997.
[Vi109] Cédric Villani. Optimal transport: Old and new, volume 338. Springer, 2009.
[VN02] Jan Van Neerven. Approximating Bochner integrals by Riemann sums. Indagationes Mathematicae, 1772) June 2002.
[WBK+22] Peter R Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et al. Outracing champion gran turimo drivers with deep reinforcement learning. Nature, 222 2022.
13

---

[WFG*24) Harley Wiltzer, Jesse Farebrother, Arthur Gretton, Yunhao Tang, André Barreto, Will Dabney, Mare G. Bellemare, and Mark Rowland. distributional analogue to the successor representation. In International Conference on Machine Learning (ICML), 2024.
[WUS23] Runzhe Wu, Masatoshi Uehara, and Wen Sun. Distributional Offliee Policy Evaluation with Predictive Error Guarantees. In International Conference on Machine Learning (ICML), 2023.
[YSTS23] Qisong Yang, Thiago D. Simão, Simon H. Tindemans, and Matthijs T. J. Spaan. Safety-constrained reinforcement learning with distributional safety critie. Machine Learning, 11(()) 2023.
[YZL+19) 19] Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie- Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[ZCZ+21] Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Tie- Yan Liu. Distribu- tional Reinforcement Learning for Multi-Dimensional Reward Functions. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
14

---

Appendices
Appendix A
In-Depth Summary of Related Work
15
Appendix B Proofs
16
B.1
Multivariate Distributional Dynamic Programming: Section 3
16
B.2 EWP Dynamic Programming: Section
18
B.3
Categorical Dynamic Programming: Section
20
B.3.1
Quality of the Categorical Fixed Point
22
B.4
Categorical TD Learning: Section 6
24
B.4.1 The Signed Measure Relaxation
24
B.4.2 Convergence of Categorical TD Learning
26
Appendix C Memory Efficiency of Randomized EWP Dynamic Programming
30
Appendix D
Nonlinearity of the Categorical MMD Projection
33
Appendix E
Experiment Details
33
Appendix Neural Multivariate Distributional TD-Learning
33
A In-Depth Summary of Related Work
In Sections and 2, we provided high-level synopsis of the state of existing work in multivariate distributional RL. In this section, we elaborate further.
Analysis techniques. Our results in this paper mostly drawn the analysis of one-dimensional distributional RL algorithms such as categorical and quantile dynamic programming, as well as their cemporal-difference learning counterparts [RBD+ 18, DRBM18, RMA +24 BDR23]. The proof techniques in these works themselves are related to contraction-basee arguments for reinforcement learning with function approximation [Tsi94, BT96, TVR97].
Multivariate distributional RL algorithms. Several prior works have contributed algorithms for multivariate distributional reinforcement learning, along with empirical demonstrations and theoretical analysis, though as we note in the main paper, the approaches proposed in this paper the first algorithms with strong theoretical guarantees and efficient tabular implementations. [FSMT19] propose a deep-learning-beee brase approach in which generative adversarial networks are used to model multivariate return distributions, and use these predictions to inform exploration strategies. [ZCZ=2) 21] propose the TD algorithm combing equally-weighted particle representations and an MMD loss that recall in Equation (3). They demonstrate the effectiveness of this algorithm in combination with deep learning function approximators, though do not analyze it. [WUS23] propose a family of algorithms for multivariate distributional RL termed fitted likelihood evaluation. These methods mirror LSTD algorithms [BB96], iteratively minimising batch objective funcction (in this case, negative log- likelioodo NLL) over growing dataset. [WUS23] demonstrate that these algorithms are performant in now-dimensione settings empirically, and provide theoretical analysis for FLE algorithms, assuming an oracle which can approximately optimise the NLL objective each algorithm step. [SFS24] also propose TD learning algorithm for one-dimensional distributional RL using categorical support and an MMD-based loss. They demonstrate strong performance of this algorithm in classic RL domains such CartPole and Mountain Car, but do not analyze the algorithm. Our analysis in this paper suggests our novel relaxation to mass-1 signed measures may be crucial to obtaining straightforwardly analyzable TD algorithm.
15

---

Finally, the concurrent work of [LK24] studied distributional Bellman operators for Banach-space valued reward functions. Their work focuses on analyzing how well the fixed point of distributional fiiit-dimenniionll multivariate Bellman equation can approximate the fixed point of a distributional Banach-space space-valuedd Bellman equation. In contrast, our work only studies fnite-dimensional reward functions, but provides explicit convergence rates and approximation bounds when the distribution representations are finite dimensional, unlike [LK24]. Moreover, [LK24] considers similar algorithm to that discussed in Theorem but for categorical representations. though its convergence is not proved. Furthermore, [LK24] did not prove convergence of any TD-learning algorithms, although they did propose some TD-learning algorithms which achieved interesting results in simulation.
B
Proofs
B.1 Multivariate Distributional Dynamic Programming: Section 3
In this section, will state some lemmas building up the proof of Theorem 2. These lemmas generalize corresponding results of [NTGV20] that were specific to the scalar reward setting. We begin with lemma that demonstrates notion of convexity for the MMD induced by conditional positive definite kernel.
Lemma 6. Let (PolaEI C. P(Y) and (qa)ez C. P()) be collections of probability measures indexed by an index set I. Suppose E P(T). Then for any characteristic kernel R, the following holds,
MMDx, [po), Ea'nT (gw:)) S sup MMD,,
aEI
Proof. is known from [Sch00] that characteristic kernels generate RKHSs H. into which probability measures be embedded. As such, it holds that
MMD.(p.g) = -111pp Myll
where || is the norm in the Hilbert space and Mp is the mean embedding of p that is, the unique element of H. such that Eysu [f() = (f. pp) for every f E H. and where (.. is the inner product in H.
Let Тр A EasT [pu] and define Tg analogously. We claim that prp = EanT [Mp.]] 2 Tip.- To see
this, let f H. and observe that
/ f(y)Tp(dy)
E y~Tp
(v)
/ (())I(da)padd))
// (d))
/ (J.Hp.)T(da)
( f, J.f(1e
(f. TPp)
16

---

where the third step invokes Fubini's 's theorem, and the penultimate steps leverages the linearity of the inner product. Notably, T acts linear operator on mean embeddings. As result, we see that
MMD,,(Pp Tq)
1pTp PTyll
[ Up Tuyl
I (Ny. - Mg (daa)
1lp.. Mga (T(aa)
sup l1,p. Maal aE2
sup MMDxe( Pga). aEI
where the penultimate inequality is due Jensen's inequality, and the final inequality holds since suPaez Pp.. Mya upper bounds the integrand, and the integral is monotone operator.
Next, we show how the MMD,, under the kernels hypothesized in Theorem 2 behave under affine transformations to random variables.
Lemma Let K. be a kernel induced by semimetric p of strong negative type defined over compuct subset y C Rd that is both shift invariant and c-homogeneous (cf. Theorem 2). Then for any y and p, q E P(Y),
MMD,((ba,)eee (ba,)) ( M//MMM,
Proof. is known [GBR+ 12] that the MMD can be expressed in terms of expected kernel evaluations, according to
MMD q) [ss( (O(YY)))+ + E
2E (x(V, Z)]
(Y,Y", (.y
2 (p(Y, yo) + pY" - ,((,
= E
[ yo) + ((Z vo) - p(Z,2') Z'))
+
2
[0 3o) + ((2.00) p( 2]]
E [pY,, 2)
1
E E 2l)(Z
2
where the last step invokes the definition of a kernel induced by semimetric, and the linearity of expectation. Then, defining Y, Y' independent samples from (ba,))p and Z, Z' independent samples from (bag)ra. we have
MMD2 (ba))2g)= E [ }
81) - (BB1::))))(((((())""))))>>)) =E(((X,,22)) - 2 (BlotX."""")++
1
E (
E [o(2,21)
2
E 7Y, + +2)
76 (B(o(Y, Y')] EE
= =7'E [DY, 2))
<PMMM)((p.)))
where the second step is change of variables, the third step invokes the shift invariance of p. and the fourth step invokes the homogeneity of
Thus, it follows that MMD, (bP, (baz):4) S M//M q).
We are now ready prove the main result of this section.
17

---

Theorem 2 (Convergent MMD dynamic programming for the multi- distribution function). Let
be kernel induced by a semimetric p on [0, (1 (1 - with strong negative rype, satisfying
1. Shift-invariance. For any z Rd, p(z + v1, 2 + 12) p(y1,
2. Homogeneity. For any Y E [0, 1), there exists c > 0 for which p(yy,,, = yf p(y1, y2).
Consider the sequence {Mk]k=1 given by mk+1 = T*
Tik- Then Tp - 17" at a geometric rate of 7t:/2
MMD x as long as MM, 11) S CC<. 00.
Proof. We begin by showing that the distributional Bellman operator T contractive in MMD We have
MT TT)
sup MMD (T), =EX
sup MMD
E [(bre))))) P (2'mPT()la)
E
[(bre)))))e((""))
=EX
We apply Lemma with I = and = Px x). yielding
MMD,(T 12) S sup sup MMD,, ((br(a))) (x'), (br(e))))na((()) x'Ex
Next, invoking Lemma 7 with the shift-invariance and c-homogeneity of n, we have
HIM,.(TTMm m, 12) S g0 sup sup MMD,, (n1 (x'), (=a), VE 2'EX
= -002 sup MMD, (m (2),
TE
MMD (71, 72).
follows that MMD, 11+1, 17*) S MM The Tn1²') // (T me of) since
17th is a fixed point of T Continuing, we see that MMD, (Tk, 17¹T S 7th /e/2 x (7o- 17t) S
ka/22 ke)k)/2) o1). Since MMD,, is metric on P([O, (1 7) Rmmax ddy for any
characteristic kernel it follows that mh approaches 17" at geometric rate.
B.2 EWP Dynamic Programming:
Section
In this section, we provide the proof of Theorem 3. To do so, prove abstract, general result, regarding any projection mapping that approximates the argmin MMD projection given in Equation (4).
Proposition (Convergence of EWP Dynamic Programming). Consider a kernel satisfying the hypotheses of Theorem 2, suppose rewards are globally bounded in each dimension in [0, Rmax), and let (TT)) h220 be a sequence of maps II: P([O, (1 - ->)=1Rmmax - 6EWP,m. satisfying
MMD
k.m () ())))))),))))))))))) n(=)) S f(d, m) <0 Vk: 2 0.
Then the iterates (Tx)k20 given by Mk+11 (k) T Th with MMD,(1o: 17 D < converge th N B(n (1 ff in MMD,,. where B denotes the closed ball in
MMD
B(n, R) {rr P(RR)* MMD. =
Proof. Let A, MMD,(m. 17 Then we have
A. MMD,, (K) T 1]6 17") k.
MMD,, (k) T 1k-1, T mm-1)) + MMD (T" m-- T7
(K)
< fd, m) / (AAAD,.(K-110 (7*-1.77"))
A S f(d, m)) +/
18

---

where the first step invokes the identity that 77" is the fixed point of 7 (which is well-defined by
Theorem 2), the second step leverages the triangle inequality, and the third step follows by the definition of [d, m.) and the contractivity of T established in Theorem 2. Unrolling the recurrence above, we have
(MK 17) = 5 2 + +[dd m)
1=0
S / D + [(d, m)
1 g0/2"
As such, as - oo, we have that
11" 1- -2ar2 )) =
f(d. m)
lim MMD, A-000
I1ki B
proving our claim.
Proposition 1, despite simplicity, reveals an interesting fact: given good enough approximate MMD projection Im in the sense that [d, m)) decays quickly with m. the dynamic programming iterates (7% 2220 will eventually be contained in (arbitrarily) small neighborhood 1y" The next result provides an example consequence of this abstract result, and establishes that poly(d) is enough for convergence to an arbitrarily small set with projected distributional dynamic programming under the EWP representation.
Finally, we can now prove Theorem 3, which we restate below for convenience.
Theorem 3. Consider a kernel induced by the semimetric p(a,, y) || - y with E (0, 2), and suppose rewards are bounded in each dimension within [0, Rmax). For any no such that MMD, (1).,7) ) S D) < cc, and any >> O, for the sequence (m=)220 defined in Equation (5), we have
with probability at least
(TK v7)) ēō (
d0.12 Rmax
log--a )):
1211-11
(11-)/2)(3 7) log
where 1 PojoProj and K-
and where Ô omits logarithmic m.
log y
Proof. For each x E X' and k [K]. denote by Ex the event given by
/d Rinax
4d 2
Rmax log 8'-1
Ea,k MMMt (x)). (x)) < (1 -1-71 - +
71 m =: f(d, m. 8")
for 8' > 0 constant to be chosen shortly. Moreover, with E = (x it holds that under& E, MMD,, (BootProj T7k mh) S [d, for all k [K]. Following the proof of Proposition 1, we have that, conditioned on E,
MMD, 17") S 7k/2 f'd, m; 5')
1 g0122
22 2 (1 /m + f(d, m2
70/2
Now, by [TSM17, Proposition A. event Ex holds with probability least - 8', since each (BootProj". is generated independently by sampling m independent draws from the distribu- tion Thp. Therefore, event & holds with probability at least |x K8. Choosing 8' = 5/ KC), have that, with probability at least 1 - 5,
2/ Rexa
2/ 2 Reax
MMDa(1..77
(1 2log(
(1 7)"
1 1- 1-70/2 (1 7)"
2d Re
2/ 1.at Reaxx
[ (1+
log m log xa
5-1 )
<
1 2 log
(1 /m
(1 21-0a/2)( /m
19

---

Following technique of [SZS+08] we establish that IT& C.r. can be computed as solution to As such, there exist universal constants Co, C1 R+ such that
d R max
log log 7 7
MMD "7) C
1 2 log
|
8-1
-1-100 (0/2)(2(1 -17)
da/2 Rmaxx
log m log
$ Co
log |22| |
log
+ log 5-1)
(14)
(0/2)(1 (1
d=/2 R
log |2
Co -0 (1- (1-2//2) 7) mm
+ loo m
log
Corollary 1. For any kernel K satisfying the hypotheses of Theorem 3, and for any no E CEWP m for which (To. I D < co, the iterates Tj++1 EWP, TT converge to set DEWP, C
'6EWP.0. ms where
6/ Ro
B 17"
max
T EWP,
/2- 7)
Proof. Proposition 1 shows that projected EWP dynamic programming converges to set with radius controlled by the quantity fd, m) that upper bounds the distance f(d, m) between (z) and that with nonzero probability,
(M) (2) at the worst state In the proof of Theorem 3,
6d"/2 R
(I (1-3)"vm Thus, there exists projection satisfying this bound. Since the projection EwP, is, by definition, the projection with the smallest possible [d, m). the claim follows directly by Proposition 1.
the randomized projections satisfy f(d, m.)
B.3 Categorical Dynamic Programming: Section
Lemma 1. Let i. be kernel induced by semimetric p on [0, (1 - (1-7)-1RR, with strong negative
rype (ef. Theorem 1). Then 1 is well bell-defined, RII C CR. and II V5 id.
Proof. Firstly, note that AR(Z) is a bounded, limte-dimensional subspace for each e X. Thus, AR(Z) is compact, and by the continuity the MMD, the infimum (7) is attained.
particular quadratic program with convex constraints.
Let K E R N(x) denote a matrix where Kú, (x((x),, {(z),). Since is a positive definite kernel when is is characteristic [GBR* 12], follows that K is a positive definite matrix. Then, for any P([O, (1 - -1 ) we have
arginf MMD, g) pEAR(E)
arrgnn MMD e) pEAR(S)
BER N(o)
N(2) N(x) Σ E =1
N(a) Σ Ps ((E()))))e((d))++))
= arginl pEAR(S)
Σ
(emo))))dd(((((() R₂
1=1
= =arginf pEAR(.))
{p Kp 2p bb.
where (K, R, g) is independent of p, so it does not influence the minimization. Now, since K is positive definite (by virtue of being characteristic) and AR(²) is closed convex subset of R NRN()) (4) it well-known that there is unique optimum, and the infimum above is attained for some p* E AR(Z)- Therefore, I is indeed well-defined, and its range is contained I confirming the first two claims. Finally, since 11R is well-defined and since MMD, is nonnegative and separates points,
20

---

I must map elements of AR()) to themselves - this is because MMD, (p. p) = 0 for the kernels we consider.
Lemma 2. Under the conditions of Lemma I, IT a nonexpansion in MMD,.- That is, for any T, E P([0, (1-7)-1 I,)x, we have MMD, (II 7h MC 72) S MMD,. 12).
Proof. Fix any x E X' and denote (x)) = {up E H: E AR()). where H is the RKHS induced by the kernel and My denotes the mean embedding of p in this RKHS. It is simple verify that p Hp is linear: for any p. 9 E P(RR) and a, B E R. for f E H with |I 1 we have
(SMoppt8a) = / [w))ap + (d)(d)) = =a / (())odd)) + B [(w))(du))
= (a, oup + Bpa);
which implies that Map+88q = 04pp Bpg. As consequence, M(x) inherits convexity from AR(
We claim that (x) is closed as a subset of H. Since p + Mp is an injection [GBR+ 12], by Lemma
since there is a unique q E AR( minimizing MMD,.(p; q). there is unique P4 E .M) attaining
the infimum || Mall over (x). Let f H \ (xr). Then there exists Mag M such that
Ilng sall = inf || - 2/ and since 1ig F ft, it follows that ve(e)(()) || μ4| = 6 <>>>> 0.
Since this is true for any μ E M (z). it follows that H. \ ((x) is open, so M (z) is closed.
We will now show that (2)) + I (2)) is nonexpansion in H. Let 77 E ER and denote by ya (3)), (2( the mean embeddings of m (I), 72(x). We slightly abuse notation and write (x; to denote the mean embedding of (1 Since M(x) is convex, for any u(x) E M(x) and
E [0, 1] we have
MMD,(( (x), (1))) = |l 1 (x)
< [1 (x) (Au(z)) + (1 - XITT MI (z)))1
[1) 11 (x) -((())())--))))))))))(()) 11 (z))
Now, by expanding the squared norms and taking À 1 O, since (xr) closed have (un (x) 11 1 (x)) - 11, (x)) < 0
(x(x), M(x)
(I (2)x) (2(x), 11 (2(x) - (x)) S0. 0,
where the second inequality follows by applying the same logic to p2(x)) Choosing (xx)) IT (z)) (2(2) = (1))) E M(x) and adding these inequalities yields
(1)(x) #2(x)) + (II (1 118 ((0a(1(())::)00 1I (z)()) - -I S..
Expanding, we see that
II me())) (100),1222))))
HI xx
S (12(1)) 11 (1(1), U - 1T (1))
[1u2(x)) 1n (x)
IT (x))
MMD,(m (m)) MMD,, MII,,(IIS (x), (2 < MMD.(m (z).
(z), 18,
confirming that nz()) + IT is non-expansion. It follows that
MMD, M, 118 sup (M) #eX
sup MMD,, (2), =
(M,,
21

---

Corollary 2. Let R. be a kernel satisfying the conditions of Theorem 2. Then for any E ER. the converge geometrically a unique fixed point.
iterates [= given by T Tx TTT
Proof. Combining Theorem and Lemma 2, see that
(T T m, IT TT S MMD, (T Th, 12) S 7 MMD (m. 72)
for some O. Thus, _ T is contraction on (6R, M.DD.. If H is the RKHS induced by K, showed in Lemma 2 that GR is isometric to a product of closed, convex subsets of H. Therefore, by the completeness H. GR is isometric to a complete subspace, and consequently xR a complete subspace under the metric MMD,.. Invoking the Banach fixed-point theorem, follows that IT T has unique fixed point x and T)s -> 12 geometrically.
B.3.1 Quality of the Categorical Fixed Point
As we saw in our analysis of multivariate DRL with EWP representations, the distance between distribution and its projection (as a function of m, d) plays major role controlling the approxi- mation error in projected distributional dynamic programming. Before proving the main results of this section, we begin by analyzing this quantity by reducing it to the largest distance between points among certain partitions of the space of returns.
Lemma Let k be kernel satisfying the conditions of Lemma 1, and for any finite R R define II P(R") - AR via 1Ip arginf gEAR. ge MMD,(pp q). Then MMD (IIp,) p) S Pe mesh (P; k).
Proof. Our proof proceeds by establishing approximation bounds of Riemann sums to the Bochner integral Up. similar to [VN02]. Let P E PR. Abusing notation to denote by 1pp the probability of the ith atom of the discrete support under 1lp, we have
MMD πp)
/ 1p
/ ((C::)[Ip(dd)) - k(1 (l(ly)
Σ m(: 2, Mp) - 1ann (n((w))(dd))
where R = {2 for some n E N. Since Hp optimizes the MMD over all probability vectors in AR. for q E Av with qi p(0.) for the ith element of have
Σ k(, 2(.2))p((.)) / k(:,))((d))
MMD mIp) -
Σ / (m((; - x(.,)))(ddy)
Σ sup
||<(- v1) x(. ((.:22)lp((0.)
sup sup ||<(; - x(1.
PEP yn an
was shown by [SSGF13] that z + k(. 2) is an isometry from (Rd 22/2) to H. where H the RKHS induced by k. Thus, we have
MM 1Ip) < sup sup p(y,, U2) mesh(P;
Since this is true for any partition MR. R. the claim follows by taking the infimum over PR.
We now move on to the main results.
22

---

Theorem 4. Let R. be kernel induced by a c-homogeneous and thift-invarian semimetric
conforming to the conditions of Theorem 2. Then the fixed point ne.x of I T satisfies
MMD, (TC. 77) 7²) golf
sup inf DEd RE)
mesh(P:k).
(9)
Proof. The proof begins in similar manner [RBD*18, Proposition 3]. Given that 11 is nonexpansion as shown in Lemma 2, we have
MMD 71) sup (
=EX
< sup MMDD, (1), NI (1)) + MMD,,
*EX'
MMD, T MMD, ,7""
113
MMD, T TO TT, T" +
MMD (T NO Tn1"") MMD (UC " "1")
MMD (NT. 7T MMD. WI
MMD (NC. "
7//2 AMI
where (a) leverages the fact that nē.x is the fixed point of 11 T and that 17" is the fixed point of T (6) follows since _ is nonexpansion by Lemma 2, and (c) follows by the contractivity of T established in Theorem 2. Finally, by Lemma 3, we have
1
MMD 7) g0/3 SE
sup MMD, ( " (x), 1- 7012 xEX PE
sup inf v meshh(P; (P:).
Finally, explicitly derive convergence rate for particular support map under the energy distance kernels.
Corollary 3. Lei R(x)= = Um. where Um is a set of m uniforrny-spaceed support points on [0, (1
7)- Rmax]. For induced by the semimetric p(z, 3) || y 1/ for a E (0, 2).
d Ron/2
(M. (1- (1- - 7)) (m - 2)a/2
Proof. We begin bounding mesh( k)) Assume m = nd for some n E N. We consider a partition P C Pu., consisting of d-dimensional hypercubes with side length (1 -) Rmax/(t (n - 1). By definition Um. clear that these hypereubes cover [0, (1 (1-7)-1 Rmax] | and each contain exactly one support point. Now, for each E P. we have
sup p(v 122 S y 91 22 E
(y (1
7) Rmaax/ (n- 1)i
where I is the vector of all ones, and is any element in A. Expanding, we have
a/2
Raaxx
sup p(y1: v2) S (1 -a
91.3 E8
n2
/ p2
( 1)a²
23

---

Since this bound holds for any 0 E P. invoking Theorem 4 yields
1
1 - -g///2 MMDx 7²)=
sup inf ZE Pm
V /mess(PP P: x))
sup DEd
mesh((:;s)
1 7a/2
1
da//2 Re
1 - ga/2
sup =E,t
(1 7) ( 2)
/ Runas
(1 (p// - 7/)0/2 2 (n 1 1///
/ Rink
(1 (1- 7)a/2 / (n 1)a/2
da/ 120/2
(1 (1-70/3](1 - 7)a/2 (ml 1/d 1)a/2"
instead m E ((n - 1) n we omit all but (n - 1)d of the support points, and achieve
d²/4 R
MMD, 77) (1 (1- - ////² ([m>/d) 1)a/2
Alternativlly, we may write
d R MT 7²) (1 (1-10/3)(1-9)) - (m 2)
B.4 Categorical TD Learning: Section 6
In this section, we prove results leading up to and including the convergence of the categorical TD- learning algorithm over mass-1 1 signed measures. First, in Section B.4.1, we show that MMD,, is in fact metric on the space mass-l signed measures, and establish that the multivariate distributional Bellman operator is contractive under these distribution representations. Subsequently, in Section B.4.2, we analyze the temporal difference learning algorithm leveraging the results from Section B.4.1.
B.4.1 The Signed Measure Relaxation
We begin by establishing that MMD, is a metric on My) for spaces y under which MMD,, is metric on P())).
Lemma 4. Let K:::: y x - R be a characteristic kernel over some space y. Then MMD,, (2) x M²))) + R+ given by (p, q) + 1|p-- pqllH defines a metric on (Y), where Up denotes the usual mean embedding of and H. is the RKHS with kernel n.
Proof. Naturally, since MMD,; given by a norm, it is non-negative, symmetric, and satisfies the triangle inequality. We must show that MMD,, (p, 9) O + g. Firstly, it is clear that MMD, = by the positive homogeneity of the norm. It remains to show that MMD, (p. 9) O p=
L (J)).. For the sake of contradiction, assume that MMD.. that this implies that MM,,(P,Q) = 0 for a pair of distinct probability measures, which is contradiction since MMD,, with characteristic k known be a metric on P())).
(p, q) O. We will show
By the Hahn decomposition theorem, we may write = pt P for neg-negative measures pt.p Therefore, for some R,,, we may express
p (a + 1)pt op
24

---

where p=,p= E P()). Likewise, there exist b E R4 and probability measure q².9 ,9 for which q (b 1)q² bq Since (P. q) by hypothesis, and by linearity of p + Up. have
0 1| - Mgllze
(aa + 1)lppt -app + bygn (b + 1)1t ll
II(a + 1)lpp
+ bbug -appp + + + 1) llp
= (a + 1-b+11 +++++1) ( + (1- (1-A))1rr
(xug- + (1-1X))
)
x a a+1
X
a + b+b+11
a a a+b+1
= + + 1) 1 Mallu,
where P= Ap + (1 A)oJ and Q = Xp + X)xx are convex combinations of probability measures, and are therefore probability measures themselves. So, we have that
Apt Xpp = (1 X)x -
101 -11- A)q
(a + 1)Apt op = (b + 1)q² - bg
:. = 9,
which contradicts our hypothesis. Therefore, MMD (p, q) =0 + p= for any P, q E M M*()) and it follows that MMD, a metric.
Next, we show that the distributional Bellman operator is contractive on the space of mass- 1 signed measure return distribution representations.
Lemma 5. Under the conditions of Corollary 2, the projected operator I T" PR PR is affine, is contractive with contraction factor 00/2, and has a unique fixed point Tisc.
Proof. Indeed, 11g,, x is, in a sense, simpler operator than 1 Since subspace of ' holds that SC. is simply Hilbertian projection, which is known to be affine and a nonexpansion [Lax02]. Moreover, since 7 acts identically on (R)) it does on P(R.") d) it immediately follows that T is /0/2 ///////ntracttionn
is an affine
M (R"), inheriting the result from
Theorem 2. Thus, we have that for any 7h, 7/2 E RR.
MMD, SC. T"
II SC.R T 12)))) MMD,, R. (TMDD,("T..... (T",,
MMMD,((
confirming that the projected operator a contraction. each x E X, follows that MMD,, metric on .R. The existence and uniqueness of the fixed point TrC follows by the Banach fixed point theorem.
Since MMD,; is metric on (R((())) for
Finally, we show that the fixed point of distributional dynamic programming with signed measure representations is roughly as accurate as ne
Theorem 5. Under the conditions of Lemma 5, we have that
MMD,(1EC.. )) S I-9072 suPxEx PeS /mesh(P; k); and
1
2. MMD, (II M 77)) S (1 + - suPpex in v me)sh(P:)..
Proof. Since 11 SC. is nonexpansion in MMD,, by Lemma 5, following the procedure of Theorem
4, we have
1
MMD (1EC.. 1 1 1- 20/2
MMD, (UD "C.."" .))).
Note that ITR "f identifies the closest point (in MMD,) to 17" in FR = ILex M and II nn identifies the closest point 17t in CR := IIeex AR(z)- Since it clear that ER C PR. follows that
MMD AMDD.nnSc.s 17) 1 20/2
M )²))
25

---

The first statement then directly follows since it was shown in Lemma 3 that MMD, (II Cxx supzEx inf PE(*R()) V mesh((P;
",",,, .1")) S
To prove the second statement, we apply the triangle inequality to yield
(T NSC. 7") S MMD CH, MSC,, IT "T MMD, C,""" ;""""")"
MMD CTMD,,TSC... 7") + MMD,, (I (1, """"", ,""),
where the second step leverages the fact that HT is nonexpansion in MMD, by Lemma 2. Applying the conclusion the first statement as well as the bound on MMD, (I,,,, ",""""") we have
(L ;s...:))) 17)) 1 70/2 DE PE
sup inf
V mmssh(P;K)
sup inf v /mess(P:x) ZE RE
1 1 - 70/2
sup inf sEx PE /mesh(P;i))
B.4.2 Convergence of Categorical TD Learning
Convergence of the proposed categorical TD-Jearing algorithm will rely on studying the iterates through an isometry to an affine subspace IL R N(2), This affine subspace that consisting of the set of state-conditioned "signed probability vectors". We define R as an affine subspace of R." for any n E N according to
R
VER"
(15)
We note that any element of R can be encoded in IL R of signed masses associated to each atom of R(x). In Lemma 8, we exhibit an isometry I between PR and ILe RN((2)
((x)
by expresing 7(x) as the sequence
Lemma 8. Let be characteristic kernel. There exists an affine isometric isomorphism I between PR and affine subspace IIrex Ry (()) (ef. (15)).
Proof. Since is characteristic, it is positive definite [GBR 12]. Thus, for each x X, define K₂ R N(x) according to
(Ka)i, n(2;, 2z)
R(x) { 1= N(2)
Then each K₂ is positive definite since x is a positive definite kernel. Let E AR(a), and define P E RNV()) and Q R N(x) such that P. p(z)) and Qi q(z)). Then, we have MMD g) = 11up
Mg 1 l
((#)
N(ar)
Σ
k(, 2()x))(()))) - Σ
x(,,
(()(())(z))
H
N()
N(x)
Σ
x(s x(<2))((2) - q(2+)),
Σ
x(- x((2z))((2z) -
j=1
(() (())
Σ
Σ (p(2)) - a(<)))(((z}) - ((,,)))((z,)2))
d=1 1=l
(P Q) (P Q)
= PPP - Q1
Since K2 is positive definite, || ||x, is a norm on R N(x) Therefore, the map I3 AAR()),MMD, + (()),, ||| (lx.) given by (p) P is a linear isometric isomorphism
26

---

onto the affine subspace of R (a)) with entries summing to 1, which we denote R (x) Moreover, since (R, (=)) Il llx,) is finite dimensional Hilbert space, it is well known that there exists linear isometric isomorphism I (R)) ((x) (()(), -I llk.)
((()) with the usual L norm. Thus, I, = 1 (AR(z), MMD,, + R ((m) is a linear isometric isomorphism. Consequently, follows that I (6R, MM),;) - ILe R((e)
(x)) given by I (TTxxx Rx(x), Ill 112,2))) linear isometric isomorphism, where ||o||2,50 = suprex l|x())]22-
Lemma Define the operator us ILLEx R R + IIEEx R R by ltt SC.xT T"I-1,
where is the isometry of Lemma 8. Let {Ur) be given by Ur = Iık, where {mete= are the dynamic programming iterates 7t+1 = "c,xT" Then U2t1 = U. Moreover, Utth is contractive whenever Il T is, and in this case, U2 - U*, where U is the unique fixed point of ur
Proof. By definition, have
UL1 Imt+1
I(I TTTT
-IT& TTITTTIIIA Ut)
U,,
which proves the first claim. Moreover, for U1 = Im and U2 Imz, we have
| - U2 | T m III T m2 12,00,00
MMDD (IT C, T m. II T Tm1)),
where the second step transforms the || 112,cc to MMD, since I is an isometry between those metric spaces. Therefore, if I T contractive with contraction factor (0, 1), we have
|ICT - U2 s BMMD 72)
= 3 [Tm I1ya 2,,0
= B |11 - U2112,00
so that L1 has the same contraction factor IT TT, Consequently, by the Banach fixed point theorem, U1T has a unique fixed point U* whenever _,, T contractive, and U. + U at the same rate 7/2 17"
The main theorem in this section is that temporal difference learning on the finite dimensional representations I(m) converges.
Proposition 2 (Convergence of categorical temporal difference learning). Let {U+) c ILE RxN(s) be given by U, = Iīi and let be a kernel satisfying the conditions of Theorem Suppose that, for each æ X, the states {Xt) and step sizes {act satisfy the Robbins-Man Munnro conditions
Σ or (x,=]] = 00 t=
Σ
( <
Then, with probability 1, Uk + U* where U is the fixed point of UT
The proof of this result natural extension of the convergence analysis of Categorical TD Learning given in [BDR23] to the multivariate return setting under the supremal MMD metric. The analysis hinges on the following general lemma.
Lemma 10 ([BDR23, Theorem 6.9]). Let O (R")& + (R")& be contractive operator with respect to || 1|2,00 with fixed point Z*, and let (R,F, (Fe)er,, P) be a filtered probability space.
Define map Ö (R")* x X Ω (R"yx such that
27

---

Ep
x x= = (0))
(16)
For stochastic process [Sk]==) adupted to [FR]k=1 with Ek Xx @ wks consider sequence {Z (R")* given by
+ +o (1-0z)2((1) wz)
x==2
(x))
(17)
(x(x)
otherwise
where {ax)k=1 is adapted to {Fx and satisfies the Robbins-Munn conditions for each X,
Oek [x== =
Σ k=1
of < 00.
Finally, for the processes {w(x) 12 k=1 where w(x); (k: (x assume the following moment condition holds,
Ep [1w( FR) S1 + C2 12
(18)
for finite universal constants C1, C2. Then, with probability 1, Zx Z*
The operator of Lemma will be substituted with Ut" governing the dynamics the encoded iterates the multi-return distribution. The stochastic operator O will be substituted with the corresponding stochastic TD operator for Lt given by
(I (bR,):I- (II) (U(z)
tn (U, (R,
T =
(19)
otherwise.
This corresponds to applying a Bellman backup from a stochastic reward R and next state 22, followed by projecting back onto PR, and applying the isometry back to II R N(a).
Proof of Proposition 2. Let max N (z). Note that for any S R can be isometrically embedded into R. by pero-padding. Thus, IL R RaN(a)) can be isometrically embedded into (RT)*; so without loss of generality, we will assume that N(x)) = n.
Define the map यल (RT)*
(R xX) (R)) X according
[IITT SC. (I-
(Ú'(U,I,, 22) (z)
21 otherwise
(20)
U((r)
Then, defining = -ç==ū Xx. (RR, Xx)) with (X,, Ak, Rr. X = Tg ~ P as in (11), we have
(x+1(x)= (IÎ + + (x)
I (1 = IT (br. (X) 1x (c))
Z (bR M (X.) + 1|x (n(x)
(o).
Note that, since 1ggh is a Hilbert projection onto an affine subspace, it is affine [Lax02]. Conse- quently, ú is an unbiased estimator of the operator 4 in the following sense,
Ep (U, Xu, (Ra, XR)) X σ Ep [III x
(bRA.)::I ( U(()]
IIT Ex [00+0))))
-I T - ((FUU)))
28

---

where the first step invokes the linearity of IS the second step invokes the linearity of the isometry Z established in Lemma and the third step due to the definition of T As result, we see that the conditions (16) and (17) of Lemma 10 are satisfied by int, the iterates {Ux} {Un]k=I> and the step sizes {o Moreover, for (x) defined by
(x(xx= = (UK, Xk. (RR, X()) (XUU)(X))) [[x6=x]
we have [=)( S Ci (z) for universal constants Ci, Chishis is shown in Lemma 11. As such, the condition of (18) satisfied.
Finally, since lr inherits contractivity from C T as shown in Lemma 5, we may invoke Lemma 10, which ensures that Us + U with probability where U unique fixed point. C
Theorem For kernel induced by a semimetric p of strong negative type, the sequence (felr= given by (11)-(13) converges 1isc.x with probability 1.
Proof. By Proposition 2, the sequence (Udr= with U2 = Inp, converges to a unique fixed point U with probability 1. Note that
U --=
U U
-IlSo T (I''
Therefore, I-10* is a fixed point of unique fixed point, it follows that 1-* năc.. Since I is an isometry, î) I - I I- with probability 1, indeed î/: Tjcc. with probability 1.
SC, T Since it was shown in Lemma that SC,s
7 has
To conclude, we prove Lemma 11, which was invoked in the proof of Proposition 2.
Lemma 11. Under the conditions of Proposition 2, it holds that for any x E X and U e ILEx R N(a)-1
E (=
+
for finite constants C1 C2 E R+.
Proof. Since I is isometry, we have that
(xF (U, (1,,
(1,21)) = | T"I (x) TTo. ((br,))z-'(r") (x')) (4))
where H is the RKHS induced by the kernel Moreover, since 11 nonexpansion in || || as argued in Lemma 5, we have that
E (-Pl(1la)
<
E
[|T"I--U(()))
(x) (( U(X) (x)
< 2 2
Ux)
+2 E x-Px(112)
Proceeding, we will bound the terms A, B. To bound A, we simply have
A - - 77" (x) + || (x)
I 1 17 (x) + || (x)
29

---

where we invoke the contraction of 24 in MMD,, from Theorem 2. Note that 17 (x) E P(OO, (1 -
- Rmax ), so it follows that || S D1,1 1 for some constant D1.1 since the kernel & is bounded compact domains. Expanding the norm of the difference above yields
A (1 + + +
= (1 (1+ + + + (r))|
for finite constant 21,2, again invoking the isometry I in the last step.
Our bound for is similar. Choose any x' E supp PT x). We consider the operator T P[[O, (1 - -1 R ) -> P([0, (1 1 (x) given by
(z) = (1(())
This operator a contraction in MMDx,, and correspondingly has a fixed point n. To see this, we note that T is simply a special case of ldth for the case (. | x) = 6g, so the contractivity and existence of the fixed point are inherited from Theorem 2. Proceeding in manner similar to the bound on A, we have
(x) =
| (
< Toa -0--000 U((x) -7aa | + Il (x)
3 7 / -U(1)))((((())))))) - Nh 11 + || (x)
$ (1 + ) + D2.2||U(x))I
where the final step mirrors the bound on A. Therefore, we have shown that
|ErU()))) Ūtth (U, (R(x)) X))()))) < 2(1 + (1,, + D2,1) + 1D1.2 + D2:2
E X-P=(1a)) (.1x)
completing the proof.
C Memory Efficiency of Randomized EWP Dynamic Programming
In Section 4, we argued for the necessity of considering projection operator in EWP dynamic programming. While we provided a randomized projection, Theorem requires that we apply only a finite amount of DP iterations. Thus, one might ask if, given that we apply only finitely many iterations, the naive unprojected EWP dynamic programming can produce accurate enough approximations of 1y" without costing too much in memory.
In this section, we demonstrate that, in fact, the algorithm described in Theorem can approximate 17 to any desired accuracy with many fewer particles. Suppose our goal is to derive some such that
MMD 7") S
for some c > 0. We will derive bounds on the number of required particles to attain such an approximation with unprojected EWP dynamic programming (denoting the number of particles Munpraj) as well as with our algorithm described in Theorem (denoting the number of particles proj- In both cases, we will compute iterates starting with some 7/o E (EWP, . with MMD, 17") S D oc. For simplicity, we will consider the energy distance kernel with a
The remainder of this section will show that the dependence of the number of atoms on both e and |X| is substantially worse in the unprojected case (that is, Mproj < Munproj for large state spaces or low error tolerance). We demonstrate this with concrete lower bounds on Munproj and upper bounds on Mprjj below; note that these bounds are not optimized for tightness generality, and are instead aimed to provide straightforward evidence of our core points above.
We will begin by bounding Munproj. In the best case, ()())) is supported on 1 particle for each x. If any state can be reached from any other state in the MDP with non-zero probability, then applying the distributional Bellman operator to no will result in 7) (x) having support on | atoms at each
30

---

state r (due to the mixture over successor states in the Bellman backup). Consequently, the iterate Th (2) will be supported on |t| atoms. Since (TMx((7) n²) S /-/²22 by Theorem 2, we require
2log(D/e) K
logy
to ensure that MMD,, 77" S Thus, we have
Munproj 2
|
log
On the other hand, the following lemma bounds Mprjj we prove the lemma at the end of this section. Lemma 12. Let Thmnss denote the output of the projected EWP algorithm described by Theorem with proj particles. Then under the assumptions of Theorem 3 and with the energy distance kernel with 1, MMD,, (Tmmag)
77") I achievable with
Mrroj c0(c7 (1 ( 7)3 polylog (
d.P max
1 3 X1, d, R Rmax, 1
(21)
For any fixed MDP with X 2 and 7 2 1/2, we have that
loge
log D log log
Munproj
exp log |.X
exp 2 log |
log 1og~-1 log D
lox [x log
exp
2og |.|
log log
since D 0 and does not depend on Meanwhile, we have proj O (oly"polyl by Lemma 12, indicating much more graceful dependence on relative to the unprojected algorithm.
On the other hand, for any fixed tolerance S DD, we immediately have
Munproj E (2
Maroj E (d polylog(d, (x))).
In the worst case, we may have d E ((22)) (any larger d will induce linearly dependent cumulants). Thus, we have
(0(181-1) d (1) (((x)-2) de O(1,,
M Mrroj
Munproj
so the projected algorithm scales much more gracefully with |. as well.
Proof of Lemma 12
Finally, we prove Lemma 12, which determines the number of atoms required achieve an e-accurate return distribution estimate with the algorithm of Theorem 3.
Lemma 12. Let Thags denote the output of the projected EWP algorithm described by Theorem with m = projj particles. Then under the assumptions of Theorem 3 and with the energy distance
kernel with 1, MMD
(MDDD(Immees (Tmmwej
Tf²)) < is achievable with
Moroj E 60(62 (1 (1 7)3 polylog
dRR max
(
1 ( 3 X(, d, Rmaxi 1 )
1
(21)
Proof. Note that, by Theorem 3, increasing Maroj can only decrease the error as long as Maroj 2 1. Therefore, as shown in (14) the proof of Theorem 3, there exists a universal constant Co > 0 such that
/ Rac R
1 log
C == Co Mipooj (1 - 72/ (1
log
+ log Maprjj
(22)
7)
31

---

Now, we write = Cocio2, = Coc1, and = v Maroj. yielding
c3
log u²
C4
u
C3
log u
2c,
u
Then, after isolating the logarithmic term and exponentiating, we see that
uc
C3
u exp
2cA
We now rearrange this expression and invoke the identity W (2))e W- ---function [CGH+96]:
where W is Lambert
ue e
exp
WE =
1
2c4
-a//2cee
e e0222
uE
ME
exp
204
204
204
2e4
eca/2e
ue
26
= 2c4
2C4
There two branches of the Lambert W-function on the reals, namely Wo and W-1- These two branches satisfy (zee=)) = 2 when x 2 -1 and W- W_1(2e") = when & S -1. our case, we know that is negative, and it is known [CGH 96] that [Wo(>)] < 1 when 1-1,0] Consequently,
when 2 2 1, we have ue 1, and substituting Miproj = u², we have 2cg
4c2 2
Mproj S
when 2 2 11
(23)
On the other hand, when 2 < 1, we have
02/2
e
= W-1
204
02/22
W-1
2c4
2c4
projj
M7
4
W2
e
2 2 -1.
2c4
Since known [CGH+96, Equations 4.19, 4.20) that W (-2) E polylog(1/2 incorporating (23), we have that
2 -1
Mpproj
W2,
c02/2c
=5
2r4
402
max (I, (Lopollog(ccee=//e=-)))
2
ACGdRR
5
(1 21 22 polylog
ē 8. 128%, d, Rmax
The upper bound given above will generally not be an integer. Howevever, increasing Mproj can only
improve the approximation error, shown in Theorem since log m/ /m/m decreases monotonically when m > 7. So, we can round Mproj up to the nearest integer (or round it down when m S 7) incurring penalty of at most atom. It follows that the randomized EWP dynamic programming
algorithm of Theorem 3 run with Mproj given by (21) produces return distribution function 77
for which (MMD,(mmaa 17")) <
32

---

Table 1: Certificate that II
is not affine
Support point 5 E
(1(E)
(2(E)
0,0)
0
0
0,
(0,2)
0
0
(0.3)
(1.0)
O
0
(1,1)
0.1999
0.2057
(1,2)
0.1999 0
0.1959 0 0.07957
(1,3)
(2,0)
0.0937
(2,1)
0.2062
0. 0.2413
(2,2)
0.1999 0
0.2026 0 0.0787
(2,3)
3.0
0.0937
(3,1)
0.0063
0
(3,2)
(3,3)
0
D Nonlinearity of the Categorical MMD Projection
Section 6, we noted that the categorical projection 1 is non-affine. Here, we provide an explicit example certifying this phenomenon.
We consider single-state MDP, since the nonlinearity issue is independent of the cardinality of the state space (the projection is applied to each state- conditoond distribution independently). We write = {0, ... 3}2, and consider the kernel induced by p(z, y) || y this resulting MMD is known as the energy distance, which is what we used in our experiments. We consider two distributions, P1 511551155 and p2 6/2.5.0]-
We consider 0.8 and compare q1 I and we note that * d2: confirming that I is not an affine map. The results are tabulated in Table with bolded entries depicting the atoms with non-negligible differences in probability under Q1
with 0 XII A. + (1-AA n.P
E
Experiment Details
TD-learning experiments
conducted NVidia A100 80G GPU parallelize experiments. Methods were implemented in Jax [BFH+18] 18], particularly with the help of JaxOpt [BBC+21] 21] for vectorizing QP solutions this helpful for computing the categorical projections discussed in this work. SGD was used for optimization, using an annealed learning rate schedule (Aa )x20 with A - 3355, satisfying the conditions of Lemma 10. Experiments with constant learning rates yielded similar results, but were less stable-this validates that the choice of learning rate schedule did not impede learning.
The dynamic programming experiments implemented in the Julia programming language [BEKS17].
In all experiments, we used the kemel induced by p(x,y) || y with reference point 0 for MMD optimization-this corresponds energy distance, and satisfies the requisite assumptions for convergent multivariate distributional dynamic programming outlined in Theorem 2.
F Neural Multivariate Distributional TD-Learning
For the sake of illustration, this section, we demonstrate that the signed categorical TD learning algorithm presented in Section can be scaled to continuous state spaces with neural networks. We will consider an environment with visual (pixel) observations of car in a parking lot, an example observation is shown in Figure 4.
33

---

Here, we consider 2-dimensional cumulants, where the first dimen- sion tracks the x coordinate of the car, and the second dimension an indicator that is if and only if the car is parked in the parking spot. We learn a multivariate return distribution function with tran- sitions sampled from trajectories that navigate around the obstacle to the parking spot. Notably, the successor features (expectation multivariate return distribution) will be zero in the first dimension, since the set of trajectories horizontally symmetric. Thus, from the successor features alone, cannot distinguish the observed policy from one that traverses straight through the obstacle!
Parking Spot
Obstadle
Figure 4: Example state in the parking environment.
Fortunately, when modeling a distribution over multivariate returns, we should see that the support of the multivariate return distribution does not include points with vanishing first dimension.
R
It
CNN
m(x+))
Ey Loss
r(xt)
T
R
R
It+1
CNN
I SC,K
17( (++++)
Figure 5: Neural architecture for modeling multi-return distributions from images.
To learn the multivariate return distribution function from images, we use convolutional neural architecture as shown in Figure
Notably, we simply use convolutional networks to model the signed masses for the fixed atomss of the categorical representation. The projection 11g computed by QP solver discussed in Section 5, and is applied only to the target distributions (thus we do not backpropagate through it).
We compared the multi-return distributions learned by our signed categorical TD method with that of [ Our results are shown Figure 6. We see that both TD-learing methods accurately estimate the distribution vanishing lateral component. Quantitatively, we see that the EWP algorithm appears to be stuck in local optimum, with some particles lying in regions of low probability mass.
multivariate returns, indicating that no multivariate return will have
Moreover, on the right side of Figure 6, we show predicted return distributions for two randomly sampled reward vectors, and quantitatively evaluate the two methods. The leftmost reward vector incentivizes the agent to take paths conservatively avoiding the obstacle on the left. The rightmost reward vector incentivizes the agent get to the parking spot as quickly possible. We see that the EWP TD learning algorithm of [ZCZ* 21] more accurately estimates the return distribution function corresponding the latter reward vector, while signed categorical TD algorithm more accurately estimates the return distribution function corresponding to the former reward vector. In both cases, both methods produce accurate estimations.
34

---

Signad Cat-TD
EWP-TD
ESR
Perood Notten
Paature
000
-29
Figure 6: Multi-return distributions learned by signed categorical TD and EWP TD, well as examples of predicted return distributions on two randomly sampled reward functions.
NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer:
[Yes]
Justification: As claimed, we provided convergence results for multivariate DRL algorithms, and tested them through simulations.
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 contribu- tions 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 settings.
It is 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: We have explicitly stated all assumptions, and our simulation section and conclusion explicitly discuss the limitations of our proposed approach.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to vi- olations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-speciffiation, asymptotic approximations only holding locally). The authors should reflect 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 with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect 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 a speech-to-ter to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 review-
grounds for rejection, worse outcome might be that reviewers discover limitations that
35

---

aren't ucknowledged 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 technical results are proved rigorously in the appendix.
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 the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but 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 in appendix 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 reproduce the main experimental results the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data provided or not)?
Answer: [Yes]
Justification: Algorithms are described explicitly, experimental setup detailed. Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
be accomplished in various ways. For example, the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 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 large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
Depending the contribution, reproducibility
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which authors welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, may be that access the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
36

---

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, as described in supplemental material? Answer: [Yes]
Justification: Code will be provided. Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (tttps://nlps.ccppbbiicc guides/ Cooeae//odeuSumt pissooooicyy for more details.
While we encourage the release of code and data, we understand that this might not possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurlPS code and data submission guidelines (https: hetlll/c/giies/Coddeee
//nips.
wwwssssiiooiiyyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts reproduce all experimental results for the new proposed method and baselines. If only 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 as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
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: Hyperparameters are disclosed.
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 level of detail that is 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: 95% confidence intervals are shown in relevant plots. 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 are 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 a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, the hypothesis of Normality of errors is not verified.
37

---

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric eπrоr bars that would yield results that are of range (e.g. negative error rates). If error bars are reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in 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:
This is discussed in the appendix.
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, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experi- mental runs as well as estimate the total compute.
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 make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, every respect, with the NeurIPS Code of Ethics https: hett::/neurppppp cc/ppbliic/tthiicccuddeelllll
Answer: [Yes]
Justification:
We adhere to the code of ethics and preserve anonymity.
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 is special consideration due 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: [NA]
Justification: The contributions are theoretical in nature, and do not have any immediate societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed. If the authors answer NA or No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.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 not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people 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.
38

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper does not pose any such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 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 (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: The paper does not use existing assets.
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 possible, include a URL.
The name of the license (e.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 twcode.coo...cc www/ com/datasets 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 of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release any new assets.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/code/mooel 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 to anonymize your assets (if applicable). You can either create an anonymized URL include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
39

---

Question: For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing 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 if the main contribution of the paper involves human subjects, then as much detail 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 Subjects
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 your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor 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 in which research is 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 adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
40


Reviews:
[{"review_id": "u8XpDa1qjV", "reviewer": "Reviewer_SWh2", "summary": "This paper studied the multivariate distributional reinforcement learning (RL) problem, in which the goal is to learn probability distribution of accumulated multi-dimensional rewards of the RL system. First, dynamic programming for multivariate distributional dynamic programming is established, then a randomized particle-based dynamic programming solution method is proposed. Furthermore, a projection-based dynamic programming algorithm is proposed for categorical multivariate distributional. Some extensions to RL setting were discussed and additional empirical observations are provided.", "questions": "* What does the right-hand-side (RHS) of the first equation in Eq. (2) mean? \n* What does the right-hand-side (RHS) of the RHS of R(x) in Line 187? Please explain all the terms including, $\\xi$, $N(x)$ and index $i$.\n* Define $\\Delta$ notation on Line 188.", "limitations": "See the weaknesses.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper studied the fundamental problem of multivariate distributional dynamic programming and RL problems. Furthermore, convergence results have been achieved for both settings under maximum mean discrepancy distance.", "weaknesses": "* Some important notations are not defined for better understanding. For example, 1) the right-hand-side (RHS) of the first equation in Eq. (2) is not clearly explained or defined; 2) the RHS of R(x) in Line 187 and the $\\Delta$ notation on Line 188 are not explained or defined. This writing style makes it difficult to understand many technical details of the paper.\n* Typo: Equation (10), $X’_r$ -> $X’_t$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback on the paper. We address the queries on notation below, and will include clarifications in the revised draft.\n\n* The RHS of equation (2) is simply the set of all empirical distributions on $m$ points. That is, the distributions obtained by picking $m$ points $\\theta_i\\in\\mathbf{R}^d$ and constructing the distributions supported on those points with equal mass. This is often referred to as a $m$-quantile representation in the $d=1$ case (see e.g. [BDR23]) and is equivalent to the model of [ZCZ+21].\n* The object $\\mathcal{R}$ in line 187 is defining a *state-dependent* support map. That is, for any input state $x$, $\\mathcal{R}(x)$ outputs a finite set of support points $\\{\\xi_i(x)\\}_{i=1}^{N(x)}\\subset\\mathbf{R}^d$. In the $d=1$ case, these would be the locations of the atoms (bins) of the categorical return distributions (see e.g. [BDM17b]) -- we generalize the notion here to have categorical supports that can depend on the state. The quantities $N(x)$ describe how many support points are in the categorical support of state $x$ -- indeed, we additionally generalize the model beyond that of [BDM17b] to accommodate state-dependent *resolution* of categorical supports. Thus, at state $x$, our model consists of $N(x)$ support points in $\\mathbf{R}^d$, and these are indexed by $i$ in the equation in question.\n* The $\\Delta$ on line 188 represents a simplex. In particular, $\\Delta_A$ for a finite set $A$ is the set of probability mass functions on $A$.\n\nThank you for pointing out these clarity issues, and we will clarify these in the revised draft. Should you have any other questions, we would be happy to discuss further.", "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": "ypkUrEVYxv", "reviewer": "Reviewer_FNo5", "summary": "In this submission, the authors combine a multivariate reward with distributional learning. They rely on a Maximum Mean Discrepancy based projection operator to obtain the first efficient and provably convergent algorithm in this setting. The key in their work is the extension to the multivariate setting of the contraction property of the Bellman operator with respect to the MMD “metric”. They use this property in two algorithms. One based on a randomized dynamic programming which maintains the number of atoms/particles through sampling and can be proved to converge in the MDD metric. The second rely on a categorical representation for which there is a MMD compatible “projection”. Finally, they show how to make these algorithms compatible with a Temporal Difference approach. All those algorithms are illustrated with numerical experiments on toy (random MDP) examples.", "questions": "- Could the authors comment why they did not observe in practice the convergence suggested by their theorem in the EWP-based technique? \n\n- In Monte Carlo approximation, the price of increasing the dimensionality is often hidden in the “variance”. Here it seems that the price is moderate with a polynomial term. Do the authors have an intuition on why such a phenomenon and why it is not related to the correlation structure between the coordinates of the reward?", "limitations": "As mentioned by the authors, the multivariate setting can only be used in the evaluation part in MDP and Reinforcement Learning and is of interest when there are more than one possible reward, which is not the most classical setting.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The results are new and the proofs seem correct .\n\n- They exploit the possible dependency between the coordinates of the multivariate reward.\n\n- The results are supported by numerical experiments.", "weaknesses": "- The experiments are made on very simple toy examples (random MDP). It would be interesting to see if their numerical results also hold with more realistic examples.\n\nTypos:\n- 614 \\widebar{MMD} -> MMD", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful assessment of our work.\n\n> Could the authors comment on why they did not observe in practice the convergence suggested by their theorem in the EWP-based technique?\n\nWe did in fact observe convergence of the EWP-based algorithms in practice -- can you please point us more specifically to where the text or results  suggest otherwise?\n\n> In Monte Carlo approximation, the price of increasing the dimensionality is often hidden in the “variance”...\n\nThis is a very interesting question! We suspect that it may be possible to provide \"instance-dependent\" bounds that sharpen with more favorable correlation structure between reward dimensions. Our results leverage the boundedness of the reward function and the structure of strongly negative-definite kernels to provide worst-case bounds. The polynomial scaling of the bound with dimensionality is inherited from the fact that the worst-case values taken on by the MMDs we consider also scale polynomially with dimension.\n\nRegarding experimental results, please see our general response, in which we demonstrate that the signed categorical TD algorithm proposed in this paper is straightforward to combine with neural network function approximation.", "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": "QJOy0ne3zh", "reviewer": "Reviewer_wmwv", "summary": "The authors propose a tractable and convergence-guaranteed method, called randomized dynamic programming, for multivariate distributional reinforcement learning (distRL). They also introduce practical algorithms, multivariate EWP-TD and signed-categorical-TD, provide an upper bound on MMD with respect to dimension $d$, and empirically compare performance using the Cramer distance in a tabular MDP environment.", "questions": "-\tIn Theorem 3, the symbol seems to be $\\leq$ instead of $\\in$.\n\n-\tI understand that the randomized DP is a sampling-based algorithm that circumvents the EWP representation in terms of MMD. However, the projection that minimizes MMD does not guarantee the uniqueness of the fixed point. How can its relaxation, the randomized DP, have a unique fixed point? Theorem 3 seems to imply a bound at $K=O(\\log m)$, suggesting that the number of atoms increases exponentially. Can Theorem 3 be analyzed with $K \\rightarrow \\infty$ for a given number of atoms $m$?\n\n-\tThe depiction of Cramer distance with respect to the number of atoms for each algorithm in Figures 2 and 3 is very interesting. While EWP-TD is hardly affected by the increase in the number of atoms, signed-Cat-TD shows a consistent decrease. When the dimension increases, will EWP-TD always have a smaller Cramer distance compared to signed-Cat-TD?\n\n\n-\tIn Line 337, can you elaborate on \"using randomized support points for the categorical algorithm\"? It is unclear whether this refers to EWP-TD or signed-Cat-TD.\n\nSince there are still some aspects I don't fully understand, I'm willing to raise the score if the authors can clearly address my questions.", "limitations": "The authors have clearly stated the assumptions and limitations of their theory.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Despite the extensive background knowledge and theoretical understanding required in this field, the authors have proficiently written the paper to be accessible to first-time readers, using standard and clear notations. \n\n- The visualizations for the proposed EWP-TD and Categorical-TD are particularly excellent and intuitive. \n\n- The authors clearly explain the issues and limitations arising from finite parameterization in distRL theory and naturally introduce the novel concept of randomized DP for EWP representation. \n\n- The theory regarding categorical representation in multivariate cases is novel and informative, and the convergence rate matches the results in univariate cases.", "weaknesses": "While the contributions of the authors are clear, it is uncertain whether the proposed theory aligns with the contributions. The proposal of a randomized DP for a tractable EWP is interesting, but its convergence and the existence of a unique fixed point do not seem to be clearly described. The detailed theory appears to focus more on the categorical representation. See the Questions section for further details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough assessment and for their interest in our work. We hope the responses below address the queries raised in the review, please let us know if you have any further questions.\n\n**Q1**: With regard to $\\leq$ vs $\\in$ in Theorem 3 -- this was a stylistic choice. We were invoking the definition of $\\widetilde{O}(\\cdot)$ as a set-valued function -- all elements of this set satisfy the upper bound that you are describing. This is the convention used, for example, in the classic CLRS *Introduction to Algorithms* text; its meaning coincides with what you are describing.\n\n**Q2**: You are correct that there is not a unique fixed point of the projected operator. The theorem is simply claiming that, after enough ($K$) iterations, we can guarantee that the resulting $\\eta_K$ will be within a very small margin of error from $\\eta^\\pi$ w.r.t. the MMD. Notably, $\\eta^\\pi$ is the unique fixed-point of the (unprojected) distributional Bellman operator, which is well-defined (by Theorem 2). Convergence results for dynamic programming algorithms without unique fixed points also exist in previous works, such as [WUS23]. Note we do not require exponentially-many atoms; the free parameter here is $m$, not $K$. Therefore, for a desired tolerance level, we require polylog-many iterations ($K$) of dynamic programming.\n\n**Q3**: This is a fantastic question. The point to note is that the convergence bound for EWP is *dimension-free*, while those of the categorical algorithms are not (they depend on $d$). However, in the case of EWP-TD, the algorithm is prone to convergence to local minima (see our discussions on lines 146 and 324). Thus, we suspect that the EWP-TD examples are likely stuck in local minima, but notably that the quality of these local minima is roughly uniformly high for enough particles (e.g., as many as we use in our experiments). As you suggest, we would expect that EWP-TD should perform favorably to the Cat-TD algorithms with high $d$ -- however this is not guaranteed. Particularly, there is no known convergence guarantee for the EWP-TD algorithm when $d > 1$. Moreover, we could potentially leverage prior knowledge of the supports of the return distributions to enhance the resolution of the Cat-TD representations in relevant areas of the space of multi-returns, which could result in improved performance.\n\n**Q4**: This was in reference to the Signed-Cat-TD algorithm. In previous experiments, the Cat-TD algorithms were representing categorical distributions on evenly-spaced points in the space of multi-returns, which necessarily requires an exponential number of support points as a function of $d$. On line 337, we meant to say that we instead fixed a number $m$ of support points and generated supports on $m$ randomly-chosen points in the space of multi-returns to avoid the exponential blowup (at the cost of lower resolution). Notably, previous theory and algorithms for categorical distributional RL do not accommodate such supports; this is a novel feature of our algorithms and theory.", "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": "WUvtKCAqVh", "reviewer": "Reviewer_9dnf", "summary": "This paper studies the theoretical foundations of multivariate distributional RL, particularly providing the convergence proof under the MMD distance. The paper first investigates the aspect of particle-based multivariate dynamic programming in Section 4 and then shifts the attention to categorical representation-based multivariate dynamic programming in Section 5 and Temporal difference learning in Section 5. Experiments are conducted mainly on the distributional successor measure setting.", "questions": "I personally believe [ZCZ+21] indeed provides many convergence guarantees, which is inconsistent with the statement in Line 30 on Page 1. Thus, I think this kind of statement may not be accurate or proper.\n\nI personally do not think the projection operator is generally necessary. Imagine both the current Bellman statistics and the TD target ones are within the real set or unbounded, e.g., the particle representation. Another concern is why we need to introduce the randomized projector, which seems not common in the existing literature.\n\nI partially disagree with the main contribution statement of the paper in lines 33 and 67: it designs (1) a computationally tractable and (2) a theoretically justified algorithm with convergence guarantee. Firstly, the theoretical part is highly reliance on existing conclusions, which may be overclaimed by this paper. Also, the authors did not extensively verify the computational efficiency by conducting large-scale experiments.\n\nHow the size of $\\eta_k$ increases exponentially with k in line 62 of Page 2.\n\nSome writings are not clear. For example, what are the two ideas mentioned in Line 24 on Page 1?\n\nIs the word cumulants referring to statistical cumulants or general statistics for a distribution? If it is the latter, a more careful choice of this word is suggested.\n\nThe equally weighted particle seems inaccurate. In non-parametric statistics, we draw some samples to characterize the empirical distribution, and some samples are assigned a signal to contribute to the mass. I am not sure why the authors emphasize the equal weight of each sample, which seems unrelated to the following analysis.\n\nMore explanation about the QR programming should be given on Page 5. The current version of the Algorithm makes it difficult for me to understand the details.\n\nIn summary, this paper gives me the feeling that 1) the theoretical contribution is within a limited scope as most conclusions and extensions seem incremental, 2) the motivation may not be practical, and it lacks sufficient experiments to demonstrate the corresponding statements.", "limitations": "yes, the authors mentioned some limitations of their work in Section 7", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "* The convergence proof of extending existing univariate DP or TD to its multivariate settings is technically sound, although the corresponding conclusions are straightforward.\n\n* The writing, including the proof, is rigorous, and relevant works are sufficiently discussed.", "weaknesses": "* **Straightforward motivation and extension**. Personally, my main concern is that the contribution of this paper is within a limited scope. Multivariate distributional RL was first studied by [ZCZ+21], where the corresponding Bellman operator with the convergence proof under the Wasserstein distance is already provided. Having this knowledge, this paper mainly investigates the convergence under the MMD distance with either particle or categorical representation. This seems incremental to the theoretical values, given that the univariate version under MMD and categorical representation (with Cramer distance) have also been studied earlier. Therefore, a similar conclusion, which applies to the multivariate version and is highly based on existing results and proof techniques, is straightforward and easy to expect, in my opinion. Additionally, I am not fully convinced by the motivation of multi-dimensional distributional RL and whether it should be valued sufficiently in practice. The authors are suggested to emphasize the practical motivations, for example, by providing concrete examples or providing more experiments in real applications. Without these, I am not sure whether investigating the foundations of this setting is really useful (in an incremental way in this paper).\n\n\n* **Less concentrated organization**. The paper first studies the particle-based MMD distributional RL, which is a straightforward extension of [NGV20]. However, the authors suddenly shift their attention to the categorical representation albeit still equipped with MMD. I am not sure the motivation behind this kind of paper organization, but it indeed made me feel that some components are unnaturally combined together, where each component seems to rely highly on the corresponding existing works, e.g., univariate MMD or categorical distributional RL. This makes it difficult to posit the contribution of this paper.\n\n\n* **Limited and less general experiments**. While I understand some of this paper's motivations come from the distributional successor representation, it would be advisable to concentrate on more acceptable experiments in distributional RL, like Atari games in [ZCZ+21] or Mujoco environments. I acknowledge that this paper is theory-oriented, but it is more suggested to provide some general experimental results like [ZCZ+21] since this paper is highly based on [ZCZ+21].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed assessment of our work.\n\nWe appreciate your comments about the scalability of our approach and the motivations for studying foundations of multivariate DRL; we have discussed these in detail in the general response.\n\n**Q1**: The reviewer claims that our work is incremental given the results of [ZCZ+21]. We respectfully disagree, and we hope the following clarify our novelty:\n\n1. [ZCZ+21] only analyzes the contractivity of the *unprojected* distributional Bellman operator. This result alone does not prove convergence of any practical algorithm. Our work is the first to provide analysis for the type of distributional operator that we can deploy in practice, and associated TD-learning algorithms, unlike [ZCZ+21].\n2. [ZCZ+21] only provides convergence analysis for dynamic programming (again, with the unprojected / intractable operator), which we also generally cannot perform in practice without knowledge of the MDP. Our work goes beyond this and provides convergence guarantees for *TD-learning*, which can be computed from samples of the MDP (and no knowledge of the transition kernel or reward function). This is a much more difficult result to achieve, and we actually require new proof techniques and models to accomplish this (e.g., we must represent distributions as signed measures). Our novel analysis inspired a new algorithm which has not yet been studied even in the $d=1$ case, and is the first provably convergent TD-learning algorithm in the $d>1$ setting.\n3. While we have done much more than extend the analysis to MMD, this itself is a bonus. There are fundamental difficulties to performing TD-learning in Wasserstein space due to biased gradients. This is another facet under which our analysis applies to a practically-relevant algorithm, unlike [ZCZ+21].\n\nTo be clear, [ZCZ+21] is fantastic and was certainly an inspiration, but our results are not simply incremental modifications of theirs.\n\n**Q2**: Projections are necessary to avoid exponential blowup in the number of particles in return distributions (see Q4) in DP methods.\nSome TD algorithms do not explicitly compute projections, just passing gradients through a fixed representation, such as equally weighted particles.\nHowever, to *prove convergence* of such algorithms, one must show that these updates track applications of a projected operator (see e.g. [RMA+23; Sec 5-6]). Crucially, convergence of TD algorithms does not follow from contractivity of the unprojected Bellman operator, and analysis of projected operators is essential. Thus, while [ZCZ+21] does not explicitly compute a projection, this is one sense in which their paper does not theoretically justify their algorithm.\n\nRandomized projections uncommon in the literature: this is a novel contribution that we introduced to tackle issues with computing the EWP projection for DP. In the case of $d>1$, the MMD projection onto the space of EWP representations is non-convex, and could be extremely expensive (if at all possible) to compute in practice. The randomized projection that we introduce allows us to design a tractable algorithm for approximating the multi-return distribution function with EWP representations, which enjoys dimension-free theoretical convergence bounds.\n\n**Q3**: As described above, our theoretical results are *not* simple corollaries of existing results. They required novel proof techniques and algorithms (Sec 4-6). Each step of the algorithms described is objectively computationally tractable. With our statements on line 33/67, we are comparing against existing analyses of multivariate distributional RL which rely on regression oracles (for which no tractable algorithm generally exists). See also our general response, which demonstrates that our proposed TD-learning algorithm can be straightforwardly scaled with neural networks.\n\n**Q4**: Imagine an MDP with $n$ states and a nonzero probability of transition between any 2 states, and $\\eta_0(x)$ is supported on 1 point for each $x$. Then, after one iteration of DP, $\\eta_1(x)$ will be supported on $n$ points (contributions from each of the $n$ successor states which have return distributions supported on $1$ point).  Now, $\\eta_2(x)$ will be supported on $n^2$ points, since it mixes contributions of $n$ return distributions each having support on $n$ points. Continuing, we see that $\\eta_k(x)$ is generally supported on $n^k$ points.\n\n**Q5**: We will make this more clear in our revision; the two ideas are modeling return distributions and modeling multivariate returns.\n\n**Q6**: Good question. Cumulant here refers to the multivariate return. While this is standard terminology in this niche of RL research (see e.g. *Bootstrapped Representations for Reinforcement Learning*, Le Lan et al. 2023), it is worth clarifying as you point out.\n\n**Q7**: Alternatively, one could model the positions of $m$ particles and their probability masses, as opposed to modeling empirical distributions. This requires more memory and is often more difficult to optimize (see [BDR23]). Can you clarify in what sense the EWP representation is \"inaccurate\"? While modeling atom probabilities provides more flexibility, EWP representations are valid probability measures, and we provide strong bounds on the quality of our EWP approximation to the true (nonparametric) multivariate return distributions (Theorem 3).\n\n**Q8**: We would be happy to explain this further, but it would help us if you could specify more precisely which part is difficult to understand. The proof of Lemma 1 explicitly constructs the QP to be solved for computing the projection. This QP is described completely in the pseudocode on the line with QPSolve -- this is a minimization of a quadratic over a convex set. We can use efficient QP solvers to solve this. It also does not preclude function approximation, since we only need to solve the QP for the target returns (which gradients are not propagated through).", "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:
# General Response 
We thank all authors for their assessments.

Reviewers praised our convergence theory (SWh2, FNo5, wmwv), rigorous proofs and discussions (FNo5, wmwv, 9dnf), and illustrative experiments (FNo5, wmwv); while the simplicity of the numerical experiments and the motivation for formalizing multivariate distributional RL were commented on by FNo5 and 9dnf.

To address comments about scalability and motivation, we illustrate in our rebuttal PDF that our Signed Categorical TD method can be scaled to large (pixel-based) state spaces, just by directly representing the multi-return distribution function with a deep neural network trained with gradient descent for the update rule of equation (12). Our illustration also demonstrates how learning distributional successor features (multivariate return distributiosn) can be very useful in practice.

Further details about our experiment setup are given below.

## Illustration Details
Henceforth, we refer to figures in the rebuttal PDF.
Our environment consists of a car navigating to a parking spot (see Figure 1 for a depiction of an observation in this environment). The reward function is two-dimensional, with dimens
1. **Lateral feature**: The $x$ coordinate of the car.
2. **Parking feature**: A sparse reward that is $1$ when the car is parked in the correct location, and $0$ otherwise.

We train the multivariate return distribution function from pixels using a convolutional neural network (architecture shown in Figure 2 of the rebuttal document) from data generated by trajectories that circumvent the obstacle in the middle of the map. Note that, by symmetry, the (expected) successor features for such data would be roughly $0$ in the lateral feature, which is impossible to distinguish from a policy that drives straight through the obstacle.

Figures 3 and 4 show the learned multivariate return distributions (e.g., distributional successor features) from the method of [ZCZ+21] and ours, respectively. In both cases, the distribution provides crucial information for distinguishing the policy from one that drives through the obstacle (unlike if we learned SFs). Both algorithms produce multimodal distributions depicting trajectories that circumvented the obstacle on either side. However, our approach based on signed categorical TD has two major advantages:
1. Unlike the local optimum found by the EWP algorithm, its probability mass is roughly contained in the support of the true multi-return distribution;
2. It is based on an algorithm whose convergence is well understood (as we prove in this paper).

Figure 5 quantitatively demonstrates the accuracy of the multivariate return distributions learned by our algorithm and that of [ZCZ+21]. Here we examine two held-out projections on the space of multivariate returns (corresponding to held-out scalar reward functions) and evaluate the Cramer error of the projected return distributions relative to Monte Carlo estimates. These two projections correspond to a diverse pair of objectives: on the left, we incentivize paths circumventing the obstacle conservatively on the left (with a large negative lateral component) and on the right we incentivize "riskier" paths that get to the parking spot quickly (with a large positive parking component and near-zero lateral component). In both cases, both algorithms achieve remarkably low Cramer error. In the left case, our method achieves lower Cramer error by a non-negligible margin.

## Utility of Multivariate Distributional RL

This illustration also motivates multivariate distributional RL as a tool for inverse RL. Many inverse RL methods aim to match successor features to a demonstration, however in this case, this would result in learning a policy that drives through the obstacle. Matching distributional SFs (using multivariate distributional RL) would prevent such behavior.

Beyond this, a formal understanding of multivariate distributional RL can provide insights for zero-shot risk-sensitive RL: the problem of predicting arbitrary statistics of return distributions for arbitrary reward functions without further training. This concept was recently explored by [WFG+24] through the distributional successor measure (DSM), but much like [ZCZ+21], the convergence of their practical algorithm was not analyzed. Our results in fact immediately provide convergence guarantees for learning the DSM in tabular MDPs, and we demonstrate this further in Section 5.1.

Meta-review: Brief summary: this paper provides convergence analysis for policy evaluation in distributional RL with multiple reward functions. There are k reward functions and the task is to predict the distribution of the returns for all k and from all states when running a fixed policy (evaluation setting). The paper derives convergence rates for distributional dynamic programming and then extends to temporal-difference methods. Theoretical results are complemented with simple experiments.

The paper was extensively debated and the point of discussion ultimately came down to the subjective point of research taste in terms of how significant the theoretical contribution of the paper is w.r.t. prior work. All reviewers agreed that the paper is sound and the contributions are novel. All reviewers also agree that the empirical study presented is not a major contribution. 

After reading the reviews, discussion, paper, and most closely related paper, I believe the paper makes a non-trivial contribution to the foundations of distributional RL in terms of convergence analysis of both dynamic programming and TD-operators. Consequently, I think it is preferable to accept the paper.
Decision: Accept (poster)