Multivariate Distributional Reinforcement Learning Using Sliced Divergences

Distributional reinforcement learning (DRL) models the full return distribution rather than expectations, but extending it to multivariate settings remains challenging. Many common metrics do not naturally generalize beyond one dimension or lose computational tractability, and the multivariate case introduces additional difficulties such as general matrix discounting, for which no contraction results are available. We introduce Sliced Distributional Reinforcement Learning (SDRL), which lifts tractable one-dimensional divergences to multivariate return distributions via projections. We prove Bellman contraction for uniform slicing under shared scalar discounting, and introduce a maximum-slicing variant with contraction under general dense discount matrices. SDRL supports a broad class of base divergences; we analyze Wasserstein, Cramér, and Maximum Mean Discrepancy (MMD), and characterize which SDRL variants suit the standard single-sample Bellman update used in distributional RL. We evaluate SDRL on a toy chain problem and a gridworld image-based environment as well as a subset of Atari games. Code is available at https://github.com/BaptisteDebes/SlicedDistributionalRL

Reinforcement learning trains agents to make decisions by learning from rewards. Most methods work with average expected rewards, but tracking the full distribution of possible outcomes can lead to richer and more robust behavior. Extending this distributional approach to high-dimensional reward settings has proven difficult: most ways of comparing distributions either become computationally intractable or lose their theoretical guarantees. We introduce Sliced Distributional Reinforcement Learning (SDRL), which addresses this by reducing multivariate distribution comparisons to a collection of one-dimensional problems through random projections. The approach is computationally lightweight and comes with theoretical guarantees: we prove Bellman contraction even under general matrix-valued discounting, a setting not covered by existing analyses. We also characterize which variants remain compatible with the single-sample training regime standard in practice, and validate SDRL on tabular, pixel-based, and Atari environments.