Offline Imitation Learning with Variational Counterfactual Reasoning
Authors: Zexu Sun
Bowei He
Jinxin Liu
Xu Chen
Chen Ma
Shuai Zhang
Venue: NeurIPS 2023 poster (2023)
Keywords: offline imitaion learning
counterfactual reasoning
data augmentation

Abstract: In offline imitation learning (IL), an agent aims to learn an optimal expert behavior policy without additional online environment interactions. However, in many real-world scenarios, such as robotics manipulation, the offline dataset is collected from suboptimal behaviors without rewards. Due to the scarce expert data, the agents usually suffer from simply memorizing poor trajectories and are vulnerable to the variations in the environments, lacking the capability of generalizing to new environments.To automatically generate high-quality expert data and improve the generalization ability of the agent, we propose a framework named \underline{O}ffline \underline{I}mitation \underline{L}earning with \underline{C}ounterfactual data \underline{A}ugmentation (OILCA) by doing counterfactual inference. In particular, we leverage identifiable variational autoencoder to generate \textit{counterfactual} samples for expert data augmentation. We theoretically analyze the influence of the generated expert data and the improvement of generalization. Moreover, we conduct extensive experiments to demonstrate that our approach significantly outperforms various baselines on both \textsc{DeepMind Control Suite} benchmark for in-distribution performance and \textsc{CausalWorld} benchmark for out-of-distribution generalization.

Offline Imitation Learning with Variational Counterfactual Reasoning
Zexu Sun Bowei He", Jinxin Liu', Xu Chen'; Chen Ma', Shuai Zhang' Gaoling School of Artificial Intelligence, Renmin University of China DDpartment of Computer Science, City University of Hong Kong *School of Engineering, Westlake University DiDi Chuxing {sunzexu21, xu chen) edu. cn, boweihe2-c@mi cityu edu, hk Liujinxin@restlake edu. cn, chenma@cityu.co edu. hk, shuai zhang@tju ju. edu. cn
Abstract
In offline imitation learning (IL., an agent aims to learn optimal expert behavior policy without additional online environment interactions. However, in many real-world scenarios, such robotics manipulation, the offline dataset is collected from suboptimal behaviors without rewards. Due to the scarce expert data, the agents usually suffer from simply memorizing poor trajectories and are vulnerable to the variations the environments, lacking the capability of generalizing to new environments. To automatically generate high-quality expert data and improve the generalization ability the agent, we propose a framework named Offline Imitation Learning with Counterfactual data Augmentation (OILCA) by doing counterfactual inference. In particular, we leverage identifiable variational autoencoder to generate counterfactual samples for expert data augmentation. We theoretically analyze the influence of the generated expert data and the improvement of generalization. Moreover, conduct extensive experiments to demonstrate that approach significantly outperforms various baselines on both DEEPMIND CONTROL SUITE benchmark for in-distribution performance and CAUSALWORLD benchmark for out-of-distribution generalization.
Introduction
By utilizing the pre-collected expert data, imitation learning (II.) allows us to circumvent the difficulty in designing proper rewards for docision-making tasks and learning expert policy. Theoretically, as long as adequate expert data are accessible, we can casily learn an imitator policy that maintains
sufficient capacity approximate the expert behaviors 255 30 29 31.. However, in practice, several challenging issues hinder its applicability to practical tasks. In particular, expert data are often limited, and due the typical requisite for online interaction with the environment, performing such online IL may be costly or unsafe in real-world scenarios such as self-driving or industrial robotics 34 37 Alternatively, in such settings, might instead have access large amounts of pre-collected unlabeled data, which are of unknown quality and may consist of both good-performing and moor-performing trajectories. For example, in self-driving tasks, number of human driving behaviors may be available; in industrial robotics domains, one may have access to large amounts of robot data. The question then arises: can we perform offline IL with only limited expert data and the previously collected unlabeled data, thus relaxing the costly online IL requirements?
Traditional behavior cloning (BC) [ directly mimics historical behaviors logged in offline data (both expert data and unlabeled data) via supervised learning. However, our above setting, BC
Corresponding author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

suffers from unstable training, as it relies on sufficient high-quality offline data, which is umre- alistic. Besides, utilizing unlabeled data indiscriminately will lead severe catastrophes: Bad trajectories mislead policy learning, and good trajectories fail to provide strong enough guidance signals for policy learning- In order better distinguish the effect of different quality trajecto- ries on policy learning, various solutions are proposed correspondingly. For example, ORIL [39 learns a reward model to relabel previously collected data by contrasting expert and unlabeled trajectories from fixed dataset; DWBC (36] introduces discriminator-weighhe the policy learning. However, such discriminator-based methods are prone to overfitting and suf- fer from poor generalization ability, especially when only very limited expert data provided.
task to assist
Indeed, the expert data play an important role in offline II. and indicate the well-performing agent's intention. Thus, how to fully understand the expert's 's be- haviors or preferences via lim- ited available expert data be- comes pretty crucial. In this paper, propose to investi- gate counterfactual techniques for interpreting such expert be- haviors. Specifically, we lever- age the variational counterfac- tual reasoning [23 to augment the expert data, as typical imita- tion learning data augmentation methods 3 easily generate noisy data and inevitably bias the learning agent. Here, we conduct the counterfactual reasoning of the expert by answering the following question:
What miihht expert take if different abserved?
Test Envirionments
Expert Data
Generaao Counterfactaal Expert data Viriational Cousterlactui Rosuming
Unlabeled Data
Unn-ef-isttbution
Agent is trained with the collected dataset containing limited expert data and large amounts of unlabeled data, and tested on both in-distribution and -isto--distribution environments.
Figure I:
"What action might the expert take if different state is observed?"
Throughout this paper, we consider the structure causal model (SCM) underlying the offline training data and introduce an exogenous variable that influences the states and yet is unobserved. This variable is utilized the minimal edits to existing expert training data so that can generate counterfactual states that cause an agent to change its behavior. Intuitively, this exogenous variable captures variations of features in the environment. By introducing additional variations in the states during training, we encourage the model to rely less on the idiosyncrasies of a given environment. In detail, we leverage an identifiable generative model to generate the counterfactual expert data, thus enhancing the agent's ' capabilities of generalization in test environments (Figure
The main contributions of this paper are summarized follows:
We propose novel learning framework OILCA for offline IL. Using both training data and the augmentation model, we can generate counterfactual expert data and improve the generalization of the learned agent.
We theoretically analyze the disentanglement identifiability of the constructed exogenous variable and the influence of augmented counterfactual expert data via sampler policy. We also guarantee the improvement of generalization ability from the perspective of error bound.
We conduct extensive experiments and provide related analysis. The empirical results about in-distribution performance on DEEPMIND CONTROL SUITE benchmark and out-of diiittibiiiiinno generalization on CAUSALWORLD benchmark both demonstrate the effectiveness of our method.
2 Related Works
Offline IL A significant inspiration for this work grows from the offline imitation learning technique on how to learn policies from the demonstrations, Most of these methods take the idea of behavior cloning (BC) 5 that utilizes supervised learning to learn to act. However, due to the presence of suboptimal demonstrations, the performance of BC is limited mediocre levels on many datasets.
To address this issue, ORIL 39 learns reward function and uses it to relabel offline trajectories.

---

However, it suffers from high computational costs and the difficulty of performing offline RL under distributional shifts. Trained on all data, BCND [26] reuses another policy learned by BC as the weight of the original BC objective, but its performance can even be worse if the suboptimal data occupies the major part of the offline dataset. LobsDICE [2 learns to imitate the expert policy via optimization in the space of stationary distributions. It solves a single convex minimization problem, which minimizes the divergence between the two-state transition distributions induced by the expert and the agent policy. CEIL []6] explicitly learns hindsight embedding function together with a contextual policy. To achieve the expert matching objective for IL, CEIL advocates for optimizing a contextual variable such that it biases the contextual policy towards mimicking expert behaviors. DWBC 36) introduces an additional discriminator to distinguish expert and unlabeled demonstrations, and the outputs of the discriminator serve the weights of the BC loss. CLUE [T7] proposes to learn an intrinsic reward that consistent with the expert intention via enforcing the embeddings of expert data to calibrated contextual representation. OILCA aims to augment the the scarce expert data to improve the performance of the learned policy.
Causal Dynamics RL Adopting this formalism allows one to cast several important problems within RL as questions of causal inference, such as off off-policy evaluation 66 (2). learning baselines for model-free RL [20), and policy transfer [TO 15). CTRL [19] applies SCM dynamics to the data augmentation in continuous sample spaces and discusses the conditions under which the generated transitions are uniquely identifiable counterfactual samples. This approach models state and action variables as unstructured vectors, emphasizing benefits in modeling action interventions for scenarios such as clinical healthcare where exploratory policies cannot be directly deployed. MOCODA 24 applies learned locally factored dynamics model to an augmented distribution of states and actions to generate counterfactual transitions for RL. FOCUS (38] can reconstruct the causal structure accurately and illustrate the feasibility of learning causal structure in offline RL. OILCA uses the identifiable generative model to infer the distribution the exogenous variable the causal MDP, then performs the counterfactual data augmentation to augment the scarce expert data in offline IL.
3 Preliminaries
3.1 Problem Definition
We consider the causal Markov Decision Process (MDP) 19 with an additive noise, our problem setting, we have an offline static dataset consisting of i..d tuples Dan = {s{, at, 11,1) " s.t. (8) .r)=
p(st, at), $1+1 ~ (st, a,, u4+1). where p(s;, az) is an offline state-action distribution resulting from some behavior policies, fe(st, ae, ut+1) represents the causal transition mechanism. M4+1 is the sample of the exogenous variable u, which is unobserved, and & is a small permutation. Let D e and Drr be the sets of expert and unlabeled demonstrations respectively, our goal is only leverage the offline batch data Dan == Dr u Du to learn an optimal policy æ without any online interaction.
3.2 Counterfactual Reasoning
We provide brief background counterfactual reasoning. Further details can be found in 23
Definition 1 (Structural Causal Model (SCM)). A structural causal model M over variables X (XI...., X,) consists of a set of independent exogenous variables = {ui...., Un) with prior distributions P (u;) and set of functions fi...., fn such that = (PA,, u;), where PA, = X parents of X,, Therefore, the distribution of the SCM, which is denoted PM is determined by the functions and the prior distributions of exogenous variables.
Inferring the exogenous random variables based the observations, we can intervene in the observa- tions and inspect the consequences.
Definition 2 (do-intervention in SCM). An intervention I = do (X, := Sa (P).,,:)) is defined as replacing some functions fi (PA, u.) with f. (PA, u.). where PA, is the intervened parents of X,. The intervened SCM is indicated as M and, consequently, its distribution is denoted as PM:j_
The counterfactual inference with which we can answer the what if questions will be obtained in the following process:
3

---

Exogenous variable
Latest representation
Observations
Counterfactuals
sample
8t
GL++
at
do-Intervention
(a) SCM of causal MDP
(b) SCM of causal MDP with do-intervention
Figure 2: SCM of causal Markov Decision Process (MDP). We incorporate an exogenous variable in the SCM that is learned and utilized for counterfactual reasoning about do-intervention.
1. Infer the posterior distribution of exogenous variable P (u: X = x), where is set observations. Replace the prior distribution P (u,)) with the posterior distribution P (u: x in the SCM.
2. We denote the resulted SCM Mx and its distribution as pMx perform intervention on My to reach pMxx!!
3. Return the output of pMail as the counterfactual inference.
SCM representation of causal MDP We encode causal MDP (under a policy т)) into an SCM M. The SCM shown in consisss of exogenous variable and a set of functions that transmit the state action a₂ and latent representation us+ of t to the next state 1 e.g. St+1 fe (St, U++1), where small perturbation, and subsequentl, perform counterfactual inference on the SCM M, we intervene in the original parents pair (st, az) to āz). Specifically, we sample the latent representation ut+1 from the distribution of exogenous variable u, and then generate the counterfactual next state $141, i.e. $t+1 fa ($s, âe, u(+1), and also the counterfactual next action â i.e. ăt+1 # 5t+1). where á the sampler policy.
the next action 0c+1, e.g. 0--1 x 8441). To
3.3 Variational Autoencoder and Identifiability
We briefly introduce the Variational Autoencoder (VAE) and present its identifiability results. The VAE can be conceptualized the amalgamation of generative latent variable model and associated inference model, both of which are parameterized by neural networks [13]. Specifically, the VAE learns the joint distribution pe (æ. z) po =))o(=),, where pa =) represents the conditional distribution of observing æ given z. Here, denotes the set of generative parameters, and (=) _1|_1 (z:) represents the factorized prior distribution of the latent variables. By
incorporating an inference model qs æ). the parameters & and / can be jointly optimized by maximizing the evidence lower bound (ELBO) on the marginal likelihood (x()).
L = (xla Po 2)] DKL (go(2 | [x)pp(z))
= =o) - DKL (2 x: (1)o x)) s log po
where DKL denotes the KL-divergence between the approximation and the true posterior, and is a lower bound of the marginal likelihood (=x) because of the non-negativity of the KL- divergence. Recently, it has been shown that VABis with unconditional prior distributions p=(2) are not identifiable [18] but the latent factors z can identified with conditionally factorized prior distribution ( z) over the latent variables to break the symmetry (9.
4 Offline Imitation Learning with Counterfactual Data Augmentation
At high level, OILCA consists of following steps: (1) data augmentation variational coun- terfactual reasoning, and (2) cooperatively learning discriminator and policy by using Dy and augmented DE. In this section, we detail these two steps and present our method's theoretical guarantees.

---

4.1 Counterfactual Data Augmentation
Our intuition that the counterfactual expert data can model deeper relations between states and actions, which can help the learned policy be generalized better. Hence, in this section, we aim to generate counterfactual expert data by answering the what if question in Section []] Consequently, counterfactual expert data augmentation is especially suitable for some real-world settings, where executing policies during learning can be too costly or slow.
As presented in Section 3.2 the counterfactual data can be generated by using the posterior of the exogenous variable obtained from the observations. Thus, we can briefly introduce the conditional VAE 113 28 to build the latent representation of the exogenous variable u. Moreover, Considering the identifiability of unsupervised disentangled representation learning [8, an additionally observed variable is needed, where could be, for example, the time index or previous data points in time series, some kind of (possibly noisy) class label, or another concurrently observed variable [9] Formally, let 0 = (f,T.A T, >) be the parameters of the following conditional generative model:
(s u. St, at, c) = Py (sc+1 $t, at, u,,a,, ().,uurrrx.. c),,
(2)
where we first define:
( st. Cr, u) Pe (s s,,a, (u)).
(3)
Equation 2 describes the generative mechanism of 8(+1 given the underlying exogenous variable u, along with (St, az). Equation (3) implies that the observed representation 8t+1 is an additive noise function, i.e., S = f s..a, (u) + E where is independent of f a..a or u. Moreover, this formulation also can be found in the SCM representation of causal MDP in Section 3.2 where we treat the function f as the parameter f of the model.
Following standard conditional VAE [28] derivation, the evidence lower bound (ELBO) for each sample for $241 of the above generative model can be written
logp(st-11 St. ae, c) 20 Φ) := log Logpy(Ssn) St. ae u) logp(c)-
DKL (4p (u) St, at St-1, e) |P: ())).
Note that the ELBO in Equation 4 contains an additional term of the p(c) that does not affect identifiability but improves the estimation for the conditional prior (21] we present the theoretical guarantees of disentanglement identifiability in Section 4.3 Moreover, the encoder Is contains all the observable variables. The reason is that, as shown in Figure 2(a) from a causal perspective, 8t, at and u form collider at 8t+1, which means that when given 8t-1, 8t and ag are related to u.
We seek to augment the expert data in DE. As shown in Figure 2(b) once the posterior of exogenous variable 9 (u St, Ge, C)) is obtained, we can perform do-intervention. In particular, we intervene in the parents (st, az) of Se41 in DE by re-sampling a different pair (5t, (+) from the collected data in Du. Moreover, we sample the latent representation u++1 from the learned posterior distribution of exogenous variable u. Then utilizing $,, âz, and u++1, we can generate the counterfactual next state according to Equation (3), which is denoted St-1- Subsequently, we pre- mre-train& sampler policy áe with the original D a to sample the counterfactual next action 0 = ft 5t+1). . Thus, for all the generated tuples (8++1, ūz+1), ), we constitute them in DE.
4.2 Offline Agent Imitation Learning
In general, the counterfactual data augmentation of our method can enhance many offline IL methods. It is worth noting that as described in 40). discriminator-base methods are easy to be over-fitted. which can more directly show the importance and effectiveness of the counterfactual augmented data. Thus, in this section, also to best leverage the unlabeled data, we use Discriminator- Discrininator-Weigheed Behavioral Cloning (DWBC) (36), a state-of-the-an discriminator-base offline IL method. This method introduces unified framework to learn the policy and discriminator cooperatively. The discriminator training gets information from the policy Tw as additional input, yielding a new discriminating task whose learning objective is as follows:
(DE, Du) =1 (DD
[-log De(Se ar log [
E (6,00))D00
I- Ds df log mw)]
17
E (er.
Iolg D (St, au T)))
(5)
5

---

Algorithm Training procedure of OILCA.
Input: Dataset D Du. Dall. pre-trained sampler policy #æ. hyperparameters 7, initial vari- ational couterfactual parameters 8, Φ. discriminator parameters ub, policy parameters w, data augmentation batch number B.
Output: Learned policy parameters w.
1: while counterfactual training do
Variational counterfactual reasoning
2: Sample (8t, (ts , ~ Dan to form training batch
3: Update 0 and Φ according to Equation
4: end while
5:
for to do
Expert data augmentation
6: Sample (S,, aes 1)- ~ Dg. (át, a.) Du to form an augmentation batch
7: Generate the counterfactual $t+1 according to Equation (2). then predict u,+1 # $1+1)
8: DEu ($t+1,0++1)
9:
end for
Learning the discriminator and policy cooperatively 11: Sample (8,, F and (5/ a() DD to form a training batch
10: while agent training do
12: Update f/ according Equation (5) every 100 training steps
D>
Discriminator learning
13: Update w according to Equation (6) every training step
D Policy learning
14:
end while
where Du is the discriminator, is called the class prior. the previous works hyperparameter often assigned as 0.5.
35
39
7 is fixed
Notice that now Tw appears in the input of Dic. which means that imitation information from log Tw will affect L.p, and further impact the learning of D Thus, inspired by the idea of adversarial training. DWBC [36] introduces new learning objective for BC Task:
] d(1 - d)
Dx =0
E (Dxx
[- log Tw (at st)]
E -D
log (a (s)
(s.,0))
(6)
1
E (s..a()DDe
log log Tw
84)
α a> 1.
where d represents a (st, at, 14,1og ww) for simplicity. a is the weight factor (a > 1). The detailed training procedure of OILCA is shown in Algorithm Moreover, also present possible combinations with other offline IL methods and the related results in Appendix
4.3 Theoretical Analysis
In this section, we theoretically analyze our method, which mainly contains three aspects: (1) disentanglement identifiability, (2) the influence of the augmented data, and (3) the generalization ability of the learned policy.
Our disentanglement identifiability extends the theory of iVAE 9.. To begin with, some assump- tions are needed. Considering the SCM in Figure 2(a) we assume the data generation process in Assumption I
Assumption 1. (a) The distribution of the exogenous variable is independent of time but dependent on the auxiliary variable (b) The prior distributions of exogenous variable are different across auxiliary variable c. (c) The trasition mechanism p($1+1 a,,a) u) are invariant across different auxiliary variable e, (d) Given the exogenous variable u, the next state 8t.1 is independent of the
auxiliary variable i.e. 8t+1 11 u.
Part of the above assumption is also used in 19]: considering the identifiability of iVAE, we also add some necessary assumptions, which are also practical. In the following discussion, we will show that when the underlying data-genen geteratennnnnngggg mechanism satisfies Assumption [[] the exogenous variable can be identified up to permutation and affine transformations if the conditional prior distribution p( (clc) belongs to general exponential family distribution.
The prior distribution of the exogenous variable p(ulc) follows general expo- nential family with its parameter specified by an arbitrary function M(e) and sufficient statis- tics T(u) = TTlu),INxx(u here T Ta(u) is defined by the concatenation of Ta(u)
Assumption 2.

---

[Ti(u') ,,,, .....Tall ) from a factorized exponential family and the outputs of neural net- work T TNN(u) with universal approximation power. The probability density can be written as:
pr Quu) (u)
T(u) (T))))(c]]]
obtain the following identifiability of the parameters in the model. For convenience, we omit the subscript of f fs,,,, as f-
Under the Assumption 2 and leveraging the ELBO in Equation
Theorem 1. Assume that we observe data sampled from generative model defined according to B and Equation ( with parameters (f. T. X), the following holds:
Equation
(i) The {s e (Ste1 =0} has measure zero, where ya is the characteristic function of the density defined in Equation (3]
(i) The function f is injective and all of its second-order cross partial derivatives exist.
(ii) The sufficient statistics T a are twice differentiable.
(iv) There exist ks distinct points c., ,...,c² such that the matrix
L= (x))))
- (
(cC))
size x k is invertible.
Then, the purameters O = (f,T.N) are identifiable up to equivalence class induced by permutation and component-wiz wise transformations.
Theorem guarantees the identifiability of Equation 2 We present its proof in Appendix A
T further implies consistent result the conditional VAE. If the variational distribution of encoder 9s is broad parametric family that includes the true posterior, we have the following results.
Moreover, Theorem
Theorem 2. Assume the following holds:
(i) There exists the (0.0) such that the family of distributions qs (u St, at 8411 contains pe St, St+1.C).
(ii) We maximize L(O, Φ) with respect to both O and Φ.
Then, given infinite data, OILCA learn the true parameters 0" == "f",T" ,*)
We present the corresponding proof in Appendix Theorem 2 is proved by assuming our conditional VAE is flexible enough to ensure the ELBO tight for some parameters and the optimization algorithm can achieve the global maximum of ELBBO.
In our framework, the current generated expert sample pairs ($ @.. are estimated based on the sampler policy ÆÆ. However, मæ may be not perfect, and its predicted results may contain noise. Thus, we would like answer: "given the noise level of the sampler policy, how many samples one need achieve sufficiently well performance?". Using K E (0, 0.5) indicates the noise level 7 tE exactly recover the true action 0c-1 (i.c., 0 ), then the generated sequences perfect without any noise. On the contrary, k = 0.5 means that îgg can only produce random results, and the generated sequences are fully noisy. Then we have the following theorem:
Theorem 3. Given a hypothesis class H. for any é (0,1) and k. (0,0.5), if E H is the pretrained policy model learned based the empirical risk minimization (ERM), and the sample complexity (i.e., number of samples) is larger than 2oog( 2
21122x22+ then the error between the model estimated and true results smaller than with probability larger than 1- 8.
The related proof details are presented in Appendix C From Theorem 3 we can see: in order to guarantee the same performance with given probability (i.e., and are fixed), one needs to generate more than 2 2X ?T1-2k)² sequences. the noise level is larger, more samples have to be generated. Extremely, if the pre-trained policy can only produce fully noisy information (i.e., k 0.5), then infinity number of samples are required, which is impossible in reality.
7

---

explains the efficacy of counterfactual augmented data by the empirical evidence. In the context of offline IL, a straightforward yet very relevant conclusion from the analysis of generalization ability is the strong dependence on the number of expert data [25 We work with finite state and action spaces ((S), |A| < oo), and for the learned policy Tws we can analyze the generalization ability from the perspective of emrrr upper bound with the optimal expert policy n².
of our For the generalization ability, g
Theorem 4. Let [Dgl be the number of empirical expert data used to train the policy and be the expected upper bound of generalization error. There exists constant h such that, if
iDEl 2
(9)
e2
and each state 8t is sampled uniformly, then, with probability least 1 1- 6, we have:
max ||x (- ((((s)) (: st)ll S.
(10)
which shows that increasing (Del drastically improves the generalization guarantee.
Note that we provide the proof details for Theorem 4 in Appendix D
Experiments
In this section, we evaluate the performance of OILCA, aiming to answer the following questions. Q1: With the synthetic toy environment of the causal MDP, can OILCA disentangle the latent representations of the exogenous variable? Q2: For an in-distribution test environment, can OILCA improve the performance? Q3: For an cut-of-disttibution test environment, the generalization? In addition, we use five baseline methods: BC-cxp (BC on the expert data D),, BC-all (BC on all demonstrations Dau), ORIL [39 BCND (26). LobsDICE [12]. DWBC (36). More details about these methods are presented in Appendix [ Furthermore, the
OILCA improve
method and baselines are all detail-tuned for better performance.
5.1 Simulations on Toy Environment (Q1)
In real-world situations, we can hardly get the actual distributions of the exogenous variable. To evaluate the disentanglement identifiability in variational counterfactual reasoning, we build a toy environment show that our method can indeed get the disentangled distributions of the exogenous variable.
Toy environment For the toy environment, we consider a simple 2D navigation task. The agent can move fixed distance in each of the four directions. States continuous and considered as the 2D position of the agent. The goal is to navigate to a specific target state within a bounding box. The reward the negative distance between the agent's state and the target state. For the initialization of the environment, we consider C' kinds of Gaussian distributions (Equation (7)) for the exogenous
variable, where C' = 3, and the conditional variable is the class label. We design the transition function by using multi-layer perceptron (MLP) with invertible activations. We present the details about data collection and model learning in Appendix E
Results We visualize the identifiability in such 2D case in Figure where we plot the sources, the posterior distributions learned by OILCA, and the posterior distributions learned by a vanilla conditional VAE, respectively. Our method recovers the original sources to trivial indeterminacies (rotation and sign flip), whereas the vanilla conditional VAE fails to separate the latent variables well. To show the effectiveness of our method in the constructed toy environment, we also conduct repeated experiments for 1K episodes (500 steps per episode) to compare OILCA with all baseline methods. The results are presented in Figure showing that OILCA achieves the highest average return. To analyze the influence of the number of augmented samples (Theorem D. we also conduct the experiments with varying numbers of counterfactual expert data; the result is shown in Figure 5 The X-axis represents the percentage of /Drl//rrl, where DE is the augmented expert data. We can observe that within a certain interval, the generalization ability of the learned policy does improve with the increase of [De].
8

---

(a)
(b)
(c)
Fenartage
Figure 3: Visualization of both observation and latent spaces of the exogenous variable. (a) Samples from the true distribution of the sources Pa- (ulc). (b) Samples from the posterior 4s (ulse, Qe, $e+1, c). (c) Samples from the posterior 9s (ulse, ae: 8t+1) with- out class label.
Figure 4: Performance of OILCA and baselines in the toy environment. We plot the mean and the standard errors of the av- five
Figure 5: Performance of OILCA with the growing percentage of [Del/Dul We plot the average re- turn's mean and standard errors over five random seeds.
eraged return random seeds.
Table 1: Results for nn-disrribution performance on DEEPMIND CONTROL SUITE. We report the average return of episodes (with the length of 1K steps) over five random seeds. The best results and second best results are bold and underlined, respectively.
Task
BC-exp
BC-all
ORIL
BCND
LobsDICE
DWBC
OILCA
195.44
269.03
221.24
243.52
292.96
382.55
608.38
CARTPOLE SWINGUP
1 7.39
7.06
14,49
a 11.33
11.05
8.95
35,54
66.59
90.01
45.08
96.06
74.53
66.87
116.05
CHEETAH RUN
11.09
31.74
+ 9.88
16.15
7.75
4.60
14.65
129.20
104.56
185.57
204.67
190.93
243.47
298.73
FINGER TURN HARD
74.59
68.87
84.90
153.28
188.84
212.39
290.28
FISH SWIM
111.73
11.93
=1.96
119.29
#11.28
+7.62
= 10.07
77.59
138.93
96.88
257.01
120.87
302.33
460,96
HUMANOID RUN
#8.33
#9.14
+10,76
+11.21
#10.66
144.53
+17.55
91.61
98.59
98.25
141.71
197.79
107.86
296.83
MANIPULATOR INSERT BALL
+3/43
92.02
119.63
105.86
220.66
299.19
238.39
305.66
MANIPULATOR INSERT PEG
115.04
16:37
$5.10
115.14
13.40
a 19..76
14.91
169.14
192.14
181.23
279.66
252.34
280.07
395.51
WALKER STAND
K.0
37.91
10.31
12.69
4.7.33
5.79
KA15
29.44
75.79
41.43
157.44
102.14
166.95
377.19
WALKER WALK
4.05
5.2
In-distribution Performance (Q2)
We use DEEPMIND CONTROL SUITE 32 to evaluate the performance of in-distribution performance. Similar to [39) also define episode positive if its episodic reward is among the top 20% episodes for the task. For the auxiliary variable c, we add three kinds of different Gaussian noise distributions into the environment (encoded detailed data collection and statistics Appendix E. We report the results Table T We observe that OILCA outperforms baselines all tasks; this shows the effectiveness of the augmented counterfactual expert data. We also find that the performance of ORIL and DWBC tends decrease in testing for some tasks (ORIL: CHEETAI RUN, WALKER WALK: DWBC: CHEBEAII RUN, MANIPULATON INSERT BALL); this "overfitting" phenomenon also occurs in experiments of previous offline RL works 14 33
{0, 1,))). Moreover,
present the
This is perhaps due to limited data size and model generalization bottleneck. DWBC gets better results on most tasks, but for some tasks, such as Manipulator Insert Ball and the Walker Walk. OILCA achieves more than twice the average return than DWBC.
5.3 Out-of-distribution Generalization (Q3)
To evaluate the out-of-distributtion generalization of OILCA, we use a benchmark named CAUSAL WORLD (2), which provides combinatorial family of tasks with common causal structure and

---

Table 2: Results for out-of-distribution generalization on WAUSALWORLD. We report the average return of episodes (length varies for different tasks) over five random seeds. All the models are trained on space A and tested on space B to show the out-of-distribution performance 2
Task Name
BC-exp
BC-all
ORIL
BCND
LobsDICE
DWBC
OILCA
281.18
176.54
339.40
228.33
243.29
479.92
976.60
REACHING
+ 16.45
+975
12.98
+ 7.144
9.84
18.75
20.13
256.64
235.58
283.91
191.23
206.44
298.09
405.08
PUSHING
12.70
10.23
19.72
12.64
15.35
14.94
24.03
270.01
258.54
388.15
221.89
337.78
366.26
491,09
PICKING
294.06
225.42
270.75
259.12
266.09
349.66
490.24
PICK AND PLACE
= 7.34
= 11.44
= 14.87
+ 8.01
= 10.31
# 7.39
11.69
496.63
394.91
388.55
339.18
362.47
481.07
831.82
STACKING2
= 7.68
16.98
10.93
+: 9.46
17.05
10.11
11.78
667.81
784.88
655,96
139.30
535.74
768. 68
994,82
TOWERS
9.27
17.17
15.14
18.22
13.76
24.77
5.76
581.91
452.88
702.15
341.97
250.44
1596.96
2617.71
STACKED BLOCKS
496.32
529.82
882.27
288.55
317.95
700.23
1348.49
CREATIVE STACKED BLOCKS
+ 26.92
= 31.01
46.79
# 19.63
# #2.00
13.71
55,05
492.78
547.32
647.95
195.06
458.27
585.98
891.14
GENERAL
217.64
28.49
244.99
29.80
=18.09
19.2
223.12
underlying factors. Furthermore, the auxiliary variable e represents the different do-interventions on the task environments (2). We collect the offline data by using three different do- to-interrentions on environment features @stage_colos
friction. floor floor_frictin fffffciioio)) to generate offline datasets, while other features set as the defaults. More data collection and statistics details is evident that OILCA
presented in Appendix E. We show the comparative results in Table 2 also achieves the best performance all the tasks; ORIL and DWBC perform the second-best results on most tasks. BCND performs poorly compared to other methods. The reason may be that the collected data in space A can be regarded low-quality data when evaluating on space B, which may Icad to the poor performance of a single BC and cause cumulative errors for the BC ensembles of BCND. LobsDICE even performs poorer than BC-exp on most tasks. This is because the KL. regularization, which regularizes the learned policy to stay close to Du is too conservative, resulting in suboptimal policy, especially when Dy contains large collection of noisy data. This indeed hurts the performance of LobsDICE for out-of-distribun generalization.
6 Conclusion
In this paper, we propose an effective and generalizable offline imitation learning framework OILCA, which can learn policy from the expert and unlabeled demonstrations. We apply a novel identifiable counterfactual expert data augmentation approach to facilitate the offline imitation learning. also analyze the influence of generated data and the generalization ability theoretically. The experimental results demonstrate that our method achieves better performance simulated and public tasks.
Acknowledgement
This work is supported in part by National Key R&D Program of China (2022ZD0120103) National Natural Science Foundation of China (No. 62102420). Beijing Outstanding Young Scientist Program Intelligent Social Governance Platform, Major Innovation
NO. BIJWZYJHO 0000000000000099 & Planning Interdisciplinary Platform for the "Double-First Class" Initiative, Renmin University of China, Public Computing Cloud, Renmin University of China, fund for building world-class universities (disciplines) of Renmin University of China, Intelligent Social Governance Platform.
References
[1] Alekh Agarwal, Nan Jiang, Sham M. Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms, 2021.
10

---

[2] Ossama Ahmed, Frederik Träuble. Anirudh Goyal, Alexander Neitz, Yoshua Bengio, Bern- hard Schölkopf, Manuel Wüthrich, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. arXiv preprint arXiv:2010.042900 2020.
[3] Dafni Antotsiou, Carlo Ciliberto, and Tae-Kyun Kim. Adversarial imitation learning with trajectorial augmentation and correction. In 202/ IEEE International Conference on Robolics and Automation (ICRA), pages 4724-4730. IEEE, 2021.
[4] Léon Bottou, Jonas Peters, Joaquin Quiñonero-Candee Denis Charles, D Max Chickering, Elon Portugaly, Dipankar Ray, Patrice Simard, and Ed Snelson. Counterfactual reasoning and learning systems: The example of computational advertising. Journal of Machine Learning Research, 14(11), 2013.
[5] Ivan Bratko, Tanja Urbančič, and Claude Sammut. Behavioural cloning: phenomena, results and problems. IFAC Proceedings Volumes, 28(21):143-145 1995.
[6] Lars Buesing, Theophane Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Jean- Baptiste Lespiau, and Nicolas Heess. Woulda, coulda, shoulda: Counterfactually-guideed policy search. arXiv preprini arXiv: 18Xiv:/8111. 18/1.100222 2018.
[7] Satoshi Hoshino, Tomoki Hisada, and Ryota Oikawa. Imitation learning based on data augmen- tation for robotic reaching. In 2021 60th Annual Conference of the Society of Instrument and Control Engineers of Japan (SICE), pages 417-424, IEEE, 2021.
[8] Divyansh Kaushik, Amrith Setlur, Eduard Hovy, and Zachary C Lipton. Explaining the efficacy of counterfactually augmented data. arXiv preprint arXiv:2010.02114 2020.
[9] Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational au- toencoders and nonlinear ica: A unifying framework. In International Conference Artificial Intelligence and Statistics, pages 2207-22 2217. PMLR, 2020.
[10] Taylor W Killian, Marzyeh Ghassemi, and Shalmali Joshi. Counterfactually guided policy transfer clinical settings. In Conference on Health, Inference, and Learning, pages 5-31, PMLR, 2022.
[11] Beomjoon Kim and Joelle Pineau. Socially adaptive path planning in human environments using inverse reinforcement learning. International Journal of Social Robotics, 8:51-66, 2016.
[12] Geon-Hyeong Kim, Jongmin Lee, Youngsoo Jang, Hongscok Yang, and Kee-Eung Kim. Lobs- dice: offline imitation learning from observation via stationary distribution correction estimation. arXiv preprint 2202:: 13536, 2022.
[13] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 13126:1312.61 14, 2013.
arXiv preprint
[14] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
[15] Yao Lai, Jinxin Liu, Zhentao Tang, Bin Wang, Jianye Hao, and Ping Luo. Chipformer: Transferable chip placement via offline decision transformer. arXiv preprint 211::2306..474,, 2023.
[16] Jinxin Liu, Li He, Yachen Kang, Zifeng Zhuang, Donglin Wang, and Huazhe Xu. Ceil: Generalized contextual imitation learning. arXiv preprint arXiv:2306.145 14534, 2023.
[17] Jinxin Liu, Lipeng Zu, Li He, and Donglin Wang. Clue: Calibrated latent guidance for offline reinforcement learning. arXiv preprint arXiv:2306. 134/2, 2023.
[18] Francesco Locatello, Stefan Bauer, Mario Lucie, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. international conference on machine learning, pages 41 14- 4124. PMLR, 2019.

---

[19] Chaochao Lu, Biwei Huang, Ke Wang, José Miguel Hernández-Lobat. Kun Zhang, and Bernhard Schölkopf. Sample-efheient reinforcement learning via counterfactual based data augmentation. arXiv preprint urXiv:2012.00000 2020.
[20] Thomas Mesnard, Théophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Harutyunyan, Will Dabney, Tom Stepleton, Nicolas Heess, Arthur Guez, et al. Counterfactual credit assignment in model-free reinforcement learning. arXiv preprint Ttt::2011.09466. 2020.
[21] Graziano Mita, Maurizio Filippone, and Pietro Michiardi. An identifiable double vae for disentangled representations. International Conference on Machine Learning, pages 7769- 7779. PMLR, 2021.
[22] Michael Oberst and David Sontag. Counterfactual off off-poliiy evaluation with gumbel-mas max structural causal models. In International Conference on Machine Learning, pages 488 -8900. PMI.R, 2019.
[23] Judea Pearl. Models, reasoning and inference. Cambridge, UK: CambriggeUiieerittypress, 19 (2), 2000.
[24] Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfactual data augmentation. arXiv preprint 2rXiv:2210.1 11287, 2022.
[25] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth interna- tional conference on artificial intelligence and statistics, pages 637-635. JMLR Workshop and Conference Proceedings, 2011.
[26] Fumihiro Sasaki and Ryota Yamashina. Behavioral cloning from noisy demonstrations, In International Conference on Learning Representations, 2021.
[27] Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[28] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015.
[29] Jonathan Spencer, Sanjiban Choudhury, Arun Venkatraman, Brian Ziebart, and J Andrew Bagnell. Feedback in imitation learning: The three regimes of covariate shift. arXiv preprint arXiv:2102.0287. 2021.
[30] Wen Sun, Anirudh Vemula, Byron Boots, and Drew Bagnell. Provably efficient imitation learning from observation alone. In International conference on machine learning, pages 4036-6045 PMLR, 2019.
[31] Gokul Swamy, Sanjiban Choudhury, Andrew Bagnell, and Steven Wu. Of moments and match- ing: game-theoretic framework for closing the imitation gap. In International Conference on Machine Learning, pages 10022-10032. PMLR, 2021.
[32] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.0000 1801.0060, 2018.
[33] Yifan Wu. George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: 111.//911111161 2019.
[34] Zheng Wu, Liting Sun, Wei 7han, Chenyu Yang, and Masayoshi Tomizuka. Efficient sampling- based maximum entropy inverse reinforcement learning with application to autonomous driving. IEEE Roboties and Automation Letters, 5(4):5355-5362, 2020.
[35] Danfei Xu and Misha Denil. Posirive-unlabeled reward learning. In Conference on Robot Learning, pages 205-219. PMLR, 2021.
12

---

[36] Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. risciieinatorwwiigheed weighted offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning, pages 24725-24742. PMIR, 2022.
[37] Zhenting Zhao, Bowei He, Wenhao Luo, and Rui Liu. Collective conditioned reflex: A bio- inspired fast emergency reaction mechanism for designing safe multi-robot systems. IEEE Robotics and Automation Letters, 7(4):10985-1099)) 2022.
[38] Zheng-Man Maeng-Maao/o Zhu, Xiong-Hui Chen, Hong-Long Tian, Kun Zhang, and Yang Yu. Offline reinforcement learning with causal structured world models. arXiv preprint arXiv:2206.01474, 2022.
[39] Konrad ZoIna, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcehre, Ziyu Wang, Yusuf Aytar, Misha Denil, Nando de Freitas, and Scott Reed. Offline learning from demonstrations and unlabeled experience. arXiv preprint 1.Xiv:2011.13885. 2020.
[40] Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. In Conference on Robot Learning, pages 247-263, PMLR, 2021.
13


Reviews:
[{"review_id": "CfAPVLYRR7", "reviewer": "Reviewer_agsz", "summary": "This paper introduces OILCA, a causality-regularized data augmentation method for offline imitation learning tasks. Overall, the idea is novel to me. The empirical results shown in the experiment section seem very promising. However, to support the claims made in the paper, more experiments are needed. Please see details in my suggestions & questions section.", "questions": "1.**(Very Important) A critical confusion is that, if you can do accurate counterfactual reasoning over other states (not included in the expert trajectory), why do not use such a model directly perform as the learned policy?**\nActually, if I am right, this is a **very strong assumption** for your theoretical proof --- if you assume the counterfactual reasoning is accurate (such that the augmentation is reliable), you have already guaranteed to achieve an improved result.\n\n2. How is your algorithm's sensitivity to hyper-parameters? This is especially important for offline RL/IL algorithms --- the application scenarios do not permit excessive hyper-parameter tuning, hence a one-for-all hyper-parameter is of great importance. I would like to see an ablation study on stress-testing not only OILCA but also other baseline methods. \nThe underlying question is: while causal augmentation is beneficial, does it add extra instability to the learning process?\n\n3. On the performance gain: Can the authors disclose what settings have been changed over previous algorithms like ORIL and DWBC?\nIf the training settings are the same, the performance gain can be attributed to the augmented expert data, then what do other algorithms perform under this augmented dataset?\n\n4. How much expert data is needed in augmentation? As it has been shown that Behavior Cloning under clean and high-quality data can perform as good as imitation learning[4], it will be great to show some results to demonstrate in a minimal effort that the data augmentation is effective. --- I'm keen to see under what specific setting, BC can be as powerful as those offline IL methods, and how your method's performance surpasses others when you gradually decrease the proportion of expert data.\n\n[4] Mandlekar, Ajay, et al. \"What matters in learning from offline human demonstrations for robot manipulation.\" arXiv preprint arXiv:2108.03298 (2021).", "limitations": "please see above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Overall, the idea is novel to me. The authors provide theoretical analysis to support their results. The empirical results shown in the experiment section seem very promising.", "weaknesses": "Missing reference:\n\nAs your work is related to causality + imitation learning, I believe the work of Causal confusion in imitation learning clearly worths a citation and discussion. Also, there is quite a lot of related work that is not discussed, e.g., with a quick Google search on causality + reinforcement learning: \n\n[1] Pim de Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. In NeurIPS, pages 11698–11709, 2019.\n\n[2] Gasse, Maxime, et al. \"Causal reinforcement learning using observational and interventional data.\" arXiv preprint arXiv:2106.14421, 2021.\n\n[3] Sun, H., & Wang, T.. Toward Causal-Aware RL: State-Wise Action-Refined Temporal Difference. arXiv preprint arXiv:2201.00354, 2022.\n\n\n\n\nSuggestions:\n\nThe flow of the introduction is a bit awkward to me: the authors mentioned the difficulty of learning from a mixed dataset containing both expert and random-quality data, but they then answer the counterfactual question of how the expert could perform under different states, this question is uncorrelated with the previous discussion. I hope the authors can update this part to better convey their idea and motivations.\n\nThe notions in Definition 2 are unclear to me. what does the \\tilde mean when it is over PA_i? This is not explained in the paper.\n\nFigure 5 is not mentioned in the text and is therefore confusing: how do you increase the percentage of expert data? Why is the proportion always smaller than 1.0? \n\nIt's sort of misleading to call a well-performing method a 'robust' method. (Q2, experiment section). To demonstrate the robustness, the authors should show far more empirical studies rather than a comparison under a single setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer asgz:\n\nThanks for your review of our paper. Here's our response to the weaknesses, questions, and limitations you have highlighted:\n\n-------\n**Weaknesses:**\n\n**W1. Missing references.**\n\n\n**A**: Thank you for the valuable suggestion. We will include and discuss the raised related works in our revision of the paper. \n\n**W2. The flow of the introduction.**\n\n**A**: Thank you for the suggestion. We will update the introduction part to further illustrate our idea and motivation. \n\nFor the limited expert data, we choose data augmentation to solve this problem, and with the advantage of counterfactual reasoning, we adopt the counterfactual data augmentation, which is to answer the counterfactual question.\n\n\n\n**W3. What does the \\tilde mean when it is over PA_i?**\n\n**A**: $\\mathbf{PA}_i$ is the parent of $X_i$. $\\tilde{\\mathbf{PA}}_i$ is the intervented parent of $X_i$. We will clarify it in our revision. \n\n\n\n\n\n**W4. Figure 5: How do you increase the percentage of expert data? Why is the proportion always smaller than 1.0?**\n\n**A**: Figure 5 presents the experiment analyzing the extent to which the number of augmented expert data influences the performance of the learned policy. This experiment is described in detail in the paper (lines 297-302).\n\nWe increase the number of augmented expert samples and concanate them into $D_E$, until the proportion $|D_E|/|D_U|=1$. So the proportion is always smaller than or equal to 1.0.\n\n\n\n\n**W5. In distribution robustness?**\n\n**A**: This robustness is to demonstrate the performance of in distribution online testing, rather than the robustness for the additive noise. \nWe apologize for any potential confusion and we will address and revise this in the final version.\n\n\n-------\n\n**Questions:**\n\n\n\n**Q1. (Very Important) Assumption: the counterfactual reasoning is accurate.**\n\n**A**: We thank the reiviewer for raising this concern. \n\nEven when the counterfactual reasoning process is not perfect,  our framework can provide a trade-off between the information brought by the augmented data and the causal model accuracy.\n\nEven when the counterfactual reasoning process is not perfect,  our framework can provide a trade-off between the information brought by the augmented data and the causal model accuracy.\n\nExperiments show that our framework can indeed bring additional performance gains.\n\n\n**Q2. Algorithm's sensitivity to hyper-parameters.**\n\n\n**A**:  Due to the space limitation here, we put our detailed rebuttal for this question into an external anonymous link (https://hackmd.io/@littleshoes/ryaYsJfhn). Please check the Q2 of Reviewer asgz in it.\n\n\n**Q3. What settings have been changed over previous algorithms like ORIL and DWBC?**\n\n**A**:  Due to the space limitation here, we put our detailed rebuttal for this question into an external anonymous link (https://hackmd.io/@littleshoes/ryaYsJfhn). Please check the Q3 of Reviewer asgz in it.\n\n\n**Q4.How much expert data is needed in augmentation?**\n\n**A**:  Due to the space limitation here, we put our detailed rebuttal for this question into an external anonymous link (https://hackmd.io/@littleshoes/ryaYsJfhn). Please check the Q4 of Reviewer asgz in it.\n\n\n"}, {"review_id": "QFwylpgJqt", "reviewer": "Reviewer_XaVg", "summary": "The paper propose a novel learning framework OILCA for offline IL, which generate counterfactual data to augment the scarce expert data. They analyze the disentanglement identifiability of the constructed exogenous variable and the counterfactual identifiability of the augmented counterfactual expert data. The experiments especially in the CausalWorld benchmark demonstrate the effectiveness of the method.", "questions": "see weaknesses", "limitations": "NAN", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. It is in time to introduce causal inference to offline reinforcement learning and offline imitation learning. This topic is valuable to the RL community;\n2. The paper is overall well-written and easy to follow for me;\n3. The way of data augmentation through counterfactual inference makes sense;\n4. The experiment is almost sufficient to demonstrate their method.\n", "weaknesses": "This paper can be improved in several areas:\n\nWriting and structure:\n\n1. In this paper, it's not clear what the spurious relations under the MDP structure are. Could you explain it based on Figure 2? Notably, it appears that there is a direct causal relationship between $u$ and $s$, which seems should not be categorized as spurious relations.\n\n2. The definition of do-intervention from line 96 to 119 is a little confusing. Specifically, the definition of do on lines 96~96 and in Figure 2(b) is replacing $f$ with another function $\\tilde f$ while keeping $u$, but on line 117, the interpretation of do changes to keeping the function $f$ unchanged and replacing $u$ with $\\tilde u$. It seems that line 117 is more in line with the actual implementation. Could the author clarify this point?\n\n3. The explanation in Section 2.3 on why IVAE solves the identifiability problem is a bit of vague. This is an important context, especially regarding the introduction of the auxiliary variable (which is c defined in Section 3.1) and its role in identifiability. It might be beneficial to move some content about c from section 3.1 to 2.3 to better inform the reader about the preliminaries.\n\nRelated work:\n\nThe authors do not mention any related work in the main body of the paper. Consider moving some of the related work from the appendix into the main text. Moreover, there have been several recent studies applying causal inference to data augmentation for RL . The authors can compare their method with these works and discuss the similarities or differences. They might also consider using these methods as baselines for comparison:\n- Sample-Efficient Reinforcement Learning via Counterfactual-Based Data Augmentation\n- MOCODA: Model-based Counterfactual Data Augmentation\n- Counterfactual data augmentation using locally factored dynamics\n- Offline Reinforcement Learning with Causal Structured World Models\n\nMethod:\n\n1. The authors rely on a pretrained expert policy in the data augmentation process. How was this expert policy obtained? If we assume that we use $\\mathcal{D}_E$ to train the expert policy, given that this paper claims $\\mathcal{D}_E$ is very limited, it should be challenging to reconstruct the expert policy adequately from $\\mathcal{D}_E$. Consequently, the constructed $\\tilde s$, $\\tilde a$ might be different from the actual data distribution in $\\mathcal{D}_E$. How can we ensure that using the augmented dataset from $\\tilde pi_e$, which is unreliable, will have a positive effect on the downstream imitation task?\n\n2. The authors state that they ultimately use Equation (2) for data augmentation, but Equation (2) uses $p(u|c)$ to generate $u$, not the posterior-based method described on line 101, which should use $q(u|s,a,s',c)$ instead. The authors need to clarify the inconsistency between these two points.\n\n3. The authors claim that their contribution is to prove that their method can improve the generalizability of imitation learning, which is a bit of overclaim. From their theoretical analysis, it only shows that increasing the amount of data can improve the accuracy of imitation learning, a relatively trivial conclusion that holds for any machine learning task. If we are conducting a theoretical analysis, we need to explain why using data generated by the counterfactual data augmentation method can better improve the algorithm's generalizability. Otherwise, it raises the question: Does data augmentation not reliant on counterfactual inference have the same improvement?\n\n Experiments:\n\n1. To my knowledge, the original DEEPMIND CONTROL SUITE environment does not have exogenous variables. How did the authors construct the experimental environment and dataset to fit the setting they proposed? Supplementary note: if this is not achievable, I don't think experiments on MuJoCo are necessary. The authors could consider designing more experiments to validate their algorithm's effectiveness on benchmarks designed for causal inference, such as:\n    - Alchemy: A benchmark and analysis toolkit for meta-reinforcement learning agents\n    - Systematic Evaluation of Causal Discovery in Visual Model Based Reinforcement Learning\n\n2. While the experimental results seem significant, the ablation study and visualizations are insufficient. Perhaps the authors could consider the differences between their data and data generated by different methods, the clustering of reconstructed exogenous variables, and whether the corresponding transitions or policy behaviors meet expectations.\n\n3. The description of the baselines is insufficient. The authors should emphasize how the baselines utilize $\\mathcal{D}_E$ and $\\mathcal{D}_U$ respectively.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer XaVg:\n\nThanks for your review of our paper. Here's our response to the weaknesses, questions, and limitations you have highlighted:\n\n-------\n\n**Writing and structure:**\n\n**1.In this paper, it's not clear what the spurious relations under the MDP structure are.**\n\n\n**A**:  We follow an standard causal MDP setting [1][2][3], where both the relationship of state and the exogenous variable are the same as ours.\n\n\n[1] Kazemi, M., & Paoletti, N. (2022). Towards Causal Temporal Reasoning for Markov Decision Processes. arXiv preprint arXiv:2212.08712.\n\n[2] Oberst, M., & Sontag, D. (2019, May). Counterfactual off-policy evaluation with gumbel-max structural causal models. In International Conference on Machine Learning (pp. 4881-4890). PMLR.\n\n[3] Tsirtsis, S., De, A., & Rodriguez, M. (2021). Counterfactual explanations in sequential decision making under uncertainty. Advances in Neural Information Processing Systems, 34, 30127-30139.\n\n**2.The definition of do-intervention from line 96 to 119 is a little confusing.**\n\n**A**:  For this definition and the related figure are confused, we revise the Figure 2(b),\nhttps://i.postimg.cc/KzWk86kG/Figure2-b.jpg\n and we will revise the related description in the final version.\n\n**3.The explanation in Section 2.3 on why IVAE solves the identifiability problem is a bit of vague.**\n\n**A**:  Thanks for the question, we will move part of the Section 3.1 into the preliminary.\n\n-------\n\n**Related Work:**\n\n**1.The authors do not mention any related work in the main body of the paper.**\n\n\n**A**:  The related work will be presented in the appendix due to space limitations. In the appendix, we will discuss the related work mentioned above [1][2][3]. Additionally, we have carefully examined [4], which proposes a method called FOCUS comprising two components: causal structure learning for a world model and offline model-based policy learning. In summary, the aforementioned related works are not suitable as baselines. [1] focuses on learning a personalized policy in a sequential decision process. [2,3] generate counterfactual samples using local causal models with factored MDP. [4] leverages causal discovery to improve the world model. However, in the final version, we will remove some related works from the main text.\n\n-------\n**Method:**\n\n**1.The authors rely on a pretrained expert policy in the data augmentation process. How was this expert policy obtained?**\n\n**A:** See the pre-trained policy answer (Q2) for Reviewer srX3.\n\n**2.Clarify the inconsistency between $p(u|c)$ and $p(u|s,a,c',c)$.**\n\n**A:** There is a mistake in the pseudocode when generating counterfactual states. Currently, the counterfactual states are generated using the posterior instead of the prior. We will correct this in the final version.\n\n**3.The authors claim that their contribution is to prove that their method can improve the generalizability of imitation learning, which is a bit of overclaim**\n\n**A:** For this question, we have presented the distribution of the counterfactual samples mentioned earlier. Drawing from the existing literature on counterfactual data augmentation [1][2], we aim to show that the samples generated by OILCA can be considered as $D_E$ with theoretical guarantees. Therefore, we analyze the outcomes from the standpoint of sample numbers.\n[1] MOCODA: Model-based Counterfactual Data Augmentation\n[2] Counterfactual data augmentation using locally factored dynamics\n\n\n-------\n**Experiment:**\n\n**1.The original DEEPMIND CONTROL SUITE environment does not have exogenous variables.**\n\n**A:** In the DEEPMIND CONTROL SUITE, different noise exists in the dynamic environment depending on the environment initialization. These variations can be considered as different distributions of the exogenous variable. Additionally, we utilize the DEEPMIND CONTROL SUITE because it is commonly employed in papers proposing these baselines. By comparing the results, we can assess the effectiveness of our training framework.\n\nAs for the causal benchmark, we evaluate the performance using the causal world dataset. This benchmark is widely used for robotics control and closely aligns with the standard offline imitation learning setting.\n\nFurthermore, we explore the above-mentioned benchmarks for causal meta-reinforcement learning and model-based reinforcement learning. However, it may require some additional time to integrate them into the standard offline imitation learning setting, and the suitability of these datasets for offline IL needs to be further discussed.\n\n**2.While the experimental results seem significant, the ablation study and visualizations are insufficient.**\n\n**A:** To address the challenge of not having access to the true distribution of the exogenous variable, we visualize the learned exogenous variable alongside the true exogenous variable using a synthetic dataset. Furthermore, rather than conducting an ablation study, we prioritize compatibility as a training framework. This is because we believe that compatibility plays a more crucial role. Considering the limited research on data augmentation in offline imitation learning, combining different data augmentation methods is seen as a valuable contribution in this field. We intend to further discuss this aspect in our future work.\n\n\n\n**3.The description of the baselines is insufficient. The authors should emphasize how the baselines utilize $\\mathcal{D}_E$ and $\\mathcal{D}_U$ respectively.**\n\n**A:** We described the baselines in the appendix, to answer this question, we will add more detailed about how these methods use the expert and unlabeled data."}, {"review_id": "NUPEbYXmXd", "reviewer": "Reviewer_Txp8", "summary": "This paper focuses on the problem of offline Imitation Learning, where an agent aims to learn an optimal expert behavior policy without additional interactions with the online environment. This setting widely exists in the real world as it usually consumes a lot of human effort and cost to collect expert data. Using sub-optimal data could be a valuable way for learning policy. The proposed method uses counterfactual data augmentation to generate high-quality data by learning the underlying data generation mechanism with unlabeled data. Then the new data is combined with expert data to train a policy model. Experiment results demonstrate that the proposed method achieves a large margin over multiple baselines.", "questions": "Points 5 and 6 of the weaknesses are the main questions I want to ask the authors.", "limitations": "No potential negative societal impact is mentioned in either the main context or the appendix. One theoretical limitation is discussed in the appendix but some empirical limitations might be ignored. The most important one is how to select and obtain variable $c$, which is not discussed in the paper.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well-written and easy to follow. Figure 1 is illustrative and helps me understand the core idea of this work.\n2. The motivation of using a small portion of expert data with a large portion of unlabeled data is important in real-world tasks. The proposed method that leverages the information in unlabeled data in an unsupervised way helps generate more expert data and therefore trains a better policy model.\n3. The proposed method is evaluated in two domains with locomotion and manipulation tasks, which of both are important and have a high impact.", "weaknesses": "1. In the abstract, there is a gap between the statement of poor generation caused by sub-optimal datasets and the elimination of spurious features. No clear connection between the sub-optimal dataset and spurious feature is mentioned. \n\n2. In line 14 of the abstract, the authors say the in-distribution robustness is improved by their method but robustness is not mentioned before. Do the authors assume that robustness is almost the same as generalization when the testing scenario is “in distribution”?\n\n3. Using counterfactual data augmentation to improve the performance of decision-making models (including RL, offline RL, IL) is already investigated by existing works. The novelty of the proposed method could be limited unless more differences between existing papers and this one are emphasized. \n\n4. In Definition 2, the notations $\\tilde{f}_i$ and $\\tilde{PA}_i$ are not defined. These notations are important for understanding this definition.\n\n5. I am a little confused about the three-step pipeline of counterfactual inference discussed after Definition 2. The second step is usually for modifying the SCM by removing the structural equations $f$ and replacing the value of $x$ with others. However, this part is not mentioned and the authors only say “denote the resulted SCM as $M_x$”. So, how do we get this $M_x$? Is the value of x changed? Figure 2 enhances my confusion since the do-intervention is only conducted to the exogenous variable $u_i$ rather than $s_t$ or $a_t$. I am not sure if this is still the standard definition of counterfactuals. It is also not consistent with the statement “What action might the expert take if a different state is observed?” since the state is not changed.\n\n 6. In section 3.1, the authors say “an additionally observed variable c is needed, where c could be, for example, the time index or previous data points in a time series, some kind of (possibly noisy) class label, or another concurrently observed variable”. However, the authors do not mention which one is considered for $c$ in this paper. I can only find some information in the experiment part, which also makes me have new questions. Does the process of selecting the variable $c$ create multiple environments with different contexts? This is usually considered a multi-task learning setting, which may not be true for a general imitation learning setting since the data collection is usually not possible to be clustered with additional context indicators. I hope the authors can provide more information about how to define and section the variable $c$ in the rebuttal.\n\n7. After reading the entire paper, I find that robustness and spurious feature is only mentioned in the introduction. The theoretical analysis and experiment do not provide any evidence of the improvement of robustness and the elimination of spurious correlation. In particular, section 4.2 tries to investigate the robustness but I cannot find the setting that supports it. Robustness should be evaluated with small disturbances to the system, including the state, action, and dynamics. After checking Appendix 5.1, I cannot find any disturbance added to the dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer Txp8:\n\nThanks for your review of our paper. Here's our response to the weaknesses, questions, and limitations you have highlighted:\n\n-------\n\n**Weaknesses**\n\n**W1.No clear connection between the sub-optimal dataset and spurious feature is mentioned.**\n\n**A**: The motivation of our method is rooted in the limitation of expert data in standard offline imitation learning. Given this constraint, we propose a method to augment the expert data, which we believe will lead to better performance of the learned policy. The objective of our method is to obtain an augmented expert dataset for training an improved offline policy. However, a crucial challenge arises: determining the most suitable data augmentation method to use. Drawing inspiration from the advancements in causal inference in RL, we opt for a counterfactual approach to augment the expert dataset.\nOne of the primary reasons for the poor generalization of the learned policy is the presence of spurious features or spurious correlations. Specifically, a distribution limitation in the training data can result in the learned policy failing to generalize effectively to unseen data. For instance, in the case of both RL and IL states, which may consist of high-dimensional vectors, some dimensions are crucial for action selection while others are not. However, the irrelevant dimensions in the state can introduce spurious correlations into the learned policy. Intervening on the states through a do intervention can modify this aspect and eliminate the spurious correlation.\n\n\n\n\n**W2.In distribution robustness**\n\n\n**A**: See our robustness explaination (W5) to Reviewer agsz.\n\n\n\n**W3.The novelty of the proposed method.**\n\n**A**: To the best of our knowledge, our method is the first to specifically address counterfactual data augmentation in offline IL. Additionally, we thoroughly reviewed relevant papers on RL and synthesized our findings in the related work section. To provide further clarity on this point, we offer additional details below:\n\n[1] Lu, C., Huang, B., Wang, K., Hernández-Lobato, J. M., Zhang, K., & Schölkopf, B. (2020). Sample-efficient reinforcement learning via counterfactual-based data augmentation. arXiv preprint arXiv:2012.09092.\nThis paper proposes counterfactual RL algorithms to learn both population-level and individual-level policies. It focuses on the sequential decision making setting and achieving an individual policy. \n\n\n[2] Pitis, S., Creager, E., & Garg, A. (2020). Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems, 33, 3976-3990.\nThis paper proposes a counterfactual data augmentation method by using the locally factored dynamics, which can generate the causally valid sample in the global model.\n\n[3] Pitis, Silviu, et al. \"Mocoda: Model-based counterfactual data augmentation.\" Advances in Neural Information Processing Systems 35 (2022): 18143-18156.\nThis paper applies a learned locally factored dynamics model to an augmented distribution of states and actions to generate counterfactual transitions for RL. \n\nOur work utilizes the variational counterfactual method to generate counterfactual samples in the standard offline IL, distinguishing it from the three most related works mentioned above. Additionally, we provide theoretical analysis for identifying both the causal effect and the exogenous variable. As far as we know, this is the first time that identifiable variational counterfactual reasoning is used in offline imitation learning.\n\n\n**W4.Definition 2 Notation.**\n\n**A**: See our definition answer (W3) to Reviewer agsz.\n\n\n**W5.Confusion about the three step pipline.**\n\n**A**: Thanks for your questions, we realize that Figure 2b causes the misleading to you, thus we revise this figure in the following:\nhttps://i.postimg.cc/KzWk86kG/Figure2-b.jpg\n\nThe SCM is typically considered as a predictive model. In our paper, Figure 2(a) presents the SCM, while Figure 2(b) illustrates the do intervention. The changed value of x corresponds to the generated $\\tilde{s}\\_{t+1}$. In conjunction with Definition 2, Figure 2(b) demonstrates that we replace the function $f\\_u$ with $\\tilde{f}\\_u$, and $u\\_{t+1}$ is changed to $\\tilde{u}\\_{t+1}$, corresponding to the change in $\\mathbf{PA}\\_i$ to $\\tilde{\\mathbf{PA}}\\_i$. Consequently, we generate the new $\\tilde{s}\\_{t+1}$, corresponding to the change in $\\mathbf{X}$. The prediction of $\\tilde{a}\\_{t+1}$ provides an answer to the counterfactual question, aligning with the statement.\n\n**W6.The obtainability of $c$.**\n\n**A**: Thanks for the thoughtful questions. This variable is readily obtainable, particularly in real-world scenarios. In the context of autonomous driving, the variable $c$ represents the diverse road conditions encountered during a single driving mission, encompassing information about various exogenous variables. In systems exhibiting periodic cycles, such as the days of the week from Monday to Sunday, the $d$-th day can be represented by the variable $c$. While these additional variables may not be prevalent in many imitation learning studies, considering the exogenous variable always allows for finding an appropriate label. In the context of the variable $c$, even in the original paper [1] proposing this additional variable, there are no strict constraints on its interpretation. This characteristic renders $c$ readily observable, particularly in numerous real-world scenarios.\n\n[1] Khemakhem, I., Kingma, D., Monti, R., & Hyvarinen, A. (2020, June). Variational autoencoders and nonlinear ica: A unifying framework. In International Conference on Artificial Intelligence and Statistics (pp. 2207-2217). PMLR.\n\n**W7.Robustness experiment**\n\n**A**: See our robustness explaination (W5) to Reviewer agsz."}, {"review_id": "HyUb55JSAL", "reviewer": "Reviewer_srX3", "summary": "This paper proposes OILCA, which addresses the scarcity of expert data in offline imitation learning by generating counterfactual samples, to imagine what the expert will do under a unobserved state. OILCA takes the perspective of Structual Causal Model (SCM). The algorithm consists of four steps: \n\n1) A heuristic expert policy is pretrained with the known expert data. \n\n2) A conditional VAE is trained on the union of expert and supplementary data with the latent variable being \"exogeneous variable\" that integrates the information from task label and transition, to rebuild the possible next state. \n\n3) Sample transitions with task labels from the expert data. Counterfactual next states are generated by applying a do-intervention on the latent variable, and expert action is then generated by the pretrained policy; the counterfactual transition is added back to the expert data. \n\n4) Arbitrary offline imitation learning method with supplementary data is applied on the augmented expert data (DWBC is empirically the best choice).\n\nOn many environments with both in-distribution and out-of-distribution dataset, OILCA is proved to be better than multiple baselines. The data augmentation technique of OILCA is orthogonal to other offline IL algorithms and can benefit many of them.\n", "questions": "Besides my concerns in the weakness section, I have two questions:\n\n1. Judging from the pseudocode, it is quite likely a generated counterfactual state is again sampled as expert data and used to generate other counterfactual states. Will this leads to wild compound error when this process goes too far from a single original expert data source? \n\n2. The pretrained expert policy is directly given in the pseudocode, but it also needs to be trained from original expert data, and thus will also suffer from out-of-distribution decision making. Does this create a chicken-and-egg dilemma, where you need to have a well-trained expert policy to generate more data, but data needs to be generated by a well-trained expert?\n\nBelow are my suggestions for the paper. I am happy to increase my score if the suggestions and concerns raised in the weakness section are addressed.\n\n1. Expand related work to include the works discussed in the weakness section.\n\n2. Give explanation for the questions proposed above, and the concern of moving parts.\n\n3. Modify the pseudocode, and possibly the main paper, to more clearly present the do-intervention process and fix the minor issue of weakness 3. a) and 3. b). ", "limitations": "The authors have discussed the limitations in the appendix line 582-587, and I think is indeed the most important problem that readers could ask: does counterfactually generated data necessarily improve the performance? Will it instead hinder the performance because of the poor quality of generated data? It is a pity that this paper does not answer the question, despite the fact that it is already a good work.\n\nAs for potential negative societal impact of the work, there is no discussion. I suggest the authors to include the discussion somewhere in the paper; though the work is purely on simulated control environment, the work on automated control would inevitably causes some concerns such as misuse of technology and potential job loss.\n\n", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**1. The proposed method is sound, and can be combined with multiple offline IL methods.** The idea of generating imaginary (counterfactual) samples to combat scarcity of expert data is a still novel direction that has been proved to be successful by several works [1, 2]. Furthermore, the algorithm is orthogonal to many offline IL algorithms and empirically proved to be an amplifier for them, which significantly increases its impact.\n\n**2. Generally, the idea of the paper is clearly conveyed through the authors' writing.** The motivation is clearly presented in the question between line 59 and 60, and the high-level process is clearly conveyed in the pseudocode. The questions asked at the beginning of the experiment section is well-answered and clearly indicates the important takeaways of the section.\n\n**3. The experiment result is solid.** Not only does OILCA works significantly and unanimously better than multiple baselines across in-data distribution and out-of-data distribution dataset on many environments, but it also improves the performance of the baselines. Moreover, the visualization on toy environment clearly shows the property and behavior pattern of the algorithm.\n\n**4. The paper also provides theoretical analysis to prove its generalization ability, which is a very important problem of IL/RL community.** Numerous methods, such as entropy regularizer [3] and pessimism [4], has been proposed to address the problem of generalization onto the vast state-action space without data coverage. Theoretical advantage on this area is a valuable contribution. \n\n**References:**\n\n[1] D. Hwang et al. Sample Generation for Reinforcement Learning via Diffusion Models. In ICLR, 2023.\n\n[2] C. Lu et al. Synthetic Experience Replay. In ArXiv, 2023.\n\n[3] J. Ho and S. Ermon. Generative Adversarial Imitation Learning. In NIPS, 2016.\n\n[4] G-H Kim et al. LobsDICE: Offline Learning from Observation via Stationary Distribution Correction Estimation. In NeurIPS, 2022.", "weaknesses": "\n**1. The related work section could be expanded to include more related work.**\n\na) The work is related to works that uses generative model and tries to do few-shot adaptation on offline IL such as FIST [1], PARROT [2] and CEIP [3] (the latter two has RL finetune, but already works decently well at the offline IL stage). While this work uses generative model in a very different way, those works should be briefly discussed due to the similarity of the problem formulation and the common use of generative model. \n\nb) The key idea of this work is to generate imaginary data to enhance the imitation learning, and the idea of data generation has already been studied by several works (see [1, 2] in the strength section). It would be great to see a discussion about the relation between this work and the existing works in data generation.\n\nThe above two parts are now missing from the related work, and poses a weakness on the investigation of the context.\n\n**2. There are many moving parts in OILCA.** While the authors claim that OILCA is a 2-step solution (line 135-136), the networks that need to be trained include two actors (pretrain & final), one discriminator, and encoder/decoder of the conditional VAE. Two pairs of them needs to be trained jointly. This could be a potential source for instability in training and extra usage of computational resource (and there is no computational resource and training time spent reported in the paper).\n\n**3. About the pseudocode:**\n\na)  \"$D_E\\leftarrow$\" should be added in front of the current line. This is a value assignment rather than an expression;\n\nb) The meaning of $c$ should be reiterated in the pseudocode for quick understanding.\n\nc) overall, I think the update process and generate process in line 3 and 7 can be clarified more clearly (for example, what is sampled in the beginning, and what is fed into a neural network with what output?) such that the do-intervention process on SCM is more clearly presented. Currently the readers can only make analogy from the definition on line 96-106 without explicit confirmation on notations, which is difficult to comprehend (especially when the parameter notation used in Eq. 2 and the pseudocode is different).\n\n\n**References:**\n\n[1] K. Hakhamaneshi et al. Hierarchical few-shot imitation with skill transition models. In ICLR, 2022.\n\n[2] A. Singh et al. Parrot: Data-driven behavioral priors for reinforcement learning. In ICLR, 2021.\n\n[3] K. Yan et al. CEIP: Combining Explicit and Implicit Priors for Reinforcement Learning with Demonstrations. In NeurIPS, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Dear reviewer srX3:**\n\nThanks for your review of our paper. Here's our response to the weaknesses, questions, and limitations you have highlighted:\n\n--------\n**Weaknesses:**\n\n**W1.The related work section could be expanded to include more related work.**\n\n**A**: Thanks for your suggestions. Firstly, we will follow your recommendations and conduct further research on the usage of generative models for few-shot adaptation in offline imitation learning. We will incorporate these papers into our revised version of the related work section. Secondly, we agree that data generation for imitation learning is highly relevant to our study. Therefore, we will include some representative works in this area in our revised related work section.\n\n**W2.Instablility analysis.**\n\n**A**:  Thanks for your insightful questions. We admit that our OILCA introduces several additional moving parts (a pretrained actor and encoder/decoder of conditional VAE) compared with vanilla DWBC. However, this does not necessarily bring the instability in training. First, the stability for joint training of final actor and discriminator has been previously proved in original paper of DWBC. Then, the pre-training of actor is just a simple supervised learning with the logged behavior data, which stability has been widely validated in various previous literacture on offline IL research. As for the joint training of encoder and decoder of conditional VAE in our counterfactual data augmentatiion module, its stability has been proved with the hyperparameter ablation study in our above response to Q2 of Reviewer agsz. \nAs for the computational resource, we apologize for nor providing related details in our submitted version considering that this is rarely addressed in the relevant literature. Actually, for the compuation device, we use the NVIDIA RTX A6000 GPU. Here, we further report the overall training time spent of our OILCA and baselines as follows. And we will supplement related analysis to the appendix in the revised version of paper.\n\nTraining time consumption for in-distribution experiments on Deepmind Control Suite (corresponding to the Table 1 in the paper).  All the methods are trained 200 epochs. \n**To save space, we put the result tables into the external anonymous link (Unit: sec.)**: https://hackmd.io/@littleshoes/ryaYsJfhn. Please check Table 10 in it.\n\nTraining time consumption for out-of-distribution experiments on CausalWorld (corresponding to the Table 2 in the paper).  All the methods are trained 200 epochs. **To save space, we put the result tables into the external anonymous link (Unit: sec.)**: https://hackmd.io/@littleshoes/ryaYsJfhn. Please check Table 11 in it.\n\nFrom the tables above, we can conclude that the additional time consumption brought by our counterfactual data augmentation is acceptable considering that offline IL component  still occupies the main part of time consumption.\n\n\n**W3. About the pseudocode:**\n\n**A**: Thanks for your suggestions. We have revised the Figure 2(b) (see https://i.postimg.cc/KzWk86kG/Figure2-b.jpg), and we will add the related descriptions in the pseudocode. \n\n--------\n**Questions:**\n\n**Q1.Counterfactual state compound error**\n\n**A**: Once we have acquired a proficiently trained counterfactual sample generation model, we randomly select $(s\\_t,a\\_t,s\\_{t+1},c)$ from the expert data collection and generate the counterfactual state $\\tilde{s}\\_{t+1}$. The generated samples are not utilized for subsequent expert data augmentation. A detailed explanation of this will be provided in the final version to prevent any misconceptions.\n\n**Q2.Detailed training procedure of pre-trained policy.**\n\n\n**A**: This question helps us to correct our presentation. We realize that we did not present the detailed training process of $\\hat{\\pi}\\_{E}$. To answer this, we provide the details here. \n One main perspective is how do we know that the interventions won’t lead to divergence, or learning use less models. \n\nIn our work, we use a pretrained policy $\\hat{\\pi}\\_E$, which can be regarded as the S-Learner [1] in causal inference. And we consider the counterfactual decision risk, the objective of trained policy as  $\\mathcal{L} = \\frac{1}{n}\\sum l(s\\_t, a\\_t) * \\frac{p(a\\_{t+1}\\mid s\\_{t+1})}{p^{do(I)}(a\\_{t+1}\\mid \\tilde{s}\\_{t+1})}$. $l$ is the policy training loss. For the sample weight $\\frac{p(a\\_{t+1}\\mid s\\_{t+1})}{p^{do(I)}(a\\_{t+1}\\mid \\tilde{s}\\_{t+1})}$ may have a big variance, we clip it with a small value $m$. Moreover, to analysis the decision risk,  we can derive the bound on the risk.\n\n[1]Zhang, W., Li, J., & Liu, L. (2021). A unified survey of treatment effect heterogeneity modelling and uplift modelling. ACM Computing Surveys (CSUR), 54(8), 1-36.\n\n\n**Lemma 1** The density of the counterfactuals based on the observations, i.e.\n$p^{do(I)}\\_{s\\_t, a\\_t}(\\tilde{s}\\_{t+1})=\\mathbb{E}\\_{s\\_{t+1} \\sim p\\_{s\\_t,a\\_t}(s\\_{t+1})}\\left[p\\_{s\\_t, a\\_t}^{{do}(I) \\mid s\\_{t+1}}(\\tilde{{s}}\\_{t+1})\\right]$\nThis result shows that the density of intervened variable $\\tilde{s}\\_{t+1}$ is the marginal of the observations.\n\n\n**Lemma 2.** We have the following lower bound on the logdensity of the counterfactuals:\n$$\\begin{aligned}  \\log \\left(p^{d o(I)}(\\tilde{s}\\_{t+1}, a\\_{t+1})\\right) \\geq& \\mathbb{E}\\_{s\\_{t+1} \\sim p\\_{s\\_t,a\\_t}(s\\_{t+1})}\\left[\\log \\left( p\\_{s\\_t, a\\_t}^{{do}(I) \\mid s\\_{t+1}}(\\tilde{{s}}\\_{t+1})\\right)\\right] \\\\\\\\ &+\\mathbb{E}\\_{{u} \\sim p({u})}\\left[\\log \\left(p\\_{s\\_t,a\\_t}^{{do}(I)}(\\tilde{s}\\_{t+1} \\mid {u})\\right)\\right]\\end{aligned}$$\n\nDue to the limited space, we only provide the brief proofs in the OpenReview. **For more detailed version, please refers to the Q2 of Reviewer srX3 in our external anonymous link (https://hackmd.io/@littleshoes/ryaYsJfhn).** And we will add them in the final revised Appendix. \n"}]

Author rebuttal:

Meta-review: This paper presents OILCA, a causality-regularized data augmentation method for offline imitation learning. The proposed framework generates counterfactual data to augment expert data and demonstrates effectiveness in experiments, particularly in the CausalWorld benchmark. The idea looks interesting and novel, and the proposed method is sound with theoretical and experimental support.
Decision: Accept (poster)