Research map / Reinforcement learning / Policy gradient and actor-critic
Foundations of Multivariate Distributional Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Gives the first oracle-free, tractable, provably convergent multivariate distributional dynamic programming and TD-learning algorithms, using MMD-projected categorical representations and a mass-1 signed-measure relaxation to obtain an affine projection.
Problem
Multivariate distributional RL (learning the joint return distribution for vector-valued rewards) lacks tractable algorithms with convergence guarantees; existing methods either lack theory or require intractable maximum-likelihood oracles, and categorical TD fails to analyze when reward dimension e
Approach
Project Bellman backups onto categorical distributions with fixed or state-dependent supports using the MMD (energy-distance kernel) via a QP; in TD learning, relax the projection to mass-1 signed measures, which makes the projection affine and yields a provably convergent stochastic approximation. Also proposes a randomized particle DP operator with high-probability guarantees.
Key results
- Proves geometric convergence of projected categorical multivariate distributional DP via MMD nonexpansion of the projection, and that the categorical fixed point approximates the true return distribution at rate sqrt(mesh) with uniform supports recovering the m^{-1/2} rate of scalar categorical DRL.
- Shows the categorical MMD projection is not affine for d>=2, and proves the signed-measure relaxed categorical TD converges with probability 1 to its fixed point.
- In simulations on random tabular MDPs (d=2, 3), categorical TD accuracy improves sharply with more atoms while EWP TD plateaus, which the authors attribute to local minima of the non-convex MMD objective.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase the number of atoms in categorical TD | Categorical TD · Random MDPs with 2-dimensional cumulants; held-out reward functions | return-distribution accuracy | Helped |
| increase the number of atoms in EWP TD | EWP TD · Random MDPs with 2-dimensional cumulants; held-out reward functions | return-distribution accuracy | No clear effect |
| use categorical TD instead of EWP TD | Categorical TD and EWP TD · Random MDPs with 2-dimensional cumulants; held-out reward functions | return-distribution accuracy | Mixed |
| use categorical TD instead of EWP TD | Categorical TD and EWP TD · Random MDPs with 3-dimensional cumulants; held-out reward functions | return-distribution accuracy | Mixed |
Surprising findings
- Standard analysis of categorical TD fails when reward dimension exceeds 1, since the categorical MMD projection is not affine in that case.
- EWP TD performs well with few atoms but does not improve much with more atoms, plausibly due to local minima.
Limitations
- Categorical representations are hard to scale to high-dimensional cumulants since support grows exponentially with d (cubic in the d=3 case with grids).
- EWP guarantees hold only with high probability and are weaker; EWP TD convergence is not established.
- Experiments are restricted to small tabular MDPs with d=2 and d=3.
Experimental setup
Datasets: random tabular MDPs (3-state MDPs; random MDPs with Dirichlet transitions and uniform-prior cumulants)
Benchmarks: random tabular MDPs with d=2 and d=3 cumulants (zero-shot scalar return distribution prediction)
Metrics: MMD / accuracy of zero-shot return distributions with 95% CI
Baselines: EWP TD (equally-weighted particle TD with MMD loss)
Builds on: categorical distributional RL (C51) energy distance / MMD kernels from semimetrics equally-weighted particle TD (Zhang et al.) distributional successor measure asynchronous stochastic approximation TD analysis
Closest papers by meaning
- Quotient-Categorical Representations for Bellman-Compatible Average-Reward Distributional Reinforcement Learning (NeurIPS 2026 · similarity 0.89)
- Categorical Distributional Reinforcement Learning with Kullback-Leibler Divergence: Convergence and Asymptotics (ICML 2025 · similarity 0.89)
- A Finite Sample Analysis of Distributional TD Learning with Linear Function Approximation (NeurIPS 2025 · similarity 0.89)
- Near-Minimax-Optimal Distributional Reinforcement Learning with a Generative Model (NeurIPS 2024 · similarity 0.87)
- Intrinsic Benefits of Categorical Distributional Loss: Uncertainty-aware Regularized Exploration in Reinforcement Learning (NeurIPS 2025 · similarity 0.87)
- Distributional Reinforcement Learning with Regularized Wasserstein Loss (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.