Paper: Rethinking Optimal Transport in Offline Reinforcement Learning (ICLR-2024)

Reviewer comment (Novelty):
$f$ seems to be learned in an adversarial training manner. Is there any connection between XMRL and ATAC [1], an offline algorithm that also uses adversarial training?

Author response:
Our method XMRL differs from ATAC by using the function $f$ for a different purpose. In XMRL, $f$ is introduced as a function to enforce constraints during the learning process, as described in Section 2, whereas in ATAC they use adversarial training to optimize the critic function $Q$ (denoted as $f$ in their paper).

Moreover, although ATAC is more closely related to the CQL approach, which we compare to in our experiments, our method can use a cost/critic function trained using ATAC's method. This synergy between XMRL and ATAC could be an interesting direction for future research, showing how adversarially trained critic functions could improve constraint-based methods like ours. In the revised version of our paper, we have updated the *related works section* to include ATAC [1] as a critic penatly method.

Reviewer comment (Experiments):
How much additional computation complexity does XMRL introduce? The authors discuss about absolute training time, but how much more time does XMRL take compared to methods like IQL or CQL?

Author response:
For the MuJoCo tasks pre-trained models were used, we trained our method for 100K steps using the pre-trained critic, the training time in the paper is reported for these settings. Comparing our method on MuJoCo Halfcheetah with training from scratch, IQL takes 10 hours, XMRL: 11 hours, CQL: 18 hours.

Regarding the AntMaze tasks, we used the JAX framework, which is known for its efficiency, and the timing was as follows: for example, for the AntMaze (large-play) ReBrac takes ~60 minutes and our takes XMRL: 75.1 minutes.

While introducing an additional discriminator into the training process could potentially increase the computational complexity, in reality XMRL introduces minimal overhead compared to other methods.

Reviewer comment (Experiments):
It seems that the larger $w$ is, the higher the performance. Can the authors further discuss the impact of $w$ on performance?

Author response:
Please, see the general response to all reviewers and revised Appendix A.2 for a detailed discussion of this topic.

Reviewer comment (Novelty):
This contribution of the paper is insignificant, as it is barely a direct application of Extremal Optimal Transport [1], without providing a new understanding of offline RL or solving/alleviating existing problems in offline RL.

Author response:
Our paper addresses the well-known "stitching problem" in offline RL, as detailed in **Section 3.2**. This problem has been recognized in the offline RL literature (Levine et al., 2020) as a significant obstacle to achieving strong performance in offline reinforcement learning. 

We believe that our contribution goes beyond a direct application of Extremal OT. We argue that without our reformulation in **Section 2**, it is unclear that Extremal OT can applicable to RL. We establish a unique connection between OT and RL domains, offering a novel perspective on the problem - **this connection has not been considered before.**

Then we demonstrates the practical applicability of this novel connection by integrating recent findings in OT, specifically Extremal OT, to effectively mitigate the snitching problem. The superior performance demonstrated in our evaluation results underscores the effectiveness of our approach compared to existing methods.

Reviewer comment (Theory):
The idea is not convincing. It is hard to see why we should consider an offline RL task an OT problem. Why do we want to preserve the distribution of $\mathcal{A}$ instead of considering the distribution $\mathcal{A}(s)$ or support $Supp(\mathcal{A}(s))$ independently for each state? According to my understanding, offline RL considers the distribution $\mathcal{A}(s)$ instead of the distribution of $\mathcal{A}$, as shown in W-BRAC (eq. 8).

Author response:
The main reason for consider an offline RL task an OT problem is because this formulation allows to solve stitching problem, and outperforms a previously proposed approaches.  

Regarding the choice between using $ \mathcal{A} $  and $  \mathcal{A}(s) $ , we'd like to clarify that our method is indeed flexible and can consider both $  \mathcal{A}$ and $  \mathcal{A}(s)$  distributions. In the conclusion section of our initial submission, we provided a rationale for our choice to prioritize $ A$  over $  \mathcal{A}(s) $ . This choice, was guided by our experimental observations. We found that focusing on $ f(a)$  rather than $ f(s,a)$  did not compromise performance, and we preferred this representation for its methodological simplicity. 

In practice we can incorporate a state-conditional information through conditional optimal transport. This simply requires adding an additional input variable to the $ f$  function to produce $f(s,a)$ .

We added a separate discussion section on this topic into the revised paper. Also we have included $ f(s,a)$  results in the revised Appendix A.3. All results with $ f(s,a)$ can be reproduced with the original code provided in the Supplementary Material. Please see the networks.py file for the state-conditional version of the potential $ f(s,a)$ .

Reviewer comment (Novelty):
XMLR does not seem better than W-BRAC. As stated in Point 2, the theoretical foundation of XMLR is not convincing. Furthermore, in practice, XMLR faces the same problems as W-BRAC: (1) both of them train an additional discriminator $f$; (2) both have a hyper-parameter to control the extreme of the policy, which is difficult to choose.

Author response:
We initially did not include W-BRAC in the direct comparisons because our experiments showed that XMLR outperformed methods that in turn showed significant improvements over W-BRAC, suggesting a clear performance hierarchy. 

Regarding the training of an additional discriminator, we highlight an important difference in **Section 2**: W-BRAC requires *Lipschitz* constraints on the discriminator. *Lipschitz* constraints often complicate W-BRAC's discriminator training, while weight clipping, gradient penalty, and other regularizations are required. In contrast, XMLR's approach does not require Lipschitz constraints on the discriminator.

The role of the parameter in W-BRAC is completely different. Please see the general response to all reviewers for a detailed discussion of this topic.

Reviewer comment (Evaluation):
The experimental results are not surprising. According to Table 2, the results of XMRL are close to the results of ReBRAC.

Author response:
We respectfully disagree. In the area of Offline RL, large jumps in performance have become increasingly uncommon. Achieving even modest improvements is often challenging, yet these improvements can still represent meaningful progress. This is particularly true in the environments we have examined, where many sophisticated methods, like IQL, ReBRAC, etc, have already set a high-performance benchmark.

Furthermore, a comparison of our XMRL method with the existing SOTA OT-based method OTR provides essential context for our claims. Table 2 in our paper shows that OTR, which uses the IQL backbone, improves on IQL in 2/6 environments, while XMRL consistently improves on CQL in every environment and improves on ReBRAC in 4/6.

Reviewer comment (Writing):
The writing should be improved. Typos, e.g., in the first line in section 2.4, should be fixed.

Author response:
Thanks for pointing this out, typos are fixed in the revised version.

Reviewer comment (Novelty):
What is the benefit of formulating an offline RL problem as an OT problem?

Author response:
Please see the answers to the Q1-4.

Reviewer comment (Theory):
I do not see the relevance of proposition 3.1 (policy improvement), since it does not appear to give any insight about the method (the proposition is not cited/used anywhere in the paper) and the proof is basically the same as the one for classic policy iteration improvement.

Author response:
Thank you for raising your concern that **Proposition 3.1** is not cited in the paper. The primary purpose of **Proposition 3.1** is to emphasize the focus of our method on policy improvement rather than behavior cloning, as stated in the Introduction and **Section 3**. We mentioned our proposition now in **Sections 1, and 3**.

Reviewer comment (Experiments):
Evaluating the performance of the method in more challenging environments or tasks would enrich the experimental section. The impact of entropy regularization on the performance of the method is not assessed.

Author response:
In response to your suggestion, we conducted experiments on the Pen and Door environments of the D4RL. Our XMRL-BC method showed performance similar to the ReBRAC method, specifically: **pen-expert: 154.9±4.2, pen-cloned: 100.1±20, pen-human: 97.5±11, door-expert: 104.9±4.5**. For the environments like door-cloned and door-human, our method achieved zero scores similar to ReBRAC. We have included these results in *Appendix A.4* and additionally provided an ablation study on parameter $w$ for these environments in **A.2**.

Reviewer comment (Experiments):
The method seems considerably sensitive to the parameter w (as seen in appendix 2), so one of the drawbacks of W-BRAC is still present in this method.

Author response:
This is not true. The drawbacks of W-BRAC are not present in our method, while the role of our parameter is completely different. The fact that a hyperparameter (even a different) is required is noted as a *limitation* in the **(Section 4.3)** of the initial submission. Please see the general response to all reviewers for a detailed discussion of this topic.

Reviewer comment (Writing):
Typos: “Fol all experiments” section 4.2.

Author response:
All typos have been corrected.