1 Introduction↩︎

Distributional reinforcement learning (DRL) models full return distributions rather than expectations, with strong empirical [1]–[4] and theoretical support [5]–[7], building on the foundational perspective of [8], [9]. In practice, DRL hinges on two coupled design choices: how to measure distributional discrepancy and the critic’s parameterization [7]. In the univariate case, several tractable combinations are available. For instance, the categorical parameterization paired with the KL divergence yields efficient updates [8], while Wasserstein-based discrepancies admit efficient estimators when combined with quantile representations [2]. This tractability is largely lost in the multivariate setting: categorical grids grow combinatorially, quantile parameterizations do not scale, and Wasserstein estimation becomes prohibitively costly, typically \(\mathcal{O}(n^3\log n)\) for general optimal transport solvers [10].

A classical approach to high-dimensional distribution comparison is slicing, which represents multivariate distributions through their one-dimensional projections and aggregates discrepancies across directions. This idea underlies the framework of Sliced Probability Divergences (SPDs) [11]–[13], where distributions are projected onto directions on the unit sphere, one-dimensional discrepancies are computed, and the results are aggregated. This projection–aggregation mechanism reduces multivariate comparison to a collection of tractable univariate problems, enabling the use of base divergences with efficient one-dimensional estimators.

Building on this slicing principle, we introduce Sliced Distributional Reinforcement Learning (SDRL), a distributional reinforcement learning framework that leverages tractable one-dimensional projections to efficiently compare multivariate return distributions. Our approach lifts base divergences with efficient one-dimensional estimators, such as Wasserstein or Cramér, to the multivariate setting.

Crucially, moving beyond scalar-valued returns introduces additional degrees of freedom, notably through more general forms of discounting. While such forms of discounting naturally arise in a variety of settings, they have so far received limited theoretical treatment in distributional reinforcement learning. Beyond uniform slicing, we consider max slicing [14] as a tool to obtain contraction guarantees in settings beyond the standard scalar-discounted case. Max slicing replaces aggregation over random directions with an optimization over the most discriminative direction. One example is the problem of Distributional Sobolev Temporal Difference, for which a concrete instantiation based on max slicing and MMD has been proposed [15]. Here we develop the general framework underlying such constructions.

Contributions. We introduce Sliced Distributional RL (SDRL), a general approach to multivariate return distributions based on sliced divergences, and establish Bellman contraction under shared scalar discounting. We further introduce a Max-Sliced variant (MSDRL) and establish contraction guarantees for matrix-discounted multivariate Bellman updates, a setting not covered by existing contraction analyses in distributional RL.

2 Background and Related Works↩︎

In the expected reinforcement learning framework, an agent interacts with an environment modeled as a Markov decision process (MDP) \((\mathcal{S}, \mathcal{A}, P, R, \{\Gamma(s,a)\})\), where rewards may be \(d\)-dimensional (\(R_t \in \mathbb{R}^d\), \(d\ge1\)). Here \(\Gamma(s,a) \in \mathbb{R}^{d\times d}\) denotes a (possibly dense) state–action dependent discount–mixing matrix. Given a policy \(\pi(a|s)\), the agent seeks to maximize the expected discounted return \[\begin{gather} Q^\pi(s,a)= \\ \mathbb{E} \!\left[ \sum_{t=0}^{\infty} \Big(\prod_{k=1}^{t} \Gamma(S_k,A_k)\Big)\, R_t \;\middle|\; S_0=s,\,A_0=a \right] \end{gather}\] with the convention \(\prod_{k=1}^{0}\Gamma(S_k,A_k) = I_d\). Classical RL methods focus on estimating \(Q^\pi(s,a)\), the expectation of the return distribution (componentwise when \(d>1\)).

The distributional perspective [8], originally developed for scalar rewards with scalar discounting, can be applied here as well: it models the full return random variable \[Z^\pi(s,a) \;=\; \sum_{t=0}^\infty \Big(\prod_{k=1}^{t} \Gamma(S_k,A_k)\Big)\, R_t ,\] whose expectation recovers \(Q^\pi(s,a) = \mathbb{E}[Z^\pi(s,a)]\) (componentwise when \(d>1\)). This viewpoint leads to the distributional Bellman operator with state–action dependent matrix discount. Let \(S' \sim P(\cdot \mid s,a)\) and \(A' \sim \pi(\cdot \mid S')\) be the next state and action; we adopt this sampling convention for all subsequent definitions of the Bellman operator. The operator is then defined as \[\label{eq:distributional95bellman95equation} (\mathcal{T}^\pi Z)(s,a) \stackrel{D}{=} R(s,a) + \Gamma(s,a)\, Z(S',A'),\tag{1}\] where \(\stackrel{D}{=}\) denotes equality in distribution.

2.0.0.1 Special cases.

1. Classical distributional RL: \(d=1\), \(\Gamma(s,a)=\gamma\in[0,1)\) [8].

2. Multivariate with shared scalar discount: \(d>1\), \(\Gamma(s,a)=\gamma I_d\) [16], [17].

3. General constant matrix: \(\Gamma(s,a)\equiv\Gamma\), where \(\Gamma\) may be possibly dense; e.g., multi-horizon return modeling with \(\Gamma=\mathrm{diag}(\gamma_1,\ldots,\gamma_d)\) that represent returns at different horizons via multiple discount factors [18].

4. State-action dependent general matrix: \(\Gamma(s,a)\) is a general (possibly dense) matrix [15].

2.1 Related works.↩︎

Wasserstein. Wasserstein metrics provide a natural divergence for distributional reinforcement learning, and in the scalar setting the distributional Bellman operator is contractive under \(\mathbf{W}_p\) [8]. However, in high dimensions Wasserstein distances suffer from dimension-dependent statistical complexity [19]. Motivated by these contraction properties, [20] reinterpret the distributional Bellman equation as an adversarial learning problem, using a learned discriminator to approximate \(\mathbf{W}_1\). In practice, such discriminators can suffer from Lipschitz violations, finite-sample bias, and optimization error, yielding objectives that may deviate substantially from true optimal-transport distances [21], [22], thus weakening contraction claims that presume exact \(\mathbf{W}_1\).

MMD. Moment matching with MMD was explored in the univariate case [23] and later extended to multivariate returns [16]. In the multivariate setting, contractivity results are available only for a narrow class of kernels [17], and identifying a kernel that is both empirically strong and contractive remains challenging [24]. Consequently, practitioners often resort to mixtures of Gaussian kernels despite their lack of contractivity guarantees in the multivariate setting.

Synthesis. Taken together, existing approaches suffer from at least one of three limitations: (1) performant methods may not be contractive, (2) theoretical guarantees do not extend to the general anisotropic discount setting we target, or (3) the estimation is too loose to support contraction claims (adversarial \(\mathbf{W}_1\)). This gap motivates our sliced approach.

3 SDRL: Distributional RL via Sliced Probability Divergences↩︎

3.1 Sliced Probability Divergences↩︎

3.1.0.1 Slicing a base divergence.

Let \(\Delta:\mathcal{P}(\mathbb{R})\times\mathcal{P}(\mathbb{R})\!\to\!\mathbb{R}_+\cup\{\infty\}\) be a divergence on one–dimensional probability laws. Fix \(p\!\ge\!1\) and let \(\sigma\) denote the uniform measure on \(\mathbb{S}^{d-1}\). For a unit direction \(\theta\!\in\!\mathbb{S}^{d-1}\), let \(P_\theta:\mathbb{R}^d\!\to\!\mathbb{R}\) denote the linear projection \(P_\theta(x)=\langle \theta,x\rangle\), and write \((P_\theta)_{\#}\mu\) for the pushforward of \(\mu\) by \(P_\theta\). Given \(\mu,\nu\in\mathcal{P}(\mathbb{R}^d)\), the associated sliced probability divergence (SPD) is \[\mathbf{S}\Delta_p^p(\mu,\nu) \;=\; \int_{\mathbb{S}^{d-1}} \Delta^p\!\bigl((P_\theta)_{\#}\mu,(P_\theta)_{\#}\nu\bigr)\,d\sigma(\theta).\] This lifts \(\Delta\) to multivariate laws by averaging the one-dimensional divergence over random projection directions [13].

3.1.0.2 Monte Carlo approximation.

In practice, this integral is estimated via Monte Carlo sampling by drawing \(L\) i.i.d.directions \(\{\theta_i\}_{i=1}^L\!\sim\!\sigma\) and computing \[\widehat{\mathbf{S}\Delta}_p^p(\mu,\nu) \;=\; \frac{1}{L}\sum_{i=1}^{L} \Delta^p\!\bigl((P_{\theta_i})_{\#}\mu,(P_{\theta_i})_{\#}\nu\bigr).\] Each projected subproblem is independent, so the \(L\) evaluations can be carried out in parallel.

3.1.0.3 Sliced Wasserstein distance.

Among sliced probability divergences, the most widely used instance is the sliced Wasserstein distance (SWD) [11], [12], where the base divergence is chosen as \(\Delta = \mathbf{W}_p\). For \(\mu,\nu\in\mathcal{P}(\mathbb{R}^d)\) and \(p\ge1\), \[\mathbf{SW}_p^p(\mu,\nu) \;=\; \int_{\mathbb{S}^{d-1}} \mathbf{W}_p^p\!\bigl((P_\theta)_{\#}\mu,(P_\theta)_{\#}\nu\bigr)\,d\sigma(\theta),\] which reduces the high-dimensional Wasserstein problem to an average of one-dimensional Wasserstein distances between the projected pushforwards \((P_\theta)_{\#}\mu\) and \((P_\theta)_{\#}\nu\). For two sets of \(n\) samples, computing each one-dimensional \(\mathbf{W}_p\) reduces to sorting the projected samples, at cost \(\mathcal{O}(n\log n)\) per direction; hence the Monte Carlo estimator with \(L\) directions costs \(\mathcal{O}(L\,n\log n)\) overall. This contrasts with computing exact OT in \(\mathbb{R}^d\) between two sets of \(n\) samples, which can scale as \(\mathcal{O}(n^3\log n)\) for standard exact solvers [25]. Further details on properties and the estimator are provided in Appendix 8.1.

3.1.0.4 Sliced Cramér distance.

In one dimension, a natural family of discrepancies is given by \(\ell_p\) distances between cumulative distribution functions: \[\ell_p^p(\mu,\nu) := \int_{\mathbb{R}}\!\big|F_\mu(u)-F_\nu(u)\big|^p\,du,\] where \(F_\mu\) and \(F_\nu\) are the univariate CDFs of \(\mu,\nu\). The special case \(p=2\) is the Cramér distance \(\mathbf{C}_2:=\ell_2^2\) [26]. The sliced Cramér distance lifts this divergence to \(\mathbb{R}^d\) via random projections: \[\mathbf{SC}_2^2(\mu,\nu) \;=\; \int_{\mathbb{S}^{d-1}} \ell_2^2\!\bigl((P_\theta)_{\#}\mu,(P_\theta)_{\#}\nu\bigr)\,d\sigma(\theta).\] This distance is also referred to as the Cramér–Wold distance and has been investigated in machine learning [27], [28]. Its estimator has the same complexity as sliced Wasserstein, since each one-dimensional Cramér evaluation costs \(\mathcal{O}(n\log n)\). The use of the Cramér distance in distributional RL has been explored in prior work [6], [29], [30]. See Appendix 8.2 for further properties and our estimator.

3.1.0.5 Sliced MMD.

Another tractable base divergence is the Maximum Mean Discrepancy (\(\mathbf{MMD}\)) [31], which has already been explored in distributional RL [24], [32], [33]. For laws \(\mu,\nu \in \mathcal{P}(\mathbb{R}^{d})\) and a kernel \(k\), the squared \(\mathbf{MMD}\) is \[\begin{align} \label{slmiafdu} \mathbf{MMD}_k^2(\mu,\nu) &= \mathbb{E}_{x,x'\sim \mu}[k(x,x')] + \mathbb{E}_{y,y'\sim \nu}[k(y,y')] \nonumber \\ &\quad - 2\,\mathbb{E}_{x\sim \mu,\,y\sim \nu}[k(x,y)]. \end{align}\tag{2}\] Lifting this discrepancy through random projections yields the sliced MMD: for \(\mu,\nu \in \mathcal{P}(\mathbb{R}^d)\), \[\label{eq:smmd95def} \begin{gather} \mathbf{SMMD}_k^2(\mu,\nu) = \\ \int_{\mathbb{S}^{d-1}} \mathbf{MMD}_k^2\!\big((P_\theta)_{\#}\mu,\,(P_\theta)_{\#}\nu\big)\,d\sigma(\theta) \end{gather}\tag{3}\] Sliced MMD was first introduced in [13]. For two sets of \(n\) samples, the base MMD estimator costs \(\mathcal{O}(n^2)\) per projection, yielding an overall sliced estimator complexity of \(\mathcal{O}(L\,n^2)\). More details on the properties of MMD and the estimator we use are provided in Appendix 8.3.

3.2 Max Sliced Probability Divergences↩︎

Max slicing was proposed in [14] to address the potential inefficiency of uniform random slicing, where many directions may be required to capture the discrepancy between two distributions. It does so by learning the most discriminative projection direction, along which the projected one-dimensional divergence is largest. It is defined as \[\label{eq:msdelta95def} \mathbf{MS}\Delta(\mu,\nu) = \sup_{\theta\in\mathbb{S}^{d-1}} \Delta\!\big((P_\theta)_{\#}\mu,\,(P_\theta)_{\#}\nu\big).\tag{4}\] This framework was originally proposed for \(\Delta=\mathbf{W}_p\), yielding the max–sliced Wasserstein distance \(\mathbf{MSW}_p\) [14]. As discussed in Section 4, we show that max slicing can provide stronger contraction guarantees than uniform slicing.

3.2.0.1 Estimation.

Since the supremum in the definition of max–sliced divergences cannot, in general, be computed exactly, it is typically approximated by iterative optimization of the projection direction on the unit sphere. At each step a gradient ascent update on the divergence is followed by renormalization onto the unit sphere, and the final direction defines the empirical estimate. The full procedure is outlined in Algorithm 7.

3.3 Problem setting and algorithmic approach↩︎

We focus on multivariate distributional policy evaluation, estimating the full return law under a fixed policy \(\pi\). We wish to model the joint vector of multivariate returns in order to capture their correlations and higher-order structure, rather than only marginal statistics. Let \(d>1\) and \(\mathcal{X}=\mathbb{R}^d\). For any policy \(\pi(\cdot\,|\,s)\), let \(\mu^\pi(s,a)\in\mathcal{P}(\mathcal{X})\) denote the law of the multivariate return \(Z^\pi(s,a)\). The distributional Bellman operator \(\mathcal{T}^{\pi}\) relates return laws across state–action pairs via \[\label{eq:bellman-mu-operator} \begin{align} (\mathcal{T}^{\pi}\mu)(s,a) := \int_{\mathcal{S}}\!\int_{\mathcal{A}}\!\int_{\mathcal{X}} (f_{\Gamma(s,a),r})_{\#}\mu(s',a') \\[3pt] \quad \times R(dr\,|\,s,a)\; \pi(da'\,|\,s')\; P(ds'\,|\,s,a). \end{align}\tag{5}\] where \(f_{\Gamma(s,a),r}(z)=r+\Gamma(s,a)\,z\) for \(z\in\mathbb{R}^d\) and \(\Gamma(s,a)\in\mathbb{R}^{d\times d}\) is a (possibly dense) discount–mixing matrix. The target in policy evaluation is the fixed point \(\mu^\pi\) of \(\mathcal{T}^{\pi}\), so \(\mathcal{T}^{\pi}\mu^\pi=\mu^\pi\).

Algorithmic approach Following the particle-based critic paradigm [16], [23], we approximate \(\mu^\pi(s,a)\) by finite sets of particles and optimize them directly. Given a transition \((s,a,r,s')\) and next action \(a'\!\sim\!\pi(\cdot|s')\), we produce a set of \(N\) predicted particles \(\{z_i = Z_\phi(s,a,i)\}_{i=1}^N\) and a set of \(N\) target particles \(\{\hat{z}_j = r + \Gamma(s,a)\,Z_\phi(s',a',j)\}_{j=1}^N\). Their discrepancy is measured by a sliced probability divergence with base \(\Delta\), using either \(L\) random projections or a single optimized direction (max–sliced), and minimizing it w.r.t.\(\phi\) yields a distributional TD update toward the matrix-discounted target (Algorithm 1).

4 Theoretical results↩︎

In this section, we provide theoretical foundations for multivariate distributional RL with sliced divergences, based on a supremum divergence and sufficient conditions for contraction of the distributional Bellman operator.

4.0.0.1 Standing notation.

Throughout this section, returns are \(\mathbb{R}^d\)–valued and we write \(\eta(s,a)\in\mathcal{P}(\mathbb{R}^d)\) for the return distribution at \((s,a)\). We write \(\Delta\) for a base divergence on \(\mathcal{P}(\mathbb{R})\), with lifted divergences \(\mathbf{S}\Delta_p\) and \(\mathbf{MS}\Delta\), and use \(\overline{\mathcal{D}}\) for the supremum lift over \((s,a)\).

We focus on the following questions:

1. Metric property: When do \(\overline{\mathbf{S}\Delta_p}\) and \(\overline{\mathbf{MS}\Delta}\) induce metrics on \(\mathcal{P}(\mathbb{R}^d)^{\mathcal{S}\times\mathcal{A}}\)?

2. Contraction property: Under what conditions on \(\Delta\) and on the discount structure \(\Gamma\) does \(\mathcal{T}^\pi\) contract in \(\overline{\mathbf{S}\Delta_p}\) or \(\overline{\mathbf{MS}\Delta}\)?

3. Sample complexity: How does the estimation error of the sliced and max–sliced divergences scale with the number of samples, and do they avoid the curse of dimensionality?

4. Stochastic training suitability: With the widely used bootstrap update that instantiates the Bellman target from a single sampled successor \((s',a')\), when do \(\mathbf{S}\Delta_p\) and \(\mathbf{MS}\Delta\) remain compatible with this training setup?

4.1 Metric property↩︎

It is known that uniform slicing preserves the metric property of a base divergence [13]. Similarly, [14] established that max–sliced Wasserstein is a metric; more generally, max–slicing preserves metricity whenever the base divergence is a metric on \(\mathcal{P}(\mathbb{R})\). Finally, supremum lifting over \((s,a)\) preserves metricity. These results are summarized in Theorem 1, with full proofs provided in Appendix 9.

4.2 Contraction property↩︎

Let \(\mathcal{D}\) be any divergence on \(\mathcal{P}(\mathbb{R}^d)\). We say that \(\mathcal{T}\) is a \(\kappa\)–contraction in \(\overline{\mathcal{D}}\) if there exists \(\kappa\in[0,1)\) such that, for all \(\eta_1,\eta_2:\mathcal{S}\times\mathcal{A}\to\mathcal{P}(\mathbb{R}^d)\), \[\label{eq:contraction95def} \overline{\mathcal{D}}\big(\mathcal{T}\eta_1,\,\mathcal{T}\eta_2\big)\;\le\;\kappa\,\overline{\mathcal{D}}(\eta_1,\eta_2).\tag{6}\]

Univariate contraction We recall sufficient conditions under which the univariate distributional Bellman operator \(\mathcal{T}^\pi\) is a \(c(\gamma)\)–contraction with respect to \(\overline{\mathcal{D}}\) (Proposition 1). This slightly generalizes a result from [34]. The proof is in Appendix 10.1.

Uniform slicing We now turn to the main multivariate contraction result. We start from the canonical setting of vector-valued returns in \(\mathbb{R}^d\) with \(d>1\), where the Bellman update uses the shared scalar discount from Section 2. This is the standard multivariate formulation considered in [16], [20]. Given \(X' \sim \eta(S',A')\), the corresponding distributional Bellman update is \[\label{eq:bellman-shared-gamma} (\mathcal{T}^\pi \eta)(s,a) = \mathrm{Law}\!\bigl(R(s,a) + \gamma I_d\, X'\bigr).\tag{7}\] The key observation is that the univariate sufficient conditions (T), (S), (M\(_p\))from Proposition 1 lift directly to the sliced objective, so 7 contracts in the sup–sliced divergence with the same factor \(c(\gamma)\) as in the univariate case. This is stated in Theorem 2; the full proof is in Appendix 10.3.

Max–slicing We now move beyond shared scalar discounting and consider anisotropic multivariate Bellman updates of the form \[\label{eq:bellman-general-gamma} (\mathcal{T}^\pi \eta)(s,a) =\mathrm{Law}\!\big(R(s,a)+\Gamma(s,a)\,X'\big),\tag{8}\] Our goal is the norm–only contraction form (A)for 8 , for some nondecreasing function \(c\): \[\begin{gather} \mathrm{\boldsymbol{\textcolor{nice_gold}{(A)}}}\qquad \overline{\mathcal{D}}\big(\mathcal{T}^\pi\eta_1,\mathcal{T}^\pi\eta_2\big) \;\le\; c(\bar L)\,\overline{\mathcal{D}}(\eta_1,\eta_2), \\ \bar L := \sup_{(s,a)\in\mathcal{S}\times\mathcal{A}}\;\|\Gamma(s,a)\|_{\mathrm{op}}. \end{gather}\] This deliberately discards any dependence on the geometry of \(\Gamma\) beyond its operator norm. Our motivation for doing so is twofold. First, this is the exact type of criterion obtained when using Wasserstein as divergence [15]. Second, it yields a simple norm–based criterion that is already useful in practice [15]).

We may wonder whether (A)is already known to hold for multivariate distributional RL under divergences that both admit tractable estimators and come with multivariate contraction results. To the best of our knowledge, there are currently two such options: MMD with contraction–admitting kernels [23] and uniform sliced divergences. The next proposition shows that, for both, (A)can fail in general; full constructions are given in Appendix 10.4.

We now turn to a sufficient condition under which (A)does hold. Theorem 3 shows that max–slicing yields (A)for 8 when the base divergence satisfies the univariate sufficient conditions of Proposition 1. The proof is given in Appendix 10.5.

4.3 Sample complexity↩︎

We summarize sample–complexity bounds for uniform slicing and max–slicing; proofs are deferred to Appendix 12.

Uniform slicing. As shown by [13] and restated in Proposition 3, uniform slicing inherits the one–dimensional rate of the base divergence, without additional dependence on the ambient dimension.

Max–slicing. For a class of base divergences, Proposition 4 gives a finite–sample bound for \(\mathbf{MS}\Delta\) under a bounded–support assumption, which is natural in RL when returns are assumed bounded. Depending on the base divergence \(\Delta\), the resulting rate scales like \((d\log n/n)^{\alpha/2}\) and therefore avoids the curse of dimensionality (polynomial dependence on \(d\)).

4.4 Suitability for stochastic training.↩︎

A general property that often matters for stochastic optimization is whether sample-based gradients match the gradient of the population objective. Let \(\mu\in\mathcal{P}(\mathbb{R}^d)\) denote a data law and let \(\nu_\phi\in\mathcal{P}(\mathbb{R}^d)\) be a parametric model. Writing \(\widehat\mu_m\) for the empirical measure of \(m\) i.i.d.samples from \(\mu\), we say that a divergence \(\mathcal{D}\) satisfies the unbiased sample gradient property if \[\label{eq:Ucond95main} \mathbb{E}_{X_{1:m}\sim \mu}\!\left[\nabla_\phi\,\mathcal{D}\!\left(\widehat{\mu}_m,\,\nu_\phi\right)\right] \;=\; \nabla_\phi\,\mathcal{D}\!\left(\mu,\,\nu_\phi\right). \qquad \mathrm{\boldsymbol{\textcolor{nice_gold}{(U)}}}\tag{9}\] When (U)fails, stochastic gradients can be biased and, in some cases, the minimizers of the expected sample objective and the population objective may differ [26].

Why this matters in distributional RL. In distributional RL, the Bellman target is the mixture law induced by transition and policy randomness, but a standard TD update typically instantiates that target from a single sampled successor \((S',A')\). As a result, training implicitly optimizes an expected sample loss rather than the population loss against the true mixture target. Wasserstein provides a canonical illustration where (U)fails (Proposition 5 in [8]), which has motivated alternatives such as quantile-regression based objectives [1], [35]. In contrast, the Cramér distance and \(\mathrm{MMD}^2\) are well-known examples that do satisfy (U) [26], [36]. This motivates asking, for each divergence used in our framework, whether it remains compatible with one-sample Bellman bootstrapping in the sense of (U).

Uniform slicing. As shown in Proposition 5, uniform slicing preserves (U)for the powered sliced objective; see Appendix 11 for the proof.

Max-slicing. Unlike uniform slicing, max-slicing introduces an additional maximization step over directions, and this changes the behavior of sample-based objectives: the direction selected by the max can depend on the same sample used to evaluate the loss. As a result, the unbiased sample gradient property (U)can fail for max-sliced objectives, even when it holds direction-wise. This means that max-slicing is generally not compatible with the usual one-sample Bellman bootstrapping setup used in distributional RL. See Appendix 11.

4.5 Instantiations↩︎

Our results are agnostic to the choice of base divergence \(\Delta\) on \(\mathcal{P}(\mathbb{R})\), provided it is a metric and satisfies (T), (S), and (M\(_p\)): the corresponding sliced and max–sliced objectives inherit metricity and yield Bellman contraction via Theorems 2 and 3. In this paper we instantiate \(\Delta\) with \(\mathbf{W}_p\), the Cramér distance \(\mathbf{C}_2\), and \(\mathbf{MMD}_k\). For \(\mathbf{MMD}\), the conclusions extend to any kernel \(k\) for which \(\mathbf{MMD}_k\) is a metric and admits a scale bound of the form (S); we use the multiquadric kernel [24], defined as \(k_h(x,y) = -\sqrt{1 + h^2 \|x-y\|^2}\), as a representative example, and use this kernel for all experiments reported in the subsequent section. A complete summary of instantiations (contraction factors and sample–complexity rates) is given in Appendix 13, Table 1.

5 Experiments↩︎

5.0.0.1 Chain environment.

We validate sliced divergences for multivariate policy evaluation on a tabular chain MDP adapted from [7], [23]. The state space is a length-\(K\) chain with stochastic transitions (Fig. 2). Rewards are positional and defined as \(R_{t+1}=e_{S_{t+1}}\in\mathbb{R}^K\), the one-hot embedding of the entered next state. Appendix 16 provides full details.

We compare two bootstrapping regimes: (i) standard one-sample distributional TD, which instantiates the Bellman target from a single successor, and (ii) a near-exact variant that explicitly constructs the transition-mixture target using the known dynamics (Appendix 16.4). This controlled comparison isolates the effect of the unbiased sample-gradient property (U): in Fig. 3, near-exact TD is largely insensitive to gradient bias, whereas under one-sample TD the objectives that violate (U)degrade substantially (notably the Wasserstein-based and max-sliced variants). Performance is measured by the empirical \(\widehat{\mathbf{W}}_2\) between Monte Carlo returns and the learned critic distribution.

5.0.0.2 Pixel-based environments.

We next evaluate the same objectives on pixel observations. All pixel-based experiments use PQN [37] as a common backbone. Its fully vectorized implementation and simplified training pipeline (notably, no target network) make it substantially faster, which lets us compare many objectives and variants under a fixed compute budget.

We consider two complementary benchmarks. Maze follows the maze environments of [16], which provide pixel observations and vector-valued rewards, and is used for policy evaluation. We assess distributional accuracy via the empirical \(\widehat{\mathbf{W}}_2\) between Monte Carlo return samples and the return distribution predicted by the critic. We include max-sliced variants for completeness: unlike the chain environment, exact TD is not available here, making this a realistic setting where the gap between objectives satisfying (U)and those that do not is practically meaningful. Figure 5 shows the environment and a representative example where the critic closely matches the Monte Carlo return distribution when trained with sliced Cramér. Figure 4 summarizes \(\widehat{\mathbf{W}}_2\) across objectives and again highlights that (U)is essential for accurate distributional learning. In particular, sliced MMD and sliced Cramér perform on par with MMD.

Atari follows [16] and considers a subset of six Atari games from the Arcade Learning Environment [38] whose primitive rewards admit a fixed decomposition into multi-dimensional components. We evaluate control performance and report aggregated normalized scores across games and variants in Figure 6. Sliced Cramér performs strongly in this benchmark.

While objectives that satisfy (U)yield accurate distributional policy evaluation from pixels, this does not necessarily translate into the strongest control performance. In particular, sliced Wasserstein–\(2\) performs remarkably well despite exhibiting biased gradients and poor distributional matching in earlier experiments. This suggests that control performance may depend primarily on accurate estimation of the return expectation even when the full return distribution is not well captured. This is consistent with the observation in Section 17.2.2, where objectives violating (U)can still accurately predict the return mean despite collapsing the full distribution. It also aligns with prior work noting that policy improvement depends on the accuracy of the expected return [7].

The full training procedure and experimental setup for the pixel-based experiments are detailed in Appendix 15. Further ablations and diagnostics are reported in Section 17, including computational cost, qualitative return distribution visualizations, and an Atari sensitivity study on the number of projection directions (Section 17.4).

6 Discussion and open question↩︎

For anisotropic updates of the form \(R+\Gamma Z(S',A')\), Wasserstein is attractive because its contraction criterion depends only on \(\|\Gamma\|_{\mathrm{op}}\). However, in our setting it faces three difficulties: (1) high computational cost in large \(d\), (2) poor statistical efficiency, and (3) incompatibility with the usual one-sample Bellman bootstrap because (U)fails. Max-slicing is designed to mitigate (1)–(2) by reducing comparisons to one-dimensional projections, but Proposition 6 shows that it can still fail (3) because of the max-selection step. Taken together, our results leave open the question: does there exist a divergence that is provably contractive under general anisotropic discounts in a norm-only sense and that also satisfies (U), while still addressing (1)–(2) through tractable estimators and dimension-friendly statistical rates?

7 Conclusion↩︎

This work studies how to extend one-dimensional distributional objectives to multivariate return distributions in a way that remains compatible with temporal-difference learning. Our main practical takeaway is that sliced Cramér provides a strong default for multivariate distributional learning. It is computationally efficient, behaves well under one-sample bootstrapping, and comes with contraction guarantees in the standard isotropic setting.

On the methodological side, we propose a principled framework to lift divergences defined on one-dimensional samples to multivariate returns through slicing. This framework identifies conditions on the one-dimensional base divergence that are sufficient to obtain contraction of the corresponding multivariate distributional Bellman operator under anisotropic updates. We further analyze max slicing as an alternative to uniform slicing, and show how it can recover contraction under the same norm-controlled Bellman updates, while also requiring care due to its sample-dependent direction selection.

In particular, our results provide tools to nuance what can go wrong in practice when optimizing sliced objectives with bootstrapped targets. We show that an objective can remain contractive at the population level while still yielding biased sample gradients in the traditional one-sample TD setting, and we characterize when this bias arises or disappears.

Empirically, our controlled policy-evaluation benchmarks confirm that satisfying the unbiased sample-gradient requirement is essential for accurate distribution matching from bootstrapped targets in the traditional one-sample regime, whereas the near-exact mixture updates largely remove these effects.

Finally, we are far from having explored the full potential of slicing, which has seen many improvements over the years [23], [39]–[41]. However, it is worth keeping in mind that improvements over uniform slicing can quickly reintroduce gradient bias, potentially making an objective unsuitable for one-sample TD and thus for traditional distributional RL.

Acknowledgements↩︎

This project has received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement n° 101021347).

Impact Statement↩︎

This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

Appendix overview↩︎

This appendix collects the technical background, proofs, and experimental details supporting the main paper. Base probability divergences (Section 8) recalls the definitions, estimators, and key structural properties of the standard discrepancies used as building blocks, namely the Wasserstein distance, the class of cumulative distribution function distances, and the Maximum Mean Discrepancy. Metric property (Section 9) proves that lifting a base metric through uniform slicing or maximization over projections preserves the metric axioms on the space of multivariate probability measures. Contraction results (Section 10) establishes sufficient conditions for the distributional Bellman operator to be a contraction, covering univariate and multivariate settings, uniform slicing, explicit counterexamples where norm control fails, and guarantees for max-sliced objectives. Mixture bias (Section 11) discusses the suitability of different divergences for stochastic gradient-based training, focusing on the unbiased sample gradient property and its interaction with slicing operations. Sample complexity (Section 12) derives statistical estimation rates, highlighting the dimension-free nature of uniform slicing and providing specific bounds for max-sliced Wasserstein and Cramer distances. Instantiations (Section 13) summarizes the specific contraction factors, sample complexity bounds, and computational costs for the considered divergence families. Pseudo-codes (Section 14) provides the algorithmic details for estimating the max-sliced divergence. Pixel-based control: maze and Atari (Section 15) describes the experimental setup for the control tasks, including environment specifications and neural network architectures. Chain environment (Section 16) details the tabular environment used for policy evaluation and the specific temporal difference algorithms employed. Finally, More results and ablations (Section 17) presents additional experimental data and ablation studies.

8 Base probability divergences↩︎

We briefly recall the base discrepancies used in our instantiations, along with the estimators and key structural properties that will be invoked later to obtain contraction factors and sample-complexity bounds. Section 8.1 reviews Wasserstein and the contraction properties we use. Section 8.2 covers the \(\ell_p\) CDF distances on \(\mathbb{R}\), including Cramér–2. Section 8.3 reviews MMD and its kernel formulations.

8.1 Wasserstein distance↩︎

The Wasserstein distance, arising from optimal transport theory [42], provides a principled way of comparing probability measures by quantifying the minimal cost of transporting mass from one distribution to another. Let \((\mathbb{R}^d, d)\) be a metric space and denote by \(\mathcal{P}_p(\mathbb{R}^d)\) the set of Borel probability measures with finite \(p\)-th moment. For \(\mu, \nu \in \mathcal{P}_p(\mathbb{R}^d)\), the \(p\)-Wasserstein distance is defined as \[\mathbf{W}_p(\mu, \nu) = \left( \inf_{\pi \in \Pi(\mu,\nu)} \int_{\mathbb{R}^d\times \mathbb{R}^d} d(x, y)^p \, d\pi(x, y) \right)^{1/p},\] where \(\Pi(\mu, \nu)\) denotes the set of couplings (or transport plans) \(\pi\) whose marginals are \(\mu\) and \(\nu\). When the underlying measures admit densities \(I_\mu\) and \(I_\nu\), we may write \(\mathbf{W}_p(I_\mu, I_\nu)\) without ambiguity.

Computing \(\mathbf{W}_p\) directly is challenging in high dimensions, but there are settings where closed-form expressions exist. In the special case where \(\mu\) and \(\nu\) are one-dimensional distributions on a normed linear space, the Wasserstein distance simplifies to \[\mathbf{W}_p(\mu, \nu) = \left( \int_0^1 \big|F_\mu^{-1}(z) - F_\nu^{-1}(z)\big|^p \, dz \right)^{1/p},\] where \(F_\mu^{-1}\) and \(F_\nu^{-1}\) are the quantile functions (inverse CDFs) of \(\mu\) and \(\nu\), respectively.

8.1.1 Estimator↩︎

For empirical measures \(\tilde{\mu} = \frac{1}{N} \sum_{n=1}^N \delta_{x_n}\) and \(\tilde{\nu} = \frac{1}{N} \sum_{n=1}^N \delta_{y_n}\) in one dimension, \(\mathbf{W}_p\) can be computed by sorting the samples and comparing corresponding order statistics [42]: \[\mathbf{W}_p(\tilde{\mu}, \tilde{\nu}) = \left( \frac{1}{N} \sum_{n=1}^N \big|x_{I_x[n]} - y_{I_y[n]}\big|^p \right)^{1/p},\] where \(I_x[n]\) and \(I_y[n]\) are the indices that sort \(\{x_n\}\) and \(\{y_n\}\) in ascending order.

8.1.2 Properties↩︎

8.1.2.1 Metric.

It is a classical result that the Wasserstein distances are genuine metrics. In particular, Proposition 2 in [43] establishes that \[\mathbf{W}_p \quad\text{is a metric on }\;\mathcal{P}_p(\mathbb{R}) \qquad\text{for every }p\in[1,\infty],\] where \(\mathcal{P}_p(\mathbb{R})=\{\mu\in\mathcal{P}(\mathbb{R}):\int |x|^p\,d\mu(x)<\infty\}\) for \(p<\infty\), and \(\mathcal{P}_\infty(\mathbb{R})=\mathcal{P}(\mathbb{R})\).

8.1.2.2 Translation invariant.

By definition

8.1.2.3 Scale contraction.

8.1.2.4 p-convexity.

8.2 The \(\ell_p\)–family of CDF distances on \(\mathbb{R}\)↩︎

Let \(\mu,\nu \in \mathcal{P}(\mathbb{R})\) be probability measures with cumulative distribution functions (CDFs) \(F_\mu,F_\nu\). For \(p \in [1,\infty)\), the \(\ell_p\) distance between \(\mu\) and \(\nu\) is defined as \[\ell_p(\mu,\nu) := \left( \int_{-\infty}^{\infty} \big| F_\mu(t) - F_\nu(t) \big|^p \, dt \right)^{1/p} \;=\; \|F_\mu - F_\nu\|_{L^p(\mathbb{R})} ,\] that is, the \(\ell_p\)–family can be seen as the \(L^p\) norm between the two CDFs [26].

8.2.0.1 Connections to other distances.

• Wasserstein distance. For \(p=1\), one recovers the \(1\)–Wasserstein distance [26]: \[\ell_1(\mu,\nu) = \mathbf{W}_1(\mu,\nu) = \int_{0}^1 \big| F_\mu^{-1}(u) - F_\nu^{-1}(u) \big| \, du .\] Thus, \(\ell_1\) coincides with the classical earth mover’s distance on \(\mathbb{R}\).

• Cramér distance. For \(p=2\), the squared \(\ell_2\) distance coincides with the Cramér distance (also denoted \(\mathbf{C}_2\)) [26]: \[\ell_2^2(\mu,\nu) = \mathbf{C}_2(\mu,\nu)= \int_{-\infty}^{\infty} \big(F_\mu(t) - F_\nu(t)\big)^2 \, dt,\] which also admits the energy distance form \[\ell_2^2(\mu,\nu) = \mathbb{E}|X-Y| - \tfrac{1}{2}\mathbb{E}|X-X'| - \tfrac{1}{2}\mathbb{E}|Y-Y'|,\] for \(X,X' \!\sim\! \mu\) and \(Y,Y'\!\sim\!\nu\) i.i.d.

• Cramér distance and the Continuous Ranked Probability Score (CRPS). The squared \(\ell_2\) distance admits a close connection to the Continuous Ranked Probability Score (CRPS) [44]. For a predictive distribution \(\mu\) with CDF \(F_\mu\) and an observation \(y\in\mathbb{R}\), the CRPS is defined as \[\mathrm{CRPS}(\mu,y) := \int_{-\infty}^{\infty} \big(F_\mu(t) - \mathbf{1}\{y \le t\}\big)^2\,dt .\] Taking expectation with respect to \(Y\sim\nu\) and exchanging expectation and integration yields \[\mathbb{E}_{Y\sim\nu}\,\mathrm{CRPS}(\mu,Y) = \int_{-\infty}^{\infty} \Big(F_\mu(t)^2 - 2F_\mu(t)F_\nu(t) + F_\nu(t)\Big)\,dt .\] Adding and subtracting \(F_\nu(t)^2\) inside the integral gives the decomposition \[\mathbb{E}_{Y\sim\nu}\,\mathrm{CRPS}(\mu,Y) = \underbrace{\int_{-\infty}^{\infty}\big(F_\mu(t)-F_\nu(t)\big)^2\,dt}_{\ell_2^2(\mu,\nu)} \;+\; \underbrace{\int_{-\infty}^{\infty}\big(F_\nu(t)-F_\nu(t)^2\big)\,dt}_{\text{constant in }\mu}.\] Consequently, minimizing the expected CRPS with respect to \(\mu\) is equivalent to minimizing the squared Cramér distance \(\ell_2^2(\mu,\nu)\), and both objectives induce identical gradients with respect to the predicted distribution, since the second term is independent of \(\mu\).

8.2.1 Estimator↩︎

8.2.1.1 Empirical CDF estimator.

For empirical measures \(\tilde{\mu} = \tfrac{1}{n} \sum_{i=1}^n \delta_{u_i}\) and \(\tilde{\nu} = \tfrac{1}{m} \sum_{j=1}^m \delta_{v_j}\) in one dimension, the \(\ell_p\) CDF distance (Cramér when \(p=2\)) admits a closed form: after merging and sorting all samples, one tracks the cumulative difference of the two empirical CDFs, which is piecewise constant between successive breakpoints. The distance then reduces to a weighted sum of gap lengths multiplied by the corresponding powers of this difference. \[\ell_p^p(\tilde{\mu},\tilde{\nu}) =\sum_{k=1}^{K-1} (t_{k+1}-t_k)\,|\Delta_k|^p.\] Algorithmically, the estimator amounts to sorting the combined samples once, tracking the cumulative difference of the two empirical CDFs, and summing the piecewise contributions. This requires \(\mathcal{O}((n{+}m)\log(n{+}m))\) time for sorting and linear time for the scan.

8.2.1.2 CRPS–based estimator for the \(\ell_2\) (Cramér–2) distance.

Building on the connection established in Section 8.2 between the squared \(\ell_2\) CDF distance (Cramér–\(2\)) and the expected Continuous Ranked Probability Score (CRPS), we construct an empirical estimator using the integral representation of the CRPS. Let \(F_\mu\) denote the CDF of \(\mu\) and \(Q_\mu(\tau)=F_\mu^{-1}(\tau)\) its quantile function. The CRPS satisfies \[\mathrm{CRPS}(\mu,y) = 2\int_0^1 \rho_\tau\!\big(y - Q_\mu(\tau)\big)\,d\tau, \qquad \rho_\tau(r)=r\big(\tau-\mathbf{1}\{r<0\}\big).\] In practice, we approximate the integral over \(\tau\) by a finite sum over uniformly spaced quantile levels \(\tau_i=(2i-1)/(2n)\) and replace \(Q_\mu(\tau_i)\) by the order statistics \(u_{(i)}\), yielding the empirical objective \[\widehat{\mathcal{L}}_{\mathrm{CRPS}} = \frac{2}{nm}\sum_{i=1}^n\sum_{j=1}^m \rho_{\tau_i}\!\big(v_j-u_{(i)}\big).\] As shown in Section 8.2, the expected CRPS differs from the squared Cramér–\(2\) distance \(\ell_2^2(\mu,\nu)\) by an additive constant independent of \(\mu\). Accordingly, we use the gradient of the estimator above as a surrogate for the gradient of the true Cramér–\(2\) objective, namely \[\nabla_\theta \widehat{\mathcal{L}}_{\mathrm{CRPS}} \;\approx\; \nabla_\theta \,\ell_2^2(\mu_\theta,\nu),\] where the approximation reflects the finite-sample and discretization effects of the estimator. Computationally, this estimator requires sorting the \(n\) samples from \(\mu\) once (\(\mathcal{O}(n\log n)\)) and evaluating \(nm\) terms, resulting in total cost \(\mathcal{O}(n\log n + nm)\).

8.2.1.3 Energy distance estimator.

A second estimator for the squared \(\ell_2\) CDF distance (Cramér–\(2\)) follows from its energy distance representation. For \(X,X'\sim\mu\) and \(Y,Y'\sim\nu\) i.i.d., one has \[\ell_2^2(\mu,\nu) = \mathbb{E}|X-Y| -\tfrac12 \mathbb{E}|X-X'| -\tfrac12 \mathbb{E}|Y-Y'|.\] Given samples \(u_1,\dots,u_n\sim\mu\) and \(v_1,\dots,v_m\sim\nu\), each expectation is estimated by its corresponding U–statistic, yielding \[\widehat{\ell}_2^2(\mu,\nu) = \frac{1}{nm}\sum_{i,j}|u_i-v_j| -\frac{1}{2n(n-1)}\sum_{i\neq i'}|u_i-u_{i'}| -\frac{1}{2m(m-1)}\sum_{j\neq j'}|v_j-v_{j'}|.\] This estimator is unbiased for \(\ell_2^2(\mu,\nu)\) and does not rely on sorting or discretization.

Equivalently, the squared Cramér–\(2\) distance coincides with the squared Maximum Mean Discrepancy \(\mathbf{MMD}_k^2(\mu,\nu)\) associated with the kernel \[k(x,y) = -|x-y|.\] Under this identification, the estimator above coincides (up to a constant factor) with the unbiased MMD U–statistic associated with the kernel \(k\), as discussed in Section 8.3. Computationally, evaluating the estimator involves \(\mathcal{O}(nm)\) cross–sample pairwise terms together with \(\mathcal{O}(n^2+m^2)\) within–sample terms, resulting in an overall cost of \(\mathcal{O}((n+m)^2)\).

8.2.2 Properties↩︎

8.2.2.1 Metric.

8.2.2.2 Translation invariant.

By non-trivial arguments (see Theorem 2 in [26] and Proposition 3.2 in [45]), the \(\ell_p\) distance is invariant under translations: for all \(\mu,\nu\in\mathcal{P}_1(\mathbb{R})\) and every \(t\in\mathbb{R}\), \[\ell_p\bigl((T_t)_{\#}\mu,(T_t)_{\#}\nu\bigr)=\ell_p(\mu,\nu), \qquad T_t(x)=x+t.\]

8.2.2.3 Scaling.

8.2.2.4 p-convexity.

8.3 MMD↩︎

The Maximum Mean Discrepancy (MMD) is a kernel-based discrepancy that measures how far apart two probability laws are in a reproducing kernel Hilbert space (RKHS) \(\mathcal{H}\). Given a positive definite kernel \(k\colon X\times X\to\mathbb{R}\) with feature map \(\phi(x)=k(x,\cdot)\), each distribution admits a mean embedding in \(\mathcal{H}\): \[\label{eq:mmd95embed} \mu_P=\int_X \phi(x)\,dP(x), \qquad \mu_Q=\int_X \phi(x)\,dQ(x).\tag{11}\] The distance between these embeddings defines \[\label{eq:mmd95def} \mathbf{MMD}_k(P,Q) \;=\; \|\mu_P-\mu_Q\|_{\mathcal{H}}.\tag{12}\] Its square admits the familiar expansion in terms of kernel expectations: \[\begin{align} \label{eq:mmd295expand} \mathbf{MMD}^2_k(P,Q) &= \|\mu_P-\mu_Q\|_{\mathcal{H}}^2 \nonumber\\ &= \iint k(x,x')\,dP(x)\,dP(x') \;+\;\iint k(y,y')\,dQ(y)\,dQ(y') \nonumber\\ &\quad-\;2\iint k(x,y)\,dP(x)\,dQ(y). \end{align}\tag{13}\]

In this paper we also use a closely related formulation that does not require explicitly starting from an RKHS. Let \(X\) be a measurable space and let \(k:X\times X\to\mathbb{R}\) be symmetric. We say that \(k\) is conditionally positive definite (CPD) if \[\label{eq:cpd95def} \iint_{X\times X} k(x,x')\,d\mu(x)\,d\mu(x') \;\geq\;0 \qquad\text{for all finite signed measures }\mu\text{ on }X\text{ with }\mu(X)=0,\tag{14}\] and conditionally strictly positive definite (CSPD) if the inequality is strict for every nonzero such \(\mu\). Under this condition, the quadratic form \[\label{eq:gamma95k95def} \gamma_k(P,Q)^2 :=\iint k(x,y)\,d(P-Q)(x)\,d(P-Q)(y)\tag{15}\] is nonnegative and behaves like an MMD. Concretely, define \[\label{eq:rho95k95def} \rho_k(x,y)\;:=\;k(x,x)+k(y,y)-2k(x,y),\tag{16}\] fix any \(z_0\in X\), and introduce the one-point centered (distance-induced) kernel \[\begin{align} \label{eq:k95circ95def} k^\circ(x,y) &:=\tfrac12\bigl[\rho_k(x,z_0)+\rho_k(y,z_0)-\rho_k(x,y)\bigr] \nonumber\\ &=\;k(x,y)-k(x,z_0)-k(z_0,y)+k(z_0,z_0). \end{align}\tag{17}\] Then \(k^\circ\) is positive definite and admits an RKHS \(\mathcal{H}_{k^\circ}\), and for any \(P,Q\) with finite integrals, \[\begin{align} \label{eq:gamma95equals95mmd} \gamma_k(P,Q)^2 &=\iint k^\circ(x,y)\,d(P-Q)(x)\,d(P-Q)(y) \nonumber\\ &=\bigl\|\mu_{k^\circ}(P)-\mu_{k^\circ}(Q)\bigr\|_{\mathcal{H}_{k^\circ}}^{2} =\mathbf{MMD}_{k^\circ}(P,Q)^2 . \end{align}\tag{18}\] This is the standard distance-induced kernel construction and equivalence result of [46].

8.3.0.1 Convention.

To keep notation light, we write \(\mathbf{MMD}_k\) to denote either the usual RKHS MMD when \(k\) is positive definite, or more generally the square root of the quadratic form \(\gamma_k(\mu,\nu)^2=\iint k\,d(\mu-\nu)\,d(\mu-\nu)\) when \(k\) is only conditionally positive definite. In the latter case, one may always replace \(k\) by its one-point centered kernel \(k^\circ\) in 17 , which is positive definite, and then \(\gamma_k(\mu,\nu)=\mathbf{MMD}_{k^\circ}(\mu,\nu)\).

8.3.1 Estimator↩︎

The \(\mathbf{MMD}\) can be approximated from samples in two standard ways, both originating from [31]. Given two sets of \(m\) samples \(\{x_i\}_{i=1}^m \sim P\) and \(\{y_i\}_{i=1}^m \sim Q\), the biased estimator is \[\label{eq:biased95estimator95main95long} \widehat{\mathbf{MMD}}_b^2 = \frac{1}{m^{2}} \sum_{i,j=1}^{m} k(x_i,x_j) + \frac{1}{m^{2}} \sum_{i,j=1}^{m} k(y_i,y_j) - \frac{2}{m^{2}} \sum_{i,j=1}^{m} k(x_i,y_j).\tag{19}\] while the unbiased estimator excludes diagonal terms: \[\label{eq:mmd95unbiased95estimator} \widehat{\mathbf{MMD}}_u^2 = \frac{1}{m(m-1)} \sum_{\substack{i,j=1 \\ i \neq j}}^m k(x_i, x_j) + \frac{1}{m(m-1)} \sum_{\substack{i,j=1 \\ i \neq j}}^m k(y_i, y_j) - \frac{2}{m^2} \sum_{i,j=1}^m k(x_i, y_j).\tag{20}\] Although the unbiased form eliminates a small finite-sample bias, the biased estimator is often preferred in practice. In particular, applications of \(\mathbf{MMD}\) to distributional RL [24], [32] consistently rely on the biased version due to its lower variance and greater numerical stability during training.

8.3.2 Properties↩︎

8.3.2.1 Metric.

8.3.2.2 Scale contraction.

Whether \(\mathbf{MMD}_k\) contracts under scaling depends on the choice of kernel. In the contraction-based distributional RL literature, only two kernel families have been instantiated to obtain a Bellman contraction: the negative-type power family introduced by [23] and the multiquadric kernel studied by [24]. We record the corresponding scale-contraction bounds in Proposition 14 (negative-type power kernels) and Proposition 13 (multiquadric kernel). The multiquadric bound is restated from [15]; the proof is given there.

8.3.2.3 p-convexity.

9 Metric property↩︎

10 Contraction property↩︎

In Section 10.1, we recall sufficient conditions for the univariate distributional Bellman operator to be a contraction. In Section 10.2, we generalize these sufficient conditions to anisotropic updates, allowing random linear maps controlled by an operator-norm bound. In Section 10.3, we apply this framework to uniform slicing and infer sufficient conditions on the underlying one-dimensional base divergence. In Section 10.4, we nuance the picture with explicit counterexamples showing that norm control alone does not ensure contraction for several multivariate objectives. Finally, in Section 10.5, we propose a solution based on max slicing and obtain sufficient conditions for a contraction guarantee under the same norm control.

10.1 Univariate case↩︎

10.2 Multivariate case↩︎

In this subsection, we keep the multivariate update as general as possible. We first treat the strong setting where the discount matrix may depend on the state action pair and may also be random, while only requiring a uniform operator norm control, as formalized in Theorem 6. We then specialize this result to the standard isotropic case where the discount is a scaled identity matrix, as stated in Corollary 1.

10.3 Uniform slicing↩︎

10.4 Negative results.↩︎

Our goal in this section is to test whether multivariate contraction can be controlled solely through the operator norm of the linear part of the Bellman update. Corollary 2 gives a simple reduction: it suffices to find one matrix \(A\) and one pair of laws \(\mu,\nu\) such that \[\label{eq:preamble95obstruction} \|A\|_{\mathrm{op}}\le \bar L \qquad\text{and}\qquad \Delta\bigl(A_{\#}\mu,\,A_{\#}\nu\bigr)>\Delta(\mu,\nu).\qquad{(5)}\] If such a triple exists, then there is an MDP/policy pair whose Bellman update respects the same norm bound \(\|A_{s,a}(C)\|_{\mathrm{op}}\le \bar L\) but for which the associated distributional Bellman operator is not nonexpansive under the supremum metric \(\bar\Delta\). In other words, ?? rules out any contraction guarantee that is meant to hold uniformly over all such updates and that depends only on \(\|A\|_{\mathrm{op}}\).

Accordingly, we now construct explicit instances of ?? for the two kernel families that have been instantiated to obtain contraction-based MMD objectives in distributional RL: the negative-type power family [23] and the multiquadric family [24], as detailed in Section 10.4.1; we then apply the same obstruction strategy to uniform sliced objectives, obtained by averaging the corresponding one-dimensional discrepancies over directions \(\theta\in\mathbb{S}^1\), as detailed in Section 10.4.2.

10.4.1 MMD.↩︎

Throughout, we use the MMD convention of Sec. 8.3. In particular, for the (possibly CPD) kernels considered below, \(\mathbf{MMD}_k^2(P,Q)\) admits the expansion \[\label{eq:mmd95expand95cross95within} \mathbf{MMD}_k^2(P,Q) = \underbrace{\mathbb{E}\,k(X,X')}_{\text{within }P} \;+\; \underbrace{\mathbb{E}\,k(Y,Y')}_{\text{within }Q} \;-\;2\, \underbrace{\mathbb{E}\,k(X,Y)}_{\text{cross}}, \qquad X,X'\!\sim P,\;Y,Y'\!\sim Q\;\text{i.i.d.}\tag{24}\]

The examples below exploit a simple mechanism visible in 24 : since the objective is a cross term minus two within terms, one can arrange for these expectations to be individually large and close, so that \(\mathbf{MMD}_k^2(P,Q)\) results from a delicate cancellation. An anisotropic contraction can then shrink the within terms much more than the cross term, breaking that cancellation and yielding an increase of \(\mathbf{MMD}_k^2\) despite the map having operator norm \(<1\).

Fix \(s\in(0,1)\) and \(\varepsilon\in(0,1)\) and consider the full-rank diagonal map on \(\mathbb{R}^2\), \[A_{s,\varepsilon}:=\mathrm{diag}(s,\;s\varepsilon), \qquad \|A_{s,\varepsilon}\|_{\mathrm{op}}=s<1.\] Fix \(D>0\) and \(M>0\) and define \[P_{D,M}:=\tfrac12\delta_{(0,M)}+\tfrac12\delta_{(0,-M)}, \qquad Q_{D,M}:=\tfrac12\delta_{(D,M)}+\tfrac12\delta_{(D,-M)}.\] Under \(A_{s,\varepsilon}\), the separation \(D\) becomes \(sD\) while the spread \(M\) becomes \(\varepsilon M\), which can be made tiny when \(\varepsilon\ll 1\).

10.4.1.1 Negative-type power family (\(\beta=1\), energy distance).

Consider the negative-type power kernel \(k_\beta(x,y):=-\|x-y\|^\beta\) on \(\mathbb{R}^2\). For \(\beta=1\), 24 becomes \[\mathbf{MMD}_{k_1}^2(P,Q) = 2\,\mathbb{E}\|X-Y\| -\mathbb{E}\|X-X'\| -\mathbb{E}\|Y-Y'\|.\] For \((P,Q)=(P_{D,M},Q_{D,M})\), the two-value enumeration yields \[\label{eq:power95beta195closedform} \mathbf{MMD}_{k_1}^2(P_{D,M},Q_{D,M}) = D+\sqrt{D^2+4M^2}-2M,\tag{25}\] and after applying \(A_{s,\varepsilon}\), \[\label{eq:power95beta195closedform95push} \mathbf{MMD}_{k_1}^2\!\big((A_{s,\varepsilon})_{\#}P_{D,M},\;(A_{s,\varepsilon})_{\#}Q_{D,M}\big) = sD+s\sqrt{D^2+4\varepsilon^2 M^2}-2s\varepsilon M.\tag{26}\]

Numerical instance. Take \(D=1\), \(M=1000\), \(s=0.75\), \(\varepsilon=10^{-6}\) and write \(A:=A_{s,\varepsilon}\). Then \(\|A\|_{\mathrm{op}}=0.75<1\), while a direct evaluation of 25 –26 gives \[\mathbf{MMD}_{k_1}^2(P_{1,1000},Q_{1,1000}) \approx 1.000249999984, \qquad \mathbf{MMD}_{k_1}^2(A_{\#}P_{1,1000},A_{\#}Q_{1,1000}) \approx 1.498501499999.\]

10.4.1.2 Negative multiquadric kernel.

Fix \(h=1\) and consider the negative multiquadric kernel \[k_{\mathrm{MQ}}(x,y):=-\sqrt{1+\|x-y\|^2}.\] With \(f(r):=\sqrt{1+r^2}\), 24 becomes \[\mathbf{MMD}_{k_{\mathrm{MQ}}}^2(P,Q) = 2\,\mathbb{E}\,f(\|X-Y\|) -\mathbb{E}\,f(\|X-X'\|) -\mathbb{E}\,f(\|Y-Y'\|).\] For \((P,Q)=(P_{D,M},Q_{D,M})\), \[\label{eq:mq95closedform} \mathbf{MMD}_{k_{\mathrm{MQ}}}^2(P_{D,M},Q_{D,M}) = f(D)+f(\sqrt{D^2+4M^2})-1-f(2M),\tag{27}\] and after applying \(A_{s,\varepsilon}\), \[\label{eq:mq95closedform95push} \mathbf{MMD}_{k_{\mathrm{MQ}}}^2\!\big((A_{s,\varepsilon})_{\#}P_{D,M},\;(A_{s,\varepsilon})_{\#}Q_{D,M}\big) = f(sD)+f\!\Big(s\sqrt{D^2+4\varepsilon^2M^2}\Big)-1-f(2s\varepsilon M).\tag{28}\]

Numerical instance. With the same choice \(D=1\), \(M=1000\), \(s=0.75\), \(\varepsilon=10^{-6}\) (so \(\|A\|_{\mathrm{op}}=0.75<1\)), a direct evaluation of 27 –28 gives \[\mathbf{MMD}_{k_{\mathrm{MQ}}}^2(P_{1,1000},Q_{1,1000}) \approx 0.414463562326, \qquad \mathbf{MMD}_{k_{\mathrm{MQ}}}^2(A_{\#}P_{1,1000},A_{\#}Q_{1,1000}) \approx 0.499999775000.\]

10.4.2 Uniform slicing.↩︎

Fix \(d=2\) and let \(\sigma\) be the uniform probability measure on \(\mathbb{S}^{1}\). Given a divergence \(\Delta\) on \(\mathcal{P}(\mathbb{R})\) and \(p\in[1,\infty)\), define its uniform sliced lift by \[\label{eq:uniform95sliced95def95p} \mathbf{S}\Delta_p(\mu,\nu) :=\Bigg(\int_{\mathbb{S}^{1}} \Delta\!\big((P_\theta)_{\#}\mu,(P_\theta)_{\#}\nu\big)^p\,d\sigma(\theta)\Bigg)^{\!1/p}, \qquad P_\theta(x):=\langle \theta,x\rangle .\tag{29}\] In this subsection we take \(\Delta=\mathbf{W}_1\) and \(p=1\), so that \[\mathbf{S} \mathbf{W}_1(\mu,\nu)=\int_{\mathbb{S}^1} \mathbf{W}_1\!\big((P_\theta)_{\#}\mu,(P_\theta)_{\#}\nu\big)\,d\sigma(\theta).\]

10.4.2.1 Wasserstein (\(\mathbf{W}_1\)).

Let \[\mu:=\tfrac12\delta_{(1,0)}+\tfrac12\delta_{(-1,0)}, \qquad \nu:=\tfrac12\delta_{(0,1)}+\tfrac12\delta_{(0,-1)}.\] Fix \(s\in(0,1)\) and \(\varepsilon\in(0,1)\) and consider the full–rank diagonal map \[A:=\mathrm{diag}(s,\;s\varepsilon), \qquad \|A\|_{\mathrm{op}}=s<1.\] Parametrize \(\theta\in\mathbb{S}^1\) as \(\theta=(\cos\varphi,\sin\varphi)\) with \(\varphi\in[0,2\pi)\). Then the one-dimensional projections are \[(P_\theta)_{\#}\mu=\tfrac12\delta_{-\cos\varphi}+\tfrac12\delta_{\cos\varphi}, \qquad (P_\theta)_{\#}\nu=\tfrac12\delta_{-\sin\varphi}+\tfrac12\delta_{\sin\varphi},\] and similarly, \[(P_\theta)_{\#}(A_{\#}\mu)=\tfrac12\delta_{-s\cos\varphi}+\tfrac12\delta_{s\cos\varphi}, \qquad (P_\theta)_{\#}(A_{\#}\nu)=\tfrac12\delta_{-s\varepsilon\sin\varphi}+\tfrac12\delta_{s\varepsilon\sin\varphi}.\] In one dimension, for \(a,b\ge 0\), \[\mathbf{W}_1\!\left(\tfrac12\delta_{-a}+\tfrac12\delta_a,\;\tfrac12\delta_{-b}+\tfrac12\delta_b\right)=|a-b|.\] Therefore, using 29 with \(p=1\) and \(d\sigma(\theta)=\frac{1}{2\pi}\,d\varphi\), we obtain \[\begin{align} \tag{30} \mathbf{S} \mathbf{W}_1(\mu,\nu) &=\frac{1}{2\pi}\int_{0}^{2\pi}\big||\cos\varphi|-|\sin\varphi|\big|\,d\varphi,\\ \tag{31} \mathbf{S} \mathbf{W}_1(A_{\#}\mu,A_{\#}\nu) &=\frac{s}{2\pi}\int_{0}^{2\pi}\big||\cos\varphi|-\varepsilon|\sin\varphi|\big|\,d\varphi. \end{align}\] We use the following identity: for any integrable \(H:[0,1]^2\to\mathbb{R}\), \[\label{eq:quadrant95reduction95H} \int_{0}^{2\pi} H\big(|\cos\varphi|,|\sin\varphi|\big)\,d\varphi = 4\int_{0}^{\pi/2} H\big(\cos\varphi,\sin\varphi\big)\,d\varphi.\tag{32}\] Applying 32 to 30 –31 (and using that \(\cos\varphi,\sin\varphi\ge 0\) on \([0,\pi/2]\)) yields \[\begin{align} \tag{33} \mathbf{S} \mathbf{W}_1(\mu,\nu) &=\frac{2}{\pi}\int_{0}^{\pi/2}|\cos\varphi-\sin\varphi|\,d\varphi,\\ \tag{34} \mathbf{S} \mathbf{W}_1(A_{\#}\mu,A_{\#}\nu) &=\frac{2s}{\pi}\int_{0}^{\pi/2}|\cos\varphi-\varepsilon\sin\varphi|\,d\varphi. \end{align}\]

We evaluate the one-dimensional integrals in 33 –34 using scipy.integrate.quad [48]. For \(s=0.9\) and \(\varepsilon=0.01\) this yields \(\|A\|_{\mathrm{op}}=0.9<1\) but \[\mathbf{S} \mathbf{W}_1(\mu,\nu)\approx 0.527393, \qquad \mathbf{S} \mathbf{W}_1(A_{\#}\mu,A_{\#}\nu)\approx 0.567286,\] so \(\mathbf{S} \mathbf{W}_1(A_{\#}\mu,A_{\#}\nu)>\mathbf{S} \mathbf{W}_1(\mu,\nu)\), which instantiates ?? with \(\Delta=\mathbf{S} \mathbf{W}_1\).

10.4.2.2 Cramér–2 (\(\ell_2^2\)).

The same construction yields a counterexample for the uniform sliced Cramér–2 divergence. For \(a,b\ge 0\), a direct computation gives \[\ell_2^2\!\left(\tfrac12\delta_{-a}+\tfrac12\delta_{a},\;\tfrac12\delta_{-b}+\tfrac12\delta_{b}\right) =\frac{1}{2}|a-b|.\] In one dimension we also have \[\mathbf{W}_1\!\left(\tfrac12\delta_{-a}+\tfrac12\delta_{a},\;\tfrac12\delta_{-b}+\tfrac12\delta_{b}\right)=|a-b|,\] hence, for every \(\theta\), \[\ell_2^2\big((P_\theta)_{\#}\mu,(P_\theta)_{\#}\nu\big)=\tfrac12\,\mathbf{W}_1\big((P_\theta)_{\#}\mu,(P_\theta)_{\#}\nu\big), \qquad \ell_2^2\big((P_\theta)_{\#}A_{\#}\mu,(P_\theta)_{\#}A_{\#}\nu\big) =\tfrac12\,\mathbf{W}_1\big((P_\theta)_{\#}A_{\#}\mu,(P_\theta)_{\#}A_{\#}\nu\big).\] Using the definition 29 with \(p=1\) and linearity of the integral, we obtain the exact identities \[\label{eq:S95c295half95SW195exact95base} \mathbf{S}\ell_2^2(\mu,\nu)=\tfrac12\,\mathbf{S} \mathbf{W}_1(\mu,\nu), \qquad \mathbf{S}\ell_2^2(A_{\#}\mu,A_{\#}\nu)=\tfrac12\,\mathbf{S} \mathbf{W}_1(A_{\#}\mu,A_{\#}\nu).\tag{35}\] In particular, for \(s=0.9\) and \(\varepsilon=0.01\) (so \(\|A\|_{\mathrm{op}}=0.9<1\)), using the values reported in the Wasserstein case we get \[\mathbf{S}\ell_2^2(\mu,\nu)\approx \tfrac12\cdot 0.527393 = 0.263697, \qquad \mathbf{S}\ell_2^2(A_{\#}\mu,A_{\#}\nu)\approx \tfrac12\cdot 0.567286 = 0.283643,\] so \(\mathbf{S}\ell_2^2(A_{\#}\mu,A_{\#}\nu)>\mathbf{S}\ell_2^2(\mu,\nu)\), which instantiates ?? with \(\Delta=\mathbf{S}\ell_2^2\).

10.4.2.3 Uniform sliced MMD with the negative-type power family.

Fix \(\beta\in(0,2]\) and consider the negative-type power kernel on \(\mathbb{R}\), \[k_\beta(x,y):=-|x-y|^\beta .\] For one-dimensional laws \(\mu,\nu\in\mathcal{P}(\mathbb{R})\), \[\mathbf{MMD}^2_{k_\beta}(\mu,\nu) = \underbrace{2\,\mathbb{E}|X-Y|^\beta}_{\text{cross}} \;-\; \underbrace{\mathbb{E}|X-X'|^\beta}_{\text{within }\mu} \;-\; \underbrace{\mathbb{E}|Y-Y'|^\beta}_{\text{within }\nu}, \qquad X,X'\!\sim\mu,\;Y,Y'\!\sim\nu\;\text{i.i.d.}\] For the two-point symmetric laws \(\mu=\tfrac12\delta_{+a}+\tfrac12\delta_{-a}\) and \(\nu=\tfrac12\delta_{+b}+\tfrac12\delta_{-b}\) with \(a,b\ge 0\), we have \[\mathbb{E}|X-Y|^\beta=\tfrac12|a-b|^\beta+\tfrac12(a+b)^\beta, \qquad \mathbb{E}|X-X'|^\beta=2^{\beta-1}a^\beta, \qquad \mathbb{E}|Y-Y'|^\beta=2^{\beta-1}b^\beta,\] and substituting into the definition yields \[\label{eq:1d95mmd295power952point} \mathbf{MMD}^2_{k_\beta}(\mu,\nu) = |a-b|^\beta+(a+b)^\beta-2^{\beta-1}\big(a^\beta+b^\beta\big).\tag{36}\]

Now consider the same pair of laws \(\mu,\nu\in\mathcal{P}(\mathbb{R}^2)\) and the full-rank anisotropic map \(A_{s,\varepsilon}=\mathrm{diag}(s,s\varepsilon)\) as above. For \(\theta=(\cos\varphi,\sin\varphi)\in\mathbb{S}^1\) with \(\varphi\in[0,2\pi)\), the projected laws are \[(P_\theta)_\#\mu=\tfrac12\delta_{+\!|\cos\varphi|}+\tfrac12\delta_{-\!|\cos\varphi|}, \qquad (P_\theta)_\#\nu=\tfrac12\delta_{+\!|\sin\varphi|}+\tfrac12\delta_{-\!|\sin\varphi|},\] and after applying \(A_{s,\varepsilon}\), \[(P_\theta)_\#(A_{s,\varepsilon})_\#\mu=\tfrac12\delta_{+\!s|\cos\varphi|}+\tfrac12\delta_{-\!s|\cos\varphi|}, \qquad (P_\theta)_\#(A_{s,\varepsilon})_\#\nu=\tfrac12\delta_{+\!s\varepsilon|\sin\varphi|}+\tfrac12\delta_{-\!s\varepsilon|\sin\varphi|}.\] Define the uniform sliced lift (with \(p=1\) and \(\sigma\) uniform on \(\mathbb{S}^1\)) by \[\mathbf{S}\mathbf{MMD}^2_{k_\beta}(\mu,\nu) :=\int_{\mathbb{S}^1}\mathbf{MMD}^2_{k_\beta}\big((P_\theta)_\#\mu,(P_\theta)_\#\nu\big)\,d\sigma(\theta).\] Using the identity 32 (as in the Wasserstein paragraph) to reduce the \(\varphi\)-integral to \([0,\pi/2]\), and then applying 36 with \(a=\cos\varphi\) and \(b=\sin\varphi\), this becomes \[\begin{align} \tag{37} \mathbf{S}\mathbf{MMD}^2_{k_\beta}(\mu,\nu) &=\frac{2}{\pi}\int_{0}^{\pi/2} \Big( |\cos\varphi-\sin\varphi|^\beta+(\cos\varphi+\sin\varphi)^\beta \nonumber\\ & -2^{\beta-1}\big(\cos^\beta\varphi+\sin^\beta\varphi\big) \Big)\,d\varphi, \\ \tag{38} \mathbf{S}\mathbf{MMD}^2_{k_\beta}\big((A_{s,\varepsilon})_\#\mu,(A_{s,\varepsilon})_\#\nu\big) &=\frac{2}{\pi}\int_{0}^{\pi/2} \Big( |s\cos\varphi-s\varepsilon\sin\varphi|^\beta+(s\cos\varphi+s\varepsilon\sin\varphi)^\beta \nonumber\\ & -2^{\beta-1}\big((s\cos\varphi)^\beta+(s\varepsilon\sin\varphi)^\beta\big) \Big)\,d\varphi . \end{align}\]

Concrete numbers (computed by the same numerical integration method as earlier). We use \(\varepsilon=0.01\) throughout, with \(s=0.9\) for \(\beta\in\{0.5,1\}\) and \(s=0.98\) for \(\beta=1.5\). When \(\beta=1\), 36 simplifies to \(\mathbf{MMD}^2_{k_1}(\tfrac12\delta_{\pm a},\tfrac12\delta_{\pm b})=|a-b|\), so the sliced objective coincides with the uniform sliced Wasserstein-1 objective from the previous paragraph; numerically this gives \[\mathbf{S}\mathbf{MMD}^2_{k_1}(\mu,\nu)\approx 0.52739, \qquad \mathbf{S}\mathbf{MMD}^2_{k_1}\big((A_{0.9,0.01})_\#\mu,(A_{0.9,0.01})_\#\nu\big)\approx 0.56729, \qquad \text{ratio}\approx 1.07564.\] For \(\beta=0.5\), \[\mathbf{S}\mathbf{MMD}^2_{k_{0.5}}(\mu,\nu)\approx 0.73545, \qquad \mathbf{S}\mathbf{MMD}^2_{k_{0.5}}\big((A_{0.9,0.01})_\#\mu,(A_{0.9,0.01})_\#\nu\big)\approx 0.88443, \qquad \text{ratio}\approx 1.20256.\] For \(\beta=1.5\), \[\mathbf{S}\mathbf{MMD}^2_{k_{1.5}}(\mu,\nu)\approx 0.29777, \qquad \mathbf{S}\mathbf{MMD}^2_{k_{1.5}}\big((A_{0.98,0.01})_\#\mu,(A_{0.98,0.01})_\#\nu\big)\approx 0.31553, \qquad \text{ratio}\approx 1.05964.\] Thus, for several \(\beta\in(0,2]\), the uniform sliced objective \(\mathbf{S}\mathbf{MMD}^2_{k_\beta}\) can increase under a full-rank anisotropic \(A_{s,\varepsilon}\) with \(\|A_{s,\varepsilon}\|_{\mathrm{op}}=s<1\).

10.4.2.4 Uniform sliced MMD with the multiquadric kernel.

Fix \(h>0\) and consider the (negative) multiquadric kernel on \(\mathbb{R}\), \[k_h(x,y):=-\sqrt{1+h^2(x-y)^2}.\] For one-dimensional laws \(\mu,\nu\in\mathcal{P}(\mathbb{R})\), with \(X,X'\sim\mu\) i.i.d.and \(Y,Y'\sim\nu\) i.i.d., \[\label{eq:mmd295mq95expand95h} \mathbf{MMD}^2_{k_h}(\mu,\nu) = \underbrace{2\,\mathbb{E}\sqrt{1+h^2(X-Y)^2}}_{\text{cross}} \;-\; \underbrace{\mathbb{E}\sqrt{1+h^2(X-X')^2}}_{\text{within }\mu} \;-\; \underbrace{\mathbb{E}\sqrt{1+h^2(Y-Y')^2}}_{\text{within }\nu}.\tag{39}\] For the two-point symmetric laws \(\mu=\tfrac12\delta_{+a}+\tfrac12\delta_{-a}\) and \(\nu=\tfrac12\delta_{+b}+\tfrac12\delta_{-b}\) with \(a,b\ge 0\), we have \[\begin{align} \mathbb{E}\sqrt{1+h^2(X-Y)^2} &=\tfrac12\sqrt{1+h^2(a-b)^2}+\tfrac12\sqrt{1+h^2(a+b)^2},\\ \mathbb{E}\sqrt{1+h^2(X-X')^2} &=\tfrac12+\tfrac12\sqrt{1+4h^2a^2},\\ \mathbb{E}\sqrt{1+h^2(Y-Y')^2} &=\tfrac12+\tfrac12\sqrt{1+4h^2b^2}, \end{align}\] and substituting into 39 yields \[\label{eq:1d95mmd295mq952point95h} \mathbf{MMD}^2_{k_h}(\mu,\nu) = \sqrt{1+h^2(a-b)^2}+\sqrt{1+h^2(a+b)^2} -1-\tfrac12\Big(\sqrt{1+4h^2a^2}+\sqrt{1+4h^2b^2}\Big).\tag{40}\]

Now consider the same \(\mu,\nu\in\mathcal{P}(\mathbb{R}^2)\) and the full-rank anisotropic map \(A_{s,\varepsilon}=\mathrm{diag}(s,s\varepsilon)\) as above. Define for \(a,b\ge 0\) \[F_{\mathrm{MQ},h}(a,b):= \sqrt{1+h^2(a-b)^2}+\sqrt{1+h^2(a+b)^2} -1-\tfrac12\Big(\sqrt{1+4h^2a^2}+\sqrt{1+4h^2b^2}\Big).\] Using the identity 32 to reduce the \(\varphi\)-integral to \([0,\pi/2]\), and applying 40 slice-wise with \(a=\cos\varphi\) and \(b=\sin\varphi\), we obtain \[\begin{align} \tag{41} \mathbf{S}\mathbf{MMD}^2_{k_h}(\mu,\nu) &=\frac{2}{\pi}\int_{0}^{\pi/2} F_{\mathrm{MQ},h}\!\big(\cos\varphi,\sin\varphi\big)\,d\varphi,\\ \tag{42} \mathbf{S}\mathbf{MMD}^2_{k_h}\big((A_{s,\varepsilon})_\#\mu,(A_{s,\varepsilon})_\#\nu\big) &=\frac{2}{\pi}\int_{0}^{\pi/2} F_{\mathrm{MQ},h}\!\big(s\cos\varphi,s\varepsilon\sin\varphi\big)\,d\varphi. \end{align}\]

Concrete numbers. We evaluate the one-dimensional integrals in 41 –42 using the same numerical integration method as earlier. Instantiating \(h=100\) (the value used in [24]), and choosing \(s=0.9\) and \(\varepsilon=0.01\) (so \(\|A_{s,\varepsilon}\|_{\mathrm{op}}=0.9<1\)), we obtain \[\mathbf{S}\mathbf{MMD}^2_{k_{100}}(\mu,\nu)\approx 51.758656, \qquad \mathbf{S}\mathbf{MMD}^2_{k_{100}}\big((A_{0.9,0.01})_\#\mu,(A_{0.9,0.01})_\#\nu\big)\approx 55.552543, \qquad \text{ratio}\approx 1.07330,\] so \(\mathbf{S}\mathbf{MMD}^2_{k_{100}}\big((A_{0.9,0.01})_\#\mu,(A_{0.9,0.01})_\#\nu\big)> \mathbf{S}\mathbf{MMD}^2_{k_{100}}(\mu,\nu)\), which instantiates ?? with \(\Delta=\mathbf{S}\mathbf{MMD}^2_{k_{100}}\).

10.5 Max slicing↩︎

11 Mixture bias↩︎

We first recall the unbiased sample gradient property (U)from Section 4.4. Let \(\mu\) denote the data law and let \(\nu_\theta\) be a parametric model law. Writing \(\widehat \mu_m\) for the empirical measure of \(m\) i.i.d.samples from \(\mu\), we say that a divergence \(\mathcal{D}\) satisfies (U)if \[\mathbb{E}_{X_{1:m}\sim \mu}\!\left[\nabla_\theta\,\mathcal{D}\!\left(\widehat{\mu}_m,\,\nu_\theta\right)\right] \;=\; \nabla_\theta\,\mathcal{D}\!\left(\mu,\,\nu_\theta\right). \qquad \mathrm{\boldsymbol{\textcolor{nice_gold}{(U)}}}\] A small but important terminology point is that divergence here only refers to a population-level separation property (nonnegativity and \(\mathcal{D}(\mu,\nu)=0\) iff \(\mu=\nu\)), so that \(\nu=\mu\) minimizes the population objective \(\nu\mapsto \mathcal{D}(\mu,\nu)\). In practice, however, learning minimizes the sample loss \(\theta\mapsto \mathcal{D}(\widehat\mu_m,\nu_\theta)\) using plug-in gradients; (U)is precisely what ensures that, in expectation over the sampling of \(\widehat\mu_m\), these gradients point in the direction of the population objective. Equivalently, when (U)holds, the expected sample loss is minimized at \(\nu=\mu\), which is why such losses are called proper in this context [26].

11.0.0.1 Why this matters in distributional RL.

In distributional RL, the Bellman target is the mixture law induced by the transition and policy randomness, as introduced by the distributional Bellman operator in Eq. 1 (Section 2). However, a standard TD update draws a single sampled \((S',A')\) (and reward) and therefore forms a one-sample based target distribution \(\mathop{\mathrm{Law}}\!\big(R(s,a)+\Gamma(s,a)\,Z_{\phi^-}(S',A')\big)\) rather than the full mixture \((\mathcal{T}^\pi Z_{\phi^-})(s,a)\). As a result, learning implicitly targets the expected sample loss, which without (U)need not coincide with the intended population objective: \[\label{eq:drl95argmin95mismatch} \arg\min_{\phi}\;\mathbb{E}_{S',A'}\!\Big[\mathcal{D}\big(\mathop{\mathrm{Law}}(R+\Gamma Z_{\phi^-}(S',A')),\;Z_\phi(s,a)\big)\Big] \;\neq\; \arg\min_{\phi}\;\mathcal{D}\big((\mathcal{T}^\pi Z_{\phi^-})(s,a),\;Z_\phi(s,a)\big),\tag{46}\] as highlighted for Wasserstein objectives in [8].

11.0.0.2 Wasserstein.

Wasserstein distances provide a canonical illustration where (U)fails in distributional RL: the Bellman target is a mixture over transition and policy randomness, yet training typically instantiates it from a single sampled next state and next action. In that regime, one generally cannot permute expectation and gradient as above, and the resulting plug-in stochastic gradients can be biased relative to the population Wasserstein objective [8], [26], [49]. It is one reason why Wasserstein-based discrepancies are delicate in stochastic training for distributional RL [8].

11.0.0.3 Energy distance and MMD.

A broad family of kernel discrepancies used in practice can be written in the quadratic “energy” form \[\label{eq:energy-form} \mathcal{E}_k(\mu,\nu) := \mathbb{E}_{x,x'\sim\mu}[k(x,x')] +\mathbb{E}_{y,y'\sim\nu}[k(y,y')] -2\mathbb{E}_{x\sim\mu,\,y\sim\nu}[k(x,y)],\tag{47}\] for a symmetric function \(k\). When \(k\) is positive definite this is exactly the usual squared MMD objective, and under our convention (Section 8.3) the same formula can also be viewed as the squared “kernel energy” \(\gamma_k(\mu,\nu)^2\) associated with a conditionally positive definite kernel [46].

For stochastic training, the key point is that the squared objective 47 satisfies (U). Indeed, \[\mathcal{E}_k(\widehat\mu_m,\nu_\theta) = \mathbb{E}_{x,x'\sim\widehat\mu_m}[k(x,x')] +\mathbb{E}_{y,y'\sim\nu_\theta}[k(y,y')] -2\mathbb{E}_{x\sim\widehat\mu_m,\,y\sim\nu_\theta}[k(x,y)],\] where the first term does not depend on \(\theta\) and thus vanishes under \(\nabla_\theta\). Linearity of expectation then yields \[\mathbb{E}_{X_{1:m}\sim\mu}\!\left[\nabla_\theta\,\mathcal{E}_k(\widehat\mu_m,\nu_\theta)\right] = \nabla_\theta\,\mathcal{E}_k(\mu,\nu_\theta),\] so (U)holds [26], [36].

11.0.0.4 Consequences for energy distance and Cramér.

Because the (squared) energy distance admits the representation 47 , it satisfies (U). Finally, in one dimension, the Cramér–2 divergence coincides with the (squared) energy distance, and the corresponding identity was already presented in Sec. 8.2: \[\ell_2^2(\mu,\nu) \;=\; \int_{\mathbb{R}}\!\bigl(F_\mu(t)-F_\nu(t)\bigr)^2\,dt \;=\; \mathbb{E}\,|X-Y| \;-\;\tfrac{1}{2}\,\mathbb{E}\,|X-X'| \;-\;\tfrac{1}{2}\,\mathbb{E}\,|Y-Y'|, \qquad X,X'\!\sim\mu,\;\;Y,Y'\!\sim\nu\;\text{i.i.d.}\]

11.1 Interaction with slicing↩︎

11.1.1 Uniform slicing and the unbiased sample gradient property↩︎

Let \(\Delta:\mathcal{P}(\mathbb{R})\times\mathcal{P}(\mathbb{R})\to[0,\infty)\) be a base divergence and let \(p\in[1,\infty)\). For \(\mu,\nu\in\mathcal{P}(\mathbb{R}^d)\), recall the uniform sliced powered objective \[\mathbf{S}\Delta_p^{\,p}(\mu,\nu) \;:=\; \int_{S^{d-1}} \Delta^p\!\big((P_\theta)_{\#}\mu,\,(P_\theta)_{\#}\nu\big)\,d\sigma(\theta), \qquad P_\theta(x)=\langle \theta,x\rangle .\]

12 Sample complexity↩︎

12.1 Uniform slicing↩︎

12.2 Max slicing↩︎

13 Instantiations↩︎

This section collects concrete instantiations of our contraction results for Wasserstein, Cramér, and MMD divergences. For each choice, we state the relevant structural properties, contraction factors under shared scalar and general matrix discounting, and we record suitability for distributional RL training and sample complexity. We cover Wasserstein in Section 13.1, Cramér in Section 13.2, and MMD in Section 13.3. We summarize all results in Table 1.

13.1 Wasserstein↩︎

Wasserstein is a metric on \(\mathcal{P}_p(\mathbb{R})\) for every \(p\in[1,\infty]\) (Proposition 2 in [43]). It satisfies (T)since it is translation invariant. It satisfies (S)with \(c(s)=s\) by linear push-forward contraction (Proposition 7). It satisfies (M\(_p\))by mixture \(p\)-convexity (Proposition 8).

Gradient bias. As discussed in Section 11, Wasserstein discrepancies generally fail (U)in the one-sample TD regime when a mixture target is instantiated from a single sampled \((S',A')\) and reward as in Eq. 46 . Therefore, gradient bias is unavoidable for one-sample TD training with Wasserstein objectives.

Contraction and sample complexity. Combining (T), Proposition 7, and Proposition 8 yields \[\mathbf{W}_p\!\bigl(T^\pi\eta_1,\,T^\pi\eta_2\bigr) \;\le\; \|\Gamma\|_{\mathop{\mathrm{op}}}\,\mathbf{W}_p\!\bigl(\eta_1,\,\eta_2\bigr),\] hence strict contraction whenever \(\|\Gamma\|_{\mathop{\mathrm{op}}}<1\). In particular, for \(\Gamma=\gamma I_d\) this reduces to the factor \(\gamma\). For empirical Wasserstein in dimension \(d\), the expected estimation error deteriorates with dimension [19], so we do not record a dimension-robust rate for \(\mathbf{W}_p\) itself.

Sliced variants. By Theorem 7 with \(\Delta=\mathbf{W}_p\), the uniform sliced Wasserstein update with \(\Gamma=\gamma I_d\) contracts with factor \(c(\gamma)<1\), \[\mathbf{SW}_p\!\bigl(T^\pi\eta_1,\,T^\pi\eta_2\bigr) \;\le\; c(\gamma)\,\mathbf{SW}_p\!\bigl(\eta_1,\,\eta_2\bigr).\]

By Theorem 8 with \(\Delta=\mathbf{W}_p\), the max-sliced Wasserstein update under a general random linear map (with envelope \(\bar L=\sup_{(s,a),C}\|A_{s,a}(C)\|_{\mathop{\mathrm{op}}}\)) contracts with factor \(c(\bar L)\), strictly so whenever \(\bar L<1\), \[\mathbf{MSW}_p\!\bigl(T^\pi\eta_1,\,T^\pi\eta_2\bigr) \;\le\; c(\bar L)\,\mathbf{MSW}_p\!\bigl(\eta_1,\,\eta_2\bigr).\]

Sample complexity (uniform slicing). Let \(p\in[1,\infty)\) and assume \(\mu\in\mathcal{P}_q(\mathbb{R}^d)\) with \(q>2p\) (finite \(q\)-th moment). Let \(\hat{\mu}_n\) be the empirical measure from \(n\) samples. Carrying the same steps as in Corollary 2 of [13] but in the one-sample setting, and plugging the 1D base bound from Theorem 1 of [19], we obtain the dimension–free rate \[\mathbb{E}\!\left[\mathbf{SW}_p(\hat{\mu}_n,\mu)\right]\;=\;\mathcal{O}\!\left(n^{-1/(2p)}\right).\] Thus, uniform slicing avoids the curse of dimensionality.

Sample complexity (max-sliced). By Theorem 10 and Corollaries 3–4, if \(\mathrm{diam}(\mathrm{supp}\,\mu)\le D\) then \[\mathbb{E}\,\mathbf{MSW}_1(\hat{\mu}_n,\mu) = \mathcal{O}\!\Big(D\,\sqrt{\tfrac{d\log n}{n}}\Big), \qquad \mathbb{E}\,\mathbf{MSW}_p(\hat{\mu}_n,\mu) = \mathcal{O}\!\Big(D\,\big(\tfrac{d\log n}{n}\big)^{\!1/(2p)}\Big) \qquad (p>1).\]

13.2 Cramér↩︎

The Cramér distance enjoys all the structural assumptions we require. By Proposition 9, it is a metric. It satisfies (T)by Proposition 2 in [26] and Proposition 3.2 in [45], and (S)with \(c(s)=s^{1/2}\) via Proposition 10. It also satisfies (M\(_p\))(Proposition 11).

Gradient bias. As discussed in Section 11, the squared energy-distance form satisfies (U), and in one dimension Cramér–2 coincides with squared energy distance. By Lemma 12, this unbiasedness lifts to uniform slicing for the powered objective, so the guarantee applies to the squared sliced Cramér objective \(\mathbf{SC}_2^{\,2}\).

Contraction factors. By Theorem 7 with \(\Delta=\mathbf{C}_2\) (so \(c(s)=s^{1/2}\)), the sliced Cramér update with \(\Gamma=\gamma I\) contracts with factor \(\gamma^{1/2}\): \[\mathbf{SC}_2\!\bigl(T^\pi\eta_1,\,T^\pi\eta_2\bigr)\;\le\;\gamma^{1/2}\,\mathbf{SC}_2\!\bigl(\eta_1,\,\eta_2\bigr).\] By Theorem 8, the max–sliced Cramér update with a general linear map \(\Gamma\) contracts with factor \(c(\bar L)=\bar L^{1/2}\), strictly so whenever \(\bar L<1\): \[\mathbf{MSC}_2\!\bigl(T^\pi\eta_1,\,T^\pi\eta_2\bigr)\;\le\;\bar L^{1/2}\,\mathbf{MSC}_2\!\bigl(\eta_1,\,\eta_2\bigr).\]

Sample complexity (uniform slicing). For the one–dimensional Cramér distance (the \(L^2\)–CDF discrepancy), Lemma 14 yields \[\mathbb{E}\,\|F_n-F\|_{L^2(F)} \;=\; \mathcal{O}\!\left(n^{-1/2}\right).\] Plugging this base rate into Theorem 9 yields the dimension–free bound \[\mathbb{E}\!\left[\mathbf{SC}_2(\hat{\mu}_n,\mu)\right]\;=\;\mathcal{O}\!\left(n^{-1/2}\right).\] Thus, uniform slicing avoids the curse of dimensionality.

Sample complexity (max–sliced). By Theorem 10 and Corollary 5, for \(\mathrm{diam}(\mathrm{supp}\,\mu)\le D\), \[\mathbb{E}\!\left[\mathbf{MSC}_2(\hat{\mu}_n,\mu)\right]\;=\;\mathcal{O}\!\Big(\sqrt{D}\,\sqrt{\tfrac{d\log n}{n}}\Big).\]

13.3 MMD↩︎

The Maximum Mean Discrepancy (MMD) with a conditionally strictly positive definite kernel [31], [46] is a valid metric on probability laws. With the multiquadric (MQ) kernel \(k_h(x,y)=-\sqrt{1+h^2\|x-y\|^2}\), it enjoys all the structural assumptions we require. By Proposition 12, it is a metric. It satisfies (T)since MMD is translation invariant for all shift–invariant kernels. It satisfies (S)with \(c(s)=s^{1/2}\) for the MQ kernel (Proposition 13), reflecting its scale–sensitivity. Finally, it satisfies (M\(_p\))by mixture convexity of RKHS embeddings (Proposition 15).

Gradient bias. As discussed in Section 11, squared kernel energy objectives of the form 47 satisfy (U). By Lemma 12, this unbiasedness lifts to uniform slicing for the powered objective \(\mathbf{S}\mathbf{MMD}_{k_h}^{\,2}\), so the guarantee applies to the squared sliced MMD.

Contraction factors. By Theorem 7 with \(\Delta=\mathbf{MMD}_{k_h}\) and the scale bound \[c(s)=s^{1/2},\] the sliced MMD update with \(\Gamma=\gamma I\) satisfies \[\mathbf{SMMD}_{k_h}\!\bigl(T^\pi\eta_1,\,T^\pi\eta_2\bigr)\;\le\;c(\gamma)\,\mathbf{SMMD}_{k_h}\!\bigl(\eta_1,\,\eta_2\bigr).\] In particular, for scalar discounts \(\gamma\in(0,1)\) we have \(c(\gamma)=\gamma^{1/2}\).

By Theorem 8, the max–sliced MMD update with a general linear map \(\Gamma\) satisfies \[\mathbf{MSMMD}_{k_h}\!\bigl(T^\pi\eta_1,\,T^\pi\eta_2\bigr)\;\le\;c(\bar L)\,\mathbf{MSMMD}_{k_h}\!\bigl(\eta_1,\,\eta_2\bigr), \qquad c(\bar L)=\bar L^{1/2}.\] In particular, under \(\bar L<1\) this is a strict contraction.

Sample complexity (uniform slicing). In one dimension, the unbiased empirical MMD (equivalently, the energy distance) is a U–statistic [31], [46], so classical U–statistic theory yields the standard rate \[\mathbb{E}\!\left[\mathbf{MMD}_{k_h}(\hat{\mu}_n,\mu)\right]\;=\;\mathcal{O}\!\left(n^{-1/2}\right).\] Plugging this into Theorem 9(i) yields the dimension–free bound \[\mathbb{E}\!\left[\mathbf{SMMD}_{k_h}(\hat{\mu}_n,\mu)\right]\;=\;\mathcal{O}\!\left(n^{-1/2}\right).\] Thus, uniform slicing avoids the curse of dimensionality.

Sample complexity (max–sliced). We were not able to establish a sharp sample complexity bound for the max–sliced MMD. Deriving such a result remains an open problem for future work.

14 Pseudo-codes↩︎

15 Pixel-based control: maze and Atari↩︎

We follow the maze environment and Atari subset of [16]. The maze environment was reimplemented purely in JAX [51] to improve parallelism. This section collects the design choices shared by both pixel-based setups. Details specific to each are provided in Section 15.4 and 15.5.

15.1 From multivariate rewards to a control signal↩︎

Our critic models the full multivariate return distribution \(Z^\pi(s,a)\in\mathbb{R}^d\). To interface with standard control algorithms, we fix a preference (weight) vector \(\alpha\in\mathbb{R}^d\) and derive a scalar action-value by taking the expected return and projecting it: \[Q^\pi_\alpha(s,a)\;:=\;\alpha^\top \mathbb{E}\!\left[Z^\pi(s,a)\right].\] This produces a standard scalar control objective while retaining a multivariate distributional critic (used for the TD updates and richer modeling). We then plug this signal into a standard control algorithm, using PQN in our experiments [37]. Other common choices include DQN [52] and actor-critic methods for continuous control [53].

15.2 Baseline algorithm: PQN↩︎

15.2.0.1 Adapting TD-\(\lambda\).

PQN uses TD-\(\lambda\) [54] to build training targets that interpolate between one-step bootstrapping and Monte Carlo returns. Given an episode (or truncated rollout) \((S_t,A_t,R_{t+1})_{t=0}^{T-1}\), a discount \(\gamma\in[0,1)\), and a bootstrap value \(V_t := V(S_t)\), define the \(n\)-step return \[\label{eq:n95step95return} G_t^{(n)} \;:=\; \sum_{k=0}^{n-1}\gamma^k\,R_{t+k+1} \;+\;\gamma^n\,V(S_{t+n}), \qquad n\ge 1.\tag{49}\] The TD-\(\lambda\) target is the geometric mixture of these \(n\)-step returns, \[\label{eq:td95lambda95forward} G_t^{\lambda} \;:=\; (1-\lambda)\sum_{n=1}^{T-t-1}\lambda^{n-1}G_t^{(n)} \;+\;\lambda^{T-t-1}G_t^{(T-t)}, \qquad \lambda\in[0,1].\tag{50}\] This canonical definition is equivalent to the backward recursion used for implementation: \[\label{eq:td95lambda95backward} G_t^{\lambda} \;=\; R_{t+1} \;+\; \gamma(1-D_{t+1}) \Big((1-\lambda)\,V(S_{t+1})+\lambda\,G_{t+1}^{\lambda}\Big),\tag{51}\] which is computed by a reverse scan over time. Here \(D_{t+1}\in\{0,1\}\) is the done flag, so the factor \((1-D_{t+1})\) disables bootstrapping after termination.

15.2.0.2 Adapting to the distributional setting.

The distributional critic outputs \(K\) particles that represent a sample-based approximation of the return distribution. In our case, these \(K\) particles are predicted jointly by a neural network and cannot be assumed to be perfectly i.i.d., so their ordering should not be used implicitly when forming bootstrapped targets.

In Q-learning, TD-\(\lambda\) combines the one-step bootstrap and the continuation through a convex combination via 51 . For distributional return learning, we reinterpret the same recursion at the level of return laws: \[\label{eq:td95lambda95backward95dist} G_t^{\lambda} \;\overset{D}{=}\; R_{t+1} + \gamma(1-D_{t+1})\Big((1-\lambda)\,Z_{t+1} + \lambda\,G_{t+1}^{\lambda}\Big), \qquad Z_{t+1}:=Z(S_{t+1},a^{\mathrm{policy}}).\tag{52}\] When approximating these laws by particles, we implement the convex combination as a mixture over the two components. Let \(\{z_{t+1}^{(k)}\}_{k=1}^K\) denote the particles predicted at \((S_{t+1},a^{\mathrm{policy}})\) and \(\{G_{t+1}^{\lambda,(k)}\}_{k=1}^K\) the particles of the continuation target. For each transition and each \(k\), sample \(b_t^{(k)}\sim\mathrm{Bernoulli}(\lambda)\) and set \[\label{eq:td95lambda95particle95mixture} c_t^{(k)} := \underbrace{\mathbf{1}\{b_t^{(k)}=0\}\,\tilde{z}_{t+1}^{(k)}}_{\text{one-step bootstrap}} \;+\; \underbrace{\mathbf{1}\{b_t^{(k)}=1\}\,\tilde{G}_{t+1}^{\lambda,(k)}}_{\text{continuation target}} .\tag{53}\]

Second, before applying this mixture we independently shuffle the \(K\) particles within each of the two sets \(\{z_{t+1}^{(k)}\}_{k=1}^K\) and \(\{G_{t+1}^{\lambda,(k)}\}_{k=1}^K\), and then define \(\tilde{z}_{t+1}^{(k)}\) and \(\tilde{G}_{t+1}^{\lambda,(k)}\) as the particles obtained after shuffling. This breaks any systematic index pairing while preserving the empirical distribution of each set.

15.3 Experimental setup↩︎

We base most of our implementation on PQN [37]. The distributional extension used in this work is summarized in Algorithm 8. Environment parameters shared by our pixel-based setups are reported in Table 2, and the default hyper-parameters for the training setup are reported in Table 5. As discussed in Section 13, we use the powered sliced objectives for all base metrics.

15.3.0.1 Neural network architecture.

We keep most of the Q-network of [37] and apply two changes: we replace ReLU with SiLU, and we change the final linear layer to output \(K\) particles per action. The resulting architecture is summarized in Table 3.

15.4 Maze↩︎

We faithfully reimplemented the maze environments introduced by [16]. The environments were rewritten entirely in JAX [51], including both the environment dynamics and the rendering pipeline. This choice was made to ensure full compatibility with PQN, which is itself implemented in JAX, and to support large-scale parallel execution with \(256\) environments in parallel. Unless stated otherwise, the training settings for the maze experiments are reported in Table 4.

15.4.0.1 Environment.

All maze variants are fixed grid environments with deterministic movement dynamics and pixel observations. The agent observation is an \(84\times84\times3\) RGB rendering of the environment. This observation fully specifies the current environment configuration relevant to control, including the maze layout, the agent position, and the reward locations. The agent occupies a single cell in the grid and can take one of four discrete actions corresponding to moving up, down, left, or right. If the target cell is free, the agent moves to that cell, otherwise the action has no effect and the agent remains in its current position. Moves also have no effect when the agent is trying to return to a previous position. Colored walls can only be passed in one direction. Yellow walls can be passed from the top or from the left. Red walls can be passed from the bottom or from the right. Episodes start from a fixed initial position shared across all maze variants and run for a fixed horizon. At each step, the environment emits a vector-valued reward.

Reward locations are indicated by colored squares in the maze. Each color corresponds to one reward component. When the agent reaches such a square, it receives a stochastic scalar reward in the corresponding component and the reward cannot be collected again within the same episode. The number of reward components and their distributions depend on the maze variant described next.

15.4.0.2 Maze variants and reward functions.

Figure 9 illustrates the three maze variants and the placement of colored reward sources. The reward dimensionality is \(d=2\) for maze-exclusive and maze-identical, and \(d=4\) for maze-multireward. The reward distributions follow [16].

• maze-exclusive and maze-identical (\(d=2\)). Each position with a green square yields a reward \(r \sim \mathcal{U}(0.2, 0.6)\) in the first reward component. Each position with a red square yields a reward \(r \sim \mathcal{U}(0.4, 0.9)\) in the second reward component.

• maze-multireward (\(d=4\)). Each orange square yields a reward \(r \sim \mathcal{U}(0.2, 0.4)\) in the first reward component. Each blue square yields a reward \(r \sim \mathcal{U}(0.8, 1.0)\) in the second reward component. Each green square yields a reward \(r \sim \mathcal{U}(0.3, 0.5)\) in the third reward component. Each red square yields a reward \(r \sim \mathcal{U}(0.5, 0.7)\) in the fourth reward component.

15.4.0.3 Evaluation protocol.

To evaluate the learned return distributions, we generate Monte Carlo (MC) rollouts from the reset distribution \(S_0 \sim P_0\) using a uniform random policy \(\pi(a\mid s)=\frac{1}{|\mathcal{A}|}\), so the rollouts capture both environment stochasticity and policy stochasticity. For a fixed horizon \(H\), each rollout yields a discounted return \[G_0^{(n)} \;:=\; \sum_{t=0}^{H-1}\gamma^t R_{t+1}^{(n)}\in\mathbb{R}^d, \qquad \widehat{\mu}_{\mathrm{MC}} := \frac{1}{N}\sum_{n=1}^N \delta_{G_0^{(n)}} .\] We then query the distributional critic at the start state for all actions and aggregate its predicted particles using the same uniform policy \(\pi\), yielding an empirical critic-implied return distribution \(\widehat{\mu}_\phi\). Finally, we compare the two empirical distributions using the exact optimal-transport estimator \(\widehat {\mathbf{W}}_2(\widehat{\mu}_{\mathrm{MC}},\widehat{\mu}_\phi)\) computed with POT [56].

15.5 Atari↩︎

15.5.0.1 Reward decomposition.

We reproduce the Atari multi-reward setup of [16]. Concretely, we consider six environments, AirRaid, Asteroids, Gopher, MsPacman, UpNDown, and Pong, and decompose each primitive reward into a fixed-dimensional vector while preserving the total reward (the sum of components equals the original scalar reward). We list the exact mappings below.

• AirRaid. Primitive rewards take values in \(\{100,75,50,25,0\}\) and are mapped to \[100\mapsto[100,0,0,0],\quad 75\mapsto[0,75,0,0],\quad 50\mapsto[0,0,50,0],\quad 25\mapsto[0,0,0,25],\quad 0\mapsto[0,0,0,0].\]

• Pong. \[-1\mapsto[-1,0],\qquad +1\mapsto[0,1],\qquad 0\mapsto[0,0].\]

• Asteroids. Let \(r\) be the primitive reward and output \([r_1,r_2,r_3]\) defined by \[r_1=\begin{cases} 20, & (r-20)\bmod 50=0,\\ 0, & \text{otherwise}, \end{cases} \qquad r_2=\begin{cases} 50, & (r-r_1-50)\bmod 100=0,\\ 0, & \text{otherwise}, \end{cases} \qquad r_3=r-r_1-r_2.\]

• MsPacman. Let \(r\) be the primitive reward and output \([r_1,r_2,r_3,r_4]\) defined by \[r_1=\begin{cases} 10, & (r-10)\bmod 50=0,\\ 0, & \text{otherwise}, \end{cases} \qquad r_2=\begin{cases} 50, & (r-r_1-50)\bmod 100=0,\\ 0, & \text{otherwise}, \end{cases}\] \[r_3=\begin{cases} 100, & (r-r_1-r_2-100)\bmod 200=0,\\ 0, & \text{otherwise}, \end{cases} \qquad r_4=r-r_1-r_2-r_3.\]

• Gopher. Let \(r\) be the primitive reward and output \([r_1,r_2]\) defined by \[r_1=\begin{cases} 20, & (r-20)\bmod 100=0,\\ 0, & \text{otherwise}, \end{cases} \qquad r_2=r-r_1.\]

• UpNDown. Let \(r\) be the primitive reward and output \([r_1,r_2,r_3]\) defined by \[r_1=\begin{cases} 10, & (r-10)\bmod 100=0,\\ 0, & \text{otherwise}, \end{cases} \qquad r_2=\begin{cases} 100, & (r-r_1-100)\bmod 200=0,\\ 0, & \text{otherwise}, \end{cases} \qquad r_3=r-r_1-r_2.\]

15.5.0.2 Scalarization rule.

Following [16], we derive a scalar control signal from the multi-dimensional return by simple summation of its components. Concretely, letting \(Z(s,a)\in\mathbb{R}^d\) denote the vector-valued return, we define the scalar action-value used for control as \[Q(s,a) \;=\; \mathbb{E}\!\left[\mathbf{1}^\top Z(s,a)\right],\] where \(\mathbf{1}\in\mathbb{R}^d\) is the all-ones vector.

15.5.0.3 Evaluation protocol.

Alongside the training environments, we run a separate pool of \(8\) evaluation environments in parallel, using the same evaluation implementation as in the public codebase of [37]. During evaluation, the agent follows a purely greedy policy (no exploration noise). This design choice is deliberate for performance: evaluation runs continuously without interrupting training. We additionally report normalized scores to aggregate results across environments. We compute \[\mathrm{Norm}(\pi)\;:=\;100 \times \frac{J(\pi)-J(\pi_{\mathrm{rand}})}{J(\pi_{\mathrm{ref}})-J(\pi_{\mathrm{rand}})},\] where \(J(\cdot)\) is the undiscounted evaluation return. The constants \(J(\pi_{\mathrm{rand}})\) and \(J(\pi_{\mathrm{ref}})\) are the random and reference scores used for normalization. We do not interpret these as human-normalized scores because our Atari setup does not use sticky actions (we follow [16]). We therefore refer to them simply as normalized scores.

16 Chain environment↩︎

Section 16.3 describes the standard one-sampled distributional TD regime (Algorithm [alg:one95sample95chain95td]). Section 16.4 describes the near-exact variant that explicitly constructs the transition-mixture Bellman target (Algorithm [alg:near95exact95chain95td]). Unless stated otherwise, distributional accuracy is evaluated at the start state using the same Monte Carlo protocol as for Maze (Section 15.4).

16.1 Environment description↩︎

We reproduce and generalize the chain environment of [7] as used by [23]. This setup is used purely for policy evaluation. The MDP has states \(\{s_0,\dots,s_{K-1}\}\) with initial state \(s_0\) and terminal (absorbing) state \(s_{K-1}\). A schematic is provided in Fig. 2. For any nonterminal state \(s_i\) with \(i\le K-2\) and action \(a\in\{\texttt{fwd},\texttt{bwd}\}\), the next state satisfies \[S_{t+1} = \begin{cases} s_{i+1}, & \text{with probability } p(a),\\ s_0, & \text{with probability } 1-p(a), \end{cases} \qquad p(\texttt{fwd})=0.9,\;\; p(\texttt{bwd})=0.1,\] and if \(S_t=s_{K-1}\) then \(S_{t+1}=s_{K-1}\). Episodes terminate upon first entering \(s_{K-1}\), and we do not emit further rewards nor bootstrap beyond termination. In the original scalar chain, the reward is \(+1\) upon entering the terminal state, \(-1\) on transitions that reset to \(s_0\), and \(0\) otherwise.

In our multivariate variant, the transition dynamics are unchanged and only the reward is modified. Here the return dimension is \(d=K\). We represent each state \(s_i\) by a one-hot vector \(e_i\in\mathbb{R}^{K}\) and emit the positional reward vector \(R_{t+1} = e_{S_{t+1}}\), which is the one-hot embedding of the entered next state. This yields a \(K\)-dimensional return that tracks discounted state occupancies in the spirit of successor features [57] while preserving the original chain’s stochasticity. In our experiments we use a scalar discount \(\gamma\) and instantiate the matrix discount as \(\Gamma=\gamma I_d\).

16.2 Empirical setup↩︎

To model the return distribution we follow [23] and, for each state–action pair, parametrize an empirical distribution with a fixed set of particles. Since the chain is tabular and small, we do not use a neural network. We initialize the particles randomly and update them using a distributional divergence, as described in the two subsequent sections. Default hyper-parameters are reported in Table 6.

16.2.0.1 Evaluation protocol.

We follow the same Monte Carlo evaluation protocol as for the Maze environments (Section 15.4). Concretely, we generate Monte Carlo (MC) rollouts starting from the initial state \(S_0=s_0\) and unroll the evaluation policy (\(\pi(s)\equiv\texttt{fwd}\)) until termination at \(s_{K-1}\). Each rollout yields a discounted return vector, and repeating this procedure produces an empirical return distribution at the start state. We then query the distributional critic at \((s_0,\texttt{fwd})\) to obtain its predicted particle distribution and compare the two, using the same distributional metric as in the Maze evaluation.

16.3 One-sampled distributional TD↩︎

In distributional RL, the (population) Bellman target at \((s,a)\) is a mixture law induced by transition and policy randomness. In most implementations, however, TD bootstrapping replaces this mixture by a one-sample instantiation based on a single sampled successor \((S_{t+1},A_{t+1})\). This is the setting discussed in Sec. 4.4.

We begin with this standard one-sampled distributional TD bootstrapping regime in our tabular chain setting. The purpose of this subsection is only to specify the target distribution produced by a single sampled successor, which we will contrast in the next subsection with a closer approximation of the full (mixture) distributional Bellman target.

For clarity, we describe this mode for the evaluation policy that always goes right, i.e.\(\pi(s)\equiv \texttt{fwd}\) for all nonterminal states. At each update, we only modify the particles at the current tabular entry \((S_t,\texttt{fwd})\) using a single observed transition.

Consider one on-policy transition \((S_t,\texttt{fwd},R_{t+1},S_{t+1})\), where \(S_{t+1}\sim P(\cdot\mid S_t,\texttt{fwd})\), and take \(A_{t+1}=\texttt{fwd}\). We form a one-sample distributional target by bootstrapping from the single successor pair \((S_{t+1},\texttt{fwd})\): \[\hat{Z}_{t+1} \;:=\; R_{t+1} \;+\; \Gamma\, Z_{\phi}(S_{t+1},\texttt{fwd}),\] where \(\Gamma\in\mathbb{R}^{d\times d}\) is the (possibly matrix) discount. In our particle implementation, \(\hat{Z}_{t+1}\) is represented by the target particle set \[\hat{z}^{(n)}_{t+1} \;:=\; R_{t+1} \;+\; \Gamma\, z^{(n)}_{\phi}(S_{t+1},\texttt{fwd}), \qquad n=1,\dots,N,\] and no gradient flows through the bootstrap particles \(z^{(n)}_{\phi}(S_{t+1},\texttt{fwd})\) when updating the particles at \((S_t,\texttt{fwd})\). The TD update at time \(t\) then matches the empirical distribution at the current tabular entry \((S_t,\texttt{fwd})\) to the empirical distribution defined by \(\{\hat{z}^{(n)}_{t+1}\}_{n=1}^N\) using a distributional divergence.

16.4 Near exact distributional TD↩︎

Because the chain dynamics are known and finite, the (population) distributional Bellman target can be approximated directly, rather than instantiated from a single sampled successor as in Sec. 16.3. We exploit the fact that, for each state–action pair \((s,a)\), the next state has only two possible outcomes, so the mixture induced by the transition kernel can be constructed explicitly and represented with a fixed particle budget.

For clarity, we describe this mode for the evaluation policy that always goes right, i.e.\(\pi(s)\equiv \texttt{fwd}\) for all nonterminal states. In this regime, all nonterminal states are updated simultaneously in a batch sweep.

Fix a nonterminal state \(s=s_i\) with \(i\le K-2\) and take \(a=\texttt{fwd}\). Then the next state is \(s_{\mathrm{succ}}=s_{i+1}\) with probability \(p=0.9\) and \(s_{\mathrm{fail}}=s_0\) with probability \(1-p=0.1\). Let \(R(s,s')\in\mathbb{R}^d\) denote the (vector) reward emitted on transition \(s\to s'\). We form the two successor-backed-up target distributions \[Z_{\mathrm{succ}} \;:=\; R(s,s_{\mathrm{succ}}) \;+\; \Gamma\, Z_{\phi}(s_{\mathrm{succ}},\texttt{fwd}), \qquad Z_{\mathrm{fail}} \;:=\; R(s,s_{\mathrm{fail}}) \;+\; \Gamma\, Z_{\phi}(s_{\mathrm{fail}},\texttt{fwd}),\] and define the corresponding mixture target \[Z_{\mathrm{targ}} \;:=\; p\, Z_{\mathrm{succ}} \;+\; (1-p)\, Z_{\mathrm{fail}}.\] In our particle implementation, \(Z_{\mathrm{targ}}\) is represented by drawing \(N\) target particles from this mixture: we pool the particles from \(Z_{\mathrm{succ}}\) and \(Z_{\mathrm{fail}}\) and resample \(N\) of them with mixture weights \(p\) and \(1-p\). As in Sec. 16.3, no gradient flows through the bootstrap particles when updating the particles at \((s,\texttt{fwd})\).

This yields a close approximation of the full mixture Bellman target while keeping the update fully tabular and using a fixed particle budget per state.

17 More results and ablations↩︎

This section collects additional results and diagnostics that complement the main experiments. Section 17.1 reports update-time scaling with the number of particles. Section 17.2 provides additional chain results, including a max-slicing sanity check and Monte Carlo versus predicted return visualizations. Section 17.3 reports a Maze ablation on the choice of Cramér estimator, together with Monte Carlo versus predicted return visualizations. Section 17.4 reports per-environment Atari learning curves and studies the effect of the number of projection directions.

17.1 Computational cost↩︎

In Figure 11 we compare the wall-clock time of a single update iteration of Algorithm [alg:near95exact95chain95td] across the objectives considered in this work, while varying the number of particles used to model the return distribution. As expected, MMD and sliced-MMD exhibit a much steeper increase in runtime as the number of particles grows, which is consistent with the quadratic complexity of standard MMD estimators (\(\mathcal{O}(n^2)\)). In contrast, the other sliced variants scale more mildly with the number of particles, which is consistent with a near \(\mathcal{O}(n\log n)\) cost when their one-dimensional base computations admit efficient sorting-based implementations.

17.2 Chain environment↩︎

17.2.1 Testing gradient bias from max slicing↩︎

We now isolate the selection bias induced by max slicing by considering a chain setting in which the multivariate return distribution is effectively one-dimensional. We do so by enforcing that all rewards are colinear: there exists a fixed direction \(u\in\mathbb{R}^d\) such that \[R_{t+1} = r_{t+1}\,u \qquad \text{for all transitions,}\] for some scalar reward process \((r_t)_{t\ge 0}\). In our experiments, we take \(u=\mathbf{1}_d\) and replicate the scalar chain reward across all \(d\) coordinates, which is a direct generalization of the scalar chain setting of [23].

17.2.1.1 Consequence for max slicing.

When the return distribution is supported on the one-dimensional subspace \(\mathrm{span}\{u\}\), the max-sliced objective reduces to a constant multiple of a one-dimensional divergence and the selection mechanism behind Lemma 13 simply doesn’t apply. Indeed, since \(Z=\sum_{t\ge 0}\gamma^t R_{t+1}=Y\,u\) for some scalar random variable \(Y\), for any slicing direction \(\theta\in S^{d-1}\) we have \[\langle \theta, Z\rangle = \langle \theta, u\rangle\,Y.\] Thus the projected distributions are scaled versions of the same scalar law. For the divergences considered in this work, this scaling induces a deterministic multiplicative factor: for Wasserstein, the standard homogeneity property gives \(\mathbf{W}_p((S_s)_{\#}\mu,(S_s)_{\#}\nu)=s\,\mathbf{W}_p(\mu,\nu)\), and for Cramér we have \(\mathbf{C}_2((S_s)_{\#}\mu,(S_s)_{\#}\nu)=s\,\mathbf{C}_2(\mu,\nu)\), shown in Proposition 10. Consequently, for all \(\theta\), \[\Delta\!\big((P_\theta)_{\#}\mu,\,(P_\theta)_{\#}\nu\big) = |\langle \theta, u\rangle|\, \Delta\!\big(\mu^{(1)},\nu^{(1)}\big),\] where \(\mu^{(1)},\nu^{(1)}\) are the corresponding scalar return distributions and \(\Delta\in\{\mathbf{W}_p,\mathbf{C}_2\}\). Taking the supremum over \(\theta\) gives \[\sup_{\theta\in S^{d-1}} \Delta\!\big((P_\theta)_{\#}\mu,\,(P_\theta)_{\#}\nu\big) = \|u\|_2\, \Delta\!\big(\mu^{(1)},\nu^{(1)}\big),\] which depends only on \(u\) and not on the sample used to form the empirical sliced losses.

17.2.1.2 Empirical verification.

Figure 12 reports, for each objective, the empirical Wasserstein–\(2\) distance between Monte Carlo returns and the learned critic distribution at the initial state, averaged over \(5\) random seeds (with \(95\%\) confidence intervals). In this degenerate setting, max-sliced-\(\mathbf{C}_2\) performs on par with the objectives that satisfy (U), consistent with the fact that the max-slicing selection mechanism (and the associated gradient bias) does not apply here.

17.2.2 Monte Carlo returns vs predicted returns↩︎

Figure 13 provides a qualitative companion to the chain results reported in Section 5 (Fig. 3). For the same evaluation setting, we directly visualize the learned return particles at the initial state alongside Monte Carlo rollouts from the true environment. Notably, although the Wasserstein–\(2\) objective can collapse to a degenerate return distribution in this setting, its predicted mean remains qualitatively comparable to that of objectives that fit the distribution well.

17.3 Maze↩︎

17.3.1 Choice of Cramér estimator↩︎

As discussed in Section 8.2.1, the Cramér–\(2\) objective admits multiple empirical estimators. In addition to our default estimator, we further tested the CRPS-based alternative while keeping the rest of the training and evaluation protocol unchanged. Figure 14 shows that this choice has no noticeable effect on distributional matching accuracy across the three environments considered, and does not change the conclusions drawn from the main maze results.

17.3.2 Monte Carlo returns vs predicted returns↩︎

Figure 15 provides a qualitative companion to the Maze policy-evaluation results reported in Section 5 (Fig. 4). For each of the three Maze variants, we visualize the critic’s learned return particles at a fixed reference state alongside Monte Carlo return samples from the true environment, plotting two return coordinates (dimensions \(0\) and \(1\)). Blue points denote the particles predicted by the critic, while orange points are Monte Carlo returns. The figure illustrates that objectives satisfying (U)(notably sliced Cramér and sliced-MMD) closely match the Monte Carlo return distribution, whereas objectives that violate (U)can exhibit visible distributional mismatch.

17.4 Atari↩︎

17.4.1 Number of projection directions↩︎

Figure 16 complements the aggregated normalized-score results in Figure 6 by showing per-environment learning curves for the same runs. We vary the number of projection directions used by the sliced objectives and visualize how this choice affects raw control performance (undiscounted evaluation returns) over training. In this experiment, performance is remarkably robust to the number of projection directions, with strong results already obtained using a single projection direction.

18 LLM Usage↩︎

We used an LLM-based assistant to support the preparation of this paper. In particular, it was employed to (i) rephrase draft paragraphs for clarity and suggest alternative framings of related work, (ii) format proofs, explore directions, and verify intermediate steps, (iii) assist in debugging code, (iv) suggest LaTeX equation formatting, and (v) help identify relevant theoretical results in preceding works. All core research contributions, including the development of theoretical results, algorithms, and experiments, were carried out by the authors.

References↩︎