An Optimal Discriminator Weighted Imitation Perspective for Reinforcement Learning

We thank the reviewer for their thoughful review of our work. We address the reviewers questions and concerns below:

I believe the performance of IDRL is not stellar. Considering that the algorithm requires multiple dataset filtering steps, the performance gains from extra computation might not necessarily suggest the significance of the results.

Note that on almost all tasks, IDRL matches or outperforms previous SOTA Primal-RL methods and greatly outperforms all Dual-RL methods. We acknowledge the additional computational cost brought by IDRL, however, we want to highlight that the core contibution of the paper is providing the first, optimal discriminator-weighted imitation view of solving offline reinforcement learning problem and showing the promising of correctly using Dual-RL methods. These new findings could potentially lead a new trend in the RL community.

Since it works with filtering, the algorithm might fail in scarcity of data. In MuJoCo, single demonstration imitation learning is a standard setting. In this case, I suspect IDRL's performance will converge with (single demonstration) behavior cloning performance.

We respectfully disagree with this argument. Although IDRL leverages imitation, it does not operate at the trajectory level. Instead, IDRL could do stitching: extracting useful parts from different trajectories—those within the optimal distribution—effectively forming a superior demonstration compared to any single existing trajectory in the offline dataset. For example, in MuJoCo, Behavior Cloning (BC) on selected high-return trajectories in the dataset (as shown by the X%-BC results in the paper) performs significantly worse than IDRL.

An analysis scalability (such as computation costs) of various offline RL tasks should be reported and experimentally validated.

We acknowledge the increased computational requirements of IDRL, as noted in the limitations section of our paper. To provide additional clarity, we present a runtime comparison of IDRL with several baseline offline RL algorithms below. On D4RL datasets, IDRL introduces only one additional iteration. During each iteration, we train $Q$, $V$ for half the training steps and $U$, $W$ for the other half. As a result, the total runtime is approximately twice that of IQL or SQL.

Is the denoising process of diffusion-based offline RL (such as Diffusion-QL) similar to filtering datasets? How is IDRL conceptually different?

The success of diffusion-based methods relies on both accurate behavior modeling and effective sampling guidance. However, both steps can introduce errors. For instance, [1] demonstrates that diffusion behavior policies may produce potential out-of-distribution (OOD) actions, leading to overestimation errors in the guided sampling process.

IDRL is related to these methods but differs significantly. IDRL is grounded in discriminator-weighted imitation learning, which avoids the use of potential OOD actions during training. This distinction ensures that IDRL does not inherit the overestimation errors associated with diffusion-based approaches.

[1] Diffusion-DICE: In-Sample Diffusion Guidance for Offline Reinforcement Learning

Please let us know if any further questions remain. We hope the reviewer can reassess our work with these clarifications.

Dear authors, I thank you for your response; now I have a better understanding of IDRL's performance.

I have a few questions on the theory.

• Theorem 1 says that we need a decoupling approach from Eq. (5) s.t. $\omega^\ast(s,a) = \omega^\ast(s) \omega^\ast(a|s)$ based on Proposition 1. Is this interpretation correct?
• Theorem 2 manifests an optimization problem for conditional visitation $\omega^\ast(a|s)$ from Theorem 1. Is this correct?
• Theorem 3 is "highly built on the Theorem 3 in Li et al. (2024)." What could be the technical novelty in this theorem proving?

Dear authors,

Thank you for your detailed response. Based on your response and the reviews from other reviewers, I would like to take some additional time to carefully reevaluate the submission and finalize my recommendation, and further discuss with the other reviewers.

We thank the reviewer for their thoughful review of our work. We address the reviewers questions and concerns below:

There should be a comparison of compute costs (e.g., run time, memory usage), given the substantial amount of modifications introduced (e.g., additional updates and iterative dataset refinement).

We acknowledge the increased computational requirements of IDRL, as noted in the limitations section of our paper. To provide additional clarity, we present a runtime comparison of IDRL with several baseline offline RL algorithms below. On D4RL datasets, IDRL introduces only one additional iteration. During each iteration, we train $Q$, $V$ for half the training steps and $U$, $W$ for the other half. As a result, the total runtime is approximately twice that of IQL or SQL.

How does correcting the visitation distribution address the fragmented trajectory problem?

The reviewer is correct in noting that states filtered with $w^*(a|s)$ are also filtered out by $w^*(s, a)$. The key distinction is that $w^*(a|s)$ does not account for the dynamics of the environment, resulting in a much higher degree of incompleteness (trajectory-level incomplete). Theoretically, filtering with $w^*(s, a)$ guarantees a valid visitation distribution (trajectory-level complete), ensuring that there are no trajectories with missing transitions. As the reviewer points out in Fig. 2e, some fragmentation is observed empirically, but the amount is much smaller compared to using $w^*(a|s)$, and it does not lead to performance collapse during training.

Line 240 states that Deep RL algorithms are prone to overestimation errors caused by fragmented trajectories. Is this conclusion based on a previous study? To the best of my knowledge, the "stitching" challenge (which is a task design factor of D4RL) requires offline RL algorithms to assemble sub-trajectories in order to solve a task [1].

We agree with the reviewer that offline RL algorithms can stitch together trajectories to solve a task. However, such algorithms are typically trained on complete trajectories with few or no missing transitions, ensuring that every transition (except the last) is supported by at least one subsequent transition and has a valid Bellman backup target. In our work, fragmented trajectories refer to unsupported transitions that lack a valid backup target due to dataset filtering.

To support our claim that fragmented trajectories lead to divergence issues, we ran IQL on randomly selected transitions from the original complete dataset. Specifically, we randomly selected 10K transitions from medium and medium-replay datasets in the Hopper and Walker environments and observed the learned value functions of IQL. The results, available at link 1, link 2, link 3, and link 4, show divergence or overestimation in these cases.

(Line 263, 283) Which equation are you referring to? I assume it is Equation 9?

Thanks for pointing this out, we have corrected the reference to Equation 9.

Without distribution correction, one might expect the algorithm to produce results similar to conventional Dual-RL methods (e.g., IQL). However, the average score in Table 2 seems to be significantly worse (56.8 vs. 77.8 of IQL on Mujoco). A more detailed ablation study may help.

The ablation study uses two iterations to isolate the effect of filtering with $w(s, a)$. As mentioned earlier, running the second iteration with transitions filtered by $w(a|s)$ tends to cause divergence and degrade performance. Results for running one iteration without distribution correction can be inferred from the performance of f-DVL in Dual-RL [1], which closely aligns with IQL:

[1] Dual RL: Unification and New Methods for Reinforcement and Imitation Learning

The proof of Theorem 1 lacks clarity for readers not familiar with Fenchel-Rockafellar duality, as the authors have omitted some details (e.g., solving for w(s)). A more detailed explanation would be helpful.

Thank you for pointing this out. The solution for $w(s)$ can be easily derived by setting $d^D$ to $d^*$ in the first term of Obj.(9). This is valid based on Lemma 2, which holds true for $d^*$. Taking the derivative and setting it to zero provides the solution for $w(s)$.

Please let us know if any further questions remain. We hope the reviewer can reassess our work with these clarifications.

We thank the reviewer for their time and effort in reviewing our paper and for the constructive comments. Below, we address the concerns in detail:

Further, this method seems to be computationally more expensive compared to other methods. It would be nice if this was discussed.

We acknowledge the increased computational requirements of IDRL, as noted in the limitations section of our paper. To provide additional clarity, we present a runtime comparison of IDRL with several baseline offline RL algorithms below. On D4RL datasets, IDRL introduces only one additional iteration. During each iteration, we train $Q$, $V$ for half the training steps and $U$, $W$ for the other half. As a result, the total runtime is approximately twice that of IQL or SQL.

How does filtering the dataset in round k, change the approximation of previously removed s,a pairs in later rounds?

Transitions filtered at iteration $k$ are no longer used in subsequent iterations. These transitions do not belong to the optimal distribution and are therefore excluded from further approximations.

How does this policy generalize to states that were filtered out?

Thank you for this insightful question. This issue is less significant if the initial state distribution during deployment remains similar to the training data, which is typically the case in standard scenarios. States filtered during training will not be visited by the policy because they lie outside the distribution of the optimal policy. However, as we noted in the limitations section, we acknowledge that generalization issues may arise when the initial state distribution during deployment deviates significantly from the offline data. Note that in this case offline RL algorithms that don't filter out dataset may also encounter issues like suboptimality of the dataset action if behavior regularization weight is strong, or overestimation caused by the value function if behavior regularization weight is weak.

Please let us know if any further questions remain. We hope the reviewer can reassess our work with these clarifications

We thank the reviewer for the effort engaged in the review phase and the constructive comments. Regarding the concerns, we provide the detailed responses separately as follows.

This paper misses some literature in RL trained with weighted loss, such as EDP, and QVPO [1, 2].

We appreciate the reviewer bringing this to our attention and will include references to these works in the revised version of the paper. However, we would like to emphasize that the primary contribution of our paper is introducing the first, optimal discriminator-weighted imitation view of solving offline reinforcement learning, which is distinct from the methodologies presented in the mentioned literature.

The reviewer believes the author should provide more explanation on how the additional variance introduced by the training of the U and W networks affects the overall stability of the algorithm.

Thank you for raising this important point. Based on our empirical evaluations, which span 24 datasets across 7 environments, we did not observe any divergence or instability during the training of $U$ and $W$. Theoretically, the learning of $U$ and $W$ is expected to be stable because their objectives are convex, which under mild assumptions guarantees convergence to the optimal solution. Additionally, the training processes for $U$ and $W$ are independent of the training of $Q$ and $V$; this can be viewed as a second phase of learning akin to training another set of $Q$ and $V$, which is similarly stable.

That said, we acknowledge that potential instabilities may arise in broader cases due to the function approximation setting, where even $Q$-learning does not have guaranteed convergence.

In Algorithm 1, the reviewer wonders whether the W network is updated based on (12) rather than (10) on line 10?

We apologize for the typo in Algorithm 1 and thank the reviewer for pointing it out. The W network is updated based on (10), not (12).

The reviewer is confused about the value of M used in the experiments, and considers further clarification is needed here.

Sorry for the confusion, $M$ is the iteration number used in IDRL, as mentioned in Algorithm 1, line 2. We use M=2 on D4RL datasets and M=3 on corrupted demonstrations, we have made it more clear in the revised version of the paper.

Please let us know if any further questions remain. We hope the reviewer can reassess our work with these clarifications

We sincerely thank the reviewer once again for their time and effort in providing valuable feedback. We respectfully note that many papers receive improved scores during the rebuttal period after addressing reviewers’ concerns. Since there appear to be no remaining issues across all reviews for this paper, and considering your assessment of its soundness as excellent and its contributions as novel, we kindly ask if you might consider revisiting your score to further support this work, if you find it appropriate. We understand that such a request may seem unusual, but given that the reviewer who gave a score of 5 has not responded, the current statistics place our paper in a highly borderline position. We truly appreciate your understanding and consideration.

Thanks for your reply. Although this work offers a novel discriminator-weighted imitation view, it looks a little bit difficult to realize and needs extra costs for U and W network. In that case, I keep my score.

Besides, I suggest the authors add a practical implementation part and instantiate the convex function $f$ as some common metrics (e.g., KL divergence). This will make it more straightforward to realize IDRL.