Title: Offline Imitation Learning with Variational Counterfactual Reasoning
Venue: NeurIPS 2023
Accepted: yes
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

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 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, we conduct extensive experiments to demonstrate that our approach significantly outperforms various baselines on both DEEPMIND CONTROL SUITE benchmark for in-distribution performance and CAUSALWORLD benchmark for out-of-distribution generalization.

1 Introduction

By utilizing the pre-collected expert data, imitation learning (IL) allows us to circumvent the difficulty in designing proper rewards for decision-making tasks and learning an expert policy. Theoretically, as long as adequate expert data are accessible, we can easily learn an imitator policy that maintains a sufficient capacity to approximate the expert behaviors [25, 30, 29, 31]. However, in practice, several challenging issues hinder its applicability to practical tasks. In particular, expert data are often limited, and due to 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 [11, 34, 37]. Alternatively, in such settings, we might instead have access to large amounts of pre-collected unlabeled data, which are of unknown quality and may consist of both good-performing and poor-performing trajectories. For example, in self-driving tasks, a 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) [5] directly mimics historical behaviors logged in offline data (both expert data and unlabeled data) via supervised learning. However, in our above setting, BC

suffers from unstable training, as it relies on sufficient high-quality offline data, which is unrealistic. Besides, utilizing unlabeled data indiscriminately will lead to severe catastrophes: Bad trajectories mislead policy learning, and good trajectories fail to provide strong enough guidance signals for policy learning. In order to better distinguish the effect of different quality trajectories 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 a fixed dataset; DWBC [36] introduces a discriminator-weighted task to assist the policy learning. However, such discriminator-based methods are prone to overfitting and suffer from poor generalization ability, especially when only very limited expert data are provided.

Indeed, the expert data play an important role in offline IL and indicate the well-performing agent's intention. Thus, how to fully understand the expert's behaviors or preferences via limited available expert data becomes pretty crucial. In this paper, we propose to investigate counterfactual techniques for interpreting such expert behaviors. Specifically, we leverage the variational counterfactual reasoning [23, 4] to augment the expert data, as typical imitation learning data augmentation

Generate Counterfactual Expert data

via Variational Counterfactual

Reasoning

What action might the

expert take if a different

state is observed?

In-distribution

Out-of-distribution

Unlabeled Data

methods [3, 7] 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 action might the expert take if a 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 as the minimal edits to existing expert training data so that we 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 1).

The main contributions of this paper are summarized as follows:

We propose a 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 a 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-distribution generalization on CAUSALWORLD benchmark both demonstrate the effectiveness of our method.

2 Related Works

Offline IL Asignificant 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 to mediocre levels on many datasets. To address this issue, ORIL [39] learns a 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 [12] 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 [16] explicitly learns a 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 as the weights of the BC loss. CLUE [17] proposes to learn an intrinsic reward that is consistent with the expert intention via enforcing the embeddings of expert data to a 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-policy evaluation [6, 22], learning baselines for model-free RL [20], and policy transfer [10, 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 a 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 of the exogenous variable in 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. In our problem setting, we have an offline static dataset consisting of i.i.d tuples D all = { s i t , a i t , s i t + 1 } n all i = 1 s.t. ( s t , a t ) ∼ ρ ( s t , a t ) , s t + 1 ∼ f ε ( s t , a t , u t + 1 ) , where ρ ( s t , a t ) is an offline state-action distribution resulting from some behavior policies, f ε ( s t , a t , u t + 1 ) represents the causal transition mechanism, u t + 1 is the sample of the exogenous variable u , which is unobserved, and ε is a small permutation. Let D E and D U be the sets of expert and unlabeled demonstrations respectively, our goal is to only leverage the offline batch data D all ∶= D E ∪D U to learn an optimal policy π without any online interaction.

3.2 Counterfactual Reasoning

We provide a brief background on counterfactual reasoning. Further details can be found in [23].

Definition 1 (Structural Causal Model (SCM)) . A structural causal model M over variables X = { X 1 , . . . , X n } consists of a set of independent exogenous variables U = { u 1 , . . . , u n } with prior distributions P ( u i ) and a set of functions f 1 , . . . , f n such that X i = f i ( PA i , u i ) , where PA i ⊂ X are parents of X i . Therefore, the distribution of the SCM, which is denoted P M , is determined by the functions and the prior distributions of exogenous variables.

Inferring the exogenous random variables based on the observations, we can intervene in the observations and inspect the consequences.

Definition 2 ( do -intervention in SCM) . An intervention I = do ( X i ∶= f i ( ˜ PA i , u i )) is defined as replacing some functions f i ( PA i , u i ) with f i ( ˜ PA i , u i ) , where ˜ PA i is the intervened parents of X i . The intervened SCM is indicated as M I , and, consequently, its distribution is denoted as P M ; I .

The counterfactual inference with which we can answer the what if questions will be obtained in the following process:

u

Exogenous variable

Latent representation

Observations

Counterfactuals

(a) SCM of causal MDP

u

sample

do

-Intervention

(b) SCM of causal MDP with

do

-intervention

Infer the posterior distribution of exogenous variable P ( u i ∣ X = x ) , where x is a set of observations. Replace the prior distribution P ( u i ) with the posterior distribution P ( u i ∣ X = x ) in the SCM.

We denote the resulted SCM as M x and its distribution as P M x , perform an intervention I on M x to reach P M x ; I .

Return the output of P M x ; I as the counterfactual inference.

SCM representation of causal MDP We encode causal MDP (under a policy π ) into an SCM M . The SCM shown in Figure 2 consists of an exogenous variable u and a set of functions that transmit the state s t , action a t and latent representation u t + 1 of u to the next state s t + 1 e.g. s t + 1 ∼ f ε ( s t , a t , u t + 1 ) , where ε is a small perturbation, and subsequently, to the next action a t + 1 , e.g. a t + 1 ∼ π (⋅ ∣ s t + 1 ) . To perform counterfactual inference on the SCM M , we intervene in the original parents pair ( s t , a t ) to ( ˜ s t , ˜ a t ) . Specifically, we sample the latent representation ˜ u t + 1 from the distribution of exogenous variable u , and then generate the counterfactual next state ˜ s t + 1 , i.e. ˜ s t + 1 ∼ f ε ( ˜ s t , ˜ a t , ˜ u t + 1 ) , and also the counterfactual next action ˜ a t + 1 , i.e. ˜ a t + 1 ∼ ˆ π (⋅ ∣ ˜ s t + 1 ) , where ˆ π is the sampler policy.

3.3 Variational Autoencoder and Identifiability

We briefly introduce the Variational Autoencoder (VAE) and present its identifiability results. The VAE can be conceptualized as the amalgamation of a generative latent variable model and an associated inference model, both of which are parameterized by neural networks [13]. Specifically, the VAE learns the joint distribution p θ ( x , z ) = p θ ( x ∣ z ) p θ ( z ) , where p θ ( x ∣ z ) represents the conditional distribution of observing x given z . Here, θ denotes the set of generative parameters, and p θ ( z ) = Π I i = 1 p θ ( z i ) represents the factorized prior distribution of the latent variables. By incorporating an inference model q ϕ ( z ∣ x ) , the parameters ϕ and θ can be jointly optimized by maximizing the evidence lower bound (ELBO) on the marginal likelihood p θ ( x ) .

where D KL denotes the KL-divergence between the approximation and the true posterior, and L is a lower bound of the marginal likelihood p θ ( x ) because of the non-negativity of the KLdivergence. Recently, it has been shown that V AEs with unconditional prior distributions p θ ( z ) are not identifiable [18], but the latent factors z can be identified with a conditionally factorized prior distribution p θ ( x ∣ z ) over the latent variables to break the symmetry [9].

4 Offline Imitation Learning with Counterfactual Data Augmentation

At a high level, OILCA consists of following steps: (1) data augmentation via variational counterfactual reasoning, and (2) cooperatively learning a discriminator and a policy by using D U and augmented D E . In this section, we detail these two steps and present our method's theoretical guarantees.

4.1 Counterfactual Data Augmentation

Our intuition is 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 1. 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 [13, 28] to build the latent representation of the exogenous variable u . Moreover, Considering the identifiability of unsupervised disentangled representation learning [18], 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 [9]. Formally, let θ = ( f , T , λ ) be the parameters of the following conditional generative model:

where we first define:

Equation (2) describes the generative mechanism of s t + 1 given the underlying exogenous variable u , along with ( s t , a t ) . Equation (3) implies that the observed representation s t + 1 is an additive noise function, i.e. , s t + 1 = f s t ,a t ( u ) + ε where ε is independent of f s t ,a t 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 s t + 1 of the above generative model can be written as:

Note that the ELBO in Equation (4) contains an additional term of the log 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 q ϕ contains all the observable variables. The reason is that, as shown in Figure 2(a), from a causal perspective, s t , a t and u form a collider at s t + 1 , which means that when given s t + 1 , s t and a t are related to u .

We seek to augment the expert data in D E . As shown in Figure 2(b), once the posterior of exogenous variable q ϕ ( u ∣ s t , a t , s t + 1 , c ) is obtained, we can perform do -intervention. In particular, we intervene in the parents ( s t , a t ) of s t + 1 in D E by re-sampling a different pair ( ˜ s t , ˜ a t ) from the collected data in D U . Moreover, we sample the latent representation ˜ u t + 1 from the learned posterior distribution of exogenous variable u . Then utilizing ˜ s t , ˜ a t , and ˜ u t + 1 , we can generate the counterfactual next state according to Equation (3), which is denoted as ˜ s t + 1 . Subsequently, we pre-train a sampler policy ˆ π E with the original D E to sample the counterfactual next action ˜ a t + 1 = ˆ π E (⋅ ∣ ˜ s t + 1 ) . Thus, for all the generated tuples ( ˜ s t + 1 , ˜ a t + 1 ) , we constitute them in D E .

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-based 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-Weighted Behavioral Cloning (DWBC) [36], a state-of-the-art discriminator-based offline IL method. This method introduces a unified framework to learn the policy and discriminator cooperatively. The discriminator training gets information from the policy π ω as additional input, yielding a new discriminating task whose learning objective is as follows:

where D ψ is the discriminator, η is called the class prior. In the previous works [35, 39], η is a fixed hyperparameter often assigned as 0.5.

Notice that now π ω appears in the input of D ψ , which means that imitation information from log π ω will affect L ψ , and further impact the learning of D ψ . Thus, inspired by the idea of adversarial training, DWBC [36] introduces a new learning objective for BC Task:

where d represents D ψ ( s t , a t , log π ω ) for simplicity, α is the weight factor ( α > 1 ) . The detailed training procedure of OILCA is shown in Algorithm 1. Moreover, we also present more possible combinations with other offline IL methods and the related results in Appendix F.

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 iV AE [9]. To begin with, some assumptions are needed. Considering the SCM in Figure 2(a), we assume the data generation process in Assumption 1.

Assumption 1. (a) The distribution of the exogenous variable u is independent of time but dependent on the auxiliary variable c . (b) The prior distributions of exogenous variable u are different across auxiliary variable c . (c) The trasition mechanism p ( s t + 1 ∣ s t , a t , u ) are invariant across different auxiliary variable c . (d) Given the exogenous variable u , the next state s t + 1 is independent of the auxiliary variable c . i.e. s t + 1 /upmodels c ∣ u .

Part of the above assumption is also used in [19]; considering the identifiability of iV AE, we also add some necessary assumptions, which are also practical. In the following discussion, we will show that when the underlying data-generating mechanism satisfies Assumption 1, the exogenous variable u can be identified up to permutation and affine transformations if the conditional prior distribution p ( u ∣ c ) belongs to a general exponential family distribution.

Assumption 2. The prior distribution of the exogenous variable p ( u ∣ c ) follows a general exponential family with its parameter specified by an arbitrary function λ ( c ) and sufficient statistics T ( u ) = [ T d ( u ) , T NN ( u )] , here T d ( u ) is defined by the concatenation of T d ( u ) =

[ T 1 ( u 1 ) T , /uni22EF , T d ( u d ) T ] T from a factorized exponential family and the outputs of a neural network T NN ( u ) with universal approximation power. The probability density can be written as:

Under the Assumption 2, and leveraging the ELBO in Equation (4), we can obtain the following identifiability of the parameters in the model. For convenience, we omit the subscript of f s t ,a t as f .

Theorem 1. Assume that we observe data sampled from a generative model defined according to Equation (2) -(3) and Equation (7) with parameters ( f , T , λ ) , the following holds:

(i) The set { s t + 1 ∈ S ∶ φ ε ( s t + 1 = 0 )} has measure zero, where φ ε is the characteristic function of the density p ε defined in Equation (3) .

(ii) The function f is injective and all of its second-order cross partial derivatives exist.

(iii) The sufficient statistics T d are twice differentiable.

(iv) There exist k + 1 distinct points c 0 , . . . , c k such that the matrix

of size k × k is invertible.

Then, the parameters θ = ( f , T , λ ) are identifiable up to an equivalence class induced by permutation and component-wise transformations.

Theorem 1 guarantees the identifiability of Equation (2). We present its proof in Appendix A.

Moreover, Theorem 1 further implies a consistent result on the conditional VAE. If the variational distribution of encoder q ϕ is a broad parametric family that includes the true posterior, we have the following results.

Theorem 2. Assume the following holds:

(i) There exists the ( θ , ϕ ) such that the family of distributions q ϕ ( u ∣ s t , a t , s t + 1 , c ) contains p θ ( u ∣ s t , a t , s t + 1 , c ) .

(ii) We maximize L( θ , ϕ ) with respect to both θ and ϕ .

Then, given infinite data, OILCA can learn the true parameters θ ∗ ∶= ( f ∗ , T ∗ , λ ∗ ) .

We present the corresponding proof in Appendix B. Theorem 2 is proved by assuming our conditional VAE is flexible enough to ensure the ELBO is tight for some parameters and the optimization algorithm can achieve the global maximum of ELBO.

In our framework, the current generated expert sample pairs ( ˜ s t + 1 , ˜ a t + 1 ) are estimated based on the sampler policy ˆ π E . However, ˆ π E may be not perfect, and its predicted results may contain noise. Thus, we would like to answer: 'given the noise level of the sampler policy, how many samples one need to achieve sufficiently well performance?'. Using κ ∈ ( 0 , 0 . 5 ) indicates the noise level of ˆ π E . If ˆ π E can exactly recover the true action a t + 1 (i.e., κ = 0 ), then the generated sequences are perfect without any noise. On the contrary, κ = 0 . 5 means that ˆ π E 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 κ ∈ ( 0 , 0 . 5 ) , if ˆ π E ∈ H is the pretrained policy model learned based on the empirical risk minimization (ERM), and the sample complexity (i.e., number of samples) is larger than 2log ( 2 H δ ) ϵ 2 ( 1 -2 κ ) 2 , then the error between the model estimated and true results is smaller than ϵ with probability larger than 1 -δ .

The related proof details are presented in Appendix C. From Theorem 3, we can see: in order to guarantee the same performance with a given probability (i.e., ϵ and δ are fixed), one needs to generate more than 2log ( 2 H δ ) ϵ 2 ( 1 -2 κ ) 2 sequences. If 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., κ = 0 . 5 ), then infinity number of samples are required, which is impossible in reality.

For the generalization ability, [8] 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∣ < ∞) , and for the learned policy π ω , we can analyze the generalization ability from the perspective of error upper bound with the optimal expert policy π ∗ .

Theorem 4. Let ∣D E ∣ 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

and each state s t is sampled uniformly, then, with probability at least 1 -δ , we have:

which shows that increasing ∣D E ∣ drastically improves the generalization guarantee.

Note that we provide the proof details for Theorem 4 in Appendix D.

5 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 out-of-distribution test environment, can OILCA improve the generalization? In addition, we use five baseline methods: BC-exp (BC on the expert data D E ), BC-all (BC on all demonstrations D all), ORIL [39], BCND [26], LobsDICE [12], DWBC [36]. More details about these methods are presented in Appendix E. Furthermore, the hyper-parameters of our 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 to 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 a fixed distance in each of the four directions. States are continuous and are considered as the 2D position of the agent. The goal is to navigate to a specific target state within a bounding box. The reward is 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 c is the class label. We design the transition function by using a 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 a 2D case in Figure 3, 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 V AE 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 4, showing that OILCA achieves the highest average return. To analyze the influence of the number of augmented samples (Theorem 4), 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 ∣D E ∣/∣D U ∣ , where D E 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 ∣D E ∣ .

6

4

2

0

-2

-4

-6

-4

-2

0

2

4

6

8

0.100

0.125

0.150

0.175

0.200

0.225

0.250

0.275

(a)

0

2

4

6

-2

0

2

4

6

8

10

0.100

0.125

0.150

0.175

0.200

0.225

0.250

0.275

(b)

-4

-2

0

2

4

6

8

-2

0

2

4

6

8

(c)

10

0.300

0.275

0.250

0.225

0.200

0.175

0.150

0.125

0.100

BC-exp

BC-all

ORIL

BCND LobsDICE DWBC

OILCA

0

2

4

6

8

10

12

14

Average Return

10%

30%

50%

70%

90%

Percentage

9.0

9.5

10.0

10.5

11.0

11.5

12.0

12.5

13.0

Average Return

5.2 In-distribution Performance (Q2)

We use DEEPMIND CONTROL SUITE [32] to evaluate the performance of in-distribution performance. Similar to [39], we also define an episode as 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 as c = { 0 , 1 , 2 } ). Moreover, we present the detailed data collection and statistics in Appendix E.1). We report the results in Table 1. We can observe that OILCA outperforms baselines on all tasks; this shows the effectiveness of the augmented counterfactual expert data. We also find that the performance of ORIL and DWBC tends to decrease in testing for some tasks (ORIL: CHEETAH RUN, WALKER WALK; DWBC: CHEETAH RUN, MANIPULATOR INSERT BALL); this 'overfitting' phenomenon also occurs in experiments of previous offline RL works [14, 33]. 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-distribution generalization of OILCA, we use a benchmark named CAUSALWORLD [2], which provides a combinatorial family of tasks with a common causal structure and

0.300

0.300

underlying factors. Furthermore, the auxiliary variable c represents the different do -interventions on the task environments [2]. We collect the offline data by using three different do -interventions on environment features ( stage_color , stage_friction . floor_friction ) to generate offline datasets, while other features are set as the defaults. More data collection and statistics details are presented in Appendix E.1. We show the comparative results in Table 2. It is evident that OILCA also achieves the best performance on 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 as low-quality data when evaluating on space B , which may lead 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 D U is too conservative, resulting in a suboptimal policy, especially when D U contains a large collection of noisy data. This indeed hurts the performance of LobsDICE for out-of-distribution generalization.

6 Conclusion

In this paper, we propose an effective and generalizable offline imitation learning framework OILCA, which can learn a policy from the expert and unlabeled demonstrations. We apply a novel identifiable counterfactual expert data augmentation approach to facilitate the offline imitation learning. We also analyze the influence of generated data and the generalization ability theoretically. The experimental results demonstrate that our method achieves better performance in simulated and public tasks.

References

Alekh Agarwal, Nan Jiang, Sham M. Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms, 2021.

Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Yoshua Bengio, Bernhard Schölkopf, Manuel Wüthrich, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. arXiv preprint arXiv:2010.04296 , 2020.

Dafni Antotsiou, Carlo Ciliberto, and Tae-Kyun Kim. Adversarial imitation learning with trajectorial augmentation and correction. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 4724-4730. IEEE, 2021.

Léon Bottou, Jonas Peters, Joaquin Quiñonero-Candela, Denis X 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.

Ivan Bratko, Tanja Urbanˇ ciˇ c, and Claude Sammut. Behavioural cloning: phenomena, results and problems. IFAC Proceedings Volumes , 28(21):143-149, 1995.

Lars Buesing, Theophane Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Jean-Baptiste Lespiau, and Nicolas Heess. Woulda, coulda, shoulda: Counterfactually-guided policy search. arXiv preprint arXiv:1811.06272 , 2018.

Satoshi Hoshino, Tomoki Hisada, and Ryota Oikawa. Imitation learning based on data augmentation for robotic reaching. In 2021 60th Annual Conference of the Society of Instrument and Control Engineers of Japan (SICE) , pages 417-424. IEEE, 2021.

Divyansh Kaushik, Amrith Setlur, Eduard Hovy, and Zachary C Lipton. Explaining the efficacy of counterfactually augmented data. arXiv preprint arXiv:2010.02114 , 2020.

Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: A unifying framework. In International Conference on Artificial Intelligence and Statistics , pages 2207-2217. PMLR, 2020.

Taylor W Killian, Marzyeh Ghassemi, and Shalmali Joshi. Counterfactually guided policy transfer in clinical settings. In Conference on Health, Inference, and Learning , pages 5-31. PMLR, 2022.

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.

Geon-Hyeong Kim, Jongmin Lee, Youngsoo Jang, Hongseok Yang, and Kee-Eung Kim. Lobsdice: offline imitation learning from observation via stationary distribution correction estimation. arXiv preprint arXiv:2202.13536 , 2022.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing offpolicy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Yao Lai, Jinxin Liu, Zhentao Tang, Bin Wang, Jianye Hao, and Ping Luo. Chipformer: Transferable chip placement via offline decision transformer. arXiv preprint arXiv:2306.14744 , 2023.

Jinxin Liu, Li He, Yachen Kang, Zifeng Zhuang, Donglin Wang, and Huazhe Xu. Ceil: Generalized contextual imitation learning. arXiv preprint arXiv:2306.14534 , 2023.

Jinxin Liu, Lipeng Zu, Li He, and Donglin Wang. Clue: Calibrated latent guidance for offline reinforcement learning. arXiv preprint arXiv:2306.13412 , 2023.

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning , pages 41144124. PMLR, 2019.

Chaochao Lu, Biwei Huang, Ke Wang, José Miguel Hernández-Lobato, Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual-based data augmentation. arXiv preprint arXiv:2012.09092 , 2020.

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 arXiv:2011.09464 , 2020.

Graziano Mita, Maurizio Filippone, and Pietro Michiardi. An identifiable double vae for disentangled representations. In International Conference on Machine Learning , pages 77697779. PMLR, 2021.

Michael Oberst and David Sontag. Counterfactual off-policy evaluation with gumbel-max structural causal models. In International Conference on Machine Learning , pages 4881-4890. PMLR, 2019.

Judea Pearl. Models, reasoning and inference. Cambridge, UK: CambridgeUniversityPress , 19 (2), 2000.

Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfactual data augmentation. arXiv preprint arXiv:2210.11287 , 2022.

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 international conference on artificial intelligence and statistics , pages 627-635. JMLR Workshop and Conference Proceedings, 2011.

Fumihiro Sasaki and Ryota Yamashina. Behavioral cloning from noisy demonstrations. In International Conference on Learning Representations , 2021.

Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms . Cambridge university press, 2014.

Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems , 28, 2015.

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.02872 , 2021.

Wen Sun, Anirudh Vemula, Byron Boots, and Drew Bagnell. Provably efficient imitation learning from observation alone. In International conference on machine learning , pages 6036-6045. PMLR, 2019.

Gokul Swamy, Sanjiban Choudhury, J Andrew Bagnell, and Steven Wu. Of moments and matching: A game-theoretic framework for closing the imitation gap. In International Conference on Machine Learning , pages 10022-10032. PMLR, 2021.

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.00690 , 2018.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Zheng Wu, Liting Sun, Wei Zhan, Chenyu Yang, and Masayoshi Tomizuka. Efficient samplingbased maximum entropy inverse reinforcement learning with application to autonomous driving. IEEE Robotics and Automation Letters , 5(4):5355-5362, 2020.

Danfei Xu and Misha Denil. Positive-unlabeled reward learning. In Conference on Robot Learning , pages 205-219. PMLR, 2021.

Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning , pages 24725-24742. PMLR, 2022.

Zhenting Zhao, Bowei He, Wenhao Luo, and Rui Liu. Collective conditioned reflex: A bioinspired fast emergency reaction mechanism for designing safe multi-robot systems. IEEE Robotics and Automation Letters , 7(4):10985-10990, 2022.

Zheng-Mao 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.

Konrad Zolna, 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 arXiv:2011.13885 , 2020.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\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.\n\nQUESTIONS:\nBesides 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).\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nPoints 5 and 6 of the weaknesses are the main questions I want to ask the authors.\n\nLIMITATIONS:\nNo 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\nThis 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.\n\nQUESTIONS:\nsee weaknesses\n\nLIMITATIONS:\nNAN", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nOverall, 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.\n\nWEAKNESSES:\nMissing 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.\n\nQUESTIONS:\n1.**(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).\n\nLIMITATIONS:\nplease see above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-16"}