Distributional Reinforcement Learning with Regularized Wasserstein Loss

Summary

This paper proposes a new RL algorithm that leverages Sinkhorn divergence, which they claimed as a regularized Wasserstein loss. Theoretically, they showed the contraction properties that align with the interpolation nature of Sinkhorn divergence between Wasserstein distance and MMD. Empirically, it outperforms or matches existing algorithms on many atari games.

Strengths

- The algorithm is clearly proposed and can be easily understood. - The authors study both theoretical (mostly contraction property) and empirical aspects of the proposed algorithm. - The experimental results seem comprehensive (on the entire atari games). They show that the algorithm performs better on complicated games.

Weaknesses

After reading the paper I am still not sure about the benefits of using sinkhorn divergence, although I agree that the authors established some theoretical results (which I also have some concerns about below) and showed experimental results. My question is more intuitive---why should we pick Sinkhorn over others at a high level of intuition? I appreciate that the authors have spent many words explaining it from various perspectives. However, I didn't totally understand them and am not convinced at this moment. Some of my specific confusion are below: - line 48-51: why does using samples instead of pre-specified statistics have better accuracy? Is there any reference? - line 55: how is the regularization aligning with the max entropy principle? This point is mentioned multiple times in the paper but not explained in detail. - line 53: "smoother" is very vague in this context. I would suggest using some different words. As for the theoretical results, I am not sure how novel they are. For instance, the authors show convergence rate and sample complexity only under limit (when epsilon is approaching zero or infinity). Considering that the algorithm reduces to either QRDQN or MMD under the limit, the results are not so surprising and are actually quite natural. In addition, I didn't find general convergence results (I mean under any value of epsilon). Hence I am doubting if there is anything significant enough in the theory. For the experiments, it generally looks good to me. I only have a minor comment for now: figure 2 seems a bit misleading--specifically, I feel plotting the ratio "Sinkhorn/QRDQN" is misleading. Think about two cases: if Sinkhorn is twice QRDQN, then the value will be 200%; on the other hand if QRDQN is twice Sinkhorn, then the ratio will be 50%. The two cases are symmetric but 200% will be visually four times higher than 50% in the plot. This may be why the Sinkhorn visually has a huge advantage over others in the figure.

Questions

Line 319: the algorithm is much better on complicated games, which is very interesting. In other words, the algorithm is probably inferior on easy games. Do you have more explanation on this?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

I didn't find any potential negative societal impact.

Summary

• This paper proposes a novel distributional RL algorithm, called SinkhornDRL, which interpolates between Wasserstein distance and MMD. The authors aim to estimate the distribution using unrestricted statistics, enhancing stability and facilitating extension to multi-dimensional reward settings. The authors also provide some theoretical guarantees on convergence of SinkhornDRL. The proposed method shows decent experimental performance in 55 Atari games and several multi-dimensional reward settings.

Strengths

• The paper reviews the relevant literatures in distributional RL and analyzes the new Sinkhorn divergence metric as an alternative to Wasserstein loss. The author provides sufficient background and advantages of Sinkhorn divergence, demonstrating a well-motivated approach. • Although the evaluation scores are presented for 40 million frames rather than 200 million, the authors ensure transparency and reproducibility by reporting raw scores for performance comparisons. The authors also present a significant amount of experiment and ablation studies to provide sufficient experimental results for the effectiveness of their algorithm.

Weaknesses

• Typos: In line 52, Sinkrhorn -> Sinkhorn • The text and figures in the paper are quite dense and difficult to read. Some parts seem unnecessary to be included in the main paper. o Algorithm 1 lacks certain technical details, and may rather be placed in Appendix. In addition, explaining the details of Algorithm 2 on Line 769 would help a lot in implementing a practical algorithm. o The text in Figure 2 is not legible. It would be better to include a table showing the mean and median of the best HNS scores, as commonly done in other distributional RL papers [1,2]. • There may be a technical error regarding Theorem 1 (3). The authors state in Eq (33) that a universal upper bound $\bar{\Delta}_{\epsilon}(a,\alpha)$ is strictly less than 1. But if $\mu$ and $\nu$ are close enough, even if non-trivial, doesn’t the scaling factor become 1 by taking the supremum? In Line 660, the authors recognize the non-expansion case, but it seems to rely on the strong assumption that the set $\{\lambda_{\epsilon}(U,V)\}$ is finite. This assumption essentially implies that the set $\{\lambda_{\epsilon}(\(mathcal{T}^{\pi}_D)^n U, (\mathcal{T}^{\pi}_D)^n V )\}$ where $n \in \mathbb{N}$ is finite. [1] Dabney, W., Rowland, M., Bellemare, M. G. & Munos, R. Distributional Reinforcement Learning with Quantile Regression. Arxiv (2017). [2] Hessel, M. et al. Rainbow: Combining Improvements in Deep Reinforcement Learning. arXiv (2017) doi:10.48550/arxiv.1710.02298.

Questions

• Could the author elaborate on the statement in Lines 129-134 that 'SinkhornDRL inherently captures the spatial and geometric layout of return distributions'? Does this imply that the KL divergence term can leverage a richer representation of data geometry? • The definition of Sinkhorn divergence for multi-dimensional reward settings is unclear. Specifically, in Line 756, the cost function applies equal weights among reward sources. Can this be generalized to a weighted sum?

Rating

5

Confidence

5

Soundness

3

Presentation

2

Contribution

3

Limitations

• In line 22, the reference [11] is cited as a risk-sensitive control, but this is close to exploration method for distributional RL.

Summary

This paper introduces Sinkhorn Distributional Reinforcement Learning (SinkhornDRL), a new algorithm designed to address the limitations of current distributional RL methods, particularly those relying on quantile regression. Existing methods often struggle with accurately capturing the characteristics of return distributions and extending to scenarios with multi-dimensional rewards. SinkhornDRL leverages Sinkhorn divergence, a regularized Wasserstein loss, to minimize the difference between current and target Bellman return distributions. This approach combines the geometric advantages of Wasserstein distance with the computational efficiency of Maximum Mean Discrepancy (MMD). The paper provides theoretical proof of SinkhornDRL's contraction properties, demonstrating its convergence behavior and relationship to other distance metrics. Empirical evaluations on the Atari games suite show that SinkhornDRL consistently outperforms or matches existing algorithms, especially in settings with multi-dimensional rewards.

Strengths

- Introduces a new family of distributional RL algorithms based on Sinkhorn divergence, expanding the toolkit for researchers and practitioners. - Provides theoretical analysis of Sinkhorn divergence in the context of distributional RL, including convergence guarantees. - Demonstrates the effectiveness of SinkhornDRL through extensive experiments on a standard benchmark, showing superior performance in many cases. - Specifically tackles the issues of inaccurate distribution capture and difficulty with multi-dimensional rewards that plague quantile regression methods.

Weaknesses

- SinkhornDRL introduces some additional computational cost compared to simpler methods like C51 and QR-DQN. - The algorithm requires tuning of additional hyperparameters (e.g., the regularization strength), which might require extra effort in practice. - The paper acknowledges that a deeper connection between theoretical properties of divergences and practical performance in specific environments remains an open question. - Some typos: - line 151 and line 647 Sinkrhon - line 163 supremal from - line 787 numeral

Questions

Could the authors address the following points: - Are there existing works on distributional reinforcement learning that use the entropic regularized Wasserstein distance? - What specific advantages does the Sinkhorn divergence offer over the entropic regularized Wasserstein distance in the context of this study?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

NA