SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zexu Sun $^{\S}$ , Bowei He $^{\dagger}$ , Jinxin Liu $^{\ddagger}$ , Xu Chen $^{\S}$ *, Chen Ma $^{\dagger}$ , Shuai Zhang $^{\mathbb{I}}$
$^{\S}$ Gaoling School of Artificial Intelligence, Renmin University of China
$^{\dagger}$ Department of Computer Science, City University of Hong Kong
$^{\ddagger}$ School of Engineering, Westlake University  $^{\ddagger}$ DiDi Chuxing
{sunzexu21, xu.chen}@ruc.edu.cn, boweihe2-c@my.cityu.edu.hk
liujinxin@westlake.edu.cn, chenma@cityu.edu.hk, shuai.zhang@tju.edu.cn
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.
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
Figure 1: Agent is trained with the collected dataset containing limited expert data and large amounts of unlabeled data, and tested on both in-distribution and out-of-distribution environments.
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).
- 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.
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 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 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.
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  $\mathcal{D}_{\mathrm{all}} = \left\{s_t^i,a_t^i,s_{t + 1}^i\right\}_{i = 1}^{n_{\mathrm{all}}}$  s.t.  $(s_t,a_t)\sim \rho (s_t,a_t),s_{t + 1}\sim f_\varepsilon (s_t,a_t,u_{t + 1})$  , where  $\rho (s_t,a_t)$  is an offline state-action distribution resulting from some behavior policies,  $f_{\varepsilon}(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  $\varepsilon$  is a small permutation. Let  $\mathcal{D}_E$  and  $\mathcal{D}_U$  be the sets of expert and unlabeled demonstrations respectively, our goal is to only leverage the offline batch data  $\mathcal{D}_{\mathrm{all}}\coloneqq \mathcal{D}_E\cup \mathcal{D}_U$  to learn an optimal policy  $\pi$  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  $\mathcal{M}$  over variables  $\mathbf{X} = \{X_1, \ldots, X_n\}$  consists of a set of independent exogenous variables  $\mathbf{U} = \{\mathbf{u}_1, \ldots, \mathbf{u}_n\}$  with prior distributions  $P(\mathbf{u}_i)$  and a set of functions  $f_1, \ldots, f_n$  such that  $X_i = f_i(\mathbf{PA}_i, \mathbf{u}_i)$ , where  $\mathbf{PA}_i \subset \mathbf{X}$  are parents of  $X_i$ . Therefore, the distribution of the SCM, which is denoted  $P^{\mathcal{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 = \text{do}\left(X_{i} := f_{i}\left(\tilde{\mathbf{PA}}_{i},\mathbf{u}_{i}\right)\right)$  is defined as replacing some functions  $f_{i}\left(\mathbf{PA}_{i},\mathbf{u}_{i}\right)$  with  $f_{i}\left(\tilde{\mathbf{PA}}_{i},\mathbf{u}_{i}\right)$ , where  $\tilde{\mathbf{PA}}_i$  is the intervened parents of  $X_{i}$ . The intervened SCM is indicated as  $\mathcal{M}^I$ , and, consequently, its distribution is denoted as  $P^{\mathcal{M};I}$ .
The counterfactual inference with which we can answer the what if questions will be obtained in the following process:
(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(\mathbf{u}_i \mid \mathbf{X} = \mathbf{x})$ , where  $\mathbf{x}$  is a set of observations. Replace the prior distribution  $P(\mathbf{u}_i)$  with the posterior distribution  $P(\mathbf{u}_i \mid \mathbf{X} = \mathbf{x})$  in the SCM.
2. We denote the resulted SCM as  $\mathcal{M}_{\mathbf{x}}$  and its distribution as  $P^{\mathcal{M}_{\mathbf{x}}}$ , perform an intervention  $I$  on  $\mathcal{M}_{\mathbf{x}}$  to reach  $P^{\mathcal{M}_{\mathbf{x}};I}$ .
3. Return the output of  $P^{\mathcal{M}_{\mathbf{x}};I}$  as the counterfactual inference.
SCM representation of causal MDP We encode causal MDP (under a policy  $\pi$ ) into an SCM  $\mathcal{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} \sim f_{\varepsilon}(s_t, a_t, u_{t+1})$ , where  $\varepsilon$  is a small perturbation, and subsequently, to the next action  $a_{t+1}$ , e.g.  $a_{t+1} \sim \pi(\cdot \mid s_{t+1})$ . To perform counterfactual inference on the SCM  $\mathcal{M}$ , we intervene in the original parents pair  $(s_t, a_t)$  to  $(\tilde{s}_t, \tilde{a}_t)$ . Specifically, we sample the latent representation  $\tilde{u}_{t+1}$  from the distribution of exogenous variable  $u$ , and then generate the counterfactual next state  $\tilde{s}_{t+1}$ , i.e.  $\tilde{s}_{t+1} \sim f_{\varepsilon}(\tilde{s}_t, \tilde{a}_t, \tilde{u}_{t+1})$ , and also the counterfactual next action  $\tilde{a}_{t+1}$ , i.e.  $\tilde{a}_{t+1} \sim \hat{\pi}(\cdot \mid \tilde{s}_{t+1})$ , where  $\hat{\pi}$  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_{\theta}(\boldsymbol{x}, \boldsymbol{z}) = p_{\theta}(\boldsymbol{x} \mid \boldsymbol{z}) p_{\theta}(\boldsymbol{z})$ , where  $p_{\theta}(\boldsymbol{x} \mid \boldsymbol{z})$  represents the conditional distribution of observing  $\boldsymbol{x}$  given  $\boldsymbol{z}$ . Here,  $\theta$  denotes the set of generative parameters, and  $p_{\theta}(\boldsymbol{z}) = \prod_{i=1}^{I} p_{\theta}(z_i)$  represents the factorized prior distribution of the latent variables. By incorporating an inference model  $q_{\phi}(\boldsymbol{z} \mid \boldsymbol{x})$ , the parameters  $\phi$  and  $\theta$  can be jointly optimized by maximizing the evidence lower bound (ELBO) on the marginal likelihood  $p_{\theta}(\boldsymbol{x})$ .
$$
\begin{array}{l} \mathcal {L} = \mathbb {E} _ {q _ {\phi} (\boldsymbol {z} | \boldsymbol {x})} \left[ \log p _ {\boldsymbol {\theta}} (\boldsymbol {x} \mid \boldsymbol {z}) \right] - D _ {\mathrm {K L}} \left(q _ {\phi} (\boldsymbol {z} \mid \boldsymbol {x}) \| p (\boldsymbol {z})\right) \\ = \log p _ {\boldsymbol {\theta}} (\boldsymbol {x}) - D _ {\mathrm {K L}} \left(q _ {\phi} (\boldsymbol {z} \mid \boldsymbol {x}) \| p _ {\boldsymbol {\theta}} (\boldsymbol {z} \mid \boldsymbol {x})\right) \leq \log p _ {\boldsymbol {\theta}} (\boldsymbol {x}), \tag {1} \\ \end{array}
$$
where  $D_{\mathrm{KL}}$  denotes the KL-divergence between the approximation and the true posterior, and  $\mathcal{L}$  is a lower bound of the marginal likelihood  $p_{\theta}(\boldsymbol{x})$  because of the non-negativity of the KL-divergence. Recently, it has been shown that VAEs with unconditional prior distributions  $p_{\theta}(z)$  are not identifiable [18], but the latent factors  $z$  can be identified with a conditionally factorized prior distribution  $p_{\theta}(\boldsymbol{x} \mid 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  $\mathcal{D}_U$  and augmented  $\mathcal{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  $\pmb{\theta} = (\pmb{f},\pmb{T},\pmb{\lambda})$  be the parameters of the following conditional generative model:
$$
p _ {\boldsymbol {\theta}} \left(s _ {t + 1}, u \mid s _ {t}, a _ {t}, c\right) = p _ {\boldsymbol {f}} \left(s _ {t + 1} \mid s _ {t}, a _ {t}, u\right) p _ {\boldsymbol {T}, \boldsymbol {\lambda}} (u \mid c), \tag {2}
$$
where we first define:
$$
p _ {\boldsymbol {f}} \left(s _ {t + 1} \mid s _ {t}, a _ {t}, u\right) = p _ {\varepsilon} \left(s _ {t + 1} - \boldsymbol {f} _ {s _ {t}, a _ {t}} (u)\right). \tag {3}
$$
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) + \varepsilon$  where  $\varepsilon$  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:
$$
\begin{array}{l} \log p \left(s _ {t + 1} \mid s _ {t}, a _ {t}, c\right) \geq \mathcal {L} (\boldsymbol {\theta}, \phi) := \log p _ {\boldsymbol {f}} \left(s _ {t + 1} \mid s _ {t}, a _ {t}, u\right) + \log p (c) - \tag {4} \\ D _ {\mathrm {K L}} \left(q _ {\phi} (u \mid s _ {t}, a _ {t}, s _ {t + 1}, c) \| p _ {\boldsymbol {T}, \lambda} (u \mid c)\right). \\ \end{array}
$$
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_{\phi}$  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  $\mathcal{D}_E$ . As shown in Figure 2(b), once the posterior of exogenous variable  $q_{\phi}\left(u\mid s_t,a_t,s_{t + 1},c\right)$  is obtained, we can perform do-intervention. In particular, we intervene in the parents  $(s_t,a_t)$  of  $s_{t + 1}$  in  $\mathcal{D}_E$  by re-sampling a different pair  $(\tilde{s}_t,\tilde{a}_t)$  from the collected data in  $\mathcal{D}_U$ . Moreover, we sample the latent representation  $\tilde{u}_{t + 1}$  from the learned posterior distribution of exogenous variable  $u$ . Then utilizing  $\tilde{s}_t,\tilde{a}_t$ , and  $\tilde{u}_{t + 1}$ , we can generate the counterfactual next state according to Equation (3), which is denoted as  $\tilde{s}_{t + 1}$ . Subsequently, we pre-train a sampler policy  $\hat{\pi}_E$  with the original  $\mathcal{D}_E$  to sample the counterfactual next action  $\tilde{a}_{t + 1} = \hat{\pi}_E(\cdot |\tilde{s}_{t + 1})$ . Thus, for all the generated tuples  $(\tilde{s}_{t + 1},\tilde{a}_{t + 1})$ , we constitute them in  $\mathcal{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  $\pi_{\omega}$  as additional input, yielding a new discriminating task whose learning objective is as follows:
$$
\begin{array}{l} \mathcal {L} _ {\boldsymbol {\psi}} (\mathcal {D} _ {E}, \mathcal {D} _ {U}) = \eta \underset {(s _ {t}, a _ {t}) \sim \mathcal {D} _ {E}} {\mathbb {E}} \left[ - \log D _ {\boldsymbol {\psi}} (s _ {t}, a _ {t}, \log \pi_ {\boldsymbol {\omega}}) \right] + \underset {(s _ {t} ^ {\prime}, a _ {t} ^ {\prime}) \sim \mathcal {D} _ {U}} {\mathbb {E}} \left[ - \log (1 - D _ {\boldsymbol {\psi}} (s _ {t} ^ {\prime}, a _ {t} ^ {\prime}, \log \pi_ {\boldsymbol {\omega}}) \right] \\ - \eta_ {\left(s _ {t}, a _ {t}\right) \sim \mathcal {D} _ {E}} \left[ - \log \left(1 - D _ {\psi} \left(s _ {t}, a _ {t}, \log \pi_ {\omega}\right)\right) \right], \tag {5} \\ \end{array}
$$
Algorithm 1 Training procedure of OILCA.
Input: Dataset  $\mathcal{D}_E$ ,  $\mathcal{D}_U$ ,  $\mathcal{D}_{\mathrm{all}}$ , pre-trained sampler policy  $\hat{\pi}_E$ , hyperparameters  $\eta$ ,  $\alpha$ , initial variational couterfactual parameters  $\theta, \phi$ , discriminator parameters  $\psi$ , policy parameters  $\omega$ , data augmentation batch number  $B$ .
Output: Learned policy parameters  $\omega$ .
1: while counterfactual training do  $\triangleright$  Variational counterfactual reasoning
2: Sample  $(s_t, a_t, s_{t+1}, c) \sim \mathcal{D}_{\mathrm{all}}$  to form a training batch
3: Update  $\theta$  and  $\phi$  according to Equation (4)
4: end while
5: for  $b = 1$  to  $B$  do  $\triangleright$  Expert data augmentation
6: Sample  $(s_t, a_t, s_{t+1}, c) \sim \mathcal{D}_E$ ,  $(\tilde{s}_t, \tilde{a}_t) \sim \mathcal{D}_U$  to form an augmentation batch
7: Generate the counterfactual  $\tilde{s}_{t + 1}$  according to Equation (2), then predict  $\tilde{a}_{t + 1} = \hat{\pi}_E\big(\cdot \big|\tilde{s}_{t + 1}\big)$
8:  $\mathcal{D}_E\cup (\tilde{s}_{t + 1},\tilde{a}_{t + 1})$
9: end for
10: while agent training do  $\triangleright$  Learning the discriminator and policy cooperatively
11: Sample  $(s_t, a_t) \sim \mathcal{D}_E$  and  $(s_t', a_t') \sim \mathcal{D}_U$  to form a training batch
12: Update  $\psi$  according Equation (5) every 100 training steps  $\triangleright$  Discriminator learning
13: Update  $\omega$  according to Equation (6) every 1 training step  $\triangleright$  Policy learning
14: end while
where  $D_{\psi}$  is the discriminator,  $\eta$  is called the class prior. In the previous works [35, 39],  $\eta$  is a fixed hyperparameter often assigned as 0.5.