Title: Foundations of Multivariate Distributional Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: distributional reinforcement learning

rl theory

dynamic programming

temporal difference learning

successor features

successor representation
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

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 mass1 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.

1 Introduction

Distributional reinforcement learning [DRL; MSK + 10, BDM17b, BDR23] focuses on 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 [YZL + 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, SLL21].

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 of multivariate cumulants, 1 rather than a scalar reward. Having learnt such a multivariate value function, it 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 a variety of applications in transfer learning, and multi-objective and constrained RL.

Multivariate distributional RL combines these two ideas, and aims to learn the full probability distribution of returns given a multivariate cumulant function. Successfully learning the multivariate

Jesse Farebrother

Mark Rowland

reward distribution opens up a 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 RL. 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 + 21], or requiring a maximum-likelihood optimisation oracle for implementation [WUS23]. Concurrently, the work of [LK24] analyzed algorithms for DRL with Banach-spacevalued rewards, and provided convergence guarantees for dynamic programming with non-parametric (intractable) distribution models.

Our central contribution in this paper is to propose algorithms for dynamic programming and temporaldifference learning in multivariate DRL which are computationally tractable and theoretically justified, with convergence guarantees. We show that reward dimensions strictly larger than 1 introduce new computational and statistical challenges. To resolve these challenges, we introduce multiple novel algorithmic techniques, including a randomized dynamic programming operator for efficiently approximating projected updates with high probability, and a novel TD-learning algorithm operating on mass1 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 → P ( X ) , and discount factor γ ∈ [0 , 1) . Unlike the standard RL setting, we consider a vector-valued reward function r : X → [0 , R max ] d , as in the literature on successor features [BDM + 17a]. Given a policy π : X → P ( A ) , we write the policy-conditioned transition kernel P π ( · | x ) = ∫ P ( · | x, a ) π (d a | x ) .

Multi-variate return distributions. We write ( X t ) t ≥ 0 for a trajectory generated by setting X 0 = x , and for each t ≥ 0 , X t +1 ∼ P π ( ·| X t ) . The return obtained along this trajectory is then defined by G π ( x ) = ∑ ∞ t =0 γ t r ( X t ) , and the (multi-)return distribution function is η π ( x ) = Law ( G π ( 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, [ZCZ + 21, CZZ + 24] show that for any reward function ˜ r : x ↦→⟨ r ( x ) , w ⟩ for some w ∈ R d , ⟨ G π ( x ) , w ⟩ = law ∑ t ≥ 0 γ t ˜ r ( X t ) for X 0 = x . Likewise, one might consider r ( x ) = δ x , in which case G π ( x ) corresponds to 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 [ZCZ + 21] that multi-return distributions obey a distributional Bellman equation, similar to the scalar case [MSK + 10, BDM17b], and defines the multivariate distributional Bellman operator T π : P ( R d ) X →P ( R d ) X by

where b y,γ ( z ) = y + γz and f ♯ µ = µ ◦ f -1 is the pushforward of a measure µ through a measurable function f . In particular, [ZCZ + 21] showed that η π satisfies the multi-variate distributional Bellman equation T π η π = η π , and that T π is a γ -contraction in W p , where W p ( η 1 , η 2 ) = sup x ∈X W p ( η 1 ( x ) , η 2 ( x )) and W p is the p -Wasserstein metric [Vil09]. This suggests a convergent scheme for approximating η π in W p by distributional dynamic programming , that is, computing the iterates η k +1 = T π η k , following Banach's fixed point theorem.

Approximating multivariate return distributions. In practice, however, the iterates η k +1 = T π η k cannot be computed efficiently, because the size of the support of η k may increase exponentially with k . A variety of alternative approaches that aim to circumvent this computational difficulty have been considered [FSMT19, ZCZ + 21, WUS23]. 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-tractable 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 [NTGV20, ZCZ + 21], is the notion of distance over probability distributions given by maximum mean discrepancies [GBR + 12, MMD]. A maximum mean discrepancy MMD κ : P ( Y ) × P ( Y ) → R + assigns a notion of distance to pairs of probability distributions, and is parametrised via a choice of kernel κ : Y × Y → R , defined by

A useful alternative perspective on MMD is that the choice of kernel κ induces a reproducing kernel Hilbert space (RKHS) of functions H , namely the closure of the span of functions of the form z ↦→ κ ( y, z ) for each y ∈ Y , with respect to the norm ∥ · ∥ H induced by the inner product ⟨ κ ( y 1 , · ) , κ ( y 2 , · ) ⟩ = κ ( y 1 , y 2 ) . With this interpretation, MMD κ ( p, q ) is equal to ∥ µ p -µ q ∥ H , where µ p = ∫ Y κ ( · , y ) p (d y ) ∈ H is the mean embedding of p (similarly for µ q ). When p ↦→ µ p is injective, the kernel κ is called characteristic , and MMD κ is then a proper metric on P ( Y ) [GBR + 12]. In the remainder of this work, we will assume that all spaces of measures will be over compact sets Y ; 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 in [SSGF13], that will be particularly useful for our analysis. These are the kernels induced by semimetrics.

̸

Definition 1. Let Y be a nonempty set, and consider a function ρ : Y × Y → R + . Then ρ is called a semimetric if it is symmetric and ρ ( y 1 , y 2 ) = 0 ⇐⇒ y 1 = y 2 . Additionally, ρ is said to have strong negative type if ∫ ρ d([ p -q ] × [ p -q ]) < 0 whenever p, q ∈ P ( Y ) with p = q .

Notably, certain semimetrics naturally induce characteristic kernels and probability metrics.

̸

Theorem 1 ([SSGF13, Proposition 29]) . Let ρ be a semimetric on a space Y have strong negative type , in the sense that ∫ ρ d([ p -q ] × [ p -q ]) < 0 whenever p = q are probability measures on a compact set Y . Moreover, let κ : Y × Y → R denote the kernel induced by ρ , that is

for some y 0 ∈ Y . Then κ is characteristic, so MMD κ is a metric.

Remark 1. An important class of semimetrics are the functions ρ α : R d × R d → R + given by ρ α ( y 1 , y 2 ) = ∥ y 1 -y 2 ∥ α 2 for α ∈ (0 , 2) . It is known that these semimetrics have strong negative type, and thus the kernels κ α induced by ρ α are characteristic [SR13, SSGF13]. The resulting metric MMD κ α is known as 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 a supremal form MMD κ of MMD, given by MMD κ ( η 1 , η 2 ) = sup x ∈X MMD κ ( η 1 ( x ) , η 2 ( 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 a kernel induced by a semimetric ρ on [0 , (1 -γ ) -1 R max ] d with strong negative type, satisfying

Shift-invariance . For any z ∈ R d , ρ ( z + y 1 , z + y 2 ) = ρ ( y 1 , y 2 ) .

Homogeneity . For any γ ∈ [0 , 1) , there exists c > 0 for which ρ ( γy 1 , γy 2 ) = γ c ρ ( y 1 , y 2 ) .

Consider the sequence { η k } ∞ k =1 given by η k +1 = T π η k . Then η k → η π at a geometric rate of γ c/ 2 in MMD κ , as long as MMD κ ( η 0 , η π ) ≤ C < ∞ .

Notably, the energy distance kernels κ α satisfy the conditions of Theorem 2, and ρ α ( γy 1 , γy 2 ) ≤ γ α ρ ( y 1 , y 2 ) by the homogeneity of the Euclidean norm, so T π is a γ α/ 2 -contraction in the energy distances. This generalizes the analogous result of [NTGV20] in the one-dimensional case.

While Theorem 2 illustrates a method for approximating η π in MMD, it leaves a 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 a tractable exact representation. As such, it will be 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 η π from samples. We provide algorithms for resolving both of these concerns-the former in Section 5 and the latter in Section 6-where we illustrate the difficulties that arise once the cumulant dimension exceeds unity.

4 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 [ZCZ + 21].

Temporal-difference learning with EWP representations. EWP representations, expressed by the class C EWP ,m , are defined by

For simplicity, we consider the case here where at each state x , the multi-return distribution is approximated by N ( x ) = m atoms. We can represent η ∈ C EWP ,m by η ( x ) = 1 m ∑ m i =1 δ θ i ( x ) for θ i : X → R d . The work of [ZCZ + 21] introduced a TD-learning algorithm for learning a C EWP ,m representation of η π , computing iterates of the particles ( θ ( k ) i ) i m =1 according to

for step sizes ( λ k ) k ≥ 0 and sampled next states X ′ ∼ P π ( · | x ) , where θ = stop-gradient ( θ ( k ) ) is a copy of θ ( 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 distributional dynamic programming [RBD + 18, RMA + 24], we have T π C EWP ,m ̸⊂ C EWP ,m ; in words, the distributional Bellman operator does not map EWP representations to themselves. Concretely, as long as there exists a state x at which the support of P π ( · | x ) is not a singleton, ( T π η )( x ) will consist of more than m atoms even when η ∈ C EWP ,m ; and secondly, if P ( · | x ) is not uniform, ( T π η )( x ) will not consist of equally-weighted particles.

Consequently, to obtain a DP algorithm over EWP representations, we must consider a projected operator of the form Π EWP T π , for a projection Π EWP : P ( R d ) X → C EWP ,m . A natural choice for this projection is the operator that minimizes the MMD of each multi-return distribution in C EWP ,m ,

Unfortunately, even in the scalar-reward ( d = 1 ) case, the operator Π m EWP ,κ is problematic; (Π m EWP ,κ η )( x ) is not uniquely defined, and Π m EWP ,κ is not a 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 Π m EWP ,κ T π has a 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 η ( x ) has finite support for each state x , the projection (Π m EWP ,κ η )( x ) is very similar to clustering, which can be 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.

Randomised dynamic programming. Towards this end, we introduce a tractable randomized dynamic programming algorithm for the EWP representation, by using a randomized proxy BootProj π κ,m for Π κ,m T π , that produces accurate return distribution estimates with high probability. The method produces the following iterates,

A similar algorithm for categorical representations was discussed in concurrent work [LK24] without convergence analysis.

The intuition is that, particularly for large m , the Monte Carlo error associated with the sample-based approximation to ( T π η )( x ) is small, and we can therefore expect the DP process, though randomised, to be accurate with high probability. This is summarised by a key theoretical result of this section; our proof of this result in the appendix provides a general approach to proving convergence for algorithms using arbitrary accurate approximations to (4) that we expect to be useful in future work.

Theorem 3. Consider a kernel κ induced by the semimetric ρ ( x, y ) = ∥ x -y ∥ α 2 with α ∈ (0 , 2) , and suppose rewards are bounded in each dimension within [0 , R max ] . For any η 0 such that MMD κ ( η 0 , η π ) ≤ D < ∞ , and any δ > 0 , for the sequence ( η k ) k ≥ 0 defined in Equation (5) , with probability at least 1 -δ we have

where η k +1 = BootProj π κ,m η k and K = ⌈ log m log γ -α ⌉ , and where ˜ O omits logarithmic factors in m .

This shows that our novel randomised DP algorithm with EWP representations can tractably compute accurate approximations to the true multivariate return distributions, with only polynomial dependence 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 one-dimensional distributional DP algorithms [RBD + 18, RMA + 24, BDR23]. Consequently, these guarantees do not provide a clear understanding of the EWP-TD method described at the beginning of this section. However, in the sequel, we introduce DP and TD algorithms based on categorical representations , for which we derive dynamic programming and TD-learning convergence bounds.

The proof of Theorem 3 is hinges on 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 1 (Convergence of EWP Dynamic Programming) . Consider a kernel satisfying the hypotheses of Theorem 2, suppose rewards are globally bounded in each dimension in [0 , R max ] , and let { Π ( k ) κ,m } k ≥ 0 be a sequence of maps Π : P ([0 , (1 -γ ) -1 R max ] d ) X → C EWP ,m satisfying

Then the iterates ( η k ) k ≥ 0 given by η k +1 = Π ( k ) κ,m T π η k with MMD κ ( η 0 , η π ) = D < ∞ converge to a set η m EWP ,κ ⊂ B ( η π , (1 -γ c/ 2 ) -1 f ( d, m )) in MMD κ , where B denotes the closed ball in MMD κ ,

As an immediate corollary of Proposition 1 and Theorem 3, we can derive an error rate for projected dynamic programming with Π m EWP ,κ as well.

Corollary 1. For any kernel κ satisfying the hypotheses of Theorem 3, and for any η 0 ∈ C EWP ,m for which MMD κ ( η 0 , η π ) ≤ D < ∞ , the iterates η k +1 = Π m EWP ,κ T π η k converge to a set η m EWP ,κ ⊂ C EWP ,m , where

5 Categorical Multivariate Distributional Dynamic Programming

Our next contribution is the introduction of a convergent multivariate distributional dynamic programming 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 to restrict each multi-return distribution to a finitely-parameterized class. In this work, we take inspiration from successful distributional RL algorithms [BDM17b, RBD + 18] and employ a categorical representation . The work of [WUS23] proposed a categorical representation for multivariate DRL, but their categorical projection was not justified theoretically, and it required a particular choice of fixed support. We propose a novel categorical representation with a finite (possibly state-dependent) support R ( x ) = { ξ ( x ) i } N ( x ) i =1 ⊂ R d , that models the multi-return distribution function η such that η ( x ) ∈ ∆ R ( x ) for each x ∈ X . The notation ξ ( x ) i simply refers to the i th support point at state x specified by R , and ∆ A denotes the probability simplex on the finite set A . We refer to the mapping R as the support map 2 and we denote the class of multi-return distribution functions under the corresponding categorical representation as C R ≜ ∏ x ∈X ∆ R ( x ) .

Categorical projection. Once again, the distributional Bellman operator is not generally closed over C R , that is, T π C R ̸⊂ C R . As such, we cannot actually employ the procedure described in Theorem 2 - rather, we need to project applications of T π back onto C R . Roughly, we need an operator Π : P ( R d ) X → C R for which Π | C R = id . Given such an operator, as in the literature on categorical distributional RL [BDM17b, RBD + 18], we will study the convergence of iterates η k +1 = Π T π η k .

Projection operators used in the scalar categorical distributional RL literature are specific to distributions over R , so we must introduce a new projection. We propose a projection similar to (4),

We will now verify that Π R C ,κ is well-defined, and that it satisfies the aforementioned properties.

Lemma 1. Let κ be a kernel induced by a semimetric ρ on [0 , (1 -γ ) -1 R max ] d with strong negative type (cf. Theorem 1). Then Π R C ,κ is well-defined, Ran (Π R C ,κ ) ⊂ C R , and Π R C ,κ | C R = id .

It is worth noting that beyond simply ensuring the well-posedness of the projection Π R C ,κ , 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 Π R C ,κ T π with a QP solver QPSolve in Algorithm 1.

Algorithm 1 Projected Categorical Dynamic Programming

Require: Support map R , kernel κ , transition kernel P π , reward function r , discount γ Require: Return distribution function η ∈ C R for x ∈ X do ( T π η ) x ← ∑ x ′ ∈X ∑ ξ ∈R ( x ′ ) P π ( x ′ | x ) η x ′ ( ξ ) δ r ( x )+ γξ K x i,j ← κ ( ξ i , ξ j ) for each ( ξ i , ξ j ) ∈ R ( x ) 2 q x j ← ∑ ξ ∈ supp ( T π η ) x ( T π η ) x ( ξ ) κ ( ξ j , ξ ) for each ξ j ∈ R ( x ) p ← QPSolve ( min p ∈ R |R ( x ) | [ p ⊤ K x p -2 p ⊤ q ] subject to p ⪰ 0 , ∑ i p i = 1 ) (Π R C ,κ T π η ) x ← ∑ ξ i ∈R ( x ) p i δ ξ i end for return Π R C ,κ T π η

Lemma 2. Under the conditions of Lemma 1, Π R C ,κ is a nonexpansion in MMD κ . That is, for any η 1 , η 2 ∈ P ([0 , (1 -γ ) -1 R max ] d ) X , we have MMD κ (Π R C ,κ η 1 , Π R C ,κ η 2 ) ≤ MMD κ ( η 1 , η 2 ) .

Categorical multivariate distributional dynamic programming. As an immediate consequence of Lemma 2, it follows that projected dynamic programming under the projection Π R C ,κ is convergent;

this is because T π is a contraction in MMD κ and Π R C ,κ is a nonexpansion in MMD κ , so the projected operator Π R C ,κ T π is a contraction in MMD κ ; a standard invokation of the Banach fixed point theorem appealing to the completenes of MMD κ certifies that repeated application of Π R C ,κ T π will result in convergence to a unique fixed point.

Corollary 2. Let κ be a kernel satisfying the conditions of Theorem 2. Then for any η 0 ∈ C R , the iterates { η k } ∞ k =1 given by η k +1 = Π R C ,κ T π η k converge geometrically to a unique fixed point.

Beyond the result of Theorem 2, Corollary 2 illustrates an algorithm for estimating η π in MMD κ provided knowledge of the transition kernel and the reward function, which is computationally tractable in tabular MDPs. Indeed, the iterates ( η k ) k ≥ 0 all lie in C R , having finitely-many degrees of freedom. Algorithm 1 outlines a computationally tractable procedure for learning η π C ,κ 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 η π C ,κ 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 a class of partitions P R ( x ) , where each P ∈ P R ( x ) satisfies

| P | = N ( x ) ;

For any θ 1 , θ 2 ∈ P , either θ 1 ∩ θ 2 = ∅ or θ 1 = θ 2 ;

∪ θ ∈ P θ = P ([0 , (1 -γ ) -1 R max ] d ) ;

Each element θ i ∈ P contains exactly one element z i ∈ R ( x ) .

For any partition P , we define a notion of mesh size relative to a kernel κ induced by a semimetric ρ ,

The following result confirms that η π C ,κ recovers η π as the mesh decreases.

Theorem 4. Let κ be a kernel induced by a c -homogeneous and shift-invariant semimetric ρ conforming to the conditions of Theorem 2. Then the fixed point η π C ,κ of Π R C ,κ T π satisfies

Thus, for any sequence of supports {R ( x ) m } m ≥ 1 for which the maximal space (in ρ ) between any two points in R ( x ) m tends to 0 as m →∞ , the fixed point η π C ,κ approximates η π to arbitrary precision for large enough m . The next corollary illustrates this in a familiar setting.

Corollary 3. Let R ( x ) = U m , where U m is a set of m uniformly-spaced support points on [0 , (1 -γ ) -1 R max ] . For κ induced by the semimetric ρ ( x, y ) = ∥ x -y ∥ α 2 for α ∈ (0 , 2) ,

With α = 1 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 -1 / 2 ) rate shown by [RBD + 18] in this setting. Thus, our results offer a new perspective on categorical DRL, and naturally generalizes the theory to the multivariate setting.

Theorem 4 relies on the following lemma about the approximation quality of the categorical MMD projection, which may be of independent interest.

Lemma3. Let κ be kernel satisfying the conditions of Lemma 1, and for any finite R⊂ R d , define Π : P ( R d ) → ∆ R via Π p = arginf q ∈ ∆ R MMD κ ( p, q ) . Then MMD 2 κ (Π p, p ) ≤ inf P ∈ P R mesh ( P ; κ ) .

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 a convergent multivariate categorical TD-learning algorithm.

5.1 Simulation: The Distributional Successor Measure

As a preliminary example, we consider 3-state MDPs with the cumulants r ( i ) = (1 -γ ) e i , i ∈ [3] for e i the i th basis vector. In this setting, η π encodes the distribution over trajectory-wise discounted state occupancies, which was discussed in the recent work of [WFG + 24] and called the distributional successor measure (DSM). Particularly, [WFG + 24] showed that x ↦→ Law ( G ⊤ x ˜ r ) for G x ∼ η π ( x ) is the return distribution function for any scalar reward function ˜ r . Figure 1 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.

6 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 a strong motivation for TD-learning algorithms [Sut88], wherein state transition data alone is used to solve the Bellman equation. In this section, we devise a TD-learning algorithm for multivariate DRL, leveraging our results on categorical dynamic programming in MMD κ .

Relaxation to signed measures. In the d = 1 setting, the categorical projection presented above is known to be affine [RBD + 18], making scalar categorical TD-learning amenable to common techniques in stochastic approximation theory. However, the projection is not affine when d ≥ 2 ; we give an explicit example in Appendix D. Thus, we relax the categorical representation to include signed measures, which will provide us with an affine projection [BRCM19]-this is crucial for proving our main result, Theorem 6. We denote by M 1 ( Y ) the set of all signed measures µ over Y with µ ( Y ) = 1 . We begin by noting that the MMD endows M 1 ( Y ) with a metric structure.

Lemma 4. Let κ : Y × Y → R be a characteristic kernel over some space Y . Then MMD κ : M 1 ( Y ) ×M 1 ( Y ) → R + given by ( p, q ) ↦→ ∥ µ p -µ q ∥ H defines a metric on M 1 ( Y ) , where µ p denotes the usual mean embedding of p and H is the RKHS with kernel κ .

We define the relaxed projection Π R SC ,κ : M 1 ([0 , (1 -γ ) -1 R max ] d ) X → ∏ x ∈X M 1 ( R ( x )) =: S R ,

Note that (10) is very similar to the definition of the categorical MMD projection in (7)-the only difference is that the optimization occurs over the larger class of signed mass-1 measures. It is 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 Π R SC ,κ T π : S R → S R is affine, is contractive with contraction factor γ c/ 2 , and has a unique fixed point η π SC ,κ .

While we have 'relaxed' the projection, the fixed point η π SC ,κ is a good approximation of η π .

Theorem 5. Under the conditions of Lemma 5, we have that

MMD κ ( η π SC ,κ , η π ) ≤ 1 1 -γ c/ 2 sup x ∈X inf P ∈ P R ( x ) √ mesh ( P ; κ ) ; and

Notably, the second statement of Theorem 5 states that projecting η π SC ,κ back onto the space of multi-return distribution functions yields approximately the same error as η π C ,κ when γ is near 1 .

In the remainder of the section, we assume access to a stream of MDP transitions { T t } ∞ t =1 ⊂ X × A × R d ×X consisting of elements T t = ( X t , A t , R t , X ′ t ) with the following structure,

where P is some probability measure and {F t } ∞ t =1 is the canonical filtration F t = σ ( ∪ t t =1 T t ) . Based on these transitions, we can define stochastic distributional Bellman backups by

which notably can be computed exactly without knowledge of P, r . Due to the stronger convergence guarantees shown for projected multivariate distributional dynamic programming, we introduce an asynchronous incremental algorithm leveraging the categorical representation, which produces iterates { ̂ η t } ∞ t =1 according to

for ̂ η 0 ∈ C R , where { α t } ∞ t =1 is any sequence of (possibly) random step sizes adapted to the filtration {F t } ∞ t =1 . 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 to establish convergence of TD-learning on S R representations.

Theorem 6. For a kernel κ induced by a semimetric ρ of strong negative type, the sequence { ̂ η t } ∞ t =1 given by (11) -(13) converges to η π SC ,κ with probability 1.

6.1 Simulations: Distributional Successor Features

To illustrate the behavior of our categorical TD algorithm, we employ it to learn the multi-return distributions for several tabular MDPs with random cumulants. We focus on the case of 2 - and 3 -dimensional cumulants, which is the setting studied in recent works regarding multivariate distributional RL [ZCZ + 21, WUS23]. Interpreting the multi-return distributions as joint distributions over successor features [BDM + 17a, SFs], we additionally evaluate the return distributions for random reward functions in the span of the cumulants. We compare our categorical TD approach with a tabular implementation of the EWP TD algorithm of [ZCZ + 21], for which no convergence bounds are known.

100

200

300

400

500

Atoms

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.45

0.50

Cramér Distance

Method

EWP-TD

Signed-Cat-TD

(a)

d

= 2

100 200 300 400 500 600 700 800 900 1000

Atoms

0.2

0.3

0.4

0.5

0.6

Cramér Distance

Method

EWP-TD

Signed-Cat-TD

(b)

d

= 3

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

1.0

0.5

0.0

0.5

1.0

Cumulant 1

1.0

0.5

0.0

0.5

1.0

Cumulant 2

Signed-Cat-TD

1.0

0.5

0.0

0.5

Return

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Probability

MC

Signed-Cat-TD

1.0

0.5

0.0

0.5

1.0

Return

0.0

0.2

0.4

0.6

0.8

1.0

Probability

MC

Signed-Cat-TD

1.0

0.5

0.0

0.5

1.0

Cumulant 1

1.0

0.5

0.0

0.5

1.0

Cumulant 2

EWP-TD

1.0

0.5

0.0

0.5

Return

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Probability

MC

EWP-TD

0.5

0.0

0.5

1.0

Return

0.0

0.2

0.4

0.6

0.8

1.0

Probability

MC

EWP-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 in 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 provably-convergent categorical TD method learns a high fidelity quantization of the true multi-return distributions.

Naturally, however, the benefits of the poly ( d ) bounds for EWP suggested by Theorem 3 become more present as we increase the cumulant dimension. Figure 2b repeats the experiment of Figure 2a with d = 3 , using randomized support points for the categorical algorithm to avoid a 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, a 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 m used in distribution representations, and interestingly motivate the use of signed measures to develop provably convergent TD algorithms.

While it is 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 d = 2 setting has important applications-indeed, efforts in safe RL depend on distinguishing a 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-dimensional cumulants also has many important applications, such as modeling close approximations to distributional successor measures [WFG + 24] for zero-shot risk-sensitive policy evaluation. In this setting, we believe EWP-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-learning algorithms is a very interesting and important open problem for future investigation.

References

[BB96] Steven J. Bradtke and Andrew G. Barto. Linear least-squares algorithms for temporal difference learning. Machine Learning , 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-Philippe Vert. Efficient and modular implicit 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 navigation 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] Marc 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 to numerical computing. SIAM Review , 59(1):65-98, 2017.

[BFH + 18] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye WandermanMilne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

[BHB + 20] 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) , 117(48):30079-30087, 2020.

[BRCM19] Marc G Bellemare, Nicolas Le Roux, Pablo Samuel Castro, and Subhodeep Moitra. Distributional reinforcement learning with linear function approximation. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2019.

[BT96] Dimitri P. Bertsekas and John N. Tsitsiklis. Neuro-dynamic programming . Athena Scientific, 1996.

[CGH + 96] 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:329359, 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):671-675, 2020.

[DRBM18] Will Dabney, Mark Rowland, Marc G. Bellemare, and Rémi Munos. Distributional Reinforcement Learning with Quantile Regression. In AAAI Conference on Artificial Intelligence , 2018.

[FSMT19] Dror Freirich, Tzahi Shimkin, Ron Meir, and Aviv Tamar. Distributional Multivariate Policy Evaluation and Exploration with the Bellman GAN. In International Conference on Machine Learning (ICML) , 2019.

[GBR + 12] Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A Kernel Two-Sample Test. Journal of Machine Learning Research (JMLR) , 13(25):723-773, 2012.

[GBSL21] Michael Gimelfarb, André Barreto, Scott Sanner, and Chi-Guhn Lee. Risk-Aware Transfer in Reinforcement Learning using Successor Features. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

[HRB + 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 Multiagent 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 Ilyas Malik. Distributional Reinforcement Learning for RiskSensitive Policies. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

[MSK + 10] Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki 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] Mark Rowland, Marc 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. A Stochastic Approximation Method. The Annals of Mathematical Statistics , 22(3):400-407, September 1951.

[RMA + 24] Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Marc G. Bellemare, and Will Dabney. An Analysis of Quantile Temporal-Difference Learning. Journal of Machine Learning Research (JMLR) , 25(163):1-47, 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 Kernel Trick for Distances. In Advances in Neural Information 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 , abs/2112.00519, 2021.

[SLL21] Wei-Fang Sun, Cheng-Kuang 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 statistics: A class of statistics based on distances. Journal of Statistical Planning and Inference , 143(8):1249-1272, August 2013.

[SSGF13] Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equivalence of distance-based and RKHS-based statistics 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.

[SZS + 08] Le Song, Xinhua Zhang, Alex Smola, Arthur Gretton, and Bernhard Schölkopf. Tailoring 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.

[Vil09] Cédric Villani. Optimal transport: Old and new , volume 338. Springer, 2009.

[VN02] Jan Van Neerven. Approximating Bochner integrals by Riemann sums. Indagationes Mathematicae , 13(2):197-208, 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 turismo drivers with deep reinforcement learning. Nature , 602(7896):223-228, 2022.

[WFG + 24] Harley Wiltzer, Jesse Farebrother, Arthur Gretton, Yunhao Tang, André Barreto, Will Dabney, Marc G. Bellemare, and Mark Rowland. A distributional analogue to the successor representation. In International Conference on Machine Learning (ICML) , 2024.

[WUS23] Runzhe Wu, Masatoshi Uehara, and Wen Sun. Distributional Offline 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 a distributional safety critic. Machine Learning , 112(3):859-887, 2023.

[YZL + 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. Distributional Reinforcement Learning for Multi-Dimensional Reward Functions. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Appendices

A In-Depth Summary of Related Work

In Sections 1 and 2, we provided a 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 on the analysis of one-dimensional distributional RL algorithms such as categorical and quantile dynamic programming, as well as their temporal-difference learning counterparts [RBD + 18, DRBM18, RMA + 24, BDR23]. The proof techniques in these works themselves are related to contraction-based 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 are the first algorithms with strong theoretical guarantees and efficient tabular implementations. [FSMT19] propose a deep-learning-based approach in which generative adversarial networks are used to model multivariate return distributions, and use these predictions to inform exploration strategies. [ZCZ + 21] propose the TD algorithm combing equally-weighted particle representations and an MMD loss that we 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 a batch objective function (in this case, a negative log-likelihood, NLL) over a growing dataset. [WUS23] demonstrate that these algorithms are performant in low-dimensional settings empirically, and provide theoretical analysis for FLE algorithms, assuming an oracle which can approximately optimise the NLL objective at each algorithm step. [SFS24] also propose a 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 as 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 a straightforwardly analyzable TD algorithm.

Finally, the concurrent work of [LK24] studied distributional Bellman operators for Banach-spacevalued reward functions. Their work focuses on analyzing how well the fixed point of a distributional finite-dimensional multivariate Bellman equation can approximate the fixed point of a distributional Banach-space-valued Bellman equation. In contrast, our work only studies finite-dimensional reward functions, but provides explicit convergence rates and approximation bounds when the distribution representations are finite dimensional, unlike [LK24]. Moreover, [LK24] considers a similar algorithm to that discussed in Theorem 3 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, we will state some lemmas building up to the proof of Theorem 2. These lemmas generalize corresponding results of [NTGV20] that were specific to the scalar reward setting. We begin with a lemma that demonstrates a notion of convexity for the MMD induced by a conditional positive definite kernel.

Lemma 6. Let ( p a ) a ∈I ⊂ P ( Y ) and ( q a ) a ∈I ⊂ P ( Y ) be collections of probability measures indexed by an index set I . Suppose T ∈ P ( I ) . Then for any characteristic kernel κ , the following holds,

Proof. It is known from [Sch00] that characteristic kernels generate RKHSs H into which probability measures can be embedded. As such, it holds that

where ∥ · ∥ is the norm in the Hilbert space H and µ p is the mean embedding of p - that is, the unique element of H such that E y ∼ p [ f ( y )] = ⟨ f, µ p ⟩ for every f ∈ H , and where ⟨· , ·⟩ is the inner product in H .

Let T p ≜ E a ∼ T [ p a ] and define T q analogously. We claim that µ Tp = E a ∼ T [ µ p a ] ≜ Tµ p . To see this, let f ∈ H , and observe that

where the third step invokes Fubini's theorem, and the penultimate steps leverages the linearity of the inner product. Notably, T acts as a linear operator on mean embeddings. As a result, we see that

where the penultimate inequality is due to Jensen's inequality, and the final inequality holds since sup a ∈I ∥ µ p a -µ q a ∥ upper bounds the integrand, and the integral is a monotone operator.

Next, we show how the MMD κ under the kernels hypothesized in Theorem 2 behave under affine transformations to random variables.

Lemma 7. Let κ be a kernel induced by a semimetric ρ of strong negative type defined over a compact subset Y ⊂ R d that is both shift invariant and c -homogeneous (cf. Theorem 2). Then for any a ∈ Y and p, q ∈ P ( Y ) ,

Proof. It is known [GBR + 12] that the MMD can be expressed in terms of expected kernel evaluations, according to

where the last step invokes the definition of a kernel induced by a semimetric, and the linearity of expectation. Then, defining ˜ Y , ˜ Y ′ as independent samples from (b a,γ ) ♯ p and ˜ Z, ˜ Z ′ as independent samples from (b a,γ ) ♯ q , we have

where the second step is a change of variables, the third step invokes the shift invariance of ρ , and the fourth step invokes the homogeneity of ρ .

We are now ready to prove the main result of this section.

Theorem 2 (Convergent MMD dynamic programming for the multi-return distribution function) . Let κ be a kernel induced by a semimetric ρ on [0 , (1 -γ ) -1 R max ] d with strong negative type, satisfying

Shift-invariance . For any z ∈ R d , ρ ( z + y 1 , z + y 2 ) = ρ ( y 1 , y 2 ) .

Homogeneity . For any γ ∈ [0 , 1) , there exists c > 0 for which ρ ( γy 1 , γy 2 ) = γ c ρ ( y 1 , y 2 ) .

Consider the sequence { η k } ∞ k =1 given by η k +1 = T π η k . Then η k → η π at a geometric rate of γ c/ 2 in MMD κ , as long as MMD κ ( η 0 , η π ) ≤ C < ∞ .

Proof. We begin by showing that the distributional Bellman operator T π is contractive in MMD κ . We have

We apply Lemma 6 with I = X and T = P π ( · | x ) , yielding

Next, invoking Lemma 7 with the shift-invariance and c -homogeneity of κ , we have

It follows that MMD κ ( η k +1 , η π ) ≤ γ c/ 2 MMD κ ( T π η k , T π η π ) = γ c/ 2 MMD κ ( T π η k , η π ) , since η π is a fixed point of T π . Continuing, we see that MMD κ ( η k , η π ) ≤ γ kc/ 2 MMD κ ( η 0 , η π ) ≤ γ kc/ 2 C ∈ O ( γ kc/ 2 ) ⊂ o (1) . Since MMD κ is a metric on P ([0 , (1 -γ ) -1 R max ] d ) X for any characteristic kernel κ , it follows that η k approaches η π at a geometric rate.

B.2 EWPDynamic Programming: Section 4

In this section, we provide the proof of Theorem 3. To do so, we prove an abstract, general result, regarding any projection mapping that approximates the argmin MMD projection given in Equation (4).

Proposition 1 (Convergence of EWP Dynamic Programming) . Consider a kernel satisfying the hypotheses of Theorem 2, suppose rewards are globally bounded in each dimension in [0 , R max ] , and let { Π ( k ) κ,m } k ≥ 0 be a sequence of maps Π : P ([0 , (1 -γ ) -1 R max ] d ) X → C EWP ,m satisfying

Then the iterates ( η k ) k ≥ 0 given by η k +1 = Π ( k ) κ,m T π η k with MMD κ ( η 0 , η π ) = D < ∞ converge to a set η m EWP ,κ ⊂ B ( η π , (1 -γ c/ 2 ) -1 f ( d, m )) in MMD κ , where B denotes the closed ball in MMD κ ,

Proof. Let ∆ k = MMD κ ( η k , η π ) . Then we have

where the first step invokes the identity that η π is the fixed point of T π (which is well-defined by Theorem 2), the second step leverages the triangle inequality, and the third step follows by the definition of f ( d, m ) and the contractivity of T π established in Theorem 2. Unrolling the recurrence above, we have

As such, as k →∞ , we have that

proving our claim.

Proposition 1, despite its simplicity, reveals an interesting fact: given a good enough approximate MMDprojection Π κ,m in the sense that f ( d, m ) decays quickly with m , the dynamic programming iterates ( η k ) k ≥ 0 will eventually be contained in a (arbitrarily) small neighborhood of η π . The next result provides an example consequence of this abstract result, and establishes that m ∈ 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 ρ ( x, y ) = ∥ x -y ∥ α 2 with α ∈ (0 , 2) , and suppose rewards are bounded in each dimension within [0 , R max ] . For any η 0 such that MMD κ ( η 0 , η π ) ≤ D < ∞ , and any δ > 0 , for the sequence ( η k ) k ≥ 0 defined in Equation (5) , with probability at least 1 -δ we have

where η k +1 = BootProj π κ,m η k and K = ⌈ log m log γ -α ⌉ , and where ˜ O omits logarithmic factors in m .

Proof. For each x ∈ X and k ∈ [ K ] , denote by E x,k the event given by

for δ ′ > 0 a constant to be chosen shortly. Moreover, with E = ∩ ( x,k ) ∈X× [ K ] E x,k , it holds that under E , MMD κ ( BootProj π κ,m η k , T π η k ) ≤ f ( d, m ; δ ′ ) for all k ∈ [ K ] . Following the proof of Proposition 1, we have that, conditioned on E ,

Now, by [TSM17, Proposition A.1], event E x,k holds with probability at least 1 -δ ′ , since each ( BootProj π κ,m η k )( x ) is generated independently by sampling m independent draws from the distribution T π η k . Therefore, event E holds with probability at least 1 -|X| Kδ ′ . Choosing δ ′ = δ/ ( |X| K ) , we have that, with probability at least 1 -δ ,

As such, there exist universal constants C 0 , C 1 ∈ R + such that

Corollary 1. For any kernel κ satisfying the hypotheses of Theorem 3, and for any η 0 ∈ C EWP ,m for which MMD κ ( η 0 , η π ) ≤ D < ∞ , the iterates η k +1 = Π m EWP ,κ T π η k converge to a set η m EWP ,κ ⊂ C EWP ,m , where

Proof. Proposition 1 shows that projected EWP dynamic programming converges to a set with radius controlled by the quantity f ( d, m ) that upper bounds the distance f ( d, m ) between η k ( x ) and Π ( k ) κ,m η k ( x ) at the worst state x . In the proof of Theorem 3, we saw that with nonzero probability, the randomized projections satisfy f ( d, m ) ≤ 6 d α/ 2 R α max (1 -γ ) α √ m . Thus, there exists a projection satisfying this bound. Since the projection Π m EWP ,κ is, by definition, the projection with the smallest possible f ( d, m ) , the claim follows directly by Proposition 1.

B.3 Categorical Dynamic Programming: Section 5

Lemma 1. Let κ be a kernel induced by a semimetric ρ on [0 , (1 -γ ) -1 R max ] d with strong negative type (cf. Theorem 1). Then Π R C ,κ is well-defined, Ran (Π R C ,κ ) ⊂ C R , and Π R C ,κ | C R = id .

Proof. Firstly, note that ∆ R ( x ) is a bounded, finite-dimensional subspace for each x ∈ X . Thus, ∆ R ( x ) is compact, and by the continuity of the MMD, the infimum in (7) is attained.

Following the technique of [SZS + 08], we establish that Π R C ,κ can be computed as the solution to a particular quadratic program with convex constraints.

Let K ∈ R N ( x ) × N ( x ) denote a matrix where K i,j = κ ( ξ ( x ) i , ξ ( x ) j ) . Since κ is a positive definite kernel when κ is characteristic [GBR + 12], it follows that K is a positive definite matrix. Then, for any ϱ ∈ P ([0 , (1 -γ ) -1 R max ] d ) , we have

where M ( κ, R , ϱ ) is independent of p , so it does not influence the minimization. Now, since K is positive definite (by virtue of κ being characteristic) and ∆ R ( x ) is a closed convex subset of R N ( x ) , it is well-known that there is unique optimum, and the infimum above is attained for some p ⋆ ∈ ∆ R ( x ) . Therefore, Π R C ,κ is indeed well-defined, and its range is contained in Π R C ,κ , confirming the first two claims. Finally, since Π R C ,κ is well-defined and since MMD κ is nonnegative and separates points,

Π R C ,κ must map elements of ∆ R ( x ) to themselves - this is because MMD κ ( p, p ) = 0 for the kernels we consider.

Lemma 2. Under the conditions of Lemma 1, Π R C ,κ is a nonexpansion in MMD κ . That is, for any η 1 , η 2 ∈ P ([0 , (1 -γ ) -1 R max ] d ) X , we have MMD κ (Π R C ,κ η 1 , Π R C ,κ η 2 ) ≤ MMD κ ( η 1 , η 2 ) .

Proof. Fix any x ∈ X and denote M ( x ) = { µ p ∈ H : p ∈ ∆ R ( x ) } , where H is the RKHS induced by the kernel κ and µ p denotes the mean embedding of p in this RKHS. It is simple to verify that p ↦→ µ p is linear: for any p, q ∈ P ( R d ) and α, β ∈ R , for all f ∈ H with ∥ f ∥ = 1 we have

which implies that µ αp + βq = αµ p + βµ q . As a consequence, M ( x ) inherits convexity from ∆ R ( x ) .

̸

We claim that M ( x ) is closed as a subset of H . Since p ↦→ µ p is an injection [GBR + 12], by Lemma 1, since there is a unique q ∈ ∆ R ( x ) minimizing MMD κ ( p, q ) , there is a unique µ q ∈ M ( x ) attaining the infimum ∥ µ p -µ q ∥ over M ( x ) . Let µ ∈ H \ M ( x ) . Then there exists µ q ∈ M ( x ) such that ∥ µ q -µ ∥ = inf ν ∈ M ( x ) ∥ µ -ν ∥ , and since µ q = µ , it follows that inf ν ∈ M ( x ) ∥ ν -µ ∥ = ϵ > 0 . Since this is true for any µ ̸∈ M ( x ) , it follows that H\ M ( x ) is open, so M ( x ) is closed.

We will now show that η ( x ) ↦→ Π R C ,κ η ( x ) is a nonexpansion in H . Let η 1 , η 2 ∈ C R , and denote by µ 1 ( x ) , µ 2 ( x ) the mean embeddings of η 1 ( x ) , η 2 ( x ) . We slightly abuse notation and write Π R C ,κ µ i ( x ) to denote the mean embedding of Π R C ,κ η i ( x ) . Since M ( x ) is convex, for any ι ( x ) ∈ M ( x ) and λ ∈ [0 , 1] we have

Now, by expanding the squared norms and taking λ ↓ 0 , since M ( x ) is closed we have

where the second inequality follows by applying the same logic to µ 2 ( x ) . Choosing ι 1 ( x ) = Π R C ,κ µ 2 ( x ) , ι 2 ( x ) = Π R C ,κ µ 1 ( x ) ∈ M ( x ) and adding these inequalities yields

Expanding, we see that

confirming that η ( x ) ↦→ Π R C ,κ η ( x ) is a non-expansion. It follows that

Corollary 2. Let κ be a kernel satisfying the conditions of Theorem 2. Then for any η 0 ∈ C R , the iterates { η k } ∞ k =1 given by η k +1 = Π R C ,κ T π η k converge geometrically to a unique fixed point.

Proof. Combining Theorem 2 and Lemma 2, we see that

for some c > 0 . Thus, Π R C ,κ T π is a contraction on ( C R , MMD κ ) . If H is the RKHS induced by κ , we showed in Lemma 2 that C R is isometric to a product of closed, convex subsets of H . Therefore, by the completeness of H , C R is isometric to a complete subspace, and consequently C R is a complete subspace under the metric MMD κ . Invoking the Banach fixed-point theorem, it follows that Π R C ,κ T π has a unique fixed point η π C ,κ , and η k → η π C ,κ 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 a distribution and its projection (as a function of m,d ) plays a major role in controlling the approximation 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.

Lemma3. Let κ be kernel satisfying the conditions of Lemma 1, and for any finite R⊂ R d , define Π : P ( R d ) → ∆ R via Π p = arginf q ∈ ∆ R MMD κ ( p, q ) . Then MMD 2 κ (Π p, p ) ≤ inf P ∈ P R mesh ( P ; κ ) .

Proof. Our proof proceeds by establishing approximation bounds of Riemann sums to the Bochner integral µ p , similar to [VN02]. Let P ∈ P R . Abusing notation to denote by Π p i the probability of the i th atom of the discrete support under Π p , we have

where R = { z i } n i =1 for some n ∈ N . Since Π p optimizes the MMD over all probability vectors in ∆ R , for q ∈ ∆ R with q i = p ( θ i ) for the i th element of P , we have

It was shown by [SSGF13] that z ↦→ κ ( · , z ) is an isometry from ( R d , ρ 1 / 2 ) to H , where H is the RKHS induced by κ . Thus, we have

Since this is true for any partition P ∈ P R , the claim follows by taking the infimum over P R .

We now move on to the main results.

Theorem 4. Let κ be a kernel induced by a c -homogeneous and shift-invariant semimetric ρ conforming to the conditions of Theorem 2. Then the fixed point η π C ,κ of Π R C ,κ T π satisfies

Proof. The proof begins in a similar manner to [RBD + 18, Proposition 3]. Given that Π R C ,κ is a nonexpansion as shown in Lemma 2, we have

where ( a ) leverages the fact that η π C ,κ is the fixed point of Π R C ,κ T π and that η π is the fixed point of T π , ( b ) follows since Π R C ,κ is a nonexpansion by Lemma 2, and ( c ) follows by the contractivity of T π established in Theorem 2. Finally, by Lemma 3, we have

Finally, we explicitly derive a convergence rate for a particular support map under the energy distance kernels.

Corollary 3. Let R ( x ) = U m , where U m is a set of m uniformly-spaced support points on [0 , (1 -γ ) -1 R max ] . For κ induced by the semimetric ρ ( x, y ) = ∥ x -y ∥ α 2 for α ∈ (0 , 2) ,

Proof. We begin bounding mesh ( P ; κ ) . Assume m = n d for some n ∈ N . We consider a partition P ⊂ P U m consisting of d -dimensional hypercubes with side length (1 -γ ) -1 R max / ( n -1) . By definition of U m , it is clear that these hypercubes cover [0 , (1 -γ ) -1 R max ] d and each contain exactly one support point. Now, for each θ ∈ P , we have

where ⃗ 1 is the vector of all ones, and y is any element in θ . Expanding, we have

Since this bound holds for any θ ∈ P , invoking Theorem 4 yields

If instead m ∈ (( n -1) d , n d ) , we omit all but ( n -1) d of the support points, and achieve

Alternatively, we may write

B.4 Categorical TD Learning: Section 6

In this section, we prove results leading up to and including the convergence of the categorical TDlearning algorithm over mass-1 signed measures. First, in Section B.4.1, we show that MMD κ is in fact a metric on the space of mass-1 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 M 1 ( Y ) for spaces Y under which MMD κ is a metric on P ( Y ) .

Lemma 4. Let κ : Y × Y → R be a characteristic kernel over some space Y . Then MMD κ : M 1 ( Y ) ×M 1 ( Y ) → R + given by ( p, q ) ↦→ ∥ µ p -µ q ∥ H defines a metric on M 1 ( Y ) , where µ p denotes the usual mean embedding of p and H is the RKHS with kernel κ .

Proof. Naturally, since MMD κ is given by a norm, it is non-negative, symmetric, and satisfies the triangle inequality. We must show that MMD κ ( p, q ) = 0 ⇐⇒ p = q . Firstly, it is clear that MMD κ ( p, p ) = 0 by the positive homogeneity of the norm. It remains to show that MMD κ ( p, q ) = 0 = ⇒ p = q .

̸

Let p = q ∈ M 1 ( Y ) . For the sake of contradiction, assume that MMD κ ( p, q ) = 0 . We will show that this implies that MMD κ ( P, Q ) = 0 for a pair of distinct probability measures, which is a contradiction since MMD κ with characteristic κ is known to be a metric on P ( Y ) .

By the Hahn decomposition theorem, we may write p = ˜ p + -˜ p -for non-negative measures ˜ p + , ˜ p -. Therefore, for some a ∈ R + , we may express

where p + , p -∈ P ( Y ) . Likewise, there exist b ∈ R + and probability measure q + , q -for which q = ( b +1) q + -bq -. Since MMD κ ( p, q ) = 0 by hypothesis, and by linearity of p ↦→ µ p , we have

where P = λp + +(1 -λ ) q -and Q = λ ′ p -+(1 -λ ′ ) q + are convex combinations of probability measures, and are therefore probability measures themselves. So, we have that

which contradicts our hypothesis. Therefore, MMD κ ( p, q ) = 0 ⇐⇒ p = q for any p, q ∈ M 1 ( Y ) , and it follows that MMD κ is 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 Π R SC ,κ T π : S R → S R is affine, is contractive with contraction factor γ c/ 2 , and has a unique fixed point η π SC ,κ .

Proof. Indeed, Π R SC ,κ is, in a sense, a simpler operator than Π R C ,κ . Since M 1 ( R ( x )) is an affine subspace of M 1 ( R d ) , it holds that Π R SC ,κ is simply a Hilbertian projection, which is known to be affine and a nonexpansion [Lax02]. Moreover, since T π acts identically on M 1 ( R d ) as it does on P ( R d ) , it immediately follows that T π is a γ c/ 2 -contraction on M 1 ( R d ) , inheriting the result from Theorem 2. Thus, we have that for any η 1 , η 2 ∈ S R ,

confirming that the projected operator is a contraction. Since MMD κ is a metric on M 1 ( R ( x )) for each x ∈ X , it follows that MMD κ is a metric on S R . The existence and uniqueness of the fixed point η π SC ,κ follows by the Banach fixed point theorem.

Finally, we show that the fixed point of distributional dynamic programming with signed measure representations is roughly as accurate as η π C ,κ .

Theorem 5. Under the conditions of Lemma 5, we have that

Proof. Since Π R SC ,κ is a nonexpansion in MMD κ by Lemma 5, following the procedure of Theorem 4, we have

Note that Π R SC ,κ η π identifies the closest point (in MMD κ ) to η π in S R := ∏ x ∈X M 1 ( R ( x )) and Π R C ,κ η π identifies the closest point to η π in C R := ∏ x ∈X ∆ R ( x ) . Since it is clear that C R ⊂ S R , it follows that

The first statement then directly follows since it was shown in Lemma 3 that MMD κ (Π R C ,κ η π , η π ) ≤ sup x ∈X inf P ∈ P R ( x ) √ mesh ( P ; κ ) .

To prove the second statement, we apply the triangle inequality to yield

where the second step leverages the fact that Π R C ,κ is a nonexpansion in MMD κ by Lemma 2. Applying the conclusion of the first statement as well as the bound on MMD κ (Π R C ,κ η π , η π ) , we have

B.4.2 Convergence of Categorical TD Learning

Convergence of the proposed categorical TD-learning algorithm will rely on studying the iterates through an isometry to an affine subspace of ∏ x ∈X R N ( x ) . This affine subspace is that consisting of the set of state-conditioned 'signed probability vectors'. We define R n 1 as an affine subspace of R n for any n ∈ N according to

We note that any element η of S R can be encoded in ∏ x ∈X R N ( x ) 1 by expresing η ( x ) as the sequence of signed masses associated to each atom of R ( x ) . In Lemma 8, we exhibit an isometry I between S R and ∏ x ∈X R N ( x ) 1 .

Lemma 8. Let κ be a characteristic kernel. There exists an affine isometric isomorphism I between S R and an affine subspace ∏ x ∈X R N ( x ) 1 (cf. (15) ).

Proof. Since κ is characteristic, it is positive definite [GBR + 12]. Thus, for each x ∈ X , define K x ∈ R N ( x ) × N ( x ) according to

Then each K x is positive definite since κ is a positive definite kernel. Let p, q ∈ ∆ R ( x ) , and define P ∈ R N ( x ) and Q ∈ R N ( x ) such that P i = p ( z i ) and Q i = q ( z i ) . Then, we have

Since K x is positive definite, ∥ · ∥ K x is a norm on R N ( x ) . Therefore, the map I 1 x : (∆ R ( x ) , MMD κ ) → ( R N ( x ) , ∥ · ∥ K x ) given by I 1 x ( p ) = P is a linear isometric isomorphism

onto the affine subspace of R N ( x ) with entries summing to 1 , which we denote R N ( x ) 1 . Moreover, since ( R N ( x ) 1 , ∥ · ∥ K x ) is a finite dimensional Hilbert space, it is well known that there exists a linear isometric isomorphism I 2 x : ( R N ( x ) 1 , ∥ · ∥ K x ) → R N ( x ) 1 with the usual L 2 norm. Thus, I x = I 2 x I 1 x : (∆ R ( x ) , MMD κ ) → R N ( x ) 1 is a linear isometric isomorphism. Consequently, it follows that I : ( C R , MMD κ ) → ∏ x ∈X R N ( x ) 1 given by I = ( ∏ x ∈X R N ( x ) , ∥ · ∥ 2 , ∞ ) is a linear isometric isomorphism, where ∥ v ∥ 2 , ∞ = sup x ∈X ∥ v ( x ) ∥ 2 .

Lemma 9. Define the operator U π : ∏ x ∈X R N ( x ) 1 → ∏ x ∈X R N ( x ) 1 by U π = I Π R SC ,κ T π I -1 , where I is the isometry of Lemma 8. Let { U t } ∞ t =1 be given by U t = I η t , where { η t } ∞ t =1 are the dynamic programming iterates η t +1 = Π R SC ,κ T π η t . Then U t +1 = U π U t . Moreover, U π is contractive whenever Π R SC ,κ T π is, and in this case, U t → U ⋆ , where U ⋆ is the unique fixed point of U π .

Proof. By definition, we have

which proves the first claim. Moreover, for U 1 = I η 1 and U 2 = I η 2 , we have

where the second step transforms the ∥ · ∥ 2 , ∞ to MMD κ since I is an isometry between those metric spaces. Therefore, if Π R C ,κ T π is contractive with contraction factor β ∈ (0 , 1) , we have

so that U π has the same contraction factor as Π R C ,κ T π . Consequently, by the Banach fixed point theorem, U π has a unique fixed point U ⋆ whenever Π R C ,κ T π is contractive, and U t → U ⋆ at the same rate as η t → η π .

The main theorem in this section is that temporal difference learning on the finite dimensional representations I ( η t ) converges.

Proposition 2 (Convergence of categorical temporal difference learning) . Let { U t } ∞ t =1 ⊂ ∏ x ∈X R N ( x ) 1 be given by U t = I ̂ η t , and let κ be a kernel satisfying the conditions of Theorem 2. Suppose that, for each x ∈ X , the states { X t } ∞ t =1 and step sizes { α t } ∞ t =1 satisfy the Robbins-Munro conditions

Then, with probability 1, U k → U ⋆ , where U ⋆ is the fixed point of U π .

The proof of this result as a 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 n ) X → ( R n ) X be a contractive operator with respect to ∥ · ∥ 2 , ∞ with fixed point Z ⋆ , and let (Ω , F , { F k } ∞ k =1 , P ) be a filtered probability space. Define a map ̂ O : ( R n ) X ×X × Ω → ( R n ) X such that

For a stochastic process { ξ k } ∞ k =1 adapted to {F k } ∞ k =1 with ξ k = X k ⊕ ω k , consider a sequence { Z k } ∞ k =1 ⊂ ( R n ) X given by

where { α k } ∞ k =1 is adapted to {F k } ∞ k =1 and satisfies the Robbins-Munro conditions for each x ∈ X ,

Finally, for the processes { w ( x ) k } ∞ k =1 where w ( x ) k = ( ̂ O ( Z k , X k , ω k ) -( O Z k )( X k )) 1 [ X k = x ] , assume the following moment condition holds,

for finite universal constants C 1 , C 2 . Then, with probability 1, Z k → Z ⋆ .

The operator O of Lemma 10 will be substituted with U π , governing the dynamics of the encoded iterates of the multi-return distribution. The stochastic operator ̂ O will be substituted with the corresponding stochastic TD operator for U π , given by

This corresponds to applying a Bellman backup from a stochastic reward R and next state x 2 , followed by projecting back onto S R , and applying the isometry back to ∏ x ∈X R N ( x ) 1 .

Proof of Proposition 2. Let n = max x ∈X N ( x ) . Note that for any m ≤ n , R m can be isometrically embedded into R n by zero-padding. Thus, ∏ x ∈X R N ( x ) 1 can be isometrically embedded into ( R n 1 ) X , so without loss of generality, we will assume that N ( x ) ≡ n .

Define the map ̂ U π : ( R n 1 ) X ×X × ( R d ×X ) → ( R n 1 ) X according to

Then, defining ̂ U π k U = ̂ U π ( U, X k , ( R k , X ′ k )) with ( X k , A k , R k , X ′ k ) = T k ∼ P as in (11), we have

̸

̸

Note that, since Π R SC ,κ is a Hilbert projection onto an affine subspace, it is affine [Lax02]. Consequently, ̂ U π is an unbiased estimator of the operator U π in the following sense,

where the first step invokes the linearity of Π R SC ,κ , the second step invokes the linearity of the isometry I established in Lemma 8 and the third step is due to the definition of T π . As a result, we see that the conditions (16) and (17) of Lemma 10 are satisfied by ̂ U π , the iterates { U k } ∞ k =1 , and the step sizes { α k } ∞ k =1 . Moreover, for w k ( x ) defined by

we have ∥ w k ( x ) ∥ 2 ≤ C 1 + C 2 ∥ U ( x ) ∥ 2 for universal constants C 1 , C 2 -this is shown in Lemma 11. As such, the condition of (18) is satisfied.

Finally, since U π inherits contractivity from Π R SC ,κ T π as shown in Lemma 5, we may invoke Lemma 10, which ensures that U k → U with probability 1, where U = U π U is a unique fixed point.

Theorem 6. For a kernel κ induced by a semimetric ρ of strong negative type, the sequence { ̂ η t } ∞ t =1 given by (11) -(13) converges to η π SC ,κ with probability 1.

Proof. By Proposition 2, the sequence { U t } ∞ t =1 with U t = I η t converges to a unique fixed point U with probability 1. Note that

Therefore, I -1 U ⋆ is a fixed point of Π R SC ,κ T π . Since it was shown in Lemma 5 that Π R SC ,κ T π has a unique fixed point, it follows that I -1 U ⋆ = η π SC ,κ . Since I is an isometry, ̂ η t = I -1 U t →I -1 U ⋆ with probability 1, so indeed ̂ η t → η π SC ,κ with probability 1.

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 ∈ X and U ∈ ∏ x ∈X R N ( x ) -1 ,

for finite constants C 1 , C 2 ∈ R + .

Proof. Since I is an isometry, we have that

where H is the RKHS induced by the kernel κ . Moreover, since Π R SC ,κ is a nonexpansion in ∥ · ∥ H as argued in Lemma 5, we have that

Proceeding, we will bound the terms A,B . To bound A , we simply have

where we invoke the contraction of U π in MMD κ from Theorem 2. Note that η π ( x ) ∈ P ([0 , (1 -γ ) -1 R max ] d ) , so it follows that ∥ η π ∥ 2 H ≤ D 1 , 1 for some constant D 1 , 1 since the kernel κ is bounded in compact domains. Expanding the norm of the difference above yields

for a finite constant D 1 , 2 , again invoking the isometry I in the last step.

Our bound for B is similar. Choose any x ′ ∈ supp P π ( · | x ) . We consider the operator ˜ T x ′ : P ([0 , (1 -γ ) -1 R max ] d ) →P ([0 , (1 -γ ) -1 R max ] d ) given by

This operator is a contraction in MMD κ , and correspondingly has a fixed point η π x ′ . To see this, we note that ˜ T x ′ is simply a special case of U π for the case P π ( · | x ) = δ x ′ , so the contractivity and existence of the fixed point are inherited from Theorem 2. Proceeding in a manner similar to the bound on A , we have

where the final step mirrors the bound on A . Therefore, we have shown that

completing the proof.

C Memory Efficiency of Randomized EWP Dynamic Programming

In Section 4, we argued for the necessity of considering a projection operator in EWP dynamic programming. While we provided a randomized projection, Theorem 3 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 η π without costing too much in memory.

In this section, we demonstrate that, in fact, the algorithm described in Theorem 3 can approximate η π to any desired accuracy with many fewer particles. Suppose our goal is to derive some η such that

for some ϵ > 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 m unproj ) as well as with our algorithm described in Theorem 3 (denoting the number of particles m proj . In both cases, we will compute iterates starting with some η 0 ∈ C EWP ,m with MMD κ ( η 0 , η π ) ≤ D < ∞ . For simplicity, we will consider the energy distance kernel with α = 1 .

The remainder of this section will show that the dependence of the number of atoms on both ϵ and |X| is substantially worse in the unprojected case (that is, m proj ≪ m unproj for large state spaces or low error tolerance). We demonstrate this with concrete lower bounds on m unproj and upper bounds on m proj below; note that these bounds are not optimized for tightness or generality, and are instead aimed to provide straightforward evidence of our core points above.

We will begin by bounding m unproj . In the best case, η 0 ( x ) 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 η 0 will result in η 1 ( x ) having support on |X| atoms at each

state x (due to the mixture over successor states in the Bellman backup). Consequently, the iterate η k ( x ) will be supported on |X| k atoms. Since MMD κ ( η k , η π ) ≤ γ 1 / 2 D by Theorem 2, we require

to ensure that MMD κ ( η K , η π ) ≤ ϵ . Thus, we have

On the other hand, the following lemma bounds m proj ; we prove the lemma at the end of this section. Lemma 12. Let η m proj denote the output of the projected EWP algorithm described by Theorem 3 with m = m proj particles. Then under the assumptions of Theorem 3 and with the energy distance kernel with α = 1 , MMD κ ( η m proj , η π ) ≤ ϵ is achievable with

For any fixed MDP with |X| ≥ 4 and γ ≥ 1 / 2 , we have that

since D > 0 and does not depend on ϵ . Meanwhile, we have m proj ∈ Θ( ϵ -2 polylog ( ϵ -1 )) by Lemma 12, indicating a much more graceful dependence on ϵ relative to the unprojected algorithm.

On the other hand, for any fixed tolerance ϵ ≤ γD , we immediately have

In the worst case, we may have d ∈ Θ( |X| ) (any larger d will induce linearly dependent cumulants). Thus, we have

so the projected algorithm scales much more gracefully with |X| as well.

Proof of Lemma 12

Finally, we prove Lemma 12, which determines the number of atoms required to achieve an ϵ -accurate return distribution estimate with the algorithm of Theorem 3.

Lemma 12. Let η m proj denote the output of the projected EWP algorithm described by Theorem 3 with m = m proj particles. Then under the assumptions of Theorem 3 and with the energy distance kernel with α = 1 , MMD κ ( η m proj , η π ) ≤ ϵ is achievable with

Proof. Note that, by Theorem 3, increasing m proj can only decrease the error ϵ as long as m proj ≥ 1 . Therefore, as shown in (14) in the proof of Theorem 3, there exists a universal constant C 0 > 0 such that

Now, we write c 3 = C 0 c 1 c 2 , c 4 = C 0 c 1 , and u := √ m proj , yielding

Then, after isolating the logarithmic term and exponentiating, we see that

We now rearrange this expression and invoke the identity W ( z ) e W ( z ) = z where W is a Lambert W -function [CGH + 96]:

There are two branches of the Lambert W -function on the reals, namely W 0 and W -1 . These two branches satisfy W 0 ( ze z ) = z when z ≥ -1 and W -1 ( ze z ) = z when z ≤ -1 . In our case, we know that z is negative, and it is known [CGH + 96] that | W 0 ( z ) | ≤ 1 when z ∈ [ -1 , 0] . Consequently, when z ≥ -1 , we have | uϵ 2 c 4 | ≤ 1 , and substituting m proj = u 2 , we have

On the other hand, when z ≤ -1 , we have

Since it is known [CGH + 96, Equations 4.19, 4.20] that W 2 -1 ( -z ) ∈ polylog (1 /z ) , incorporating (23), we have that

The upper bound given above will generally not be an integer. Howevever, increasing m proj can only improve the approximation error, as shown in Theorem 3 since log m/ √ m decreases monotonically when m > 7 . So, we can round m proj up to the nearest integer (or round it down when m ≤ 7 ) incurring a penalty of at most one atom. It follows that the randomized EWP dynamic programming algorithm of Theorem 3 run with m proj given by (21) produces a return distribution function η m proj for which MMD κ ( η m proj , η π ) ≤ ϵ .

D Nonlinearity of the Categorical MMD Projection

In Section 6, we noted that the categorical projection Π R C ,κ is non-affine. Here, we provide an explicit example certifying this phenomenon.

We consider a single-state MDP, since the nonlinearity issue is independent of the cardinality of the state space (the projection is applied to each state-conditioned distribution independently). We write R = { 0 , . . . , 3 } 2 , and consider the kernel κ induced by ρ ( x, y ) = ∥ x -y ∥ 2 -this resulting MMDis known as the energy distance, which is what we used in our experiments. We consider two distributions, p 1 = δ [1 . 5 , 1 . 5] and p 2 = δ [2 . 5 , 0] .

̸

Weconsider λ = 0 . 8 and compare q 1 = Π R C ,κ ( λp 1 +(1 -λ ) p 2 ) with q 2 = λ Π R C ,κ p 1 +(1 -λ )Π R C ,κ p 2 , and we note that q 1 = q 2 ; confirming that Π R C ,κ is not an affine map. The results are tabulated in Table 1, with bolded entries depicting the atoms with non-negligible differences in probability under q 1 , q 2 .

E Experiment Details

TD-learning experiments were conducted on a NVidia A100 80G GPU to parallelize experiments. Methods were implemented in Jax [BFH + 18], particularly with the help of JaxOpt [BBC + 21] for vectorizing QP solutions - this was helpful for computing the categorical projections discussed in this work. SGD was used for optimization, using an annealed learning rate schedule ( λ k ) k ≥ 0 with λ k = k -3 / 5 , 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 were implemented in the Julia programming language [BEKS17].

In all experiments, we used the kernel induced by ρ ( x, y ) = ∥ x -y ∥ 2 with reference point 0 for MMDoptimization-this corresponds to the 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, in this section, we demonstrate that the signed categorical TD learning algorithm presented in Section 6 can be scaled to continuous state spaces with neural networks. We will consider an environment with visual (pixel) observations of a car in a parking lot, an example observation is shown in Figure 4.

Here, we consider 2-dimensional cumulants, where the first dimension tracks the x coordinate of the car, and the second dimension is an indicator that is 1 if and only if the car is parked in the parking spot. We learn a multivariate return distribution function with transitions sampled from trajectories that navigate around the obstacle to the parking spot. Notably, the successor features (expectation of multivariate return distribution) will be zero in the first dimension, since the set of trajectories is horizontally symmetric. Thus, from the successor features alone, one cannot distinguish the observed policy from one that traverses straight through the obstacle!

Car

Obstacle

Parking Spot

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.

CNN

CNN

To learn the multivariate return distribution function from images, we use a convolutional neural architecture as shown in Figure 5.

Notably, we simply use convolutional networks to model the signed masses for the fixed atoms of the categorical representation. The projection Π R SC ,κ is computed by a QP solver as 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 [ZCZ + 21]. Our results are shown in Figure 6. We see that both TD-learning methods accurately estimate the distribution over multivariate returns, indicating that no multivariate return will have a vanishing lateral component. Quantitatively, we see that the EWP algorithm appears to be stuck in a local optimum, with some particles lying in regions of low probability mass.

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 to get to the parking spot as quickly as possible. We see that the EWP TD learning algorithm of [ZCZ + 21] more accurately estimates the return distribution function corresponding to the latter reward vector, while our 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.

0.75

0.50

0.25 0.00

0.25

0.50

0.75

Lateral Feature

0.0

0.2

0.4

0.6

0.8

1.0

Parked Feature

0.6

0.4

0.2

0.0

0.2

0.4

0.6

Lateral Feature

0.0

0.2

0.4

0.6

0.8

1.0

Parked Feature

Deep EWP-TD

30

20

10

0

10

20

0.00

0.01

0.02

0.03

0.04

0.05

Density

EWP, Cramer = 0.0823

Cat, Cramer = 0.0806

0.1

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0

1

2

3

4

5

6

Density

EWP, Cramer = 0.0088

Cat, Cramer = 0.0091

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: 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 contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other 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 no limitation while the answer No means that the paper has limitations, but those are not discussed in 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 violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-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 reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that

aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All 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 in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Algorithms are described explicitly, experimental setup is 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 is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, 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 NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

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 a 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, or as 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 be 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables 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 experiments?

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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental 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, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [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 no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is 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, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

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 a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

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 a 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, paperswithcode.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 new 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/model as part of their submissions 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 or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: 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.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human 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 of 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 to adhere to 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 as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* **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].\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nyes, the authors mentioned some limitations of their work in Section 7", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\n-\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.\n\nLIMITATIONS:\nThe authors have clearly stated the assumptions and limitations of their theory.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

comments: 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.

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.
* Typo: Equation (10), $X’_r$ -> $X’_t$

QUESTIONS:
* What does the right-hand-side (RHS) of the first equation in Eq. (2) mean? 
* 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$.
* Define $\Delta$ notation on Line 188.

LIMITATIONS:
See the weaknesses.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13