Foundations of Multivariate Distributional Reinforcement Learning

Distributional reinforcement learning [MSK⁺10, BDM17b, BDR23, DRL] 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,¹¹1Cumulants refer to accumulated quantities in RL (e.g., rewards or multivariate rewards)—not to be confused with statistical cumulants. 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 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-space-valued 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 temporal-difference 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 mass-$1$ signed measures. These new techniques recover existing bounds even in the scalar reward case, and provide new insights into the behavior of distributional RL algorithms as a function of the reward dimension.

We consider a Markov decision process with Polish state space $\mathcal{X}$, action space $\mathcal{A}$, transition kernel $P:\mathcal{X}\times\mathcal{A}\to\mathscr{P}(\mathcal{X})$, and discount factor $\gamma\in[0,1)$. Unlike the standard RL setting, we consider a vector-valued reward function $r:\mathcal{X}\to[0,R_{\max}]^{d}$, as in the literature on successor features [BDM⁺18]. Given a policy $\pi:\mathcal{X}\to\mathcal{P}(\mathcal{A})$, we write the policy-conditioned transition kernel $P^{\pi}(\cdot\mid x)=\int P(\cdot\mid x,a)\pi(\mathrm{d}a\mid x)$.

Multi-variate return distributions. We write $(X_{t})_{t\geq 0}$ for a trajectory generated by setting $X_{0}=x$, and for each $t\geq 0$, $X_{t+1}\sim P^{\pi}(\cdot|X_{t})$. The return obtained along this trajectory is then defined by $G^{\pi}(x)=\sum_{t=0}^{\infty}\gamma^{t}r(X_{t})$, and the (multi-)return distribution function is $\eta^{\pi}(x)=\text{Law}\left(G^{\pi}(x)\right)$.

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 $\tilde{r}:x\mapsto\langle r(x),w\rangle$ for some $w\in\operatorname{\mathbf{R}}^{d}$, $\langle G^{\pi}(x),w\rangle\operatorname{=_{\rm law}}\sum_{t\geq 0}\gamma^{t}\tilde{r}(X_{t})$ for $X_{0}=x$. Likewise, one might consider $r(x)=\delta_{x}$, in which case $G^{\pi}(x)$ corresponds to the per-trajectory discounted state visitation measure, and [WFG⁺24] demonstrated methods for learning the distribution of $G^{\pi}$ 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 $\mathcal{T}^{\pi}:\mathcal{P}(\operatorname{\mathbf{R}}^{d})^{\mathcal{X}}\to\mathcal{P}(\operatorname{\mathbf{R}}^{d})^{\mathcal{X}}$ by

where $\mathrm{b}_{y,\gamma}(z)=y+\gamma z$ and $f_{\sharp}\mu=\mu\circ f^{-1}$ is the pushforward of a measure $\mu$ through a measurable function $f$. In particular, [ZCZ⁺21] showed that $\eta^{\pi}$ satisfies the multi-variate distributional Bellman equation $\mathcal{T}^{\pi}\eta^{\pi}=\eta^{\pi}$, and that $\mathcal{T}^{\pi}$ is a $\gamma$-contraction in $\overline{W}_{p}$, where $\overline{W}_{p}(\eta_{1},\eta_{2})=\sup_{x\in\mathcal{X}}W_{p}(\eta_{1}(x),\eta_{2}(x))$ and $W_{p}$ is the $p$-Wasserstein metric [Vil09]. This suggests a convergent scheme for approximating $\eta^{\pi}$ in $\overline{W}_{p}$ by distributional dynamic programming, that is, computing the iterates $\eta_{k+1}=\mathcal{T}^{\pi}\eta_{k}$, following Banach’s fixed point theorem.

Approximating multivariate return distributions. In practice, however, the iterates $\eta_{k+1}=\mathcal{T}^{\pi}\eta_{k}$ cannot be computed efficiently, because the size of the support of $\eta_{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 [NGV20, ZCZ⁺21], is the notion of distance over probability distributions given by maximum mean discrepancies [GBR⁺12, MMD]. A maximum mean discrepancy $\mathrm{MMD}_{\kappa}:\mathscr{P}(\mathcal{Y})\times\mathscr{P}(\mathcal{Y})\to\operatorname{\mathbf{R}}_{+}$ assigns a notion of distance to pairs of probability distributions, and is parametrised via a choice of kernel $\kappa:\mathcal{Y}\times\mathcal{Y}\rightarrow\operatorname{\mathbf{R}}$, defined by

A useful alternative perspective on MMD is that the choice of kernel $\kappa$ induces a reproducing kernel Hilbert space (RKHS) of functions $\mathcal{H}$, namely the closure of the span of functions of the form $z\mapsto\kappa(y,z)$ for each $y\in\mathcal{Y}$, with respect to the norm $\|\cdot\|_{\mathcal{H}}$ induced by the inner product $\langle\kappa(y_{1},\cdot),\kappa(y_{2},\cdot)\rangle=\kappa(y_{1},y_{2})$. With this interpretation, $\mathrm{MMD}_{\kappa}(p,q)$ is equal to $\|\mu_{p}-\mu_{q}\|_{\mathcal{H}}$, where $\mu_{p}=\int_{\mathcal{Y}}\kappa(\cdot,y)p(\mathrm{d}y)\in\mathcal{H}$ is the mean embedding of $p$ (similarly for $\mu_{q}$). When $p\mapsto\mu_{p}$ is injective, the kernel $\kappa$ is called characteristic, and $\mathrm{MMD}_{\kappa}$ is then a proper metric on $\mathcal{P}(\mathcal{Y})$ [GBR⁺12]. In the remainder of this work, we will assume that all spaces of measures will be over compact sets $\mathcal{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.

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

To warm up, we begin by demonstrating that indeed the (multivariate) distributional Bellman operator is contractive in a supremal form $\overline{\mathrm{MMD}_{\kappa}}$ of MMD, given by $\overline{\mathrm{MMD}_{\kappa}}(\eta_{1},\eta_{2})=\sup_{x\in\mathcal{X}}\mathrm{MMD}_{\kappa}(\eta_{1}(x),\eta_{2}(x))$, for a particular class of kernels $\kappa$. Our first theorem generalizes the analogous results of [NGV20] in the scalar case to multivariate cumulants. The proof of Theorem 3, as well as proofs of all remaining results, are deferred to Appendix B.

[Convergent MMD dynamic programming for the multi-return distribution function]theoremdpmmd Let $\kappa$ be a kernel induced by a semimetric $\rho$ on $[0,(1-\gamma)^{-1}R_{\max}]^{d}$ with strong negative type, satisfying

1. 1.

   Shift-invariance. For any $z\in\operatorname{\mathbf{R}}^{d}$, $\rho(z+y_{1},z+y_{2})=\rho(y_{1},y_{2})$.

2. 2.

   Homogeneity. For any $\gamma\in[0,1)$, there exists $c>0$ for which $\rho(\gamma y_{1},\gamma y_{2})=\gamma^{c}\rho(y_{1},y_{2})$.

Consider the sequence $\left\{{\eta}_{k}\right\}_{k=1}^{\infty}$ given by $\eta_{k+1}=\mathcal{T}^{\pi}\eta_{k}$. Then $\eta_{k}\to\eta^{\pi}$ at a geometric rate of $\gamma^{c/2}$ in $\overline{\mathrm{MMD}_{\kappa}}$, as long as $\overline{\mathrm{MMD}_{\kappa}}(\eta_{0},\eta^{\pi})\leq C<\infty$.

Notably, the energy distance kernels $\kappa_{\alpha}$ satisfy the conditions of Theorem 3, and $\rho_{\alpha}(\gamma y_{1},\gamma y_{2})\leq\gamma^{\alpha}\rho(y_{1},y_{2})$ by the homogeneity of the Euclidean norm, so $\mathcal{T}^{\pi}$ is a $\gamma^{\alpha/2}$-contraction in the energy distances. This generalizes the analogous result of [NGV20] in the one-dimensional case.

While Theorem 3 illustrates a method for approximating $\eta^{\pi}$ 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, $\mathcal{T}^{\pi}$ cannot be represented or computed exactly, and instead we must study algorithms for approximating $\eta^{\pi}$ 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.

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 $\mathscr{C}_{\mathrm{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 $\eta\in\mathscr{C}_{\mathrm{EWP},m}$ by $\eta(x)=\frac{1}{m}\sum_{i=1}^{m}\delta_{\theta_{i}(x)}$ for $\theta_{i}:\mathcal{X}\to\operatorname{\mathbf{R}}^{d}$. The work of [ZCZ⁺21] introduced a TD-learning algorithm for learning a $\mathscr{C}_{\mathrm{EWP},m}$ representation of $\eta^{\pi}$, computing iterates of the particles $(\theta^{(k)}_{i})_{i=1}^{m}$ according to

for step sizes $(\lambda_{k})_{k\geq 0}$ and sampled next states $X^{\prime}\sim P^{\pi}(\cdot\mid x)$, where $\overline{\theta}=\texttt{stop-gradient}(\theta^{(k)})$ is a copy of $\theta^{(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⁺23], we have $\mathcal{T}^{\pi}\mathscr{C}_{\mathrm{EWP},m}\not\subset\mathscr{C}_{\mathrm{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^{\pi}(\cdot\mid x)$ is not a singleton, $(\mathcal{T}^{\pi}\eta)(x)$ will consist of more than $m$ atoms even when $\eta\in\mathscr{C}_{\mathrm{EWP},m}$; and secondly, if $P(\cdot\mid x)$ is not uniform, $(\mathcal{T}^{\pi}\eta)(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 $\Pi_{\mathrm{EWP}}\mathcal{T}^{\pi}$, for a projection $\Pi_{\mathrm{EWP}}:\mathcal{P}(\operatorname{\mathbf{R}}^{d})^{\mathcal{X}}\to\mathscr{C}_{\mathrm{EWP},m}$. A natural choice for this projection is the operator that minimizes the MMD of each multi-return distribution in $\mathscr{C}_{\mathrm{EWP},m}$,

Unfortunately, even in the scalar-reward ($d=1$) case, the operator $\Pi_{\mathrm{EWP},\kappa}^{m}$ is problematic; $(\Pi_{\mathrm{EWP},\kappa}^{m}\eta)(x)$ is not uniquely defined, and $\Pi_{\mathrm{EWP},\kappa}^{m}$ is not a non-expansion in $\overline{\mathrm{MMD}_{\kappa}}$ [LB22, RMA⁺23]. 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 $\Pi_{\mathrm{EWP},\kappa}^{m}\mathcal{T}^{\pi}$ 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 $\eta(x)$ has finite support for each state $x$, the projection $(\Pi_{\mathrm{EWP},\kappa}^{m}\eta)(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 $\mathsf{BootProj}^{\pi}_{\kappa,m}$ for $\Pi_{\kappa,m}\mathcal{T}^{\pi}$, 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 $(\mathcal{T}^{\pi}\eta)(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.

theoremewpmmddpconvergence Consider a kernel $\kappa$ induced by the semimetric $\rho(x,y)=\|x-y\|_{2}^{\alpha}$ with $\alpha\in(0,2)$, and suppose rewards are bounded in each dimension within $[0,R_{\max}]$. For any $\eta_{0}$ such that $\overline{\mathrm{MMD}_{\kappa}}(\eta_{0},\eta^{\pi})\leq D<\infty$, and any $\delta>0$, for the sequence $(\eta_{k})_{k\geq 0}$ defined in Equation (5), with probability at least $1-\delta$ we have

where $\eta_{k+1}=\mathsf{BootProj}^{\pi}_{\kappa,m}\eta_{k}$ and $K=\lceil\frac{\log m}{\log\gamma^{-\alpha}}\rceil$, and where $\widetilde{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⁺23, 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 4 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.

[Convergence of EWP Dynamic Programming]propositiondpewp Consider a kernel satisfying the hypotheses of Theorem 3, suppose rewards are globally bounded in each dimension in $[0,R_{\max}]$, and let $\{\Pi^{(k)}_{\kappa,m}\}_{k\geq 0}$ be a sequence of maps $\Pi:\mathcal{P}([0,(1-\gamma)^{-1}R_{\max}]^{d})^{\mathcal{X}}\to\mathscr{C}_{\mathrm{EWP},m}$ satisfying

Then the iterates $(\eta_{k})_{k\geq 0}$ given by $\eta_{k+1}=\Pi^{(k)}_{\kappa,m}\mathcal{T}^{\pi}\eta_{k}$ with $\overline{\mathrm{MMD}_{\kappa}}(\eta_{0},\eta^{\pi})=D<\infty$ converge to a set $\boldsymbol{\eta}_{\mathrm{EWP},\kappa}^{m}\subset\overline{B}(\eta^{\pi},(1-\gamma^{c/2})^{-1}f(d,m))$ in $\overline{\mathrm{MMD}_{\kappa}}$, where $\overline{B}$ denotes the closed ball in $\overline{\mathrm{MMD}_{\kappa}}$,

As an immediate corollary of Proposition 4 and Theorem 4, we can derive an error rate for projected dynamic programming with $\Pi_{\mathrm{EWP},\kappa}^{m}$ as well.

corollarydpewporacle For any kernel $\kappa$ satisfying the hypotheses of Theorem 4, and for any $\eta_{0}\in\mathscr{C}_{\mathrm{EWP},m}$ for which $\overline{\mathrm{MMD}_{\kappa}}(\eta_{0},\eta^{\pi})\leq D<\infty$, the iterates $\eta_{k+1}=\Pi_{\mathrm{EWP},\kappa}^{m}\mathcal{T}^{\pi}\eta_{k}$ converge to a set $\boldsymbol{\eta}_{\mathrm{EWP},\kappa}^{m}\subset\mathscr{C}_{\mathrm{EWP},m}$, where

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 $\mathcal{R}(x)=\left\{{\xi(x)}_{i}\right\}_{i=1}^{N(x)}\subset\operatorname{\mathbf{R}}^{d}$, that models the multi-return distribution function $\eta$ such that $\eta(x)\in\Delta_{\mathcal{R}(x)}$ for each $x\in\mathcal{X}$. The notation $\xi(x)_{i}$ simply refers to the $i$th support point at state $x$ specified by $\mathcal{R}$, and $\Delta_{A}$ denotes the probability simplex on the finite set $A$. We refer to the mapping $\mathcal{R}$ as the support map²²2In many applications, the most natural support map is constant across the state space. and we denote the class of multi-return distribution functions under the corresponding categorical representation as $\mathscr{C}_{\mathcal{R}}\triangleq\prod_{x\in\mathcal{X}}\Delta_{\mathcal{R}(x)}$.

Categorical projection. Once again, the distributional Bellman operator is not generally closed over $\mathscr{C}_{\mathcal{R}}$, that is, $\mathcal{T}^{\pi}\mathscr{C}_{\mathcal{R}}\not\subset\mathscr{C}_{\mathcal{R}}$. As such, we cannot actually employ the procedure described in Theorem 3 – rather, we need to project applications of $\mathcal{T}^{\pi}$ back onto $\mathscr{C}_{\mathcal{R}}$. Roughly, we need an operator $\Pi:\mathcal{P}(\operatorname{\mathbf{R}}^{d})^{\mathcal{X}}\to\mathscr{C}_{\mathcal{R}}$ for which $\Pi\rvert_{\mathscr{C}_{\mathcal{R}}}=\operatorname{\mathsf{id}}$. Given such an operator, as in the literature on categorical distributional RL [BDM17b, RBD⁺18], we will study the convergence of iterates $\eta_{k+1}=\Pi\mathcal{T}^{\pi}\eta_{k}$.

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

We will now verify that $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}$ is well-defined, and that it satisfies the aforementioned properties. {restatable}lemmacatprojwelldefined Let $\kappa$ be a kernel induced by a semimetric $\rho$ on $[0,(1-\gamma)^{-1}R_{\max}]^{d}$ with strong negative type (cf. Theorem 1). Then $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}$ is well-defined, $\operatorname{\mathsf{Ran}}(\Pi_{\mathrm{C},\kappa}^{\mathcal{R}})\subset\mathscr{C}_{\mathcal{R}}$, and $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\rvert_{\mathscr{C}_{\mathcal{R}}}=\operatorname{\mathsf{id}}$. It is worth noting that beyond simply ensuring the well-posedness of the projection $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}$, Lemma 7 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 $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\mathcal{T}^{\pi}$ with a QP solver $\mathsf{QPSolve}$ in Algorithm 1.

lemmaprojorthogonal Under the conditions of Lemma  7, $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}$ is a nonexpansion in $\overline{\mathrm{MMD}_{\kappa}}$. That is, for any $\eta_{1},\eta_{2}\in\mathcal{P}([0,(1-\gamma)^{-1}R_{\max}]^{d})^{\mathcal{X}}$, we have $\overline{\mathrm{MMD}_{\kappa}}(\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\eta_{1},\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\eta_{2})\leq\overline{\mathrm{MMD}_{\kappa}}(\eta_{1},\eta_{2}).$

Categorical multivariate distributional dynamic programming. As an immediate consequence of Lemma 1, it follows that projected dynamic programming under the projection $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}$ is convergent; this is because $\mathcal{T}^{\pi}$ is a contraction in $\overline{\mathrm{MMD}_{\kappa}}$ and $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}$ is a nonexpansion in $\overline{\mathrm{MMD}_{\kappa}}$, so the projected operator $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\mathcal{T}^{\pi}$ is a contraction in $\overline{\mathrm{MMD}_{\kappa}}$; a standard invokation of the Banach fixed point theorem appealing to the completenes of $\overline{\mathrm{MMD}_{\kappa}}$ certifies that repeated application of $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\mathcal{T}^{\pi}$ will result in convergence to a unique fixed point.

corollarydpprojfixedpoint Let $\kappa$ be a kernel satisfying the conditions of Theorem 3. Then for any $\eta_{0}\in\mathscr{C}_{\mathcal{R}}$, the iterates $\left\{{\eta}_{k}\right\}_{k=1}^{\infty}$ given by $\eta_{k+1}=\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\mathcal{T}^{\pi}\eta_{k}$ converge geometrically to a unique fixed point.

Beyond the result of Theorem 3, Corollary 1 illustrates an algorithm for estimating $\eta^{\pi}$ in $\overline{\mathrm{MMD}_{\kappa}}$ provided knowledge of the transition kernel and the reward function, which is computationally tractable in tabular MDPs. Indeed, the iterates $(\eta_{k})_{k\geq 0}$ all lie in $\mathscr{C}_{\mathcal{R}}$, having finitely-many degrees of freedom. Algorithm 1 outlines a computationally tractable procedure for learning $\eta_{\mathrm{C},\kappa}^{\pi}$ 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 $\eta_{\mathrm{C},\kappa}^{\pi}$ from Corollary 1 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 $\mathscr{P}_{\mathcal{R}(x)}$, where each $P\in\mathscr{P}_{\mathcal{R}(x)}$ satisfies

1. 1.

   $|P|=N(x)$;

2. 2.

   For any $\theta_{1},\theta_{2}\in P$, either $\theta_{1}\cap\theta_{2}=\emptyset$ or $\theta_{1}=\theta_{2}$;

3. 3.

   $\cup_{\theta\in P}\theta=\mathcal{P}([0,(1-\gamma)^{-1}R_{\max}]^{d})$;

4. 4.

   Each element $\theta_{i}\in P$ contains exactly one element $z_{i}\in\mathcal{R}(x)$.

For any partition $P$, we define a notion of mesh size relative to a kernel $\kappa$ induced by a semimetric $\rho$,

The following result confirms that $\eta_{\mathrm{C},\kappa}^{\pi}$ recovers $\eta^{\pi}$ as the mesh decreases.

theoremdpconsistent Let $\kappa$ be a kernel induced by a $c$-homogeneous and shift-invariant semimetric $\rho$ conforming to the conditions of Theorem 3. Then the fixed point $\eta_{\mathrm{C},\kappa}^{\pi}$ of $\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\mathcal{T}^{\pi}$ satisfies

Thus, for any sequence of supports $\{\mathcal{R}(x)_{m}\}_{m\geq 1}$ for which the maximal space (in $\rho$) between any two points in $\mathcal{R}(x)_{m}$ tends to $0$ as $m\to\infty$, the fixed point $\eta_{\mathrm{C},\kappa}^{\pi}$ approximates $\eta^{\pi}$ to arbitrary precision for large enough $m$. The next corollary illustrates this in a familiar setting.

corollarymeshcovering Let $\mathcal{R}(x)=U_{m}$, where $U_{m}$ is a set of $m$ uniformly-spaced support points on $[0,(1-\gamma)^{-1}R_{\max}]$. For $\kappa$ induced by the semimetric $\rho(x,y)=\|x-y\|_{2}^{\alpha}$ for $\alpha\in(0,2)$,

With $\alpha=1$ and $d=1$, the MMD in Corollary 1 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 $\Theta(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 1 relies on the following lemma about the approximation quality of the categorical MMD projection, which may be of independent interest.

lemmammdprojapprox Let $\kappa$ be kernel satisfying the conditions of Lemma 7, and for any finite $\mathcal{R}\subset\operatorname{\mathbf{R}}^{d}$, define $\Pi:\mathcal{P}(\operatorname{\mathbf{R}}^{d})\to\Delta_{\mathcal{R}}$ via $\Pi p=\operatorname*{arginf}_{q\in\Delta_{\mathcal{R}}}\mathrm{MMD}_{\kappa}(p,q)$. Then $\mathrm{MMD}_{\kappa}^{2}(\Pi p,p)\leq\inf_{P\in\mathscr{P}_{\mathcal{R}}}\mathsf{mesh}(P;\kappa)$.

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.

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 $\overline{\mathrm{MMD}_{\kappa}}$.

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\geq 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 $\mathcal{M}^{1}(\mathcal{Y})$ the set of all signed measures $\mu$ over $\mathcal{Y}$ with $\mu(\mathcal{Y})=1$. We begin by noting that the MMD endows $\mathcal{M}^{1}(\mathcal{Y})$ with a metric structure.

lemmammdsigned Let $\kappa:\mathcal{Y}\times\mathcal{Y}\to\operatorname{\mathbf{R}}$ be a characteristic kernel over some space $\mathcal{Y}$. Then $\mathrm{MMD}_{\kappa}:\mathcal{M}^{1}(\mathcal{Y})\times\mathcal{M}^{1}(\mathcal{Y})\to\operatorname{\mathbf{R}}_{+}$ given by $(p,q)\mapsto\|\mu_{p}-\mu_{q}\|_{\mathcal{H}}$ defines a metric on $\mathcal{M}^{1}(\mathcal{Y})$, where $\mu_{p}$ denotes the usual mean embedding of $p$ and $\mathcal{H}$ is the RKHS with kernel $\kappa$.

We define the relaxed projection $\Pi_{\mathrm{SC},\kappa}^{\mathcal{R}}:\mathcal{M}^{1}([0,(1-\gamma)^{-1}R_{\max}]^{d})^{\mathcal{X}}\to\prod_{x\in\mathcal{X}}\mathcal{M}^{1}(\mathcal{R}(x))=:\mathscr{S}_{\mathcal{R}}$,

Note that $\eqref{eq:catproj:sm}$ 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.

lemmasmprojlikecatproj Under the conditions of Corollary 1, the projected operator $\Pi_{\mathrm{SC},\kappa}^{\mathcal{R}}\mathcal{T}^{\pi}:\mathscr{S}_{\mathcal{R}}\to\mathscr{S}_{\mathcal{R}}$ is affine, is contractive with contraction factor $\gamma^{c/2}$, and has a unique fixed point $\eta_{\mathrm{SC},\kappa}^{\pi}$.

While we have “relaxed” the projection, the fixed point $\eta_{\mathrm{SC},\kappa}^{\pi}$ is a good approximation of $\eta^{\pi}$.

theoremsmqualityfixedpoint Under the conditions of Lemma 6, we have that

1. 1.

   $\overline{\mathrm{MMD}_{\kappa}}(\eta_{\mathrm{SC},\kappa}^{\pi},\eta^{\pi})\leq\frac{1}{1-\gamma^{c/2}}\sup_{x\in\mathcal{X}}\inf_{P\in\mathscr{P}_{\mathcal{R}(x)}}\sqrt{\mathsf{mesh}(P;\kappa)}$; and

2. 2.

   $\overline{\mathrm{MMD}_{\kappa}}(\Pi_{\mathrm{C},\kappa}^{\mathcal{R}}\eta_{\mathrm{SC},\kappa}^{\pi},\eta^{\pi})\leq(1+\frac{1}{1-\gamma^{c/2}})\sup_{x\in\mathcal{X}}\inf_{P\in\mathscr{P}_{\mathcal{R}(x)}}\sqrt{\mathsf{mesh}(P;\kappa)}$.

Notably, the second statement of Theorem 6 states that projecting $\eta_{\mathrm{SC},\kappa}^{\pi}$ back onto the space of multi-return distribution functions yields approximately the same error as $\eta_{\mathrm{C},\kappa}^{\pi}$ when $\gamma$ is near $1$.

In the remainder of the section, we assume access to a stream of MDP transitions $\left\{{T}_{t}\right\}_{t=1}^{\infty}\subset\mathcal{X}\times\mathcal{A}\times\operatorname{\mathbf{R}}^{d}\times\mathcal{X}$ consisting of elements $T_{t}=(X_{t},A_{t},R_{t},X^{\prime}_{t})$ with the following structure,

where $\mathbf{P}$ is some probability measure and $\left\{{\mathcal{F}}_{t}\right\}_{t=1}^{\infty}$ is the canonical filtration $\mathcal{F}_{t}=\sigma(\cup_{t=1}^{t}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 $\left\{{\widehat{\eta}}_{t}\right\}_{t=1}^{\infty}$ according to

for $\widehat{\eta}_{0}\in\mathscr{C}_{\mathcal{R}}$, where $\left\{{\alpha}_{t}\right\}_{t=1}^{\infty}$ is any sequence of (possibly) random step sizes adapted to the filtration $\left\{{\mathcal{F}}_{t}\right\}_{t=1}^{\infty}$. 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 $\mathscr{S}_{\mathcal{R}}$ representations.

theoremcattdconvergence For a kernel $\kappa$ induced by a semimetric $\rho$ of strong negative type, the sequence $\left\{{\widehat{\eta}}_{t}\right\}_{t=1}^{\infty}$ given by (11)-(13) converges to $\eta_{\mathrm{SC},\kappa}^{\pi}$ with probability 1.

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.

The authors would like to thank Yunhao Tang, Tyler Kastner, Arnav Jain, Yash Jhaveri, Will Dabney, David Meger, and Marc Bellemare for their helpful feedback, as well as insightful suggestions from anonymous reviewers. This work was supported by the Fonds de Recherche du Québec, the National Sciences and Engineering Research Council of Canada, and the compute resources provided by Mila (mila.quebec).