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
Venue: NeurIPS 2024 poster (2024)
Keywords: Imitation Learning
Behavioral Cloning
Robust Learning
Primary area: reinforcement_learning
TL;DR: We propose DrilDICE, a method that leverages a distributionally robust objective with using DICE technique, to mitigate the covariate shift in offline imitation learning with arbitrary data state distributions.

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.

Mitigating Covariate Shift in Behavioral Cloning via Robust Stationary Distribution Correction
Seokin Seo', Byung-JJn Lee 2.3 Jongmin Lee", HyeongJoo Hwang',
Hongseok Yang', Kee-Eung Kim
KAIST, "Korea University. "Gasss Labs Inc., +UC Berkeley
siseoßai kaist. ac. kr, byungjunlee@korea.co jongmin. leedberkeley. hjhwangdai kaist. kr. hongseok, bangseok.ynggeeee yangekaist. ac.kr kr, kekimokaist.c ac
kr,
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 collected from an arbitrary state distrihution 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 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 DriIDICE, which utilizes a distributionally robust BC objective by employing stationary distribution correction ratio estimation (DICE) derive 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.
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 Despite recent works on offline IL that explore scenarios involving supplementary datasets, such suboptimal demonstrations (5) [4). behavioral cloning (BC) remains a compelling option in practice since BC does not require additional datasets except expert demonstrations, However, the efficacy 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, substantial body of IL research is dedicated to addressing this covariate shift issue 5 24
infeasible or expensive [5] 13 28
A common approach for solving offline IL problems is to the so-called distribution matching 14 [5] 31), which aims to align the stationary distribution of the
objective or its variant [1] 12
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
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

assumption may break, that is, sampled expert demonstrations may be from shifted version of the stationary distribution. For instance, expert demonstrations may be collected by first sampling states from 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 result, the offline II. algorithms based on the distribution-matchin 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) 251. 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 offine 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 Work
2.1 Offline Imitation Learning with Arbitrary State Distributions
We consider a Markov decision process (MDP) without rewards, which is defined as a tuple of (8. A, I", po. 7) where S is a state space, AA is an action space, T : A + A(8) is transition
distribution, Po E (SS) is an initial state distribution and 7 E 00, 1) is a discounted factor. We focus on the class of deterministic policies 11 = {at S + A}. Given a policy π E 11, the
stationary distribution d, of π is defined as d=(s,a) = (1 7) Ei= 7a Pr(s; s, a a) where
so ~ Po: a = (st)), St+1 ~ TTss, az). For convenience, we use gg to denote the policy of an expert and write simply de for the stationary distribution d=æ of TT.
Assuming the existence of expert policy ाæ II, the goal of imitation learning (IL) is to recover Te by utilizing some demonstrations of the expert. Specifically, offline II. prohibits online interactions and relies solely given offinn dataset that consists of demonstrations of the expert. We consider that the demonstration dataset ' is the collection of (s, s') triplets where is sampled from an arbitrary state distribution ds(s). a is determined by (s)) and s/ is sampled from T [s, a). Note that conventional IL approaches assume that dp is close to the expert's 's state distribution dg. However, unlike previous studies, we focus the imitation learning with an arbitrary state distribution dp. without assuming that dp to be a state stationary distribution of any policy.
As mentioned in the introduction, one common approach for offline IL is use the distribution matching objective or its variant [ 12 14 15) and to find an imitator policy whose stationary distribution matches that of the expert policy. But we will show later in the paper, the algorithms based on this approach do not perform well when there is a covariate shift in the given demonstration dataset d (s, a. in a way explained in the previous paragraph. In the paper, we propose an offline II. algorithm that achieves good performance in the presence of such a covariate shift by using a version of distributionally robust optimization.
2.2 Covariate Shift in Imitation Learning
Covariate shift is a widely used term in machine learning [ 71. Traditionally, it refers to the phenomenon where the distribution of covariates (inputs) during training differs from that during testing, i.e. given a covariate variable X E X' and a response variable Y E y,
Prein [Y|X- x) = Prest (Y = x)) Ex,
Purin (X) A Ples(
The distribution shift that is considered in this paper has this form where state and action correspond to X and and Purin and Prest correspond to the distribution dn of the given demonstration dataset and the stationary distribution dg of the expert policy πg.
2

---

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 worst- case supremum by using kernel methods 255 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 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 dp of the dataset being a noisy version of dg 6 [2). These issues are usually addressed by using robust optimization techniques [6 19,, which are also utilized in our approach. However, important to note that the types of distribution shifts considered in our work and in prior researches differ within the context of offlin II.
To emphasize, we are not focused on addressing covariate shifts induced by transition shifts noisy demonstrations. Instead, our work investigates 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 (s) of the samples in the given demonstration dataset differs from dg(s) but both distributions share the expert policy given states, i.e., TD(s) = Tg(ss) for all s E S.
3 Mitigating Covariate Shift in BC
3.1 Robustness to Covariate Shift
Assume that we have a free access to Ty. such that we can optimize our trained policy π to close to it on any state. We can write the BC objective for optimizing as:
min Es-dg ((T(s)) Tx())]
for some supervised learning loss (e.g. squared error, cross-entropy. ...). Previous studies on imitation learning have assumed that dp dg. Under this assumption, BC approach of:
min Eg-dn (((T(s)))))ss))))))))))){<)))[ (E))
does show strong performance 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 stationary distribution dDM that is close to data distribution dp with the following objective:
dDM == arg max - D(d(ss dEA(S)
s. dd(ssa (1- + [[(sss Vs
3,a
where DD is any divergence between two probability distributions. The solution of the optimization problem dDM can be used for the state distribution for BC instead of dg, 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.c., when the state distribution of dp diverges from that of dg, both two approaches above can fail miserably, as there no guarantee on improving policy performance when dg and 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 know Egwdg [[()(s), ((T(s) (T)) & maxdea(s) ( (s))], which holds
with dy as well, and we can instead optimize for
min max Ee<d(((T(s)) TE()))). dE
which corresponds to an bouer- minimization for optimizing ग with respect to some uncertainty set @ C (S)). However, when @ is choosed as (SS), 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, choice of crucial to obtain an appropriate solution of distributionally robust optimization.
3

---

3.2 Towards Less Pessimistic Robust Objective
For notational convenience, denote (+)) := {(T(s), (s)) In the objective (3) the worst-case policy loss considered among an uncertainty Q. However, when Q = A(S). would be too pessimistic and the objective can be improved when we consider a smaller choice of Q. To keep the fact that is an upper-bound minimization for Eq. [ which is the loss with dg, the expert state distribution dg should still be contained in Q even if we aim to use smaller set.
To facilitate this, impose two constraints Q with respect to state distribution d: (1) d should satisfy Bellman flow constraint, (2) should be close enough to dp to prevent d diverging too far from _ The first constraint constraint that should be satisfied by dr, and thus it effectively reduces the potential distribution set for d if we impose it properly. The second constraint may not seem necessary, but
estimate the policy loss from finite samples from dp, diverging too far from dp will reduce effective number of samples. Then, by introducing f-divergence regularization enforce the second constraint, we consider the following constrained optimization problem:
min max dEA(S)
[ (s)] (d(()d())))
(4)
s.t. Ed(s,ar -1-7)) 7))p(()) - T(sBs, A& E
where m denotes /-divergence with convex function f. We develop a practical algorithm in stationary distribution correction ratio estimation (DICE) framework style. From Eq. 4 we can take Lagrangian to solve the inner constrained optimization, i.c.
max min Σ d(s, (s) - a& aD (d(s al)dssa)) dEA(S)
8,a
Σ v(s)[(1 ())) + Σ a(s(5, - Ed(s,a
(5)
3,
max min(1 7) Br [s)] + ((s,,,)))d af (w(s, a)) w (s,
where 6x, = (C+(8) E. T("')s, v(s)), w(s,u) = d(s.a) and a class of
ap(aaa)
functions w : S A R R22-
By Slater's condition 3 since the strong duality holds on Eq. 5 we can change the minimax into maximin optimization as follows:
min max E [v(s]] + + (s..,,'')'d a-aff (w(s,a)) w(s, (s,
(6)
WEW
and the closed-form solution w² can obtained by solving for the inner maximization.
Proposition 1. (Lee et al. (T6)) The closed-form solution for the inner maximization of Eq. 0 is
w " a) = max (P) (Enp(),a) )) Vs, a
After obtaining t the minimax optimization 6 is converted into minimization problem with respect to v. As result, each objective for can be summarized as follows:
min (1 - (1-7) [Mss]] + B(s,a, [-af (s, + w w (s,
(8)
min E (s,a)/dd [w] (s, (s)))
Hence, the problem can be practically addressed by alternatively optimizing v and æ following recent developments in DICE approaches [6 [7]. We call this distributionally robust optimization approach to covariate shift in offline II. as DriIDICE (Distributionally Robust Imitation Learning via DICE).
Note that this constraint is analogous the of robustness radius used in 19

---

Sch-TV
0.2
0.0
-0000
-75
-5.0
-25
0.0
2.5
3.0
7.5
10.0
-1.00
-7.
55.0
-25
0.0
2.5
5.0
7.5
10.0
Figure 1: Illustration of soft TV-distance. (left) f functions, (right) corresponding derivatives f".
3.3 Soft TV-Distance
To facilliate fair comparison with the uncertainty defined by total variation (TV) distance, as utilized in [19], we introduce the soft-TV distance as specific choice of f-divergence of DriIDICE. It important to note that Proposition I requires (f")=1, the invertable derivative of f. However, generator function of TV-distance STV(#) (z) = 2 1/, lacks invertible derivative as illustrated in Figure rendering the direct application TV distance in DrilDICE challenging.
We technically overcome this limitation by relaxing the derivative function of TV-distance. Given that the derivative function TV-distance manifests step function, we choose relax this function by employing the tanh function. Specifically, we utilize the log-cosh function [2] to f as follows:
1 fSuNTV (a) log(cosh (x: 1)), (f'ORTV) (c) tanh (2y) +1.
By plugging SorrTy into Eq. 7 obtain the closed formed-form solution w of the inner maximization objective (6) tailored to a specific choice of fsont.V as follows:
(207 (8, (8, (207m((5a
(20T
w (s, a) ReLU
(tanh-
(10)
a
where ReLU(z) max 0). This choice of f enables DriIDICE to obtain closed form solution of the inner maximization problem while maintaining similar properties of TV-distance. For other possible choice of /-divergence DriIDICE, see Section B in the supplementary material.
4 Experiments
4.1 Comparison on Baselines Objectives
To provide clear descriptions about the baselines used for comparison in our experiments, we summarize the relevant objectives in Table II
Adversarial Weighted BC (AW-BC) We refer to adversarial weighting by following the termi- nology from (27 as the minimax objective without constraints. As we discussed in Section 3.22 the adversarial weighting objective is also an upper-bound overly pessimistic because it considers the entire state distribution space, including distributions that do not correspond to the stationary distribution of any policy.
the target objective. However, it tends to be
Distributionally Robust BC (DR-BC) To adjust the level of the robustness, DR-BC 19 can be employed in this context. Despite originally designed address transition shifts, DR-BC remains relevant to our scenario since it considers the uncertainty set over state distributions by adopting robustness radius hyperparameter p. Notably, f-divergence constraint of DriIDICE functionally analogous to the robustness radius constraint Drv Drv(d))))) p. However, our approach, which incorporates Bellman flow constraints, addresses more restricted set when constrained by the equivalent radius level, thereby offering tighter upper-bound to Eq. [

---

Functionality
Objective
Distribution matching
min Bag [Cx(s)]
S.L.
ddm arg max [Dy(d(s,a)ldd((s..) dEQ
AW-BC objective
min max deAS
E [
DR-BC objective
min dPPv(dld
max
Exna[Cr(s]]
Best-case weighting
min max de
E. Exvel C+)) -DDD(((s..allld((s..))
Worst-case weighting (Ours)
min de
- (a0p
Table 1: Objective comparisons for related approaches. Denote stationary distributions class Q ie., Q = {d E A(S) d(s) = (1 - (s)) + E (s]8,, a)d(s, a)vs E S)
Best-case Weighting If the sign of the cost function in the worst-case weighting objective is converted, the objective seeks to find a cost-minimizing stationary distribution that is close to the data collection policy. We call this objective the best-case weighting. When dg Q, the best-case weighting minimizes the lower bound of the target objective since minde Egl [E(C((s)) S. Edg (C (s)], which is not relevant minimize the target objective. Despite there no direct connection, to ensure that performance of our method is not due to side effects of algorithm's 's implementation, we include this method for performance comparison in the following experiment section.
4.2 Toy Domain Experiment: Four Rooms domain with Imbalanced Datasets
4.2.1
Experiment Settings
Room
Room
Four Rooms a grid- world environment which aims to find path from a starting state to goal state. Fach cell in the grid represents state and the starting state and the goal state are marked by the orange and the green box in Figure S respectively. The agent can choose one of 4 actions from each state: UP, UDW,, LEFT, RIGHT, and the rooms numbered shown in Figure 2 We pre-collect expert dataset DE by using a determinis- tic expert policy ag and online interactions.
Four Rooms environment
Covariate shift scenarios To investigate that approach can effectively addresses the covariate shift, need the dataset that notably deviates from the expert stationary distribution dg. 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 data collection device (e.g. cameras sensors) operates different recording frequencies in each room, the frequency of observing a room in the dataset will not match dg. To simulate this scenario, we design datasets where the marginal room action) distribution of dataset deviates from dg. By manipulating the marginal proportion of 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 marginal distribution a certain fixed probability and resample transitions to
Room 3
Room 4
Figure 2: Four Rooms environment and deterministic expert (red arrows).
Pa( Pa(u) of the variable be manipulated configure the dataset D, through the following process. (See Section Cin in the supplementary material)
Dj = {(s. a)lu pi(u), s Dg(s(u), a = TT(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 zero (i.e. max,es (s) 0), we conduct experiments using function approximators instead of tabular setting. We initialize xo for the cost function C, as BC policy. To solve optimization problems, each algorithm is implemented

---

Scenario
Metries
BC
DemoDICE
AW-BC
DR-BC
OpHDICE-BC
DrIDICE (Ours)
iordalized score
90.84 0.60
91.62
50.84
91.38 0.73
94.30+
98.04 0.48
Rocen Marginal Manipulation
Woest-25%
1.64
1
Target 0-1 loss (x 10"))
8.69 4 0.06
8.28 0.06
7.92
4.52::0.55
7,85 a 0.05
$,31 0.07
Normalized score
89.16 1.07
89.72 1.00
89.20
89.28 1.08
94.06 0.65
94.44 0.62
Room
Wonst-5%%
$7.84 3.90
59.36 3.83
57.28 2.88
58.32 3,92
76.24
77.76
Target D-1 loss (x (x10P)
0.17
10.51 0.14
9.34 0.00
10.72 1,03
6,85 0.08
6.46 0.07
Notualiool score
88.50
89.66
1.21
88.68
1.16
88.70 4: 1,26
94,20 0.98
95.04
Room
Woest-25%
55.28 +4.74
59.60
55.52 4:.40
55.92 4,75
76,80 3.92
80.16 3.44
Targtt 0-1 loss (x
1..6
13.23 0.15
11.75
13.74:4 1.11
9.64
8,51
Notmalized score
0.75
Room 4
Woest 25%
64.00
65.36 2.81
62.72 2
64.00 # 2.89
2.115
79.68 1.49
Target loss 10")
9.30
8,20
Normal
84
Action Marginal Manipulation
Action UP
Worst-25%
13.9214.43
1
43:44 4.18
13.92 1. 1,43
68.32 15.32.22:33
72.88 2.43
Target loss
Normalized
0.96
91,28
0.95
89.96 0.96
93.62 0.62
94,60 0.39
Action DOWN
Worst-25%
60.80
65.36 3.27
02.24
60.80 3.45
74.48 2.47
78.40
Target loss 10%)
0..45
9.13
8.79 0.11
10.45 0.89
7.73
0.06
Normalized score
90.18+1111
90.32 1,10
90.80 1:11
90.18 1,11
91.86 1.03
92.62 0.95
Action LEFT
Worst-25%
61.60
2,166
64.40
4:002
(0 4.15
68.32
71,36
Target 0-1 loss (x
1.04
11116.0.15
10.56 0.115
11.12 1.02
9.91
9,16
Normalized score
93.04 2.663
93.12
92.76.0..51
93.44 0.60
94,46 0.44
94.52 0.44
Worst.25%
72.48
73.76
Target D-1 loss (xP)
7,07 20006
7.04
0.06
6.58
6.62 4.04
6,48 0.00
6.56 0.05
Table 2: Comparison of different methods across manipulated datasets in Four Rooms environment. (p(u) 0.4) Each experiment repeated with 50 times and the average values with their standard errors are reported. The best average values are highlighted in bold.
dp(s)
dp(s) W
dram (8)
Figure 3: Visualizations of policy behaviors and weights. (a) dd(s) dataset state distribution, (b) du... behavior of BC policy, (c) ((s)w" (s): corrected state distribution that maximizes Cx. (d) behavior of dz... (s): DrilDICE policy (BC weighted with w (s)).
with a two-step procedure: (1) optimize each objective by using the cost C., determined by an
initial policy and obtain wap: (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] C For choices of f-divergence for OptiDICE-BC and DriIDICE, we select KI diveeieeeee instead of soft TV-distance for simplicity of convex optimization solver implementation. We consider following baselines:
BC: a standard behavioral cloning without any regularization.
DemoDICE [4 representative distribution matching approach offline IL. We compare special case of DemoDICE that does not exploits supplementary datasets.
AW-BC: adversarial weighting method without Bellman flow constraints and robustness radius.
DR-BC
[9
distributionally robust BC method without Bellman flow constraints.
OptiDICE-BC [16]
representrative method the best-case weighting.
Evaluation metrics
The following metries 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-1 loss (i.e. E,, ([( + {
'For convenience, we evaluate the expected loss over d., instead of the original target dg. It is important to note that DRO objectives also upper bounds for this expected loss if dz E Q.

---

Task
p(D.)
BC
DemoDICE
AW-BC
DR-BC
OptiDICE-BC
DriIDICE (Ours)
0.1
4.2
3.2
17.9
2,30
1.3
52.2 +
Rebalanced by state
hopper
0.5
35.4
3 4
15.1
2..6
36.7
8.6
67.1 = 8.2
0.9
2.5
13.4
27,9
10.4 + 1.8
36.4 6.1
0.1
18.9
2..1
1.2
3.3
4.9
51.6 8.2
walker2d
0.5
3.1
3.7
2..
45.1
8.1 4.4
13.7
7.4
2.36
3.08
4..0
1.5
7.5
0.1
5.2
3..1
4.3
3.89
1.1
$2.5
halfchectah
0.5
38.0#3.1
4..9
32.6
26,2
6.1
55.0 3.3
0.9
3.1
1.0
7.8
9,0
1.0 1.1
22.3 + 2.9
0.1
4.0
3.0
2..6
2.5
11.72211
56.6 11.9
Rehalanced by action
hopper
2,,44499
3..3
30.8
35,1
11.8
73.8
0.9
3.6
16.5 1.4
1.93
3.36
19.4
49.0
0.1
23.6
6.56
2..4
31,2
7.3
70.6
walker2d
0.5
6.7
3.5
25,7
3,99
6.4
72.1 8.7
0.9
16.9
12.0
16,9
37.6
4.6
69.5 + 8.5
4..9
3.19
4..7
2..5
8.4
56.4 4.6
halfcheetah
0.5
4.5
3.0
3..7
33.4
4.6
60.8 1.6
0.9
2..9
8.8
1.72
12.1
0.7
29.2 + 4.6
Table 3:
Performance comparison on Scenario (rebalanced dataset). p(D,) determines the propor- tion of dataset D1, which is close to the representative point. Each experiment is repeated with times and the average normalized scores with their standard errors reported. The highest mean performance scores are highlighted in bold.
4.3 Results
As depicted Table 2 DriIDICE 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 DE 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, DriIDICE 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.
Figure jillustrates the behaviors of BC and DriIDICE in the Room 3 scenario. With increased dataset proportion visiting Room 3, as shown in Figure The performs accurately in Room 3, predicting the correct expert actions in nearly all states (26 out of 27) in Room 3. However, its performance remains suboptimal in other rooms. DriIDICE, leveraging the cost derived from Tbe- computes the corrected state distribution through the worst-case weighting, denoted as whe as depicted in Figure (c). Notably, states mispredicted by Toe or those with high probabilities for the dataset distribution tend to obtain relatively higher weight values, contributing more to loss optimization. Consequently, DriIDICE effectively trains the agent by correcting suboptimal behaviors.
4.4
D4RL Dataset with Covariate-Shifted Expert Demonstrations
4.4.1 Covariate Shift Scenarios
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-works and simulate those scenarios by restructuring the given dataset to simulate these dataset deviation scenarios (For experiments on standard scenarios, re- datasets
walker2d-expert, half included in D4RI. benchmark (9] as original datasets. We utilize the soft TV-distance defined in Section 3 3.3as as f-divergence for OptiDICE-BC. and DriIDICE for these scenarios (For experiments with another choice of fidieergence, see Section E.4) See Section D for implementation details.
fer to Section Eb. We utilize dopper-expen
Rebalanced dataset From a practical perspective, data collector may encounter scenarios in which the costs associated with taking an action or visiting specific state exhibit substantial costs varying across different states 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 to 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 significantly different, hence their proportion shifted from the original datasct.
Scenario 1:

---

400
Trajectory index
01.50
05.50
Counts
200
15
20
200
400
600
800
1000
10
30
40
Timesteps
Timesteps
(a) Timesteps subsampled for Scenario 2.
(b) Segments collected for Scenario 3.
Figure 4: Illustrative examples for generating datasets for Scenario and
Task
(a.b)
BC
DemoDICE
AW-BC
DR-BC
OptiDICE-BC
DriIDICE (Ours)
(1. 1)
28.9
2,84
18.0
21.1
2,9
4.14
(1.5)
3.9
25.7
1.6
2.7
19.3 1.2
4.66
hopper
(5
2.1
2.4
23.2
47.5
2,07
34.7 9.0
Timestep dependency
(5
6.7
11.7
27.7
23.2
1..1
2.06
1)
2.3
21.7.4.11
27.6
45.7
1.0
81,2
(1
5.2
6..0
34.7+5.9
57.3
17.6+2.7
walker2d
15 I)
8.1
7.4008
8.8
18.0
4.4
48.2
(5,5)
6.2
7.2
10.6
1.85
0.7
5.2
1)
33.7
3.2
33.4
17.1
9.9
5.6
(1,5)
2.6
1.42
71.1
61.8
2.3
77.1 + 2.4
halfcheetah
1)
2.4
4.6
4.240.4
3.8
1.
5.7
(5.5)
2.0
2.922.
0.9
2.9111
-1220.4
5.5
Table 4: Performance comparison on Scenario 2 (time- seine-deeendently collected dataset). Each experi-
ment is repeated times, and the average normalized scores with their standard errors are reported. The highest mean performance scores are highlighted in bold.
In order to implement this scenario, we measure the distance between each state and the represen- tative point of the states, then split the dataset into two groups D1. D2 based on the statistics of distances: Dj comprises states close to the point. while D2 is consists of states farther from the point. Subsequently, we resample datasets according predetermined proportion of p(D1) 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 dependenedd variables [8 [8 26). In real-worke scenarios, data collection agents affected by seasonality (e.g., temperature, humidity) can cause the data to be collected far from the expert agent's s full demonstration. For instance, consider sensor with the collecting frequency is sensitive the temperature, making it prone to frequent breakdowns the summer, and this sensor should collect expert data throughout the entire year. However, due its frequent failures during the summer, the collected dataset will deviate from dg-
Motivated by these scenarios, simulate time-dependent covariate shift scenarios by subsampling timesteps of the trajectory in extensive manners. To model the timestep sampling distribution, we utilize Beta distribution B(a, b). Since the support of ((a, b) is [0, 1]. with multiplying with the maximum timestep and discretization, we can sample timesteps in various shapes of distributions by adjusting parameters b. We conduct experiments four parameter combinations, (a, 1), (1,5), (5, 1), (5) and the timestep distribution notably varies as depicted in Figure 4a
Scenario 3: Segmented trajectory dataset As similarly explored in (30), when decision-making horizons are extensively long, it 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] this scenario, we adjust the number of segments used for training to investigate the relationship between performance and the number of segments.

---

hopper
walker2d
halfcheetah
Worst25N Nommlizd score
75
25
20
60-
20
20
Target MSE
0.9
1.2
2.0
1.0
0.7
0.0
0.6
10
10
100
200
Number Segments
Number Segments
Number Sogments
BC
DeMoDICE
AW-BC
DR-BC
OptDICE-BC
DriiDICE (Ours)
Figure 5: Performance comparison on Scenario 3 (segmented dataset) along the number of segments. The points and shaded areas indicate the means and standard errors measured over 5 repetitions.
4.4.2 Results
Scenario 1 Table summarizes the performance comparison on the rebalanced datasets. From the table, we observe that DriIDICE 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 robustness of DriIDICE 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, DriIDICE 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 As depicted in Table 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 DriIDICE consistently demonstrates exceptional robustness. This emphasizes that the critical role of incorporating Bellman flow constraint to enhance robustness in scenarios involving this type covariate shift.
Scenario Figure 5 presents a performance comparison across three metrics while varying the number of segments. As illustrated, DrilDICE's normalized score consistently increases 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 metries 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 in the supplementary material).
In summary, by constraining the uncertainty set to a plausible set, DriIDICE effectively minimizes target MSE and enables BC to imitate the agent robust various covariate shift scenarios.
5 Conclusions and Limitations
We propose DriIDICE, 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-flowe constrained worst-case objective, our approach effectively minimize the surrogate loss of the expected error w.r. non-shifted target distribution. We also suggests an extensive set of practical covariate shift scenarios of our interest and empirically show that DriIDICE 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 real-world application of our method is also attractive extend of the future work.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by IITP grant funded by MSIT of Korea (No. RS-2020-112009H4 No. RS- 2022-11220311 I, No. RS--019--III 111907555, No. RS-2019-II190077 No. 2022022-00003 No. RS- 2024-00343989, No. RS-2022-00455888), LITP-ITRC (LITP-2024-RS) 20--2222-22222--0--33333333 NRF of Korea (No. RS-2024-00451 162, No. RS-2023-002796600, BK21 Four project, KAIST-NAVER Hypercreative Al Center, and NSF A14OPT AI Centre.
References
[1] Niall Adams. Dataset shift in machine learning. Journal of the Royal Statistical Society Series
A: Statistics in Society, 2009.
[2] Mohammad Ali Bashiri, Brian Ziebart, and Xinhua Zhang. Distributionally robust imitation learning. Advances in Neural Information Processing Systems, 2021.
[3] Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
[4] Jongseong Chae, Seungyul Han, Whiyoung Jung, Myungsik Cho, Sungho Choi, and Youngchul Sung. Robust imitation learning against variations in environment dynamics. In Proceedings of the 39th International Conference Machine Learning. PMLR, 2022.
[5] Jonathan Chang, Masatoshi Uehara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift in imitation learning via offline data with partial coverage. Advances in Neural Information Processing Systems, 2021.
[6] Ruidi Chen and Ioannis Ch. Paschalidis. Distributionally robust learning. Foundations and Trends in Optimization, 2020.
[7] Xiangli Chen, Mathew Monfort, Anqi Liu, and Brian D Ziebart. Robust covariate shift regression. In Proceedings of the 18th International Conference on Artificial Intelligence and Statistics. PMLR, 2016.
[8] Yasaman Ensafi, Saman Hassanzadeh Amin, Guoqing Zhang, and Bharat Shah. Time-series fore- casting of seasonal items sales using machine learning a comparative analysis. International Journal of Information Management Data Insights, 2022.
[9] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.002211 2020.
[10] Rui Gao, Xi Chen, and Anton J. Kleywegt. Wasserstein distributionally robust optimization and variation regularization. Operations Research, 2020.
[11] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in Neural Information Processing Systems, 2016.
[12] Daniel Jarrett, Ioana Bica, and Mihaela van der Schaar. Strictly batch imitation learning by energy-based distribution matching. Advances in Neural Information Processing Systems, 2020.
[13] Shengyi Jiang, Jingeheng Pang, and Yang Yu. Offline imitation learning with misspecified simulator. Advances in Neural Information Processing Systems, 2020.
[14] Geon-Hyeong Kim, Seokin Seo, Jongmin Lee, Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang. and Kee-Eung Kim. DemoDICE: Offline imitation learning with supplementary imperfect demonstrations. International Conference on Learning Representations, 2022.
[15] Ilya Kostrikov, Ofr Nachum, and Jonathan Tompson. Imitation learning via off aff-policy distribu- tion matching. arXiv preprint arXiv: 1912.090223 2019.
[16] Jongmin Lee, Wonseok Jeon, Byungjun Lee, Joelle Pineau, and Kee-Eung Kim. OptiDICE: Offline policy optimization via stationary distribution correction estimation. In Proceedings of the 24th International Conference Machine Learning. PMLR, 2021.

---

[17] Ofir Nachum, Yinlam Chow, Bo Dai. and Lihong Li. DualDICE: Behavior-agnoste aggnosiic estimation of discounted stationary distribution corrections. Advances in Neural Information Processing Systems, 2019.
[18] Richard Paap, Philip Hans Franses, and Henk Hoek. Mean shifts, unit roots and forecasting seasonal time series. International Journal of Forecasting, 1997.
[19] Kishan Panaganti, Zaiyan Xu, Dileep Kalathil, and Mohammad Ghavamzadeh. Distributionally robust behavioral cloning for robust imitation learning. In 62nd IEEE Conference Decision and Control, 2023.
[20] Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review. arXiv preprint arXiv; 1908. 05659, 2019.
[21] Stephane Ross and Drew Bagnell. Efficient reductions for imitation learning. Proceedings of the 13th International Conference on Artificial Intelligence and Statistics. PMLR, 2010.
[22] Resve A. Saleh and A.K.Md. Ehsanes Saleh. Statistical properties of the log-cosh loss function used in machine learning. arXiv preprint arXiv:2208.0455555 2022.
[23] Fumihiro Sasaki and Ryota Yamashina. Behavioral cloning from noisy demonstrations. In International Conference on Learning Representations, 2020.
[24] Jonathan Spencer, Sanjiban Choudhury, Arun Venkatraman, Brian Ziebart, and J Andrew Bagnell. Feedback in imitation learning: The three regimes of covariate shift. arXiv preprint 0rXiv:2102.02872. 2021.
[25] Matthew Staib and Stefanie Jegelka. Distributionally robust optimization and generalization in kernel methods. Advances in Neural Information Processing Systems, 2019.
[26] Xuezhi Wang, Tzu-Kuo Huang, and Jeff Schneider. Active transfer learning under model shift. In Proceedings of the 31st International Conference on Machine Learning. PMLR, 2014.
[27] Junfeng Wen, Chun-Nam Yu, and Russell Greiner. Robust learning under uncertain test distributions: Relating covariate shift to model misspecification. In Proceedings of the s1 International Conference on Machine Learning. PMLR, 2014.
[28] Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. seetiiinntto-weighheee offline imitation learning from suboptimal demonstrations. In Proceedings of the 39th International Conference on Machine Learning. PMLR, 2022.
[29] Haoran Xu, Li Jiang, Jianxiong Li, Zhuoran Yang, Zhaoran Wang, Victor Wai Kin Chan, and Xianyuan Zhan. Offline RL with no OOD actions: In-sample learning via implicit value regularization. In International Conference Learning Representations, 2023.
[30] Kai Yan, Alex Schwing, and Yu-Xiong Wang. A simple solution for offline imitation from ob- servations and examples with possibly incomplete trajectories. Advances in Neural Information Processing Systems, 2024.
[31] Xin Zhang, Yanhua Li, Ziming Zhang, and Zhi-Li Zhang. f-GAIL: Learning f --iieergence for generative adversarial imitation learning. Advances in Neural Information Processing Systems, 2020.
12

---

Supplementary Material
Contents
Suboptimality Bounds with Arbitrary State Distributions
14
Summary of --ivergencc Choices
15
Experimental Settings for Four Rooms Environment
16
Marginal probability the original dataset
16
Implementation details
16
Hyperparameters
16
Experimental Settings for D4RL Benchmark
16
D.I Implementation details
16
D.2 Hyperparameters
16
D.3 Number of tub-trajectories
17
Additional Experimental Results
17
E,1
Performance comparison complete trajectories)
17
Performance comparison on lower quality segments
17
E.3
Weak correlation between target MSE and normalized score
18
E
Performance comparison with a different fddivergence choice
18
icenses
19
Computation Resources
19
Broader Impacts
19
13

---

A Suboptimality Bounds with Arbitrary State Distributions
In this section, we present an example of BC where the performance guarantee fails when the data state distribution dp differs from expert stationary dg- For simplicity, here we focus on the finit-hhorizon average state distribution, but this can be straightforwardly extended to the infinite- horizon discounted state distribution. Consider finite-horizon setting with a horizon H 2 and denote the state distribution d' of т at timestep i. Define the H-horizon state distribution d of π (s) == H Σ d' (s) for all 8 S. Denote the performance of л as J(r) := ElT Ex; [R(s, at)] = HE,, [R(s.a) a)] where R x A 7 [0, 1] is a true reward function. Denote surrogate loss function É, which measures how far the learner policy is from the expert policy for the state 8 E S. Here, consider 0-1 loss, i.c. ((T(s), TT())) = [[[(() + [g)s)].
Proposition 2. (Ross and Bagnell (211) The suboptimality of an imitator policy ft is bounded by
J(TE) J(R)= < H (((*(s), TT()))
(11)
In our problem setting, we assume the lack of access to dataset sampled from (s(s), instead deal with dataset sampled from an arbitrary distribution (s). When applying BC approach to minimize Eg-dp [s)) (s)))] it is possible to construct an example where this expected loss fails to upper bound the policy suboptimality, (TT))- - (m.). To illustrate, we revisit [20] and consider an example that tightens the inequality
Assume 0 < < 1/H, < 1/H..
$2
a1
TB (3a)
FE
d1
02
01
dg
H
0
$1
01
(*(s))
[1
with prob. - eH with prob. cH
a
02
$3
[22
d+
(1 cH)
CH
11
a22
dp.
5
Figure A: Example of 21
Table A: Policies and their corresponding state distributions.
Consider an MDP illustrated in Figure where S = (s,, $2, 53). A {a1, 02). the transition T is deterministic with T($1, .1) 82 (51, = $3, T (s2, a1) = 83, Ts2,, 42) = $2, ( a1) = 82, (s33, a2) = 83 and 81 the initial state. We define the expert policy mæ and an imitator policy à and derive their corresponding H-horizon state distributions dg, dz in Table A Define (ss,a) = I[a Ts() for all E. A. Then, (TT)), J(r) would be:
H = H
J(TE) = H Σ ds(s)) Σ
Te(a|s)R(s,a)) H
H
H
8ES
aEA
J(ft) H Σ (s) Σ ((s = H
1
(1
cH)
+
H -
-1-cH)
H - eH²
H
H
aEAA
Hence, J ((TT) J(a) eHf Note that Eg [[(f, TE) # -el + H-1 .0 0+0. 0-1 = therefore
= the both left-hand and right-hand sides of the equality [1] equal to cH²
However, considering the expect loss with respect to d'p,, where dp, (81) = d, (x2) =
6, dp, (*a) = 0 as described in Table A we can calculate Eg-do, [f(A, TF)] = de H. Then,
J(TB) J(+) cH² > 3 = H 18( ne)
where < 1/H. As 5 decreases, the right-hand side also decreases, allowing 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.
14

---

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 uncertainty set Q C A(S) with de E 2 we can guarantee that ((FE ) & H2 supaeo (xsall(<()) (p))())))))))))))))))))))))))))) Furthermore, given Q1 C 22 S A(S) with de E 21. the upper-bound of 21 would be lighter than 22, i.e.
J(TE) - ((x) = 2²2 sup Es-al(T(S), S Hf2 sup (e<al((((ss,TEEE)) de
dE
Let O, == {d E A(S) Drv(d, dp) S p} and Q (d A(S) Drv(d, dp) S p. satisfies Bellman flow constraints). Since de satisfies Bellman flow constraints, dg 2p. then dg E QDriDICE with QDriDICE C Qp. Then, the uncertainty set of Q QDriDICE provides a more tighter upper-bound for the suboptimality compared to that of Qp.
B
Summary of f-Divergence Choices
Define
x lolog x +1 (x 1)2
if0 <
(ep((y)
if <
[sotexx (x)
and ()bot()))+
i 2
y y++1
if 2 0'
ReLU(x) max(0, z). ELU(x) = (exp(x)= -1 if x < 0 be summarized in Table B
Then, f-divergence choices can
if x 2 0'
Divergence
f(x)
(f))-(v)
W W,
KL Divergence
x lloggg x 1 1)2 xsun
exp(y - 1)
exp
x --Dveergeecc
(x+1)
ReLU
(auls,a,s')
a Ga (siant) +11
Soft 22--ivergenne
ELU
TV-Distance
1/ 1x 1|
Soft TV-Distance
2 log(cosh( 1))
tanh -2y)
ReLU
(tanh- -1 (20,u(aaadd)) +1)
Table B: Summary of fddivergences and their associated functions.
1.02
5.75
3.75
1.50
0.25
1.00
-1.23
0.75
TV
3.2
Sotty"
0.25
-3.5
-3.02
Figure B: Visualization of f functions and derivatives f" corresponding to f-divergence choices. We also visualize corresponding w max(((') 0), which is closed form solution of the inner maximization in Eq.
15

---

C
Experimental Settings for Four Rooms Environment
C.1
Marginal probability of the original dataset
Rooms
Room
Room 2
Room
Room
(u)
0.342
0.291
0.089
0.278
Actions
"UP"
"DOWN"
"LEFT"
"RIGHT"
(u(
0.080
0.453
0.016
0.451
Table C: Marginal probability DE over rooms and actions.
C.2 Implementation details
The numher of transitions: 1000 + |
RBF feature bandwidth: 10
The number of representative points for RBF feature: 5
Cx: 1-0 loss
C.3 Hyperparameters
a E {100, 50, 20, 10, 0.2, (for OptiDICE-BC, DriIDICE)
p {100, 50, 20, 10, 2, 0.5, 0.2, 0.13 (for DR-BC) Logit margin maximum: log(100)
D Experimental Settings for D4RL Benchmark
D.1 Implementation details
Since the number of so is much smaller than s, we follow heuristics to estimate Eawdp [u(s)]
instead of Eu [v[s]] in the objective (8).
In Scenario 3, determine starting timestep of each segment, sample timesteps by using Geometric distribution with p = 5 x 10-3
Cgi Mcan Squared Error (MSE) loss
D.2
Hyperparameters
Hyperparameter Policy distribution
DemoDICE AW-BC
DR-BC
OptDICEEBC
DRiDICE (Ours)
Tanh Normal 512
Barch size
Policy learning rale
3 x 10-5
hidden units
2256 256
Training iteration
500K
10
10
10
10
[10-1 1000,,000, 10-2]
(or we) learning rale
3x
Table D: Summary of hyperparameters used in D4RL benchmark experiments.
"To prevent support mismatch, data point for each S is added to dataset.
16

---

D.3
Number of sub-trajectories
Problem setting
hopper
walker2d
halfcheetah
Scenario Scenario 2
100
100
250
100
100
250
Scenario 3
[10,20,30,40550] [10,20,30,40.5]]
[50.100.155.220.255)
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.
E
Additional Experimental Results
E.1 Performance comparison on complete trajectories
To compare performance of IL methods on standard scenarios, we have conducted additional ex- periments using complete expert trajectories using D4RL expert-v2 dataset varying the number of trajectories in f1, 5, 10, 50). The results are detailed in Figure C As demonstrated, DriIDICE can deal with sampling errors of small datasets, showing superior data efficiency compared other methods.
hopper
walker2d
halfcheetah
Nomeited venre
100
100
75-
75
59
25-
Number
Number Trajectories
Number Tajacturies
BC
DemoDICE
AW-BC
DR-BC
OptiDICE-BC
DrilDICE (Ours)
Figure C: Performance comparison on complete trajectory scenarios. The points/shaded areas indicate the means/standard errors measured repetitions.
E.2 Performance comparison on 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 original dataset, D4RL, expert-v2, we use D4RI. medimm-d-eee quality demonstrations as our imitation standard for comparison. The results are depicted in Figure [ 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.
hopper-medium
walker2d-mes 2d-mmdium
halfcheetah-medium
Target rett Nomalred S0OW
29
0
1.5
0.8
1.0
0.5
3.2
Segments
Segments
Segments
DeMODICE
AW-BC
OptIDICE-BC
DriIDICE (Ours)
Figure D: Performance comparison Scenario 3 (segmented dataset) with D4RL medium-v2 datasets. Each black dashed-dot line expresses the averaged normalized score of the used dataset. The points/shaded areas indicate the means/standard errors measured over 5 repetitions.
17

---

E.3
Weak correlation between target MSE and normalized score
hopper
walker2d
halfcheetah
100
Normallzer aaaallzzdd
60-
40
6.8
1.0
1.2
0.6
0.8
1.
1.2
0.5
1.0
1.5
2.0
2.5
3.0
Target MSE
Target MSE
Target MSE
Figure E: Correlation between the target MSE and normalized scores.
E,4 Performance comparison with different f-divergence choice
To evaluate the effect of f f-divergence choice on imitation performance, we additionally evaluate OptiDICE-BC and DriIDICE with the soft xi=divergence, as originally utilized in [16]. across all considered scenarios. The results for scenarios 1-3 are presented in Table F G and Figure
respectively. These results indicate that the choice of the soft TV-distance significantly enhances the performance DriIDICE compared to using soft-x?. To explain this performance gain, we hypothesize that the soft-TV distance provides more discriminative weighting of samples based their long-term policy errors. As shown in Figure B, conventional f-divergences (e.g. KI, soft-x?, ...) make less pronounced to the long-term policy error, yet the soft-TV distance responds sensitively to changes in e, resulting in more pronounced weight w. This enables BC loss to more selectively focus critical samples with significant magnitude of long-term policy errors, thereby effectively enhancing performance, akin to the benefits observed in Sparse Q-Learning 229
Task
p(D.)
BC
DR-BC (TV)
OptIDICE-BC
OptDDICE-BC (Sof--TV)
DrlIDICE
DriIDICE
(Soft-x²)
(Soft-x")
(SoftTVV)
0.1
24.7 4.2
2.3
13.3
12.7
58.9
5.6
Rehalanced by state
hopper
3..4
3..7
2..4
8.6200
60.9
67.1
0.9
11.242.5
2..4
2..
1..4
2.0
36.4
0.1
18.9
14.7
0.9
4.9
31.9
51.68. 8.2
walker2d
0.5
2.19
8.1
11,11A
8.1004
3..5
73.7 5,4
0.9
30.4
7.7
1..5
1.7
43.3
77.6
0.1
4.2
32.9
33.7
1.111.
5.2
52.5 3.6
halfchectah
0.5
3.10
26.2
3300
6,1
41.1 41.1
55.0 3,3
0.9
1..5
9.0
2.2 1.2
10
7.3 1.6
22.3 + 2.9
0.1
29.7
2.5
1.14
11.742.1
4..3
56.6 11.9
Rehalanced by action
hopper
0.5
4.99
35.1
30.0
11.8111
9.54
73.8 3.6
0.9
3.6
3.,6
38.9
19,4
63.1
49,0 12.3
0.1
2..6
31.2
12.1
7.3
4..7
7.26
walker2d
0.5
6.7
30.5
1..4
6.4
47.9 12.4
72.1
0.9
16.9
3..6
1..6
4.6
43.7
6
41.9
27.5
halfcheetah
0.5
4..8
3.54
4..5
0.9
2..9
12.1
4.3
0.7
9.6
29.2
Table Performance comparison on Scenario (rebalanced dataset) including different choice of -diveergence 2 (i-vieergence) for OptiDICE-BC and DriIDICE.
18

---

Task
p(D:)
BC
DR-BC
OptiDICE-BC
OptiDICE-BC
DriIDICE
DriIDICE (Son-TV)
(TV)
(Sofl-x²)
(Sof--TV)
(Snft-x")
(1
2.8
21.1
50.3 + 6.6
22.8 +
7.7
5.1
(1.5)
0.9
1.7
39.5 4.00
1.2
3.92
45.6 4.6
hopper
(5. 1)
2.1
17.5
48.4 13.0
25.7
9.2
34.7 9.0
Timestep Timestep dependency
(5.5)
6.7
2.3
32.5 10.8
1:.1
4..2
25.6 6.0
1)
2..0
45.7
1..
1.0
51.9
5,42
(1.5)
6..5
57.3
37.3
17.6
64.5
84.3
walker2d
(5.
8.1
18.0
4.4
4.4
23.3
4 8.3
(5.5)
6.2
1..5
8.3
5.5
14.4
$2.6
(1,1)
33.7
17.1
33.7
9,9
33,4+3.7
44.2 5.6
(1,5)
72.7
$2.9 4.00
24.4
69.6
77.1
halfcheetah
(5,
3.8
1.3
5.7
(5.5)
2.0
2.9
1.6
-1.2
4.6
5.3
Table G: Performance comparison on Scenario (time-dependently collected dataset) including a different choice of f-divergence (soft x==divergence) for OptiDICE-BC and DriIDICE.
hopper
walker2d
halfcheetah
Normalizdd score
75
25
Worze 2336.
40-
1.2
Target MBE
0.7
1.0
0.8
1.0
10
20
10
100
200
250
Number
Number Segments
Number Segmantss
DR-BC
OpUDICE-BC (Softx²
OptIDICE-BC (Soft-T)
DriIDICE (5oft-x²)
DrIDICE (SoftTVV
Figure F: Performance comparison on Scenario (segmented dataset) including different choice of f-divergence (soft x" (-ddivergence) for OptiDICE-BC and DriIDICE.
F
Licenses
For all experiments in Four Rooms environment, CVXPY and convex optimization solver MOSEK with academic licenses. We also use Mujoco free licenses. Our code has been developed upon MIT licenses.
G Computation Resources
We used Google cloud computing engine with 100 c2 c2-standard-4 instances that have the following system specification:
Series C2
CPU Base Frequency 3.1 GHz
Family: compute-optimizz
CPU Turbo Frequency: 3.8 GHz
vCPU :4
CPU Max. Turbo Frequency 3.9 GHz
Memory: 16 GB
Network Bandwidth 10 Gbps
CPU Manufacturer Intel
Max. Disk Size 257 TB
CPU Platform Intel Cascade Lake
Max. Number of Disks 128
H
Broader Impacts
Since we study distributionally robust learning, our approach can used 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.
19

---

NeurIPS Paper Checklist
1. Claims
Question: 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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, please see Section
Guidelines:
The answer NA means that the paper has 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 the scope of the claims made, e.g. if the approach was only tested on few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texx system might not be used reliably to provide closed captions for online lectures because it fails 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 faimess.
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 impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification:
We does not claim any theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide hyperparameters and code in the supplementary materials.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

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.
Please see the NeurlPS code and data submission guidelines (https://nips.cm ec/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips. //nips.cc/pil wwwppbblccgguiicuceccccccoe.oopp
http://oooSuuuiss..ooeee htt.SSbbbiiiiooooiiiiiooooi for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, 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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
22

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 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 figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 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).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Yes, read and conducted research with Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. 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 socictal impact 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.
23

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Experiments of our research grounded only on simulators.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses for existing assets
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?
Question:
Answer: [Yes]
Justification:
We describe License in Appendix
Guidelines:
The answer NA means that the paper does not 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: does not release new assets
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Our research does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Our research does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "BvOhWju3S3", "reviewer": "Reviewer_zQPx", "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.", "questions": "- Section 3.2 requires further details and refined writing. For example, \n  - Slater's condition/assumption must be explicitly mentioned since the drildice algorithm relies on it. \n  - 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$.) \n  - 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\n  - These details will also help in providing some theoretical justifications from previous DICE-related works, which this paper currently lacks.\n\n- Comparison with previous works must be further expanded. Especially with [18].\n  - 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$. \n  - 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!\n  - Based on the above points, [18] also considers distribution matching in their inner problem w.r.t transition distribution T shifts.\n  - 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.\n  - [18] provides both theoretical guarantees and empirical evaluations and this paper provides only empirical evaluations.\n\n- Regarding experiments\n  - 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?\n  - 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].\n  - 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.\n\nMy score reflects this review provided here.", "limitations": "na", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The design of drildice in Section 3.2 and its evaluation on Mujoco benchmark are the strengths of this work.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First and foremost, we sincerely appreciate your insightful feedback. \n\n**1. Justification on Alternative Optimization**\n\nTo 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.\n\n**2. Clarification on $d_D$ and Our Problem Setting**\n\nIt 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.\n\nAdditionally, 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. \n\nAlthough 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. \n\n**3. Comparison with DR-BC [18]**\n\nIn 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:\n\n- **Choice of Uncertainty Set**:  To define its uncertainty set, DR-BC employs TV distance while DrilDICE uses $f$-divergence, which is more general.\n- **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.\n\n**4. Technical Adjustment for DrilDICE and Summary of $f$-divergence**\n\nWe initially attempted to align the $f$-divergence to the TV distance used in DR-BC for a fair comparison. \nOur formulation requires the inverse of the derivative of $f$ for a closed form solution of the weight $w$. \nHowever, since the derivative $f$ of TV distance is a step function and is not invertible, we cannot directly use TV distance. \nWe 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. \nWe refer to this as **the soft-TV distance** and utilized this in Mujoco-based experiments. (See **Figure A** for visualization)\n\nIn 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. \nAs the reviewer commented, we will include this table in the main manuscript.\n\n| Divergence | $f(x)$ | $(f')^{-1}(y)$ | $w^*_{\\pi,\\nu}$ |\n| --- | --- | --- | --- |\n| KL Divergence | $x\\log x$ | $\\exp(y-1)$ | $\\exp\\left(\\frac{e_{\\pi,\\nu}(s,a, s')}{\\alpha}-1\\right)$ |\n| $\\chi^2$-Divergence | $\\frac{1}{2}(x-1)^2$ | $(y+1)$ | $\\mathrm{ReLU}(\\frac{e_{\\pi,\\nu}(s,a, s')}{\\alpha} + 1)$ |\n| Soft-$\\chi^2$ Divergence | $f_{\\text{soft}-\\chi^2} (x) $ | $(f^{'})^{-1}_{\\text{soft}-\\chi^2} (y) $ | $\\mathrm{ELU}\\left(\\frac{e_{\\pi,\\nu}(s,a, s')}{\\alpha}\\right) + 1$ |\n| TV Distance | $\\frac{1}{2}\\|x-1\\|$ | - | - |\n| Soft-TV Distance | $\\frac{1}{2} \\log (\\cosh( x-1))$ | $\\tanh^{-1}(2y)+1$ | $\\mathrm{ReLU}(\\tanh^{-1}(\\frac{2e_{\\pi,\\nu}(s,a, s')}{\\alpha})+1)$ |\n\nwhere $f_{\\text{soft}-\\chi^2} (x) := x\\log x -x +1$ if $0<x<1$ and $(x-1)^2$ if $x \\ge 1$, \n$(f^{'})^{-1}_{\\text{soft}-\\chi^2} (y) := \\exp(y)$ if $y<0$ and $y+1$ if $y \\ge 0$,\n$\\mathrm{ReLU}(x):=\\max(0,x)$, $\\mathrm{ELU}(x):= \\exp(x)-1$ if $x<0$ and $x$ for $x\\ge0$.\n\n[A] Saleh et al., \"Statistical Properties of the log-cosh Loss Function used in Machine Learning.\" arXiv preprint (2022).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p4zMf1JovM", "reviewer": "Reviewer_oZLc", "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.", "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": "The authors have provided a discussion about the limitations and broader societal impacts.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed solution is technically reasonable and theoretically sound.", "weaknesses": "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.\n2. The experiments are conducted on simulated covariate shifts. More evaluation in real-world cases will further improve this paper. \n3. Some notations have not been well defined, the $\\Delta S$ in Equation 2, and the $\\mathcal{W}$ in Equation 5.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback.\n\n**1. Experimental Design and Real-World Applicability Concerns**\n\nOur 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].\n\nBeyond 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.\n\n[A] Yan et al., “A Simple Solution for Offline Imitation from Observations and Examples with Possibly Incomplete Trajectories,” NeurIPS 2023.\n\n\n**2. Clarification on Unclear Notations**\n\nWe 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.\n\n**3. Applicability in Partial Coverage Scenarios**\n\nDrilDICE, 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$. \nThis assumption presents a challenge in ensuring the robustness of DICE approaches when the support of $d$ does not fully cover that of $d_E$.\n\nDespite this, DrilDICE has shown strong performance in scenarios with limited data coverage. \nIn 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.\n\nWhile 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.\n\n\n**4. Expanded Evaluation Results on Main Scenarios**\n\nWe 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.\n\nBased 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.\n\nThe 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.\n\nThe 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].\n\n[B] Saleh et al., \"Statistical Properties of the log-cosh Loss Function used in Machine Learning.\" arXiv preprint (2022).\n\n[C] Xu et al., “Offline RL with No OOD Actions: In-sample Learning via Implicit Value Regularization”, ICLR 2023.\n\n\nThank 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Vt9rMmTRa9", "reviewer": "Reviewer_k6J7", "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.", "questions": "1.\tThe advantage over OptiDICE-BC seems marginal in Table 2 and Table 4. What could be the reason?\n2.\tWhat does “marginal room (or action) distribution of dataset” mean in line 189?\n3.\tFor mujoco experiments, the authors only use the expert split. How does the method perform in other splits that have lower quality?", "limitations": "One limitation mentioned by the authors in the last section is that: don’t consider uncertainty from transition shifts or noisy demonstrations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThis paper is well-written and well-organized.\n2. The proposed method seems promising.\n3.\tA lot of covariate shift scenarios are investigated in the experiments, including discrete navigation tasks and mujoco control tasks.", "weaknesses": "1.\tThe experimental results do not fully support the advantages over baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful feedback.\n\n**1. Addressing Concerns on Marginal Performance Improvements Compared to Baselines**\n\nIn 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.\n\nBased 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. \nThe 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.\n\nThe 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.\n\nThe 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].\n\nIn 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. \n\n[A] Saleh et al., \"Statistical Properties of the log-cosh Loss Function used in Machine Learning.\" arXiv preprint (2022).\n\n[B] Xu et al., “Offline RL with No OOD Actions: In-sample Learning via Implicit Value Regularization”, ICLR 2023.\n\n\n**2. Clarification on the Terminology \"Marginal Room (or Action) Distribution of Dataset\"**\n\nIn line 189, the term \"marginal room (or action) distribution of dataset\" specifically refers to $p(u)$, \nthe 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. \n\nFor 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$.\n\nHere, 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.\n\n\n**3. Performance Comparison with Lower Quality Segments**\n\nDespite our primary focus on expert-quality datasets, we conducted experiments with additional datasets on Scenario 3 (segment datasets) to ensure consistency across different datasets. \nRather 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.\n\nGiven 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.\n\n\nThank 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# General Response

We are grateful for the insightful and detailed feedback provided by all reviewers. Below, we summarize our response to the main concerns raised. Should any points require further clarification or detailed discussion, we are fully prepared to engage in discussions during the author-reviewer discussion period.

## **1. Comparison with DR-BC [18]**
Thanks to reviewer zQPx’s suggestion, we found that DR-BC’s objective is also related to our problem setting. 
In short, DR-BC considers the uncertainty set of arbitrary state distributions by utilizing TV distance. We evaluated DR-BC as a baseline for all scenarios. 

## **2. Soft-TV Distance: Technical Adjustment of $f$ for DrilDICE**
To provide a fair comparison between DrilDICE and DR-BC, we introduced **the soft-TV distance** into DrilDICE. This relaxed version of the TV distance has an invertible derivative of $f$, enabling DrilDICE to obtain a closed form solution of $w^*_{\pi,\nu}$ while maintaining similar properties of TV distance. See **Figure A** in the PDF to compare with other $f$-divergence. 

## **3. Additional Experimental Results**
We have expanded our main experiments by incorporating three additional methods: DR-BC, OptiDICE-BC (w\ soft-TV) and DrilDICE (w\ soft-TV) for our main scenarios. Due to space constraints, we omitted the results for DemoDICE, AW-BC, worst-25% score, target 0-1 loss in tables. For the expanded main experiments, please refer to the following:
- Four Rooms : **Table A** in the comment to the rebuttal
- Scenario 1 (rebalanced dataset) : **Table B** in the comment to the rebuttal
- Scenario 2 (time-dependently subsampled dataset) : **Table C** in the comment to the rebuttal
- Scenario 3 (segmented trajectory dataset) : **Figure B** in the PDF

Additionally, we conducted experiments in the following additional problem settings: 
- Complete trajectory setting to address reviewer zQPx’s concern : **Figure C** in the PDF
- Medium-quality segment setting as commented by reviewer k6J7 : **Figure D** in the PDF

In summary, the results of additional experiments consistently demonstrate that using the soft-TV distance as an $f$-divergence significantly improves the performance of DrilDICE, outperforming baselines including DR-BC.

Meta-review: 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)
Decision: Accept (poster)