Four Rooms Environment

Table A. Expanded performance comparison of normalized scores on Four Rooms environment. (corresponds to Table 2)

Scenario 1: Rebalanced Dataset

Table B. Expanded performance comparison on Scenario 1 (rebalanced dataset). (corresponds to Table 3)

Scenario 2: Time-dependently Subsampled Dataset

Table C. Expanded performance comparison on Scenario 2 (time-dependently collected dataset). (corresponds to Table 4)

Thank you for your insightful feedback.

1. Addressing Concerns on Marginal Performance Improvements Compared to Baselines

In response, we have updated our main experiments to include the DRO baseline DR-BC [18] and variation of $f$-divergence for our approach accordingly: inspired by the choice of $f$ in DR-BC, we introduced the soft-TV distance, where the $f$ function is a log-cosh function and its derivative is the tanh function. [A] This enables us to obtain a closed form solution of $w$ by using Proposition 1.

Based on this, we have expanded our analysis by incorporating three additional methods: DR-BC for all scenarios, OptiDICE-BC (with soft-TV), and DrilDICE (with soft-TV) for Mujoco scenarios. The expanded results for Four Rooms environment and Scenario 1,2 can be found in Table A, B, C in the following comment, and the result for Scenario 3 can be found in Figure B in the PDF.

The results consistently show that DrilDICE with the soft-TV outperforms baselines including DR-BC across most scenarios. Notably, DrilDICE with soft-TV, utilizing $f$-divergence similar to the one used in DR-BC, consistently outperforms DR-BC in most evaluation scenarios. We attribute this performance improvement to the remaining key difference: the inclusion of a Bellman flow constraint, which is a key contribution of our work.

The results also demonstrate that the choice of the soft-TV significantly enhances the performance of DrilDICE compared to using soft-$\chi^2$. To explain this performance gain, we hypothesize that the soft-TV distance provides more discriminative weighting $w$ of samples based on their long-term policy errors. As shown in Figure A in the PDF, conventional $f$-divergences (e.g. KL, soft-$\chi^2$, …) make $w$ less pronounced to the long-term policy error $e$, yet the soft-TV distance responds sensitively to changes in $e$, resulting in a more pronounced $w$. This enables BC loss to more selectively focus on critical samples with significant magnitude of long-term policy errors, thereby effectively enhancing performance, akin to the benefits observed in Sparse Q-Learning [C].

In summary, the choice of the $f$-divergence for DrilDICE is critical for performance. The use of soft-TV distance enhanced DrilDICE's performance across most problem settings, directly addressing the reviewer's concerns regarding its comparative performance. We hope these updates can address your concerns about the initial marginal improvements.

[A] Saleh et al., "Statistical Properties of the log-cosh Loss Function used in Machine Learning." arXiv preprint (2022).

[B] Xu et al., “Offline RL with No OOD Actions: In-sample Learning via Implicit Value Regularization”, ICLR 2023.

2. Clarification on the Terminology "Marginal Room (or Action) Distribution of Dataset"

In line 189, the term "marginal room (or action) distribution of dataset" specifically refers to $p(u)$, the proportion of transitions containing target factors (e.g. room visitation, action) within the manipulated dataset. To illustrate, these factors in our experiment include (1) the room visitation and (2) the action of the transition. By intentionally manipulating frequencies of these factors, we designed covariate shift scenarios.

For example, in the Room 1 manipulation scenario, we initially split the expert dataset into two subsets $D_A,D_B$ based on whether each transition’s state was associated with Room 1 or not. We then subsampled transitions from subsets $D_A, D_B$ using the predetermined proportion $p(u), 1-p(u)$, respectively and combined them to construct the shifted dataset $D_i$.

Here, the term "a marginal room (or action) distribution" refers to these proportions $p(u)$, representing how frequently transitions with a target factor (e.g. room visiting, action) are sampled relative to others in the dataset. We will revise this into a more clear term such as “a proportion of transitions containing a target factor” and spend our effort to improve presentations.

3. Performance Comparison with Lower Quality Segments

Despite our primary focus on expert-quality datasets, we conducted experiments with additional datasets on Scenario 3 (segment datasets) to ensure consistency across different datasets. Rather than employing the D4RL expert-v2, we use medium-v2 quality demonstrations as our imitation standard for comparison. The results are depicted in Figure D of the PDF.

Given that relying solely on the normalized scores may not accurately reflect the fidelity of imitation for a medium-quality policy, we primarily measured the target MSE for our comparisons. The results indicate that both DrilDICE and DR-BC demonstrate competitive imitation performance compared to other baselines, with a consistent decrease in target MSE as the number of segments increases.

Thank you once more for your valuable feedback. If there are any more questions or concerns, please feel free to respond, and we will address them quickly.

We respectfully remind you that less than 48 hours remain in our discussion period. We are dedicated to addressing any remaining concerns you may have. In short, we addressed your key concerns as follows:

• Concerns on Marignal Improvements:

  • We found that performance of DrilDICE depends on a choice of $f$-divergence. When similarly aligned with TV distance, DrilDICE consistently outperforms all baselines.

• Comparison on Lower-Quality Datasets:

  • We conduct additional experiments on segmented trajectory scenario (Scenario 3) with a medium-quality dataset, instead of the expert dataset.

Please refer to our detailed rebuttal for more comprehensive information. If you have any further questions or concerns, please do not hesitate to share your comments.

We sincerely thank you again for your reviews.

Four Rooms Environment

Table A. Expanded performance comparison of normalized scores on Four Rooms environment. (corresponds to Table 2)

Scenario 1: Rebalanced Dataset

Table B. Expanded performance comparison on Scenario 1 (rebalanced dataset). (corresponds to Table 3)

Scenario 2: Time-dependently Subsampled Dataset

Table C. Expanded performance comparison on Scenario 2 (time-dependently collected dataset). (corresponds to Table 4)

Thank you for your valuable feedback.

1. Experimental Design and Real-World Applicability Concerns

Our experimental setup, though based on simulated data, is intended to rigorously assess the adaptability and robustness of our approach under realistic conditions. This strategy also ensures reproducibility and is consistent with established methodologies in the field, as demonstrated by similar studies like Yan et al., which also employs segments of the D4RL dataset to test hypotheses in learning from observations (LfO) problem setting [A].

Beyond the simulated shifts, we also have conducted tests on complete trajectories with limited data sizes—a scenario frequently encountered in real-world applications. The results, illustrated in Figure C, demonstrate that DrilDICE adeptly handles sampling errors in small datasets and surpasses existing approaches in performance.

[A] Yan et al., “A Simple Solution for Offline Imitation from Observations and Examples with Possibly Incomplete Trajectories,” NeurIPS 2023.

2. Clarification on Unclear Notations

We clarify that $\Delta S$ represents a set of arbitrary state distributions, and $\mathcal{W}$ denotes a function space of $w(s,a)$ with $w:\mathcal{S}\times\mathcal{A} \to \mathbb{R}$. These updates will be reflected in the final version of the paper.

3. Applicability in Partial Coverage Scenarios

DrilDICE, like other DICE approaches, relies on the assumption that the support of the target state distribution $d$ should encompass the support of the expert's stationary distribution $d_E$. This assumption presents a challenge in ensuring the robustness of DICE approaches when the support of $d$ does not fully cover that of $d_E$.

Despite this, DrilDICE has shown strong performance in scenarios with limited data coverage. In our experiments, particularly in Scenario 3 and the complete trajectories experiment (see Figure B, C in the PDF respectively), we varied dataset sizes to adjust data coverage. The result suggests that DRO approaches possibly handle sampling errors in small datasets.

While these results are encouraging, the comprehensive applicability of our method when $d$ partially covers $d_E$ is yet to be fully established. We believe that further exploration of this critical aspect represents a promising research topic.

4. Expanded Evaluation Results on Main Scenarios

We have updated our main experiments to include a baseline DR-BC [18] and a variation of $f$-divergence for our approach accordingly: Inspired by the choice of $f$ in DR-BC, we introduced the soft-TV distance, where the $f$ function is a log-cosh function and its derivative is the tanh function. [B] This provides a relaxed and invertible version of TV distance, and enables us to obtain a closed form solution of $w$ by using Proposition 1.

Based on this, we have expanded our analysis by incorporating three additional methods: DR-BC for all scenarios, OptiDICE-BC (with soft-TV), and DrilDICE (with soft-TV) for Mujoco scenarios. The expanded results for Four Rooms environment and Scenario 1,2 can be found in Table A, B and C in the comment, and the result for Scenario 3 can be found in Figure B in the PDF.

The results consistently show that DrilDICE with the soft-TV outperforms baselines including DR-BC across most scenarios. Notably, DrilDICE with soft-TV, utilizing $f$-divergence similar to the one used in DR-BC, consistently outperforms DR-BC in most evaluation scenarios. We attribute this performance improvement to the remaining key difference: an inclusion of a Bellman flow constraint, which is a key contribution of our work.

The results also demonstrate that the choice of the soft-TV significantly enhances the performance of DrilDICE compared to using soft-$\chi^2$. To explain this performance gain, we hypothesize that the soft-TV distance provides more discriminative weighting $w$ of samples based on their long-term policy errors. As shown in Figure A in the PDF, conventional $f$-divergences (e.g. KL, soft-$\chi^2$, …) make $w$ less pronounced to the long-term policy error $e$, yet the soft-TV distance responds sensitively to changes in $e$, resulting in a more pronounced $w$. This enables BC loss to more selectively focus on critical samples with significant magnitude of long-term policy errors, thereby effectively enhancing performance, akin to the benefits observed in Sparse Q-Learning [C].

[B] Saleh et al., "Statistical Properties of the log-cosh Loss Function used in Machine Learning." arXiv preprint (2022).

[C] Xu et al., “Offline RL with No OOD Actions: In-sample Learning via Implicit Value Regularization”, ICLR 2023.

Thank you again for your insightful comments. If you have any further concerns or questions about this, please feel free to reply and we will address them promptly.

We gently remind you that less than 48 hours remain in our discussion period. We are committed to thoroughly addressing the reviewer’s remaining concerns. In summary, our responses are as follows:

• Concerns on Experiments on Simulated Shift:

  • Our covariate shift experiment setting ensures robustness and reproducibliity of comparison. Similar simulated shifts have been utilized in exisiting literature.
  • Additionally, we have included natural complete trajectory scenarios where DrilDICE shows significant data-efficiency.

• Concerns on Partial Coverage:

  • Although DrilDICE is not specifically designed to address partial coverage issues, it has shown empirical robustness in such scenarios.

• Expanded Main Experiments:

  • We introduced DR-BC as a new baseline across all experiments; DrilDICE consistently outperforms DR-BC when selecting a similar $f$-divergence.

Please see our rebuttal for more information. If you have any further questions or concerns, do not hesitate to share your comments.

Again, thank you for your thoughtful reviews.

Four Rooms Environment

Table A. Expanded performance comparison of normalized scores on Four Rooms environment. (corresponds to Table 2)

Scenario 1: Rebalanced Dataset

Table B. Expanded performance comparison on Scenario 1 (rebalanced dataset). (corresponds to Table 3)

Scenario 2: Time-dependently Subsampled Dataset

Table C. Expanded performance comparison on Scenario 2 (time-dependently collected dataset). (corresponds to Table 4)

5. Expanded Results on Main Experiments

We added three additional methods: DR-BC for all scenarios, OptiDICE-BC (with soft-TV), DrilDICE (with soft-TV) for Mujoco scenarios. The expanded results for Four Rooms environment and Scenario 1,2 are presented respectively in Table A, B, and C respectively. The results for Scenario 3 are detailed in Figure B of the PDF.

The results consistently show that DrilDICE with the soft-TV outperforms baselines including DR-BC across most scenarios. Notably, DrilDICE with soft-TV, utilizing $f$-divergence similar to the one used in DR-BC, consistently outperforms DR-BC in most evaluation scenarios. We attribute this performance improvement to the remaining key difference: an inclusion of a Bellman flow constraint, which is a key contribution of our work.

The results also demonstrate that the choice of the soft-TV significantly enhances the performance of DrilDICE compared to using soft-$\chi^2$. To explain this performance gain, we hypothesize that the soft-TV distance provides more discriminative weighting $w$ of samples based on their long-term policy errors. As shown in Figure A in the PDF, conventional $f$-divergences (e.g. KL, soft-$\chi^2$, …) make $w$ less pronounced to the long-term policy error $e$, yet the soft-TV distance responds sensitively to changes in $e$, resulting in a more pronounced $w$. This enables BC loss to selectively focus on critical samples with significant magnitude of long-term policy errors, thereby effectively enhancing performance, akin to the benefits observed in Sparse Q-Learning [B].

[B] Xu et al., “Offline RL with No OOD Actions: In-sample Learning via Implicit Value Regularization”, ICLR 2023.

6. Clarification on Datasets Used in Four Rooms Experiments

In our Four Rooms experiment, we designed the covariate shift scenarios caused by unobserved factors affecting data curation, specifically considering such factors as (1) room visitation or (2) action. In the Room 1 manipulation scenario, we initially split the expert dataset into two subsets $D_A,D_B$ based on whether each transition’s state was associated with Room 1 or not. We then subsampled transitions from subsets $D_A, D_B$ using the predetermined proportion $p(u), 1-p(u)$, respectively and combined them to construct the shifted dataset $D_i$.

Specifically, we utilized 100 original episodes, comprising a total of 3994 transitions with a maximum episode length of 50, for $D_E$. From these, after setting $p(u) = 0.4$, we subsampled 1000 transitions to construct $D_i$, which includes 40% (400 transitions) from $D_A$, and 60% (600 transitions) from $D_B$ (e.g., transitions in Rooms 2, 3, 4). To avoid potential issues with support coverage—beyond the focus of our study—we ensured that all states ($|S| = 11 \times 11 = 121$) in $D_i$ appear at least once, as detailed in the Appendix.

7. Experiments on Standard Scenarios

We acknowledge that our initial experiments did not include comparisons on standard scenarios. To address this, we have conducted additional experiments using complete expert trajectories using D4RL expert-v2 dataset varying the number of trajectories in {1, 5, 10, 50}. The results are detailed in Figure C in the PDF. We also observed that DrilDICE can deal with sampling errors of small datasets, showing a superior performance compared to other methods.

Once again, we extend our sincere thanks for your valuable feedback. We believe that addressing your concerns has substantially enhanced the quality of our research. We plan to incorporate all our discussions into the final version. If you have any remaining concerns or questions, please do not hesitate to comment and we will respond as soon as possible.

First and foremost, we sincerely appreciate your insightful feedback.

1. Justification on Alternative Optimization

To clarify, our initial statement suggested that alternating optimization provides a practical approach to approximating the solution of the problem as outlined in Equation (7). In response, we will clearly revise this statement to “... the problem can be practically addressed by alternately optimizing $\nu$ and $\pi$, following recent developments in DICE approaches. [citations]” with related literature including suggested DualDICE and OptiDICE. Additionally, we will provide a more detailed exposition on the derivation (e.g. details on Slater’s condition, assumptions) in the forthcoming revision.

2. Clarification on $d_D$ and Our Problem Setting

It is important to note that we did not assume that $d_D$ is a stationary distribution of any policy; rather, it is an arbitrary state distribution. Hence, even if the transition dynamics $T$ and the policy $\pi_E$ remain unchanged, $d_D$ can deviate from $d_E$ since we do not assume that $d_D$ should be induced from MDPs. We clarify this point in the revision.

Additionally, we also clarify that our problem setting concerns shifts in stationary distribution (say $d$-shifts), not transition shifts ($T$-shifts). Specifically, we assumed that the expert dataset is a collection of transitions $(s,a,s’)$, where $s$ sampled from $d_D$, an expert action $a$ decided by the deterministic expert $\pi_E$, and $s’$ sampled from $T(s’|s,a)$ (which will not be shifted in the testing phase). Even when $T$ and $\pi_E$ remain unchanged, practical constraints such as high costs or limited recording frequencies often result in dataset with sparsely collected transitions. This sparsity leads to incomplete trajectories, causing a deviation of $d_D(s)$ from $d_E(s)$. We do not focus on delayed transitions such as ($s_t, a_t, s_{t+k}$) with $k \ge 2$, which may arise from not collecting immediate subsequent states.

Although we maintain that our setting does not directly involve $T$-shifts, we compare our approach with DR-BC, as their objective also pertains to $d$-shifts.

3. Comparison with DR-BC [18]

In response to the reviewer’s suggestion, we have revisited DR-BC [18] and investigated the method more thoroughly. While DR-BC primarily aims to address $T$-shift scenarios, its objective considers the uncertainty set of arbitrary state distributions $d(s)$, which is potentially applicable to our scenario. Hence, we evaluate DR-BC as a baseline for all of our scenarios. To compare, the key differences between DR-BC and our approach can be summarized as:

• Choice of Uncertainty Set: To define its uncertainty set, DR-BC employs TV distance while DrilDICE uses $f$-divergence, which is more general.
• Adoption of Bellman Flow Constraints: For an uncertainty set, DR-BC considers arbitrary state distributions, including non-stationary ones, while DrilDICE considers state stationary distributions by enforcing Bellman flow constraints.

4. Technical Adjustment for DrilDICE and Summary of $f$-divergence

We initially attempted to align the $f$-divergence to the TV distance used in DR-BC for a fair comparison. Our formulation requires the inverse of the derivative of $f$ for a closed form solution of the weight $w$. However, since the derivative $f$ of TV distance is a step function and is not invertible, we cannot directly use TV distance. We address this issue by adopting the log-cosh function [A] for $f$, which has a tanh function as $f’$; a relaxed version of the step function while ensuring invertibility. We refer to this as the soft-TV distance and utilized this in Mujoco-based experiments. (See Figure A for visualization)

In summary, the following table compares the choice of $f$-divergence for DrilDICE, highlighting our method's adaptability in obtaining practical solutions through refined mathematical approaches. As the reviewer commented, we will include this table in the main manuscript.

where $f_{\text{soft}-\chi^2} (x) := x\log x -x +1$ if $0<x<1$ and $(x-1)^2$ if $x \ge 1$, $(f^{'})^{-1}_{\text{soft}-\chi^2} (y) := \exp(y)$ if $y<0$ and $y+1$ if $y \ge 0$, $\mathrm{ReLU}(x):=\max(0,x)$, $\mathrm{ELU}(x):= \exp(x)-1$ if $x<0$ and $x$ for $x\ge0$.

[A] Saleh et al., "Statistical Properties of the log-cosh Loss Function used in Machine Learning." arXiv preprint (2022).

We kindly remind you that less than 48 hours remain in our discussion period. We are committed to addressing any remaining concerns.

In summary, we have addressed your key concerns as follows:

• Clarification of Problem Setting: (1) we do not assume that the data distribution $d_D$ is not stationary, (2) we consider shifts of $d$, not $T$.
• Comparison with DR-BC [18]: (1) a choice of $f$-divergence for the uncertainty set, (2) our method includes Bellman flow constraints.
• Additional Experiments:
  • DR-BC has been added as a baseline across all experiments. DrilDICE consistently outperforms DR-BC under similar choices of $f$-divergence.
  • We have incorporated complete trajectory scenarios, demonstrating significant data efficiency in DrilDICE.

For further information, please see our rebuttal. If you have any additional questions or concerns, we encourage you to provide your comments.

Thank you again for your insightful reviews.