Research map / Reinforcement learning / Policy gradient and actor-critic
Distributional Reinforcement Learning with Regularized Wasserstein Loss
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.
Proposes SinkhornDRL, a distributional RL algorithm that minimizes the Sinkhorn divergence between particle-based return distributions and Bellman targets, with contraction theory and Atari experiments.
Problem
Quantile-regression distributional RL suffers from quantile non-crossing and is hard to extend to multi-dimensional rewards, since computing multi-dimensional Wasserstein distances is intractable.
Approach
Use the entropically regularized Sinkhorn divergence (interpolating between Wasserstein and MMD) as the loss over unrestricted particle representations of return distributions, with an unrectified kernel cost -||x-y||^2, and prove contraction of the distributional Bellman operator for all regularization strengths.
Key results
- SinkhornDRL achieves state-of-the-art mean, median, and IQM (5%) human-normalized scores across most training phases on 55 Atari games (3 seeds, 40M frames).
- SinkhornDRL outperforms MMD-DQN in most of six multi-dimensional reward Atari games.
- SinkhornDRL adds roughly 20% computational cost over MMD-DQN.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace MMD divergence with Sinkhorn divergence | SinkhornDRL vs MMD-DQN · Atari 55 games | Mean, median, and IQM (5%) HNS | Mixed |
| Replace MMD divergence with Sinkhorn divergence | SinkhornDRL vs MMD-DQN · Atari games with multi-dimensional rewards (six games) | performance | Helped |
| Replace MMD divergence with Sinkhorn divergence | SinkhornDRL vs MMD-DQN · Atari 55 games | computational cost: around 20% | Hurt |
| Use only N=2 generated samples (ablation) | SinkhornDRL · Seaquest | training stability | Hurt |
| Increase the number of generated samples excessively (ablation) | SinkhornDRL · Atari games | performance | Hurt |
| Use an overly large or small smoothing hyperparameter ε (ablation) | SinkhornDRL · Breakout and Seaquest | performance | Hurt |
Surprising findings
- SinkhornDRL converges more slowly early in training than MMD-DQN, attributed to its larger contraction factor.
- Too large or too small regularization epsilon, or too few particles, causes numerical instability and degraded performance.
Limitations
- Increased computational cost relative to MMD-DQN and other baselines.
- Requires tuning additional hyperparameters (epsilon, L, N), sensitive to them.
- No deeper connection established between divergence properties and practical performance in a given environment.
Experimental setup
Datasets: Atari 2600 (55 games)
Benchmarks: Atari 2600 (55 games)
Metrics: mean human-normalized score median human-normalized score IQM human-normalized score per-game return ratio improvement
Baselines: DQN C51 QR-DQN MMD-DQN
Builds on: QR-DQN MMD-DQN C51 Sinkhorn iterations
Closest papers by meaning
- Intrinsic Benefits of Categorical Distributional Loss: Uncertainty-aware Regularized Exploration in Reinforcement Learning (NeurIPS 2025 · similarity 0.88)
- Categorical Distributional Reinforcement Learning with Kullback-Leibler Divergence: Convergence and Asymptotics (ICML 2025 · similarity 0.88)
- Near-Minimax-Optimal Distributional Reinforcement Learning with a Generative Model (NeurIPS 2024 · similarity 0.88)
- Foundations of Multivariate Distributional Reinforcement Learning (NeurIPS 2024 · similarity 0.87)
- Bellman Unbiasedness: Toward Provably Efficient Distributional Reinforcement Learning with General Value Function Approximation (ICML 2025 · similarity 0.86)
- More Benefits of Being Distributional: Second-Order Bounds for Reinforcement Learning (ICML 2024 · similarity 0.86)
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.