Mitigating Covariate Shift in Behavioral Cloning via Robust Stationary Distribution Correction

Authors: seokin  seo
byung-jun  lee
jongmin  lee
hyeongjoo  hwang
hongseok  yang
kee-eung  kim
  kaist
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 12:30:28

Abstract: We consider offline imitation learning (IL), which aims to train an agent to imitate from the dataset of expert demonstrations without online interaction with the environment. Behavioral Cloning (BC) has been a simple yet effective approach to offline IL, but it is also well-known to be vulnerable to the covariate shift resulting from the mismatch between the state distributions induced by the learned policy and the expert policy. Moreover, as often occurs in practice, when expert datasets are collected from an arbitrary state distribution instead of a stationary one, these shifts become more pronounced, potentially leading to substantial failures in existing IL methods. Specifically, we focus on covariate shift resulting from arbitrary state data distributions, such as biased data collection or incomplete trajectories, rather than shifts induced by changes in dynamics or noisy expert actions. In this paper, to mitigate the effect of the covariate shifts in BC, we propose DrilDICE, which utilizes a distributionally robust BC objective by employing a stationary distribution correction ratio estimation (DICE) to derive a feasible solution. We evaluate the effectiveness of our method through an extensive set of experiments covering diverse covariate shift scenarios. The results demonstrate the efficacy of the proposed approach in improving the robustness against the shifts, outperforming existing offline IL methods in such scenarios.

Main paper: 
 introduction: 
 Imitation learning (IL) aims to recover the expert behavior from the dataset of demonstrations. The standard IL setting assumes that the imitator is allowed to interact with the environment during training, as it provides valuable information regarding state transitions. On the other hand, offline IL requires training without online interactions, reflecting scenarios where the interactions are either infeasible or expensive [5, 12, 13, 28] . Despite recent works on offline IL that explore scenarios involving supplementary datasets, such as suboptimal demonstrations [5, 14] , behavioral cloning (BC) remains a compelling option in practice since BC does not require additional datasets except expert demonstrations. However, the efficacy of BC can be compromised when the imitator policy's behavior deviates from the underlying data distribution. This phenomenon, known as covariate shift, presents a significant challenge, often resulting in performance degradation. Consequently, a substantial body of IL research is dedicated to addressing this covariate shift issue [5, 24] .A common approach for solving offline IL problems is to use the so-called distribution matching objective or its variant [11, 12, 14, 15, 31] , which aims to align the stationary distribution of the imitator policy with that of the expert policy. However, this approach crucially assumes that expert demonstrations are sampled from the stationary distribution of the expert policy. In practice, this assumption may break, that is, sampled expert demonstrations may be from a shifted version of the stationary distribution. For instance, expert demonstrations may be collected by first sampling states from a distribution that is different from the stationary distribution of the expert policy, and then labeling sampled states with expert actions. In this case, expert demonstrations are not samples from the stationary distribution. As a result, the offline IL algorithms based on the distribution-matching approach might perform badly in this covariate shift case.In this paper, we address the issue of the covariate shift in the dataset of expert demonstrations, particularly when the data distribution does not match the stationary distribution of the expert policy. We begin by arguing that BC objective is more natural to consider the shift than the distribution matching objective. Then, inspired by the principle of distributionally robust optimization [2, 6, 20, 25] , we propose an adversarial objective for offline IL that addresses covariate shifts in BC training. Instead of simply considering the worst-case with respect to all possible distributions on state-action pairs, our objective considers only those distributions that arise as the stationary distributions of policies, i.e., distributions that satisfy the Bellman flow constraint. Leveraging the techniques from the stationary distribution correction ratio (DICE) algorithm, we introduce DrilDICE, which efficiently solves our optimization problem and computes an imitator robust to covariate shifts of the dataset. In addition, we suggest the practical covariate shift scenarios that may arise in offline IL applications. Under those problem settings, we compare our approach with baselines and demonstrate that our approach can imitate the agent robust to covariate shift of our interests.2 Background and Related WorkQuestion: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: Yes, we clearly reflect our scope and contribution in abstract and introduction.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
offline imitation learning with arbitrary state distributions: 
 We consider a Markov decision process (MDP) without rewards, which is defined as a tuple of ⟨S, A, T, p 0 , γ⟩ where S is a state space, A is an action space, T : S × A → ∆(S) is a transition distribution, p 0 ∈ ∆(S) is an initial state distribution and γ ∈ [0, 1) is a discounted factor. We focus on the class of deterministic policies Π = {π : S → A}. Given a policy π ∈ Π, the stationary distribution d π of π is defined as d π (s, a) = (1 -γ) One significant branch of research addressing the issue of covariate shift is distributionally robust optimization (DRO). DRO involves considering an uncertainty set of the distribution and optimizing for the worst-case loss within this set [6, 20] . Recent studies have proposed method to attain the worstcase supremum by using kernel methods [25] or the integral probability metric such as Wasserstein metric [10] . Additionally, in the field of imitation learning, DRO has been applied to achieve robust learning in the presence of noisy experts [2] or transition shifts [19] .A closely related problem studied in IL is robust imitation learning, which focuses on learning a policy that can tolerate various shifts in the transition distribution T of the MDP [4, 19] . Another related scenario is the noisy expert scenario, where the demonstration dataset is corrupted by noise, resulting in the distribution d D of the dataset being a noisy version of d E [6, 23] . These issues are usually addressed by using robust optimization techniques [6, 19] , which are also utilized in our approach. However, it is important to note that the types of distribution shifts considered in our work and in prior researches differ within the context of offline IL.To emphasize, we are not focused on addressing covariate shifts induced by transition shifts or noisy demonstrations. Instead, our work investigates a type of covariate shift where the dataset still consists of expert actions, but states are not sampled from the stationary distribution of the expert policy: the distribution d D (s) of the samples in the given demonstration dataset differs from d E (s), but both distributions share the expert policy given states, i.e., π D (s) = π E (s) for all s ∈ S.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: Yes, please see Section 5Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
robustness to covariate shift: 
 Assume that we have a free access to π E , such that we can optimize our trained policy π to be close to it on any state. We can write the BC objective for optimizing π as:EQUATIONfor some supervised learning loss ℓ (e.g. squared error, cross-entropy, ...). Previous studies on imitation learning have assumed that d D ≈ d E . Under this assumption, BC approach of:EQUATIONdoes show strong performance if we have sufficient data. Otherwise, if the expert data is not sufficient, we can also incorporate transition information and benefit from the distribution matching (DM) approach, which aims to get a stationary distribution d DM that is close to data distribution d D with the following objective:d DM := arg max d∈∆(S) -D(d(s, a)∥d D (s, a)) s. t. a d(s, a) = (1 -γ)ρ 0 (s) + γ s,ā T (s|s, ā)d(s, ā) ∀s ∈ Swhere D is any divergence between two probability distributions. The solution of the optimization problem d DM can be used for the state distribution for BC instead of d E , and it is widely known that this can complement lack of data with additional transition information, leading to improved performance. Nevertheless, under harsh covariate shift, i.e., when the state distribution of d D diverges from that of d E , both two approaches above can fail miserably, as there is no guarantee on improving policy performance when d E and d D are significantly different. See Section A for an example of failure cases.To overcome the covariate shift, we adopt a distributionally robust objective, i.e., for an arbitrary distribution d, we knowE s∼d E [ℓ(π(s), π E (s))] ≤ max d∈∆(S) E s∼d [ℓ(π E (s), π(s))],which holds with d π as well, and we can instead optimize for minEQUATIONwhich corresponds to an upper-bound minimization for optimizing π with respect to some uncertainty set Q ⊆ ∆(S). However, when Q is choosed as ∆(S), this objective can easily lead to overly pessimistic solutions, as we seek for any extreme state distribution d that maximizes the policy loss, and it will easily place all its probability on a single state. Hence, a choice of Q is crucial to obtain an appropriate solution of distributionally robust optimization.For notational convenience, denote C π (s) := ℓ(π(s), π E (s)). In the objective (3), the worst-case policy loss is considered among an uncertainty set Q. However, when Q = ∆(S), it would be too pessimistic and the objective can be improved when we consider a smaller choice of Q. To keep the fact that it is an upper-bound minimization for Eq. 1, which is the loss with d E , the expert state distribution d E should still be contained in Q, even if we aim to use a smaller set.To facilitate this, we impose two constraints on Q with respect to a state distribution d: . The first constraint is a constraint that should be satisfied by d E , and thus it effectively reduces the potential distribution set for d if we impose it properly. The second constraint may not seem necessary, but as we estimate the policy loss from finite samples from d D , diverging too far from d D will reduce effective number of samples. Then, by introducing f -divergence regularization to enforce the second constraint, we consider the following constrained optimization problem:EQUATIONwhere D f denotes f -divergence with a convex function f . We develop a practical algorithm in a stationary distribution correction ratio estimation (DICE) framework style. From Eq. 4, we can take Lagrangian ν to solve the inner constrained optimization, i.e.EQUATIONwhere e π,ν := (C π (s) + γ s ′ T (s ′ |s, a)ν(s ′ ) -ν(s)), w(s, a) := d(s,a) d D (s,a) and W is a class of functions w : S × A → R ≥0 . By Slater's condition [3] , since the strong duality holds on Eq. 5, we can change the minimax into a maximin optimization as follows:EQUATIONand the closed-form solution w * can be obtained by solving for the inner maximization.Proposition 1. (Lee et al. [16] ) The closed-form solution for the inner maximization of Eq. ( 6) isEQUATIONAfter obtaining w * π,ν , the minimax optimization ( 6) is converted into a minimization problem with respect to ν. As a result, each objective for ν, π can be summarized as follows:EQUATIONHence, the problem can be practically addressed by alternatively optimizing ν and π following recent developments in DICE approaches [16, 17] . We call this distributionally robust optimization approach to covariate shift in offline IL as DrilDICE (Distributionally Robust Imitation Learning via DICE). To facilliate a fair comparison with the uncertainty set defined by total variation (TV) distance, as utilized in [19] , we introduce the soft-TV distance as a specific choice of f -divergence of DrilDICE.It is important to note that Proposition 1 requires (f ′ ) -1 , the invertable derivative of f . However, the generator function of TV-distance f TV (x) := 1 2 |x -1|, lacks an invertible derivative as illustrated in Figure 1 , rendering the direct application of TV distance in DrilDICE challenging.We technically overcome this limitation by relaxing the derivative function of TV-distance. Given that the derivative function of TV-distance manifests as a step-function, we choose to relax this function by employing the tanh function. Specifically, we utilize the log-cosh function [22] to f as follows:EQUATIONBy plugging f ′ SoftTV into Eq. 7, we can obtain the closed-form solution w * π,ν of the inner maximization objective (6) tailored to a specific choice of f Soft-TV as follows:EQUATIONwhere ReLU(x) := max(x, 0). This choice of f enables DrilDICE to obtain a closed form solution of inner maximization problem while maintaining similar properties of TV-distance. For other possible choice of f -divergence in DrilDICE, see Section B in the supplementary material.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA]Answer: [Yes] Justification: We used D4RL dataset, which is public, and provide toy domain code to generate datasets.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
comparison on baselines objectives: 
 To provide clear descriptions about the baselines used for comparison in our experiments, we summarize the relevant objectives in Table 1 .Adversarial Weighted BC (AW-BC) We refer to adversarial weighting by following the terminology from [27] as the minimax objective without constraints. As we discussed in Section 3.2, the adversarial weighting objective is also an upper-bound of the target objective. However, it tends to be overly pessimistic because it considers the entire state distribution space, including distributions that do not correspond to the stationary distribution of any policy.To adjust the level of the robustness, DR-BC [19] E d [C π (s)] ≤ E d E [C π (s)],which is not relevant to minimize the target objective. Despite there is no direct connection, to ensure that performance of our method are not due to side effects of algorithm's implementation, we include this method for performance comparison in the following experiment section. Four Rooms environment Four Rooms is a gridworld environment which aims to find a path from a starting state to a goal state. Each cell in the grid represents a state and the starting state and the goal state are marked by the orange and the green box in Figure 2 respectively. The agent can choose one of 4 actions from each state: UP,DOWN, LEFT, RIGHT, and the rooms are numbered as shown in Figure 2 . We pre-collect expert dataset D E by using a deterministic expert policy π E and online interactions.To investigate that our approach can effectively addresses the covariate shift, we need the dataset that notably deviates from the expert stationary distribution d E . A possible realistic scenario of the dataset deviation occurs when data collectors gather data at different frequencies for each state or action. For example, if a data collection device (e.g. cameras or sensors) operates at different recording frequencies in each room, the frequency of observing a room in the dataset will not match d E . To simulate this scenario, we design datasets where the marginal room (or action) distribution of dataset deviates from d E . By manipulating the marginal proportion of a predetermined variable in dataset, we generate a deviated dataset from the original dataset. We consider a total of 8 problem settings by manipulating the marginal distribution of four rooms and four actions. In each scenario, we set a marginal distribution p i (u) of the variable to be manipulated u to a certain fixed probability and resample transitions to configure the dataset D i through the following process. (See Section C in the supplementary material)D i = {(s, a)|u ∼ p i (u), s ∼ D E (s|u), a = π E (s)}Implementations and baselines Since we are not interested in trivial scenarios where the policy can easily minimize the supervised loss over entire state space to zero (i.e. max s∈S C π (s) = 0), we conduct experiments using function approximators instead of a tabular setting. We initialize π 0 for the cost function C π as BC policy. To solve optimization problems, each algorithm is implemented with a two-step procedure: (1) optimize each objective by using the cost C π0 determined by an initial policy π 0 and obtain w * π0 , (2) train π with the weighted BC. For policy and weight models, we use linear approximators with RBF features, which exploit distances from representative points. For more details, see Section C. For choices of f -divergence for OptiDICE-BC and DrilDICE, we select KL-divergence instead of the soft TV-distance for a simplicity of convex optimization solver implementation. We consider following baselines:• BC: a standard behavioral cloning without any regularization.• DemoDICE [14] : a representative distribution matching approach to offline IL. We compare a special case of DemoDICE that does not exploits supplementary datasets.• AW-BC: adversarial weighting method without Bellman flow constraints and robustness radius.• DR-BC [19] : distributionally robust BC method without Bellman flow constraints.• OptiDICE-BC [16] : a representrative method as the best-case weighting.The following metrics are measured with 100 episodes for each trained policy.• Normalized score: a normalized episode return that scales from 0 (random score) up to 100 (expert score) averaged by 100 episodes.• Worst-25% performance : the normalized scores averaged by the worst 25% episodes.• Target 0-1 loss: the averaged 0- As depicted Table 2 , DrilDICE outperformed baselines in the Four Rooms environment under various covariate shift scenarios. Our approach achieved the highest normalized scores across all scenarios and consistently demonstrated superior performance in the worst-25% performance. In scenarios involving Room 3, where the original dataset D E had only 8.9% coverage, the probability of observing states in Room 3 was increased to 40%, causing significant covariate shifts. This led to a significant degradation in the robust performance of BC (worst-25%). However, DrilDICE successfully improved the worst-25% performance in Room 3 to levels comparable to other scenarios, demonstrating its robustness to covariate shifts induced by dataset deviations. predicting the correct expert actions in nearly all states (26 out of 27) in Room 3. However, its performance remains suboptimal in other rooms. DrilDICE, leveraging the cost derived from π bc , computes the corrected state distribution through the worst-case weighting, denoted as w * bc , as depicted in Figure 3-(c ). Notably, states mispredicted by π bc or those with high probabilities for the dataset distribution tend to obtain relatively higher weight values, contributing more to loss optimization. Consequently, DrilDICE effectively trains the agent by correcting suboptimal behaviors.To investigate our approach also benefits distributionally robust training in more complex tasks, we conduct experiments on continuous control domains. In addition, we devise three practical problem settings that could possibly occur in the real-world and simulate those scenarios by restructuring the given dataset to simulate these dataset deviation scenarios (For experiments on standard scenarios, refer to Section E). We utilize hopper-expert, walker2d-expert, halfcheetah-expert datasets included in D4RL benchmark [9] as original datasets. We utilize the soft TV-distance defined in Section 3.3 as f -divergence for OptiDICE-BC and DrilDICE for these scenarios (For experiments with another choice of f -divergence, see Section E.4). See Section D for implementation details.Scenario 1: Rebalanced dataset From a practical perspective, a data collector may encounter scenarios in which the costs associated with taking an action or visiting a specific state exhibit substantial costs varying across different states or actions. In such scenarios, the data collected in states or actions with lower costs is likely to be observed, while those with higher costs tend 2.0 ± 0.9 2.9 ± 2.2 0.9 ± 0.7 2.9 ± 1.1 -1.2 ± 0.4 5.5 ± 0.8Table 4 : Performance comparison on Scenario 2 (time-dependently collected dataset). Each experiment is repeated 5 times, and the average normalized scores with their standard errors are reported. The highest mean performance scores are highlighted in bold.to be underrepresented. To simulate these circumstances, we partition the state or action space and manipulate their mixture ratio to generate an imbalanced dataset, under assuming the costs of collecting data points in each group is significantly different, hence their proportion is shifted from the original dataset.In order to implement this scenario, we measure the distance between each state and the representative point of the states, then split the dataset into two groups D 1 , D 2 based on the statistics of distances: D 1 comprises states close to the point, while D 2 is consists of states farther from the point. Subsequently, we resample datasets according to a predetermined proportion of p(D 1 ) in {0.1, 0.5, 0.9} by uniformly sampling different numbers of samples from each group.Scenario 2: Time-dependently collected dataset One of well-known practical issues of covariate shift in supervised learning and time-series forecasting communities is dataset shift caused by seasonal or time-dependent variables [8, 18, 26] . In real-world scenarios, data collection agents affected by seasonality (e.g., temperature, humidity) can cause the data to be collected far from the expert agent's full demonstration. For instance, consider a sensor with the collecting frequency is sensitive to the temperature, making it prone to frequent breakdowns in the summer, and this sensor should collect expert data throughout the entire year. However, due to its frequent failures during the summer, the collected dataset will deviate from d E .Motivated by these scenarios, we simulate time-dependent covariate shift scenarios by subsampling timesteps of the trajectory in extensive manners. Scenario 3: Segmented trajectory dataset As similarly explored in [30] , when decision-making horizons are extensively long, it is not feasible to gather multiple long trajectories keep tracking from initial states to ensure the stationarity of the expert dataset. Instead, it is more practical to gather shorter segments of expert demonstrations in such scenarios. We simulate this scenario by utilizing short segments of the pre-collected trajectories. To collect this, for each trajectory in the original dataset, we sample the starting timestep of the segment by using a subsampling method similar to Scenario 2. Then, we extract consequent segments with fixed-length timesteps and configure the segmented trajectory dataset similar to Figure 4b . In this scenario, we adjust the number of segments used for training to investigate the relationship between performance and the number of segments.Scenario 1 Table 3 summarizes the performance comparison on the rebalanced datasets. From the table, we observe that DrilDICE overally outperforms other methods across different proportions and tasks and outperforms 14 out of 15 problem settings with signficant margins. DemoDICE, AW-BC and OptiDICE-BC fail to outperform BC more than half of the problem settings, which demonstrates the distribution matching and the best-case weighting are not robust to these shift scenarios. Moreover, we can conclude a robustness of DrilDICE is not gifted by side effects from the implementation of DICE algorithms at least this scenario. DR-BC, which does not incorporate Bellman flow constraints, also exhibits improved robustness, outperforming BC in 9 out of 15 problem settings. However, DrilDICE consistently surpasses DR-BC across all settings, illustrating that the inclusion of Bellman flow constraints is crucial for effectively addressing the covariate shift of our interest.Scenario 2 As depicted in Table 4 , DrilDICE shows robust overall performance, achieving the highest mean performance in 11 out of 12 settings. Remarkably, while all baseline methods, including DR-BC, fail to surpass BC more than half of the problem settings, only DrilDICE consistently demonstrates exceptional robustness. This emphasizes that the critical role of incorporating Bellman flow constraint to enhance robustness in scenarios involving this type of covariate shift.Scenario 3 Figure 5 presents a performance comparison across three metrics while varying the number of segments. As illustrated, DrilDICE's normalized score consistently increases as the number of segments grows, and the target MSE exhibits a steady decreases maintaining levels that are lower or comparable to those of other approaches. The robust performance metrics also display consistent behavior, with an exception of walker2d task. Interestingly, we observed that a weak correlation between the target MSE and the episode return (see Figure E in the supplementary material).In summary, by constraining the uncertainty set to a plausible set, DrilDICE effectively minimizes target MSE and enables BC to imitate the agent robust to various covariate shift scenarios. 
conclusions and limitations: 
 We propose DrilDICE, a offline IL approach that is robust to the covariate shift caused by the data distribution deviated from the stationary distribution of the expert. By optimizing the Bellman-flowconstrained worst-case objective, our approach effectively minimize the surrogate loss of the expected error w.r.t non-shifted target distribution. We also suggests an extensive set of practical covariate shift scenarios of our interest and empirically show that DrilDICE successfully imitate the expert robust to those shift scenarios. For a limitation, we don't consider uncertainty from transition shift or noisy demonstrations. We expect that extending our approach into such scenario will be beneficial on many practical scenarios, such as Sim2Real or transfer learning tasks. Moreover, we believe exploring a real-world application of our method is also an attractive extend of the future work. = H-1 t=0 E d t π [R(s t , a t )] = HE dπ [R(s, a)]where R : S × A → [0, 1] is a true reward function. Denote a surrogate loss function ℓ, which measures how far the learner policy π is from the expert policy for the state s ∈ S. Here, consider 0-1 loss, i.e. ℓ(π(s), π E (s)) = I[π(s) ̸ = π E (s)]. Proposition 2. (Ross and Bagnell [21] ) The suboptimality of an imitator policy π is bounded byEQUATIONIn our problem setting, we assume the lack of access to dataset sampled from d E (s), instead deal with dataset sampled from an arbitrary distribution d D (s). When applying BC approach to minimizeE s∼d D [ℓ(π(s), π E (s))], it is possible to construct an example where this expected loss fails to upper bound the policy suboptimality, J(π E ) -J(π). To illustrate, we revisit [21] and consider an example that tightens the inequality (11) . Assume 0 < ϵ ≤ 1/H, 0 < δ < 1/H. EQUATIONTable A: Policies and their corresponding state distributions.Consider an MDP illustrated in Figure A , where S = {s 1 , s 2 , s 3 }, A = {a 1 , a 2 }, the transition T is deterministic with T (s 1 , a 1 ) = s 2 , T (s 1 , a 2 ) = s 3 , T (s 2 , a 1 ) = s 3 , T (s 2 , a 2 ) = s 2 , T (s 3 , a 1 ) = s 2 , T (s 3 , a 2 ) = s 3 and s 1 is the initial state. We define the expert policy π E and an imitator policy π and derive their corresponding H-horizon state distributions d E , d π in Table A . Define R(s, a) = I[a = π E (s)] for all s ∈ S, a ∈ A. Then, J(π E ), J(π) would be:J(π E ) = H s∈S d E (s) a∈A π E (a|s)R(s, a) = H 1 H + H -1 H = H J(π) = H s∈S d π (s) a∈A π(a|s)R(s, a) = H 1 H • (1 -ϵH) + H -1 H • (1 -ϵH) = H -ϵH 2 Hence, J(π E ) -J(π) = ϵH 2 . Note that E s∼d E [ℓ(π, π E )] = 1 H • ϵH + H-1 H • 0 + 0 • 1 = ϵ,therefore the both left-hand and right-hand sides of the equality (11) are equal to ϵH 2 . However, considering the expect loss with respect to d D1 , whered D1 (s 1 ) = δ, d D1 (s 2 ) = 1 - δ, d D1 (s 3 ) = 0 as described in Table A, we can calculate E s∼d D 1 [ℓ(π, π E )] = δϵH. Then, J(π E ) -J(π) = ϵH 2 > δϵH 3 = H 2 E s∼d D 1 [ℓ(π, π E )]where δ < 1/H. As δ decreases, the right-hand side also decreases, allowing it to be made arbitrarily small regardless of the left-hand side. This implies that there exist a scenario in which minimizing BC loss with an arbitrary data state distribution does not guarantee a reduction in the performance gap between the expert and the imitator.Although straightforward, we emphasize that distributionally robust optimization (DRO) approaches can minimize the suboptimality under the realizability assumption, i.e, d E ∈ Q. Remark 1. Given an uncertainty set Q ⊆ ∆(S) with d E ∈ Q, we can guarantee that J(π E ) -J(π) ≤ H 2 sup d∈Q E s∼d [l(π(s), π E (s))]. Furthermore, given Q 1 ⊆ Q 2 ⊆ ∆(S) with d E ∈ Q 1 , the upper-bound of Q 1 would be tighter than Q 2 , i.e. J(π E ) -J(π) ≤ H 2 sup d∈Q1 E s∼d [l(π(s), π E (s))] ≤ H 2 sup d∈Q2 E s∼d [l(π(s), π E (s))] Let Q ρ := {d ∈ ∆(S) : D TV (d, d D ) ≤ ρ} and Q DrilDICE ρ = {d ∈ ∆(S) : D TV (d, d D ) ≤ ρ, d satisfies Bellman flow constraints}. Since d E satisfies Bellman flow constraints, if d E ∈ Q ρ , then d E ∈ Q DrilDICE ρ with Q DrilDICE ρ ⊆ Q ρ .Then, the uncertainty set of Q DrilDICE ρ provides a more tighter upper-bound for the suboptimality compared to that of Q ρ . 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Yes, we include training and test details in Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Yes, we report standard errors for all experiments we did.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: Yes, we provide resource information in the appendix Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: We describe broader impacts in the appendix Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Experiments of our research grounded only on simulators.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We describe License in Appendix F.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Definef Soft-χ 2 (x) := x log x -x + 1 if 0 < x < 1 (x -1) 2 if x ≥ 1 and (f ′ Soft-χ 2 ) -1 (y) := exp(y) if y < 0 y + 1 if y ≥ 0 , ReLU(x) := max(0, x), ELU(x) := exp(x) -1 if x < 0 x if x ≥ 0 . Then, f -divergence choices can be summarized in Table B. Divergence f (x) (f ′ ) -1 (y) w * π,ν KL Divergence x log x exp(y -1) exp eπ,ν (s,a,s ′ ) α -1 χ 2 -Divergence 1 2 (x -1) 2 (y + 1) ReLU eπ,ν (s,a,s ′ ) α + 1 Soft χ 2 -Divergence f Soft-χ 2 (x) (f ′ ) -1 Soft-χ 2 (y) ELU eπ,ν (s,a,s ′ ) α + 1 TV-Distance 1 2 |x -1| - - Soft TV-Distance 1 2 log(cosh(x -1)) tanh -1 (2y) ReLU tanh -1 2eπ,ν (s,a,s ′ ) α + 1Table B: Summary of f -divergences and their associated functions. Table E : Number of sub-trajectories used in each scenario. Each sub-trajectory consists of one initial transition and 50 subsampled transitions from a complete trajectory. To enhance dataset support, one complete trajectory is appended for each specified count of sub-trajectories.To compare performance of IL methods on standard scenarios, 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 To evaluate the effect of f -divergence choice on imitation performance, we additionally evaluate OptiDICE-BC and DrilDICE with the soft χ 2 -divergence, as originally utilized in [16] , across all considered scenarios. The results for scenarios 1-3 are presented in Table F, G and Figure F respectively. These results indicate that the choice of the soft TV-distance significantly enhances the performance of DrilDICE compared to using soft-χ 2 . To explain this performance gain, we hypothesize that the soft-TV distance provides more discriminative weighting of samples based on their long-term policy errors. As shown in Figure B , conventional f -divergences (e.g. KL, soft-χ 2 , ...) make less pronounced to the long-term policy error, yet the soft-TV distance responds sensitively to changes in e, resulting in a more pronounced weight 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 [29] . Task p(D1) BC DR-BC (TV) OptiDICE-BC (Soft-χ 2 ) OptiDICE-BC (Soft-TV) DrilDICE (Soft-χ 2 ) DrilDICE (Soft-For all experiments in Four Rooms environment, we use CVXPY and convex optimization solver MOSEK with academic licenses. We also use Mujoco free licenses. Our code has been developed upon MIT licenses.We used Google cloud computing engine with 100 c2-standard-4 instances that have the following system specification:• Series : C2 Since we study distributionally robust learning, our approach can be used in fairness or bias reduction. As we discussed in experiment section, our algorithm is also beneficial when the training agent in the cost-constrained environment. Further, we expect our research contributes to build more trustworthy and robust AI systems.NeurIPS Paper ChecklistJustification: Yes, we read and conducted research with Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).38th Conference on Neural Information Processing Systems (NeurIPS 2024).Note that this constraint is analogous to the use of a robustness radius used in[19].For convenience, we evaluate the expected loss over dπ instead of the original target dE. It is important to note that DRO objectives are also upper-bounds for this expected loss if dπ ∈ Q. 


Decision: As there is a consensus on the acceptance, I would recommend it. It would be great if the author could add some more theoretical analysis (finite sample analysis) to see how we really address covariate shifts in a fundamental way. 
(Nitpick: I would recommend following the standard rule (not using official comments) in the rebuttal phase from next time)

Reviews: review: paper_summary: paper_summary: This paper devices distribution correction ratio estimation (DICE)-based optimization to mitigate the covariate shift issue in the Behavior cloning algorithm. They test their heuristic loss on Mujoco benchmarks.
strength_weakness: strength_weakness: The results in this work are preliminary and need to be evaluated carefully. I move all my concerns to the Questions section below instead of treating these as weaknesses. I will also rely on the author-reviewer and reviewer-reviewer discussion periods for updating my decisions.
questions: questions: - Section 3.2 requires further details and refined writing. For example, 
  - Slater's condition/assumption must be explicitly mentioned since the drildice algorithm relies on it. 
  - Example of $f$ and $f'$ should be mentioned in this section to make sense of the experiments section (they are only mentioned in Appendix as KL and $\chi^2$.) 
  - More details are needed to support this statement: "the problem can be solved by alternatively optimizing $\nu$ and $\pi$." Such as citing previous DICE works and so on. E.g reference https://proceedings.neurips.cc/paper_files/paper/2019/file/cf9a242b70f45317ffd281241fa66502-Paper.pdf
  - These details will also help in providing some theoretical justifications from previous DICE-related works, which this paper currently lacks.

- Comparison with previous works must be further expanded. Especially with [18].
  - Both this paper Section 3.1 and [18] share similar motivation to come up with loss functions to mitigate covariate shift created by $d_D\neq d_E$. 
  - At line 92, it is mentioned that $d_D\neq d_E$ but the underlying policy (which is $\pi_E$) is the same for the two state-action visitation distributions. Now, considering the definition of state-action visitation distribution at line 54, it is easy to see that $d_D$ and $d_E$ in fact only differ in the shifts corresponding to transition distribution T of the underlying MDP. Thus, this paper is considering transition distribution shifts as opposed to the claim at line 89!
  - Based on the above points, [18] also considers distribution matching in their inner problem w.r.t transition distribution T shifts.
  - Despite this similarity between [18] and this work, both take different approaches to design the loss. This work takes the Bellman flow constraints route whereas [18] considers direct DRO on the transition distribution T shifts.
  - [18] provides both theoretical guarantees and empirical evaluations and this paper provides only empirical evaluations.

- Regarding experiments
  - Can you please explain how $D_i$ at line 194 differs in size of the dataset compared to $D_E$? With $p(u)=0.4$, is it the case that 40% of states are sampled from $D_E$ and the rest are random states?
  - The analogy of sampling in line 187 "if a data collection device (e.g. cameras or sensors) operates at different recording frequencies in each room" can also be considered as shifts in transition distribution T, thus sharing the similarities as discussed before. Similar observation can be made for the Scenarios 1 and 3. Connection to transition shifts is more apparent in Scenario 2. So I think it is worthwhile adding more benchmark algorithms [2, 18].
  - Can you also include another scenario, training without covariate shift for instance? BC and other non-robust versions should outperform in such ablation-type tests.

My score reflects this review provided here.
limitations: limitations: na
paper_summary: paper_summary: This paper studies imitation learning when the offline dataset does not come from the stationary expert distribution. To address this problem, the authors introduce the objective of distributionally robust optimization into behavioral cloning. To avoid overly pessimistic solutions, the authors further incorporate the Bellman flow constraint and empirical distribution divergence into the min-max optimization objective. Based on the DICE technology, the solution can be derived. The experimental results have shown the proposal's effectiveness on corrupted datasets.
strength_weakness: strength_weakness: 1. I agree this paper is well-contained. This work successfully introduces the min-max formulation of DRO to the imitation learning community. However, the new setting proposed in the paper is not well-supported. The shift between the offline dataset and the stationary expert distribution is a reasonable setting. However, all experiments are validated based on simulated shifts, which somewhat makes me feel less excited.
2. The experiments are conducted on simulated covariate shifts. More evaluation in real-world cases will further improve this paper. 
3. Some notations have not been well defined, the $\Delta S$ in Equation 2, and the $\mathcal{W}$ in Equation 5.
questions: questions: In the navigation task like the current experiments in the Four Rooms environment, limited expert trajectories from some start points, especially, have only partial coverage on the state space, is there also a covariate shift with the stationary distribution? I personally think this might be a more real and direct case in real-world applications. Is the method proposed still effective in this case?
limitations: limitations: The authors have provided a discussion about the limitations and broader societal impacts.
paper_summary: paper_summary: This paper is motivated by the observation that BC is well-known to be vulnerable to the covariate shift resulting from the mismatch between the state distributions induced by the learned policy and the data collection policy. To solve this problem, the authors formulate a robust BC training objective and employ a stationary distribution correction ratio estimation (DICE) to derive a feasible solution. They evaluate the effectiveness of our method through an extensive set of experiments covering diverse covariate shift scenarios.
strength_weakness: strength_weakness: 1.	The experimental results do not fully support the advantages over baselines.
questions: questions: 1.	The advantage over OptiDICE-BC seems marginal in Table 2 and Table 4. What could be the reason?
2.	What does “marginal room (or action) distribution of dataset” mean in line 189?
3.	For mujoco experiments, the authors only use the expert split. How does the method perform in other splits that have lower quality?
limitations: limitations: One limitation mentioned by the authors in the last section is that: don’t consider uncertainty from transition shifts or noisy demonstrations.

Comments: title: Dear Reviewer oZLc
Reply to authors
Dear Reviewer k6J7
Dear Reviewer oZLc
Dear Reviewer zQPx
Additional Experimental Results
Additional Experimental Results
Additional Experimental Results
Rebuttal (continued)
comment: We are awaiting your feedback on any remaining concerns, and we would be very happy to resolve these issues. 

We are eager to address remaining concerns to enhance our manuscript, so please do not hesitate to let us know.

Thank you once again for your invaluable reviews.
We're glad to hear that some of your concerns have been addressed.

We truly believe that your insightful feedback has notably improved the quality of our manuscript.

If you have any remaining concerns, please let us know so that we can further enhance the quality of our research.
Thanks for your response. I agree that segmented trajectories have simulated a type of real-world covariate shift.
We're glad to hear that your concerns have been addressed, and we sincerely appreciate your positive review! 

We will incorporate all the findings discussed with you into the revised manuscript.
Thanks to the authors for adding new experimental results to address my concerns. The new results provide new evidence and support for the proposed method. I don't have further questions and I raised my score to 6 to lean to accept this paper.
We sincerely thank the reviewer for the thoughtful review and positive assessment of our work! 

We believe that your insights clearly enhance the quality of our manuscript. We will incorporate your valuable feedback and suggestions into the next revision.
Thank you for reflecting on the reviews. I have updated my score from 4 to 6. Good luck.
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.
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.
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.
## **Four Rooms Environment**

| Scenario | BC | OptiDICE-BC | DR-BC | DrilDICE (Ours) |
| --- | --- | --- | --- | --- |
| Room 1 | 90.84 ± 0.69 | 94.30 ± 0.41 | 91.38 ± 0.73 | **95.04 ± 0.48** |
| Room 2 | 89.16 ± 1.07 | 94.06 ± 0.65 | 89.28 ± 1.08 | **94.44 ± 0.62** |
| Room 3 | 88.50 ± 1.27 | 94.20 ± 0.98 | 88.70 ± 1.26 | **95.04 ± 0.86** |
| Room 4 | 90.94 ± 0.75 | 94.26 ± 0.54 | 90.94 ± 0.75 | **94.92 ± 0.37** |
| Action UP | 84.96 ± 1.33 | 92.06 ± 0.69 | 84.96 ± 1.33 | **93.22 ± 0.61** |
| Action DOWN | 89.96 ± 0.96 | 93.62 ± 0.62 | 89.96 ± 0.96 | **94.60 ± 0.39** |
| Action LEFT | 90.18 ± 1.11 | 91.86 ± 1.03 | 90.18 ± 1.11 | **92.62 ± 0.95** |
| Action RIGHT | 93.04 ± 0.63 | 94.46 ± 0.44 | 93.44 ± 0.60 | **94.52 ± 0.44** |

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

## **Scenario 1: Rebalanced Dataset**

| Scenario | Task | $p(D_1)$ | BC | OptiDICE-BC (Soft-$\chi^2$) | DrilDICE (Soft-$\chi^2$) | DR-BC | OptiDICE-BC (Soft-TV) | DrilDICE (Soft-TV) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Rebalanced by state | hopper | 0.1 | 24.65 ± 4.15 | 35.26 ± 3.19 | **58.92 ± 4.30** | 27.02 ± 4.31 | 12.72 ± 1.27 | 52.22 ± 5.57 |
|  |  | 0.5 | 35.38 ± 4.08 | 24.41 ± 5.12 | 60.91 ± 7.61 | 36.71 ± 3.56 | 8.61 ± 1.97 | **67.12 ± 8.18** |
|  |  | 0.9 | 11.23 ± 2.49 | 17.26 ± 2.10 | 28.29 ± 3.03 | 27.37 ± 4.89 | 10.44 ± 1.76 | **36.39 ± 6.10** |
|  | walker2d | 0.1 | 18.85 ± 4.05 | 5.91 ± 0.59 | 31.92 ± 5.94 | 14.69 ± 3.26 | 4.91 ± 1.08 | **51.55 ± 8.16** |
|  |  | 0.5 | 22.91 ± 3.13 | 11.12 ± 1.36 | 30.53 ± 3.47 | 45.08 ± 9.98 | 8.09 ± 0.42 | **73.74 ± 5.37** |
|  |  | 0.9 | 30.42 ± 7.05 | 17.51 ± 2.61 | 43.31 ± 4.54 | 46.04 ± 8.35 | 7.69 ± 0.45 | **77.60 ± 5.45** |
|  | halfcheetah | 0.1 | 49.31 ± 5.16 | 33.42 ± 4.70 | 44.79 ± 5.16 | 32.85 ± 3.82 | 7.05 ± 1.47 | **52.45 ± 3.62** |
|  |  | 0.5 | 37.98 ± 3.07 | 33.36 ± 3.04 | 41.13 ± 3.53 | 26.16 ± 4.91 | 6.14 ± 1.21 | **55.04 ± 3.27** |
|  |  | 0.9 | 15.54 ± 3.06 | 2.21 ± 1.16 | 7.28 ± 1.59 | 8.96 ± 3.26 | 1.02 ± 1.12 | **22.28 ± 2.88** |
| Rebalanced by action | hopper | 0.1 | 29.71 ± 4.00 | 28.37 ± 1.11 | 42.29 ± 6.39 | 25.92 ± 2.45 | 11.71 ± 2.12 | **56.60 ± 11.90** |
|  |  | 0.5 | 26.35 ± 4.88 | 30.03 ± 4.23 | 53.37 ± 9.50 | 35.13 ± 5.41 | 11.79 ± 1.14 | **73.80 ± 3.63** |
|  |  | 0.9 | 30.50 ± 3.60 | 38.92 ± 4.78 | **63.14 ± 7.12** | 36.56 ± 2.29 | 19.42 ± 2.82 | 48.99 ± 12.27 |
|  | walker2d | 0.1 | 23.61 ± 5.10 | 12.06 ± 1.20 | 40.72 ± 1.13 | 31.18 ± 4.23 | 7.27 ± 0.46 | **70.60 ± 3.21** |
|  |  | 0.5 | 32.29 ± 6.74 | 16.40 ± 1.70 | 47.93 ± 12.36 | 30.52 ± 3.89 | 6.37 ± 0.97 | **72.09 ± 8.70** |
|  |  | 0.9 | 16.87 ± 2.80 | 15.64 ± 3.57 | 43.68 ± 11.50 | 37.55 ± 8.97 | 4.60 ± 0.97 | **69.51 ± 8.54** |
|  | halfcheetah | 0.1 | 41.91 ± 4.80 | 26.62 ± 2.54 | 32.76 ± 1.87 | 27.50 ± 1.01 | 8.41 ± 3.38 | **56.42 ± 4.57** |
|  |  | 0.5 | 45.80 ± 4.45 | 45.52 ± 3.24 | 48.08 ± 5.50 | 33.39 ± 6.52 | 4.64 ± 0.84 | **60.81 ± 1.56** |
|  |  | 0.9 | 25.91 ± 3.35 | 4.28 ± 1.52 | 9.57 ± 2.34 | 12.08 ± 2.01 | 0.59 ± 0.68 | **29.19 ± 4.58** |

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


## **Scenario 2: Time-dependently Subsampled Dataset**

| Task | (a, b) | BC | OptiDICE-BC (Soft-$\chi^2$) | DrilDICE (Soft-$\chi^2$) | DR-BC | OptiDICE-BC (Soft-TV) | DrilDICE  (Soft-TV) |
| --- | ------- | --- | --- | --- | --- | --- | --- |
| hopper | (1, 1) | 28.89 ± 3.77 | 50.33 ± 6.60 | **54.83 ± 7.66** | 21.10 ± 2.26 | 22.77 ± 3.94 | 45.44 ± 5.11 |
|  | (1, 5) | 31.03 ± 0.90 | 39.54 ± 4.02 | 37.18 ± 7.92 | 25.00 ± 1.66 | 19.25 ± 1.21 | **45.60 ± 4.63** |
|  | (5, 1) | 26.75 ± 7.12 | **48.40 ± 12.98** | 39.91 ± 9.20 | 17.51 ± 3.38 | 25.68 ± 6.01 | 34.71 ± 9.00 |
|  | (5, 5) | 27.65 ± 6.71 | 32.46 ± 10.79 | **40.24 ± 7.41** | 23.20 ± 6.32 | 14.12 ± 3.59 | 25.61 ± 6.03 |
| walker2d | (1, 1) | 28.95 ± 5.34 | 17.42 ± 3.00 | 51.85 ± 5.30 | 45.66 ± 9.92 | 6.13 ± 1.03 | **81.21 ± 5.40**|
|  | (1, 5) | 61.48 ± 5.19 | 37.25 ± 5.66 | 64.46 ± 7.92 | 57.29 ± 4.79 | 17.55 ± 2.65 | **84.28 ± 4.89** |
|  | (5, 1) | 8.13 ± 0.72 | 4.43 ± 0.91 | 23.31 ± 3.44 | 17.97 ± 2.58 | 4.37 ± 0.52 | **48.23 ± 8.30** |
|  | (5, 5) | 6.65 ± 1.20 | 8.50 ± 2.27 | 14.40 ± 3.21 | 12.45 ± 1.84 | 5.54 ± 0.68 | **52.57 ± 6.24** |
| halfcheetah | (1, 1) | 33.74 ± 2.99 | 33.65 ± 5.49 | 33.43 ± 3.73 | 17.09 ± 2.43 | 9.93 ± 3.56 | **44.17 ± 5.62** |
|  | (1, 5) | 72.72 ± 2.60 | 52.94 ± 3.98 | 69.63 ± 3.86 | 61.81 ± 2.50 | 24.42 ± 3.27 | **77.12 ± 2.42** |
|  | (5, 1) | 2.35 ± 0.51 | 2.46 ± 1.05 | 3.97 ± 1.18 | 3.81 ± 1.26 | 1.29 ± 1.13 | **5.68 ± 1.50** |
|  | (5, 5) | 2.01 ± 0.91 | 1.61 ± 1.30 | 4.61 ± 1.55 | 2.85 ± 1.05 | -1.19 ± 0.37 | **5.50 ± 0.83** |

**Table C.** Expanded performance comparison on Scenario 2 (time-dependently collected dataset). (corresponds to Table 4)
## **Four Rooms Environment**

| Scenario | BC | OptiDICE-BC | DR-BC | DrilDICE (Ours) |
| --- | --- | --- | --- | --- |
| Room 1 | 90.84 ± 0.69 | 94.30 ± 0.41 | 91.38 ± 0.73 | **95.04 ± 0.48** |
| Room 2 | 89.16 ± 1.07 | 94.06 ± 0.65 | 89.28 ± 1.08 | **94.44 ± 0.62** |
| Room 3 | 88.50 ± 1.27 | 94.20 ± 0.98 | 88.70 ± 1.26 | **95.04 ± 0.86** |
| Room 4 | 90.94 ± 0.75 | 94.26 ± 0.54 | 90.94 ± 0.75 | **94.92 ± 0.37** |
| Action UP | 84.96 ± 1.33 | 92.06 ± 0.69 | 84.96 ± 1.33 | **93.22 ± 0.61** |
| Action DOWN | 89.96 ± 0.96 | 93.62 ± 0.62 | 89.96 ± 0.96 | **94.60 ± 0.39** |
| Action LEFT | 90.18 ± 1.11 | 91.86 ± 1.03 | 90.18 ± 1.11 | **92.62 ± 0.95** |
| Action RIGHT | 93.04 ± 0.63 | 94.46 ± 0.44 | 93.44 ± 0.60 | **94.52 ± 0.44** |

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

## **Scenario 1: Rebalanced Dataset**

| Scenario | Task | $p(D_1)$ | BC | OptiDICE-BC (Soft-$\chi^2$) | DrilDICE (Soft-$\chi^2$) | DR-BC | OptiDICE-BC (Soft-TV) | DrilDICE (Soft-TV) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Rebalanced by state | hopper | 0.1 | 24.65 ± 4.15 | 35.26 ± 3.19 | **58.92 ± 4.30** | 27.02 ± 4.31 | 12.72 ± 1.27 | 52.22 ± 5.57 |
|  |  | 0.5 | 35.38 ± 4.08 | 24.41 ± 5.12 | 60.91 ± 7.61 | 36.71 ± 3.56 | 8.61 ± 1.97 | **67.12 ± 8.18** |
|  |  | 0.9 | 11.23 ± 2.49 | 17.26 ± 2.10 | 28.29 ± 3.03 | 27.37 ± 4.89 | 10.44 ± 1.76 | **36.39 ± 6.10** |
|  | walker2d | 0.1 | 18.85 ± 4.05 | 5.91 ± 0.59 | 31.92 ± 5.94 | 14.69 ± 3.26 | 4.91 ± 1.08 | **51.55 ± 8.16** |
|  |  | 0.5 | 22.91 ± 3.13 | 11.12 ± 1.36 | 30.53 ± 3.47 | 45.08 ± 9.98 | 8.09 ± 0.42 | **73.74 ± 5.37** |
|  |  | 0.9 | 30.42 ± 7.05 | 17.51 ± 2.61 | 43.31 ± 4.54 | 46.04 ± 8.35 | 7.69 ± 0.45 | **77.60 ± 5.45** |
|  | halfcheetah | 0.1 | 49.31 ± 5.16 | 33.42 ± 4.70 | 44.79 ± 5.16 | 32.85 ± 3.82 | 7.05 ± 1.47 | **52.45 ± 3.62** |
|  |  | 0.5 | 37.98 ± 3.07 | 33.36 ± 3.04 | 41.13 ± 3.53 | 26.16 ± 4.91 | 6.14 ± 1.21 | **55.04 ± 3.27** |
|  |  | 0.9 | 15.54 ± 3.06 | 2.21 ± 1.16 | 7.28 ± 1.59 | 8.96 ± 3.26 | 1.02 ± 1.12 | **22.28 ± 2.88** |
| Rebalanced by action | hopper | 0.1 | 29.71 ± 4.00 | 28.37 ± 1.11 | 42.29 ± 6.39 | 25.92 ± 2.45 | 11.71 ± 2.12 | **56.60 ± 11.90** |
|  |  | 0.5 | 26.35 ± 4.88 | 30.03 ± 4.23 | 53.37 ± 9.50 | 35.13 ± 5.41 | 11.79 ± 1.14 | **73.80 ± 3.63** |
|  |  | 0.9 | 30.50 ± 3.60 | 38.92 ± 4.78 | **63.14 ± 7.12** | 36.56 ± 2.29 | 19.42 ± 2.82 | 48.99 ± 12.27 |
|  | walker2d | 0.1 | 23.61 ± 5.10 | 12.06 ± 1.20 | 40.72 ± 1.13 | 31.18 ± 4.23 | 7.27 ± 0.46 | **70.60 ± 3.21** |
|  |  | 0.5 | 32.29 ± 6.74 | 16.40 ± 1.70 | 47.93 ± 12.36 | 30.52 ± 3.89 | 6.37 ± 0.97 | **72.09 ± 8.70** |
|  |  | 0.9 | 16.87 ± 2.80 | 15.64 ± 3.57 | 43.68 ± 11.50 | 37.55 ± 8.97 | 4.60 ± 0.97 | **69.51 ± 8.54** |
|  | halfcheetah | 0.1 | 41.91 ± 4.80 | 26.62 ± 2.54 | 32.76 ± 1.87 | 27.50 ± 1.01 | 8.41 ± 3.38 | **56.42 ± 4.57** |
|  |  | 0.5 | 45.80 ± 4.45 | 45.52 ± 3.24 | 48.08 ± 5.50 | 33.39 ± 6.52 | 4.64 ± 0.84 | **60.81 ± 1.56** |
|  |  | 0.9 | 25.91 ± 3.35 | 4.28 ± 1.52 | 9.57 ± 2.34 | 12.08 ± 2.01 | 0.59 ± 0.68 | **29.19 ± 4.58** |

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


## **Scenario 2: Time-dependently Subsampled Dataset**

| Task | (a, b) | BC | OptiDICE-BC (Soft-$\chi^2$) | DrilDICE (Soft-$\chi^2$) | DR-BC | OptiDICE-BC (Soft-TV) | DrilDICE  (Soft-TV) |
| --- | ------- | --- | --- | --- | --- | --- | --- |
| hopper | (1, 1) | 28.89 ± 3.77 | 50.33 ± 6.60 | **54.83 ± 7.66** | 21.10 ± 2.26 | 22.77 ± 3.94 | 45.44 ± 5.11 |
|  | (1, 5) | 31.03 ± 0.90 | 39.54 ± 4.02 | 37.18 ± 7.92 | 25.00 ± 1.66 | 19.25 ± 1.21 | **45.60 ± 4.63** |
|  | (5, 1) | 26.75 ± 7.12 | **48.40 ± 12.98** | 39.91 ± 9.20 | 17.51 ± 3.38 | 25.68 ± 6.01 | 34.71 ± 9.00 |
|  | (5, 5) | 27.65 ± 6.71 | 32.46 ± 10.79 | **40.24 ± 7.41** | 23.20 ± 6.32 | 14.12 ± 3.59 | 25.61 ± 6.03 |
| walker2d | (1, 1) | 28.95 ± 5.34 | 17.42 ± 3.00 | 51.85 ± 5.30 | 45.66 ± 9.92 | 6.13 ± 1.03 | **81.21 ± 5.40**|
|  | (1, 5) | 61.48 ± 5.19 | 37.25 ± 5.66 | 64.46 ± 7.92 | 57.29 ± 4.79 | 17.55 ± 2.65 | **84.28 ± 4.89** |
|  | (5, 1) | 8.13 ± 0.72 | 4.43 ± 0.91 | 23.31 ± 3.44 | 17.97 ± 2.58 | 4.37 ± 0.52 | **48.23 ± 8.30** |
|  | (5, 5) | 6.65 ± 1.20 | 8.50 ± 2.27 | 14.40 ± 3.21 | 12.45 ± 1.84 | 5.54 ± 0.68 | **52.57 ± 6.24** |
| halfcheetah | (1, 1) | 33.74 ± 2.99 | 33.65 ± 5.49 | 33.43 ± 3.73 | 17.09 ± 2.43 | 9.93 ± 3.56 | **44.17 ± 5.62** |
|  | (1, 5) | 72.72 ± 2.60 | 52.94 ± 3.98 | 69.63 ± 3.86 | 61.81 ± 2.50 | 24.42 ± 3.27 | **77.12 ± 2.42** |
|  | (5, 1) | 2.35 ± 0.51 | 2.46 ± 1.05 | 3.97 ± 1.18 | 3.81 ± 1.26 | 1.29 ± 1.13 | **5.68 ± 1.50** |
|  | (5, 5) | 2.01 ± 0.91 | 1.61 ± 1.30 | 4.61 ± 1.55 | 2.85 ± 1.05 | -1.19 ± 0.37 | **5.50 ± 0.83** |

**Table C.** Expanded performance comparison on Scenario 2 (time-dependently collected dataset). (corresponds to Table 4)
## **Four Rooms Environment**

| Scenario | BC | OptiDICE-BC | DR-BC | DrilDICE (Ours) |
| --- | --- | --- | --- | --- |
| Room 1 | 90.84 ± 0.69 | 94.30 ± 0.41 | 91.38 ± 0.73 | **95.04 ± 0.48** |
| Room 2 | 89.16 ± 1.07 | 94.06 ± 0.65 | 89.28 ± 1.08 | **94.44 ± 0.62** |
| Room 3 | 88.50 ± 1.27 | 94.20 ± 0.98 | 88.70 ± 1.26 | **95.04 ± 0.86** |
| Room 4 | 90.94 ± 0.75 | 94.26 ± 0.54 | 90.94 ± 0.75 | **94.92 ± 0.37** |
| Action UP | 84.96 ± 1.33 | 92.06 ± 0.69 | 84.96 ± 1.33 | **93.22 ± 0.61** |
| Action DOWN | 89.96 ± 0.96 | 93.62 ± 0.62 | 89.96 ± 0.96 | **94.60 ± 0.39** |
| Action LEFT | 90.18 ± 1.11 | 91.86 ± 1.03 | 90.18 ± 1.11 | **92.62 ± 0.95** |
| Action RIGHT | 93.04 ± 0.63 | 94.46 ± 0.44 | 93.44 ± 0.60 | **94.52 ± 0.44** |

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

## **Scenario 1: Rebalanced Dataset**

| Scenario | Task | $p(D_1)$ | BC | OptiDICE-BC (Soft-$\chi^2$) | DrilDICE (Soft-$\chi^2$) | DR-BC | OptiDICE-BC (Soft-TV) | DrilDICE (Soft-TV) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Rebalanced by state | hopper | 0.1 | 24.65 ± 4.15 | 35.26 ± 3.19 | **58.92 ± 4.30** | 27.02 ± 4.31 | 12.72 ± 1.27 | 52.22 ± 5.57 |
|  |  | 0.5 | 35.38 ± 4.08 | 24.41 ± 5.12 | 60.91 ± 7.61 | 36.71 ± 3.56 | 8.61 ± 1.97 | **67.12 ± 8.18** |
|  |  | 0.9 | 11.23 ± 2.49 | 17.26 ± 2.10 | 28.29 ± 3.03 | 27.37 ± 4.89 | 10.44 ± 1.76 | **36.39 ± 6.10** |
|  | walker2d | 0.1 | 18.85 ± 4.05 | 5.91 ± 0.59 | 31.92 ± 5.94 | 14.69 ± 3.26 | 4.91 ± 1.08 | **51.55 ± 8.16** |
|  |  | 0.5 | 22.91 ± 3.13 | 11.12 ± 1.36 | 30.53 ± 3.47 | 45.08 ± 9.98 | 8.09 ± 0.42 | **73.74 ± 5.37** |
|  |  | 0.9 | 30.42 ± 7.05 | 17.51 ± 2.61 | 43.31 ± 4.54 | 46.04 ± 8.35 | 7.69 ± 0.45 | **77.60 ± 5.45** |
|  | halfcheetah | 0.1 | 49.31 ± 5.16 | 33.42 ± 4.70 | 44.79 ± 5.16 | 32.85 ± 3.82 | 7.05 ± 1.47 | **52.45 ± 3.62** |
|  |  | 0.5 | 37.98 ± 3.07 | 33.36 ± 3.04 | 41.13 ± 3.53 | 26.16 ± 4.91 | 6.14 ± 1.21 | **55.04 ± 3.27** |
|  |  | 0.9 | 15.54 ± 3.06 | 2.21 ± 1.16 | 7.28 ± 1.59 | 8.96 ± 3.26 | 1.02 ± 1.12 | **22.28 ± 2.88** |
| Rebalanced by action | hopper | 0.1 | 29.71 ± 4.00 | 28.37 ± 1.11 | 42.29 ± 6.39 | 25.92 ± 2.45 | 11.71 ± 2.12 | **56.60 ± 11.90** |
|  |  | 0.5 | 26.35 ± 4.88 | 30.03 ± 4.23 | 53.37 ± 9.50 | 35.13 ± 5.41 | 11.79 ± 1.14 | **73.80 ± 3.63** |
|  |  | 0.9 | 30.50 ± 3.60 | 38.92 ± 4.78 | **63.14 ± 7.12** | 36.56 ± 2.29 | 19.42 ± 2.82 | 48.99 ± 12.27 |
|  | walker2d | 0.1 | 23.61 ± 5.10 | 12.06 ± 1.20 | 40.72 ± 1.13 | 31.18 ± 4.23 | 7.27 ± 0.46 | **70.60 ± 3.21** |
|  |  | 0.5 | 32.29 ± 6.74 | 16.40 ± 1.70 | 47.93 ± 12.36 | 30.52 ± 3.89 | 6.37 ± 0.97 | **72.09 ± 8.70** |
|  |  | 0.9 | 16.87 ± 2.80 | 15.64 ± 3.57 | 43.68 ± 11.50 | 37.55 ± 8.97 | 4.60 ± 0.97 | **69.51 ± 8.54** |
|  | halfcheetah | 0.1 | 41.91 ± 4.80 | 26.62 ± 2.54 | 32.76 ± 1.87 | 27.50 ± 1.01 | 8.41 ± 3.38 | **56.42 ± 4.57** |
|  |  | 0.5 | 45.80 ± 4.45 | 45.52 ± 3.24 | 48.08 ± 5.50 | 33.39 ± 6.52 | 4.64 ± 0.84 | **60.81 ± 1.56** |
|  |  | 0.9 | 25.91 ± 3.35 | 4.28 ± 1.52 | 9.57 ± 2.34 | 12.08 ± 2.01 | 0.59 ± 0.68 | **29.19 ± 4.58** |

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


## **Scenario 2: Time-dependently Subsampled Dataset**

| Task | (a, b) | BC | OptiDICE-BC (Soft-$\chi^2$) | DrilDICE (Soft-$\chi^2$) | DR-BC | OptiDICE-BC (Soft-TV) | DrilDICE  (Soft-TV) |
| --- | ------- | --- | --- | --- | --- | --- | --- |
| hopper | (1, 1) | 28.89 ± 3.77 | 50.33 ± 6.60 | **54.83 ± 7.66** | 21.10 ± 2.26 | 22.77 ± 3.94 | 45.44 ± 5.11 |
|  | (1, 5) | 31.03 ± 0.90 | 39.54 ± 4.02 | 37.18 ± 7.92 | 25.00 ± 1.66 | 19.25 ± 1.21 | **45.60 ± 4.63** |
|  | (5, 1) | 26.75 ± 7.12 | **48.40 ± 12.98** | 39.91 ± 9.20 | 17.51 ± 3.38 | 25.68 ± 6.01 | 34.71 ± 9.00 |
|  | (5, 5) | 27.65 ± 6.71 | 32.46 ± 10.79 | **40.24 ± 7.41** | 23.20 ± 6.32 | 14.12 ± 3.59 | 25.61 ± 6.03 |
| walker2d | (1, 1) | 28.95 ± 5.34 | 17.42 ± 3.00 | 51.85 ± 5.30 | 45.66 ± 9.92 | 6.13 ± 1.03 | **81.21 ± 5.40**|
|  | (1, 5) | 61.48 ± 5.19 | 37.25 ± 5.66 | 64.46 ± 7.92 | 57.29 ± 4.79 | 17.55 ± 2.65 | **84.28 ± 4.89** |
|  | (5, 1) | 8.13 ± 0.72 | 4.43 ± 0.91 | 23.31 ± 3.44 | 17.97 ± 2.58 | 4.37 ± 0.52 | **48.23 ± 8.30** |
|  | (5, 5) | 6.65 ± 1.20 | 8.50 ± 2.27 | 14.40 ± 3.21 | 12.45 ± 1.84 | 5.54 ± 0.68 | **52.57 ± 6.24** |
| halfcheetah | (1, 1) | 33.74 ± 2.99 | 33.65 ± 5.49 | 33.43 ± 3.73 | 17.09 ± 2.43 | 9.93 ± 3.56 | **44.17 ± 5.62** |
|  | (1, 5) | 72.72 ± 2.60 | 52.94 ± 3.98 | 69.63 ± 3.86 | 61.81 ± 2.50 | 24.42 ± 3.27 | **77.12 ± 2.42** |
|  | (5, 1) | 2.35 ± 0.51 | 2.46 ± 1.05 | 3.97 ± 1.18 | 3.81 ± 1.26 | 1.29 ± 1.13 | **5.68 ± 1.50** |
|  | (5, 5) | 2.01 ± 0.91 | 1.61 ± 1.30 | 4.61 ± 1.55 | 2.85 ± 1.05 | -1.19 ± 0.37 | **5.50 ± 0.83** |

**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.