Technical Summary: Multivariate Distributional Reinforcement Learning Using Sliced Divergences

Problem Statement

Distributional Reinforcement Learning (DRL) models the full return distribution rather than just its expectation, offering theoretical and empirical advantages over classical expected RL. However, extending DRL to multivariate settings (where rewards are $d$-dimensional vectors, $d > 1$) remains a significant challenge.

Existing approaches face three primary limitations:

1. Computational Intractability: Common metrics like the Wasserstein distance suffer from high computational costs in high dimensions (typically $O(n^3 \log n)$ for optimal transport solvers) and poor statistical rates that depend on the ambient dimension.
2. Lack of Contraction Guarantees: While scalar DRL has well-established contraction results under the Bellman operator, multivariate settings with general matrix discounting (where the discount factor is a state-action dependent matrix $\Gamma(s, a) \in \mathbb{R}^{d \times d}$) lack rigorous contraction guarantees. Existing methods often fail to contract under anisotropic (non-uniform) discounting.
3. Gradient Bias in Stochastic Training: Many divergences, particularly Wasserstein, fail the unbiased sample gradient property (U) when the Bellman target is instantiated from a single sampled successor (standard TD learning). This leads to biased gradients and suboptimal distributional matching.

Methodology: Sliced Distributional RL (SDRL)

The authors introduce Sliced Distributional Reinforcement Learning (SDRL), a framework that lifts tractable one-dimensional divergences to multivariate return distributions via random projections (slicing).

Core Mechanism

Instead of computing a high-dimensional divergence directly, SDRL projects the multivariate distributions onto random directions $\theta \in S^{d-1}$, computes a base divergence $\Delta$ in one dimension, and aggregates the results.

• Uniform Slicing: Averages the base divergence over $L$ random directions:
  $$S\Delta_p^p(\mu, \nu) = \int_{S^{d-1}} \Delta^p((P_\theta)_\#\mu, (P_\theta)_\#\nu) d\sigma(\theta)$$
• Max-Slicing (MSDRL): Optimizes over the most discriminative direction to obtain stronger contraction guarantees for general matrix discounts:
  $$MS\Delta(\mu, \nu) = \sup_{\theta \in S^{d-1}} \Delta((P_\theta)_\#\mu, (P_\theta)_\#\nu)$$

The framework supports a broad class of base divergences, specifically analyzing Wasserstein ($W_p$), Cramér distance ($C_2$), and Maximum Mean Discrepancy (MMD).

Theoretical Contributions

The paper establishes rigorous theoretical foundations for SDRL:

1. Metric Properties: Proves that both uniform and max-sliced lifts preserve the metric axioms if the base divergence is a metric.
2. Contraction Guarantees:
   • Uniform Slicing: Proves Bellman contraction under shared scalar discounting ($\Gamma = \gamma I$) for any base divergence satisfying translation non-expansion, scale contraction, and mixture $p$-convexity.
   • Max-Slicing: Introduces a novel contraction result for general dense matrix discounting (anisotropic updates). It proves that max-slicing yields contraction under the operator norm of the discount matrix, a setting where uniform slicing and standard MMD often fail to contract.
3. Sample Complexity: Demonstrates that uniform slicing inherits dimension-free convergence rates from the base 1D divergence, avoiding the curse of dimensionality. Max-slicing rates scale polynomially with dimension ($O(d \log n / n)$), which is significantly better than exact high-dimensional OT.
4. Gradient Bias Analysis:
   • Uniform Slicing: Preserves the unbiased sample gradient property (U) if the base divergence satisfies it. This makes Sliced Cramér and Sliced MMD compatible with standard one-sample TD bootstrapping.
   • Max-Slicing: Fails property (U) due to the selection bias introduced by the maximization step. The direction selected depends on the specific sample batch, leading to biased gradients in standard TD settings.

Experimental Results

The authors evaluate SDRL on three benchmarks: a tabular chain MDP, a pixel-based maze environment, and a subset of Atari games.

1. Chain Environment (Policy Evaluation):

   • Compares one-sample TD (standard bootstrapping) vs. near-exact TD (explicit mixture construction).
   • Finding: Objectives satisfying property (U) (e.g., Sliced Cramér, Sliced MMD) achieve accurate distributional matching in the one-sample regime. Objectives violating (U) (e.g., Sliced Wasserstein, Max-Sliced variants) degrade significantly unless the full mixture target is used.
   • Max-Slicing: Performs well in the degenerate one-dimensional case (where selection bias vanishes) but fails in general multivariate settings under one-sample TD.

2. Pixel-Based Maze (Policy Evaluation):

   • Evaluates distributional accuracy using the empirical Wasserstein-2 distance between Monte Carlo returns and the critic's prediction.
   • Finding: Sliced Cramér and Sliced MMD perform on par with their non-sliced counterparts and significantly outperform Sliced Wasserstein in the one-sample TD regime, confirming the importance of property (U).

3. Atari Games (Control):

   • Evaluates control performance on decomposed multi-dimensional rewards.
   • Finding: While objectives satisfying (U) provide accurate distributional estimates, Sliced Wasserstein-2 surprisingly achieves strong control performance despite poor distributional matching and biased gradients. This suggests that for control, accurate estimation of the return expectation may be more critical than the fidelity of the full distribution, even if the distributional objective is biased.

Significance and Claims

The paper claims to provide a principled framework for extending DRL to multivariate settings while maintaining computational tractability and theoretical guarantees.

• Practical Recommendation: The authors identify Sliced Cramér as a strong default for multivariate distributional learning. It is computationally efficient ($O(n \log n)$), satisfies the unbiased gradient property (enabling standard TD training), and offers contraction guarantees in the standard isotropic setting.
• Theoretical Insight: The work clarifies the trade-offs between different divergence choices. It highlights that while Max-Slicing solves the contraction problem for general matrix discounts, it introduces gradient bias that makes it unsuitable for standard one-sample TD without modification.
• Gap Identification: The paper concludes that a divergence that is simultaneously provably contractive under general anisotropic discounts (in a norm-only sense), satisfies the unbiased sample gradient property, and offers tractable estimation remains an open question.

In summary, SDRL offers a scalable and theoretically grounded approach to multivariate DRL, with Sliced Cramér emerging as the most robust practical choice for standard reinforcement learning pipelines.