Distributional Reinforcement Learning with Regularized Wasserstein Loss

Thank you for taking the time to review our paper. We appreciate your positive assessment and insightful feedback, and we would like to address the concerns you raised in your review.

My question is more intuitive---why should we pick Sinkhorn over others at a high level of intuition?

We summarize the explanations in the introduction to provide a high-level intuition. Exiting distributional RL algorithms that rely on Wasserstein distance often struggle with 1) inaccuracies in representing distribution using pre-specified statistics, i.e., quantiles, and 2) difficulties in extending to the multi-dimensional reward setting. Sinkhorn divergence stands out for efficiently approximating multi-dimensional Wasserstein distance (by introducing an entropic regularization), simultaneously overcoming the two limitations. Additionally, the smoother transport plan and induced smoothness by the entropic regularization can also benefit the optimization in the RL learning process, making it a preferred choice.

line 48-51: why does using samples instead of pre-specified statistics have better accuracy? Is there any reference?

Yes, using samples provides a more direct and accurate representation of return distributions. The MMD-DQN paper [1] discussed this extensively, and a more recent study [2] also adopts sample-based representation. These approaches, rooted in kernel methods, contrast with our method's foundation in optimal transport, offering a different approach to compare the distribution difference.

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.

According to the Sinkhorn divergence literature, the KL regularization in essence amounts to vanilla entropy regularization in terms of the joint distribution / optimal coupling. This is because:

where $\mathcal{H}$ is the entropy. Thus, the objective function is equivalent to

where the $\mathcal{H}(\mu)$ and $\mathcal{H}(\nu)$ are constant for fixed marginal distributions, which are independent of the minimizer. In this equivalent objective function form, minimizing the Sinkhorn divergence also encourages to maximize the entropy $\mathcal{H}(\Pi)$, leading to a more uniformly distributed optimal transport plan / joint distribution. This aligns well with the maximum entropy principle.

line 53: "smoother" is very vague in this context. I would suggest using some different words.

Thanks for bringing this to our attention. The term "smoother" refers to the effect of the regularization in Sinkhorn divergence, which encourages a more uniformly distributed transport plan, as analyzed above. This more uniformly distributed transport plan is in contrast to the potentially sparse plan resulting from optimizing the unregularized Wasserstein distance. We will refine this terminology and provide a clearer explanation in the revised version.

As for the theoretical results, I am not sure how novel they are. ... 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) ...

We want to clarify that Theorem 1(3) is precisely the contraction conclusion under the general $\epsilon \in (0, +\infty)$, which is the main theoretical contribution of this paper. We encourage a further review of Theorem 1(3) and the subsequent proof sketch on Page 5 of our paper.

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.

We acknowledge the asymmetry of the ratio improvement/percentage increase metric, i.e., $(A-B)/B$, in the visual representation, but it is intuitively interpretable and thus more commonly used in practice.

Another metric is the log difference metric, i.e., $\log(A) - \log(B)$, which is symmetric as $\log(A/B)=\log(k)$ vs $\log(A/B)=\log(1/k)=-\log(k)$. However, the log differences are undefined for zero or negative values, which is its main limitation. As the raw score in each Atari game can be either positive or negative, the ratio improvement/percentage increase metric is more broadly applicable and particularly preferable in evaluating the RL algorithm across a wide range of Atari games than other metrics like the log difference.

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?

The finding that the algorithm is much better on complicated games does not imply that it is inferior on easy games. After checking the raw score table across 55 Atari games, we find our algorithm is also on par with or even performs better than other baselines on easy games, such as Atlantis, Enduro, FishingDerby, and MsPacman.

We appreciate your insights and suggestions, which will guide our revisions to improve the clarity and accuracy of our manuscript. Please feel free to let us know if you have any further questions.

Reference

[1] Distributional Reinforcement Learning via Moment Matching (AAAI 2021)

[2] Distributional Bellman Operators over Mean Embeddings (ICML 2024)

Hello Reviewer uu9w: The authors have responded to your comments. I would expect that you would respond in kind.

Thank you for taking the time to review our paper. We appreciate your positive assessment and insightful feedback and would like to address the concerns you raised in the Weakness and Question parts of your review.

Weakness 1: 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 ....

• We acknowledge that Algorithm 1 does not involve too many technical details as it is aimed at outlining a generic update procedure of SinkhornDRL. As mentioned in Line 264, we have also provided a full version of SinkhornDRL in Algorithm 3 of Appendix G, in which we add more necessary technical details. In terms of Algorithm 2, given that Sinkhorn iteration is a well-established algorithm with guaranteed convergence, we thus leave it in the Appendix. We appreciate your suggestion and will add more explanations about Algorithm 2 in the appendix.

• Thanks for this great advice. Following your suggestion, we have also included a table to summarize the results in the global response. It shows that our proposed SinkhornDRL achieves the best performance in terms of Mean and IQM(5%) HNS as well as the "> DQN" metric, and is on par with MMD-DQN in terms of Median of HNS.

Weakness 2: There may be a technical error regarding Theorem 1 (3). ...uniform upper bound and finite set assumption.

The contraction conclusion indeed requires a uniform upper bound (also after taking the supremum) over all possible return distributions of $\Delta_{\epsilon}^{U, V}$ to remain strictly below 1, which can be satisfied by straightforwardly assuming a finite (yet arbitrarily large in the state and action spaces) MDP. The finite MDP ensures a finite return distribution set, further implying a finite set of $\lambda_{\epsilon}(U, V) \in (0, 1)$. Consequently, it is straightforward to establish that $\sup_{U, V} \Delta_{\epsilon}^{U, V} < 1$ as shown in Eq.21 under this finite condition. We argue that the finite MDP is not a strong condition but a common setting as most RL theory is established on it. In addition, we can also relax this finite MDP assumption by ruling out the extreme case, where there exists a series of elements in this set that can be arbitrarily close to 1. We have provided an example of this extreme case in Line 666. Since expressing this extreme case mathematically is complicated, we, therefore, retain the more common (and mild) finite MDP condition.

Question 1: 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?

By the nature of optimal transport distances, Sinkhorn divergence inherently captures the spatial and geometric layout, which can be demonstrated by its definition. In particular, Sinkhorn divergence and other optimal transports are defined by measuring the cost of moving mass from one point to another in a space, where the cost directly depends on the spatial distances between points. This is because the optimal transport plan is not just about matching quantiles but also considers where mass is located and where it needs to go, which inherently reflects the spatial relationships and geometric layout of distributions. This is also the primary advantage of optimal transport distances over others, e.g., MMD. In addition, the KL divergence also involves the joint distribution (transport plan), which retains the properties of optimal transport distances.

Question 2: 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?

Yes, it is feasible to generalize the cost function to a weighted form, akin to an extension from Euclidean distance to Mahalanobis distance. By using the weighted cost function, the optimal transport distances can be more closely aligned with specific applications. The weighted cost function can potentially lead to benefits, provided that we have access to accurate or relevant weights, which typically require prior knowledge.

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

We acknowledge this citation issue and will revise it accordingly. Thanks for bringing this to our attention.

We appreciate your insights and suggestions, which will guide our revisions to improve the clarity and accuracy of our manuscript. Please feel free to let us know if you have any further questions.

Hello Reviewer Fdqh: The authors have responded to your comments. I would expect that you would respond in kind.

Thank you for taking the time to review our paper. We appreciate your positive assessment and insightful feedback and would like to address the concerns you raised in your review.

Question 1: Are there existing works on distributional reinforcement learning that use the entropic regularized Wasserstein distance?

To our knowledge, our study is the first to investigate the entropic regularized Wasserstein distance in the context of distributional RL.

Question 2: What specific advantages does the Sinkhorn divergence offer over the entropic regularized Wasserstein distance in the context of this study?

As highlighted in our introduction, the specific advantages of Sinkhorn divergence, which is an entropic regularized Wasserstein distance, over the vanilla Wasserstein distance in distributional RL can be summarized as follows.

• Enhanced accuracy in representing return distributions Sinkhorn divergence uses the samples to depict return distributions, offering more flexibility than the pre-specified statistics, e.g., quantiles, used in vanilla Wasserstein distance. Consequently, our approach naturally circumvents the non-crossing issues of the learned quantile curves.

• Applicability to multi-dimensional rewards Many RL tasks involve multi-dimensional reward structures, but it remains elusive how to use quantile regressions to approximate a multi-dimensional Wasserstein distance. In contrast, Sinkhorn divergence can effectively approximate a multi-dimensional Wasserstein distance, thus providing an efficient solution to address these RL tasks,

• Increased robustness to noises via smoother transport plans The entropic regularization incorporated in Sinkhorn divergence fosters smoother transport plans compared with those derived from unregularized Wasserstein distance. This makes it less sensitive to noises in the learning process.

• Stable optimization in RL learning The entropic regularization turns the objective into a strongly convex problem, and the induced smoothness facilitates faster and more stable convergence in RL learning.

We acknowledge the typos you pointed out and will correct them. Thank you once again for the time and effort you dedicated to reviewing our work.

Thank you for your responses! It seems there was a misunderstanding. For both questions 1 and 2, I was referring to the entropic regularized Wasserstein distance W_{c, \varepsilon}. In the paper, you adopted the Sinkhorn divergence \bar{W}{c, \varepsilon}. I was curious why the entropic regularized Wasserstein distance W{c, \varepsilon} was not used instead.