Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributional Reinforcement Learning with Regularized Wasserstein Loss

Ke Sun, Yingnan Zhao, Wulong Liu, Bei Jiang, Linglong Kong

TL;DR

The introduced SinkhornDRL enriches the family of distributional RL algorithms, contributing to interpreting the algorithm behaviors compared with existing approaches by contributing to interpreting the algorithm behaviors compared with existing approaches by investigation into their relationships.

Abstract

The empirical success of distributional reinforcement learning (RL) highly relies on the choice of distribution divergence equipped with an appropriate distribution representation. In this paper, we propose \textit{Sinkhorn distributional RL (SinkhornDRL)}, which leverages Sinkhorn divergence, a regularized Wasserstein loss, to minimize the difference between current and target Bellman return distributions. Theoretically, we prove the contraction properties of SinkhornDRL, aligning with the interpolation nature of Sinkhorn divergence between Wasserstein distance and Maximum Mean Discrepancy (MMD). The introduced SinkhornDRL enriches the family of distributional RL algorithms, contributing to interpreting the algorithm behaviors compared with existing approaches by our investigation into their relationships. Empirically, we show that SinkhornDRL consistently outperforms or matches existing algorithms on the Atari games suite and particularly stands out in the multi-dimensional reward setting. \thanks{Code is available in \url{https://github.com/datake/SinkhornDistRL}.}.

Distributional Reinforcement Learning with Regularized Wasserstein Loss

TL;DR

The introduced SinkhornDRL enriches the family of distributional RL algorithms, contributing to interpreting the algorithm behaviors compared with existing approaches by contributing to interpreting the algorithm behaviors compared with existing approaches by investigation into their relationships.

Abstract

The empirical success of distributional reinforcement learning (RL) highly relies on the choice of distribution divergence equipped with an appropriate distribution representation. In this paper, we propose \textit{Sinkhorn distributional RL (SinkhornDRL)}, which leverages Sinkhorn divergence, a regularized Wasserstein loss, to minimize the difference between current and target Bellman return distributions. Theoretically, we prove the contraction properties of SinkhornDRL, aligning with the interpolation nature of Sinkhorn divergence between Wasserstein distance and Maximum Mean Discrepancy (MMD). The introduced SinkhornDRL enriches the family of distributional RL algorithms, contributing to interpreting the algorithm behaviors compared with existing approaches by our investigation into their relationships. Empirically, we show that SinkhornDRL consistently outperforms or matches existing algorithms on the Atari games suite and particularly stands out in the multi-dimensional reward setting. \thanks{Code is available in \url{https://github.com/datake/SinkhornDistRL}.}.
Paper Structure (39 sections, 4 theorems, 42 equations, 11 figures, 4 tables, 3 algorithms)

This paper contains 39 sections, 4 theorems, 42 equations, 11 figures, 4 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Preliminary Knowledge
  3. Distributional Reinforcement Learning
  4. Divergences between Measures
  5. Related Work
  6. Sinkhorn Distributional RL (SinkhornDRL)
  7. Sinkhorn Divergence and New Convergence Properties in Distributional RL
  8. Extension to Multi-dimensional Return Distributions
  9. SinkhornDRL Algorithm and Approximation
  10. Experiments
  11. Performance of SinkhornDRL
  12. Sensitivity Analysis and Computational Cost
  13. Modeling Joint Return Distribution for Multi-Dimensional Reward Functions
  14. Conclusion, Limitations and Future Work
  15. Appendix
  16. ...and 24 more sections

Key Result

Proposition 1

$\mathfrak{T}^{\pi}$ is non-expansive under $\text{MI}_\Pi^\infty$ for any non-trivial joint distribution $\Pi$.

Figures (11)

  • Figure 1: Mean (left), Median (middle), and IQM (5$\%$) (right) of Human-Normalized Scores (HNS) summarized over 55 Atari games. We run 3 seeds for each algorithm.
  • Figure 2: Ratio improvement of return for SinkhornDRL over QR-DQN (left) and MMD-DQN (right) averaged over 3 seeds. The ratio improvement is calculated by (SinkhornDRL - QR-DQN) / QR-DQN in (a) and (SinkhornDRL - MMD-DQN) / MMD-DQN in (b), respectively.
  • Figure 3: Sensitivity analysis of SinkhornDRL on Breakout and Seaquest in terms of $\varepsilon$, number of samples, and number of iteration $L$. Learning curves are reported over three seeds.
  • Figure 4: Performance of SinkhornDRL on six Atari games with multi-dimensional reward functions.
  • Figure 5: Optimal transport plans for via Sinkhorn Iterations in SinkhornDRL on three Atari games. The first row denotes the (two-dimensional) spatial transport plans across different data points, while the second row represents the heat map of the obtained transport plan (optimal coupling).
  • ...and 6 more figures

Theorems & Definitions (6)

  • Proposition 1
  • Proposition 2
  • Theorem 1
  • Corollary 1
  • proof
  • proof