SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Notice that now  $\pi_{\omega}$  appears in the input of  $D_{\psi}$ , which means that imitation information from  $\log \pi_{\omega}$  will affect  $\mathcal{L}_{\psi}$ , and further impact the learning of  $D_{\psi}$ . Thus, inspired by the idea of adversarial training, DWBC [36] introduces a new learning objective for BC Task:
$$
\begin{array}{l} \mathcal {L} _ {\pi} = \alpha \underset {(s _ {t}, a _ {t}) \sim \mathcal {D} _ {E}} {\mathbb {E}} \left[ - \log \pi_ {\omega} \left(a _ {t} \mid s _ {t}\right) \right] - \underset {(s _ {t}, a _ {t}) \sim \mathcal {D} _ {E}} {\mathbb {E}} \left[ - \log \pi_ {\omega} \left(a _ {t} \mid s _ {t}\right) \cdot \frac {\eta}{d (1 - d)} \right] \tag {6} \\ + \underset {(s _ {t} ^ {\prime}, a _ {t} ^ {\prime}) \sim \mathcal {D} _ {U}} {\mathbb {E}} \left[ - \log \pi_ {\omega} (a _ {t} ^ {\prime} \mid s _ {t} ^ {\prime}) \cdot \frac {1}{1 - d} \right], \quad \alpha > 1. \\ \end{array}
$$
where  $d$  represents  $D_{\psi}(s_t, a_t, \log \pi_{\omega})$  for simplicity,  $\alpha$  is the weight factor ( $\alpha > 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 iVAE [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 transition mechanism  $p(s_{t+1} \mid 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} \perp c \mid 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-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  $\lambda(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) =$
$[\pmb{T}_1(u^1)^T,\dots ,\pmb{T}_d(u^d)^T ]^T$  from a factorized exponential family and the outputs of a neural network  $\pmb{T}_{NN}(u)$  with universal approximation power. The probability density can be written as:
$$
p _ {\boldsymbol {T}, \boldsymbol {\lambda}} (u \mid c) = \frac {\boldsymbol {Q} (u)}{\boldsymbol {Z} (u)} \exp \left[ \boldsymbol {T} (u) ^ {T} \boldsymbol {\lambda} (c) \right]. \tag {7}
$$
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  $\pmb{f}_{s_t,a_t}$  as  $\pmb{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, \lambda)$ , the following holds:
(i) The set  $\{s_{t + 1}\in S:\varphi_{\varepsilon}\big(s_{t + 1} = 0\big)\}$  has measure zero, where  $\varphi_{\varepsilon}$  is the characteristic function of the density  $p_{\varepsilon}$  defined in Equation (3).
(ii) The function  $\pmb{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, \ldots, c^k$  such that the matrix
$$
L = \left(\boldsymbol {\lambda} \left(c ^ {1}\right) - \boldsymbol {\lambda} \left(c ^ {0}\right), \dots , \boldsymbol {\lambda} \left(c ^ {k}\right) - \boldsymbol {\lambda} \left(c ^ {0}\right)\right) \tag {8}
$$
of size  $k\times k$  is invertible.
Then, the parameters  $\pmb{\theta} = (\pmb{f},\pmb{T},\lambda)$  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_{\phi}$  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  $(\theta, \phi)$  such that the family of distributions  $q_{\phi}(u \mid s_t, a_t, s_{t+1}, c)$  contains  $p_{\theta}(u \mid s_t, a_t, s_{t+1}, c)$ .
(ii) We maximize  $\mathcal{L}(\theta, \phi)$  with respect to both  $\theta$  and  $\phi$ .
Then, given infinite data, OILCA can learn the true parameters  $\pmb{\theta}^{*} \coloneqq (\pmb{f}^{*}, \pmb{T}^{*}, \pmb{\lambda}^{*})$ .
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  $(\tilde{s}_{t + 1},\tilde{a}_{t + 1})$  are estimated based on the sampler policy  $\hat{\pi}_E$ . However,  $\hat{\pi}_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  $\kappa \in (0,0.5)$  indicates the noise level of  $\hat{\pi}_E$ . If  $\hat{\pi}_E$  can exactly recover the true action  $a_{t + 1}$  (i.e.,  $\kappa = 0$ ), then the generated sequences are perfect without any noise. On the contrary,  $\kappa = 0.5$  means that  $\hat{\pi}_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  $\mathcal{H}$ , for any  $\epsilon, \delta \in (0,1)$  and  $\kappa \in (0,0.5)$ , if  $\hat{\pi}_E \in \mathcal{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  $\frac{2\log\left(\frac{2\mathcal{H}}{\delta}\right)}{\epsilon^2(1 - 2\kappa)^2}$ , then the error between the model estimated and true results is smaller than  $\epsilon$  with probability larger than  $1 - \delta$ .
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.,  $\epsilon$  and  $\delta$  are fixed), one needs to generate more than  $\frac{2\log\left(\frac{2\mathcal{H}}{\delta}\right)}{\epsilon^2(1 - 2\kappa)^2}$  sequences. If the noise level  $\kappa$  is larger, more samples have to be generated. Extremely, if the pre-trained policy can only produce fully noisy information (i.e.,  $\kappa = 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|, |\mathcal{A}| < \infty)$ , and for the learned policy  $\pi_{\omega}$ , we can analyze the generalization ability from the perspective of error upper bound with the optimal expert policy  $\pi^{*}$ .
Theorem 4. Let  $|\mathcal{D}_E|$  be the number of empirical expert data used to train the policy and  $\epsilon$  be the expected upper bound of generalization error. There exists constant  $h$  such that, if
$$
\left| \mathcal {D} _ {E} \right| \geq \frac {h | \mathcal {S} | | \mathcal {A} | \log (| \mathcal {S} | / \delta)}{\epsilon^ {2}}, \tag {9}
$$
and each state  $s_t$  is sampled uniformly, then, with probability at least  $1 - \delta$ , we have:
$$
\max  _ {s _ {t}} \left\| \pi^ {*} (\cdot | s _ {t}) - \pi_ {\omega} (\cdot | s _ {t}) \right\| _ {1} \leq \epsilon . \tag {10}
$$
which shows that increasing  $|\mathcal{D}_E|$  drastically improves the generalization guarantee.
Note that we provide the proof details for Theorem 4 in Appendix D.
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  $\mathcal{D}_E$ ), BC-all (BC on all demonstrations  $\mathcal{D}_{\mathrm{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 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 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  $|\mathcal{D}_E| / |\mathcal{D}_U|$ , where  $\mathcal{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  $|\mathcal{D}_E|$ .
Figure 3: Visualization of both observation and latent spaces of the exogenous variable. (a) Samples from the true distribution of the sources  $p_{\theta^*}(u|c)$ . (b) Samples from the posterior  $q_{\phi}(u|s_t, a_t, s_{t+1}, c)$ . (c) Samples from the posterior  $q_{\phi'}(u|s_t, a_t, s_{t+1})$  without class label.
Figure 4: Performance of OILCA and baselines in the toy environment. We plot the mean and the standard errors of the averaged return over five random seeds.
Figure 5: Performance of OILCA with the growing percentage of  $|\mathcal{D}_E| / |\mathcal{D}_U|$ . We plot the average return's mean and standard errors over five random seeds.
Table 1: Results for in-distribution 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 Name	BC-exp	BC-all	ORIL	BCND	LobsDICE	DWBC	OILCA
CARTPOLE SWINGUP	195.44 ± 7.39	269.03 ± 7.06	221.24 ± 14.49	243.52 ± 11.33	292.96 ± 11.05	382.55 ± 8.95	608.38 ± 35.54
CHEETAH RUN	66.59 ± 11.09	90.01 ± 31.74	45.08 ± 9.88	96.06 ± 16.15	74.53 ± 7.75	66.87 ± 4.60	116.05 ± 14.65
FINGER TURN HARD	129.20 ± 4.51	104.56 ± 8.32	185.57 ± 26.75	204.67 ± 13.18	190.93 ± 12.19	243.47 ± 17.12	298.73 ± 5.11
FISH SWIM	74.59 ±11.73	68.87 ± 11.93	84.90 ±1.96	153.28 ±19.29	188.84 ±11.28	212.39 ±7.62	290.28 ± 10.07
HUMANOID RUN	77.59 ±8.63	138.93 ±9.14	96.88 ±10.76	257.01 ±11.21	120.87 ±10.66	302.33 ±14.53	460.96 ±17.55
MANIPULATOR INSERT BALL	91.61 ±7.49	98.59 ±1.62	98.25 ±2.68	141.71 ±15.30	197.79 ±1.98	107.86 ±15.01	296.83 ±3.43
MANIPULATOR INSERT PEG	92.02 ±15.04	119.63 ±6.37	105.86 ±5.10	220.66 ±15.14	299.19 ±3.40	238.39 ±19.76	305.66 ±4.91
WALKER STAND	169.14 ±8.00	192.14 ±37.91	181.23 ±10.31	279.66 ±12.69	252.34 ±7.73	280.07 ±5.79	395.51 ±8.05
WALKER WALK	29.44 ±2.18	75.79 ±6.34	41.43 ±4.05	157.44 ±9.31	102.14 ±5.94	166.95 ±10.68	377.19 ±15.87
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 CAUSAL-WORLD [2], which provides a combinatorial family of tasks with a common causal structure and
Table 2: Results for out-of-distribution generalization on CAUSALWORLD. 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
REACHING	281.18 ± 16.45	176.54 ± 9.75	339.40 ± 12.98	228.33 ± 7.14	243.29 ± 9.84	479.92 ± 18.75	976.60 ± 20.13
PUSHING	256.64 ± 12.70	235.58 ± 10.23	283.91 ± 19.72	191.23 ± 12.64	206.44 ± 15.35	298.09 ± 14.94	405.08 ± 24.03
PICKING	270.01 ± 13.13	258.54 ± 16.53	388.15 ± 19.21	221.89 ± 7.68	337.78 ± 12.09	366.26 ± 8.77	491.09 ± 6.44
PICK AND PLACE	294.06 ± 7.34	225.42 ± 12.44	270.75 ± 14.87	259.12 ± 8.01	266.09 ± 10.31	349.66 ± 7.39	490.24 ± 11.69
STACKING2	496.63 ± 7.68	394.91 ± 16.98	388.55 ± 10.93	339.18 ± 9.46	362.47 ± 17.05	481.07 ± 10.11	831.82 ± 11.78
TOWERS	667.81 ± 9.27	784.88 ± 17.17	655.96 ± 15.14	139.30 ± 18.22	535.74 ± 13.76	768.68 ± 24.77	994.82 ± 5.76
STACKED BLOCKS	581.91 ± 26.92	452.88 ± 18.78	702.15 ± 15.30	341.97 ± 33.69	250.44 ± 14.08	1596.96 ± 81.84	2617.71 ± 88.07
CREATIVE STACKED BLOCKS	496.32 ± 26.92	529.82 ± 31.01	882.27 ± 46.79	288.55 ± 19.63	317.95 ± 32.03	700.23 ± 13.71	1348.49 ± 55.05
GENERAL	492.78 ±17.64	547.32 ±8.49	647.95 ±24.39	195.06 ±9.80	458.27 ±18.69	585.98 ±19.25	891.14 ±23.12
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  $\mathcal{D}_U$  is too conservative, resulting in a suboptimal policy, especially when  $\mathcal{D}_U$  contains a large collection of noisy data. This indeed hurts the performance of LobsDICE for out-of-distribution generalization.
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.
[1] Alekh Agarwal, Nan Jiang, Sham M. Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms, 2021.
[36] 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.
[37] 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.
[38] 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.
[39] 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.
[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.