Rethinking Optimal Transport in Offline Reinforcement Learning

Thank you for spending time reviewing our paper and providing valuable feedback that will help us improve the manuscript. Please find below the answers to your questions.

Q1: $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?.

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.

---

Q2: 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?.

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.

---

Q3: It seems that the larger w is, the higher the performance. Can the authors further discuss the impact of w on performance??

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

---

Concluding remarks We truly value your reviews. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. If the responses above are sufficient, we kindly ask that you consider raising your score.

Thank you for taking the time to review our paper and provide useful feedback. Your questions will help us improve the manuscript. Below are the answers to your questions.

Q1: 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.

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.

---

Q2: 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 A instead of considering the distribution or support A(s) independently for each state? According to my understanding, offline RL considers the distribution A(s) instead of the distribution of A, as shown in W-BRAC (eq. 8)?

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)$ .

---

Q3.1: XMLR does not seem better than W-BRAC. As stated in 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.

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.

---

Q3.2 both methods have a hyper-parameter to control the extreme of the policy, which is difficult to choose.

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.

Thank you for your detailed analysis of our paper and your thoughtful suggestions for improvement. All typos have been corrected. Below is our response to your comments.

Q1: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.

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.

---

Q2:Evaluating the performance of the method in more challenging environments or tasks would enrich the experimental section.

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.

---

Q3:The impact of entropy regularization on the performance of the method is not assessed?

We did not perform an ablation study on entropy regularization as it is not part of our contribution. Entropy regularization was only applied in the context of CQL-based experiments performed on the Antmaze environments, while the CQL-backbone includes entropy regularization by design. Conversely, for the simpler MuJoCo environments, we opted for a version of CQL Eq.(14) without any regularization.

---

Q4: 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.

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.

---

Concluding remarks: Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. If you finds the responses above are sufficient, we kindly ask that you consider raising your score.