SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(Repeat from the main paper.) To gain more insight into Equation 4 that captures the quality of IL (the degree of similarity to the expert data), we define  $D(\cdot, \cdot)$  as the sum of reverse KL and forward KL divergence, i.e.,  $D(q, p) = D_{\mathrm{KL}}(q \| p) + D_{\mathrm{KL}}(p \| q)$ , and derive an alternative form for Equation 4:
$$
\arg \min  _ {\mathbf {z} ^ {*}} D (\pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*}), \pi_ {E} (\boldsymbol {\tau})) = \arg \max  _ {\mathbf {z} ^ {*}} \underbrace {\mathcal {I} (\mathbf {z} ^ {*} ; \boldsymbol {\tau} _ {E}) - \mathcal {I} (\mathbf {z} ^ {*} ; \boldsymbol {\tau} _ {\theta})} _ {\mathcal {I} _ {\mathrm {M I}}} - \underbrace {D _ {\mathrm {K L}} (\pi_ {\theta} (\boldsymbol {\tau}) , \pi_ {E} (\boldsymbol {\tau}))} _ {\mathcal {I} _ {D}},
$$
Below is our derivation:
$$
\begin{array}{l} \min  _ {\mathbf {z} ^ {*}} D \left(\pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*}), \pi_ {E} (\boldsymbol {\tau})\right) \\ = \min _ {\mathbf {z} ^ {*}} \mathbb {E} _ {p (\mathbf {z} ^ {*})} \left[ D _ {\mathrm {K L}} \left(\pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*}) \| \pi_ {E} (\boldsymbol {\tau})\right) + D _ {\mathrm {K L}} \left(\pi_ {E} (\boldsymbol {\tau}) \| \pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*})\right) \right] \\ = \min _ {\mathbf {z} ^ {*}} \mathbb {E} _ {p (\mathbf {z} ^ {*}) \pi_ {\theta} (\pmb {\tau} | \mathbf {z} ^ {*})} \left[ \log \pi_ {\theta} (\pmb {\tau} | \mathbf {z} ^ {*}) - \log \pi_ {E} (\pmb {\tau}) \right] \\ + \mathbb {E} _ {p (\mathbf {z} ^ {*}) \pi_ {E} (\boldsymbol {\tau})} \left[ \log \pi_ {E} (\boldsymbol {\tau}) - \log \pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*}) \right] \\ = \min  _ {\mathbf {z} ^ {*}} \mathbb {E} _ {p (\mathbf {z} ^ {*}) \pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*})} \left[ \log \frac {p (\mathbf {z} ^ {*} | \boldsymbol {\tau}) \pi_ {\theta} (\boldsymbol {\tau})}{p (\mathbf {z} ^ {*})} - \log \pi_ {E} (\boldsymbol {\tau}) \right] \\ + \mathbb {E} _ {p (\mathbf {z} ^ {*}) \pi_ {E} (\boldsymbol {\tau})} \left[ \log \pi_ {E} (\boldsymbol {\tau}) - \log \frac {p (\mathbf {z} ^ {*} | \boldsymbol {\tau}) \pi_ {\theta} (\boldsymbol {\tau})}{p (\mathbf {z} ^ {*})} \right] \\ = \min _ {\mathbf {z} ^ {*}} \mathbb {E} _ {p (\mathbf {z} ^ {*}) \pi_ {\theta} (\boldsymbol {\tau} | \mathbf {z} ^ {*})} \left[ \log \frac {p (\mathbf {z} ^ {*} | \boldsymbol {\tau})}{p (\mathbf {z} ^ {*})} + \log \frac {\pi_ {\theta} (\boldsymbol {\tau})}{\pi_ {E} (\boldsymbol {\tau})} \right] - \mathbb {E} _ {p (\mathbf {z} ^ {*}) \pi_ {E} (\boldsymbol {\tau})} \left[ \log \frac {p (\mathbf {z} ^ {*} | \boldsymbol {\tau})}{p (\mathbf {z} ^ {*})} + \log \frac {\pi_ {\theta} (\boldsymbol {\tau})}{\pi_ {E} (\boldsymbol {\tau})} \right] \\ = \max  _ {\mathbf {z} ^ {*}} \mathcal {I} (\mathbf {z} ^ {*}; \boldsymbol {\tau} _ {E}) - \mathcal {I} (\mathbf {z} ^ {*}; \boldsymbol {\tau} _ {\theta}) - D (\pi_ {\theta} (\boldsymbol {\tau}), \pi_ {E} (\boldsymbol {\tau})), \\ \end{array}
$$
where  $\pmb{\tau}_{E}:=\pmb{\tau}\sim\pi_{E}(\pmb{\tau}),\pmb{\tau}_{\theta}:=\pmb{\tau}\sim p(\mathbf{z}^{*})\pi_{\theta}(\pmb{\tau}|\mathbf{z}^{*})$
8 More Comparisons and Ablation Studies
8.1 Offline Comparison on D4RL Expert Domain Dataset
In Table 6, we provide the normalized return of our method and baseline methods on the reward-free D4RL [22] expert dataset. Consistently, we can observe that CEIL achieves a significant improvement over the baseline methods in both  $S$ -off-  $LfD$  and  $S$ -off-  $LfO$  settings. Compared to the state-of-the-art offline IL baselines, CEIL also shows competitive results on the challenging cross-domain offline IL settings ( $C$ -off-  $LfD$  and  $C$ -off-  $LfO$ ).
8.2 Generalizability on Cross-domain Offline IL Settings
In the standard cross-domain IL setting, the goal is to extract expert-relevant information from the mismatched expert demonstrations/observations (expert domain) and to mimic such expert behaviors in the training environment (training domain). Thus, we validate the performance of the learned policy in the training environment (i.e., the environment where the offline data was collected). Here, we also study the generalizability of the learned policy by evaluating the learned policy in the expert environment (i.e., the environment where the mismatched expert data was collected). We provide the normalized scores (evaluated in the expert domain) in Table 7. We can find that across a range of cross-domain offline IL tasks, CEIL consistently demonstrates better (zero-shot) generalizability compared to baselines.
8.3 Ablating the Cross-domain Regularization
We now conduct ablation studies to evaluate the importance of cross-domain regularization in Equation 9 (in the main paper). In Figure 5, we provide the performance improvement when we
Table 6: Normalized scores (averaged over 30 trails for each task) on D4RL expert dataset. Scores within two points of the maximum score are highlighted. hop: Hopper-v2. hal: HalfCheetah-v2. wal: Walker2d-v2. ant: Ant-v2.
hop expert	hal expert	wal expert	ant expert	sum
S-off-LfD	ORIL (TD3+BC)	97.5	91.8	14.5	76.8	280.6
SQIL (TD3+BC)	25.5	14.4	8.0	44.3	92.1
IQ-Learn	37.3	9.9	46.6	85.9	179.7
ValueDICE	65.6	2.9	28.2	90.5	187.1
DemoDICE	107.3	87.1	104.8	114.2	413.3
SMODICE	111.0	93.5	108.2	122.0	434.7
CEIL	106.0	96.0	115.6	117.8	435.4
S-off-LfO	ORIL (TD3+BC)	64.2	92.1	12.2	44.3	212.8
SMODICE	111.3	93.7	108.0	122.0	435.0
CEIL	103.3	96.8	110.0	126.4	436.5
C-off-LfD	ORIL (TD3+BC)	24.4	78.3	29.3	32.1	164.1
SQIL (TD3+BC)	12.2	19.9	8.8	21.2	62.0
IQ-Learn	25.9	31.2	31.7	55.8	144.6
ValueDICE	18.6	9.8	8.3	22.3	59.0
DemoDICE	111.5	88.7	107.9	122.5	430.6
SMODICE	111.1	93.8	108.2	120.9	434.0
CEIL	105.8	97.1	108.6	112.2	423.7
C-off-LfO	ORIL (TD3+BC)	22.5	76.6	11.2	28.2	138.6
SMODICE	111.2	93.7	108.1	117.7	430.7
CEIL	113.0	90.1	108.7	125.2	437.0
Table 7: Normalized scores (evaluated on the expert dataset over 30 trails for each task) on 2 cross-domain offline IL settings:  $C$ -off-LfD and  $C$ -off-LfO. Scores within two points of the maximum score are highlighted. m: medium. mr: medium-replay. me: medium-expert. e: expert.
Hopper-v2	HalfCheetah-v2	sum
m	mr	me	e	m	mr	me	e
C-off-LfD	ORIL (TD3+BC)	74.7	16.7	45.0	21.4	2.2	0.8	-0.3	-2.2	158.3
SQIL (TD3+BC)	33.6	21.6	14.5	14.5	18.2	7.5	20.9	20.9	151.8
IQ-Learn	11.8	9.7	17.1	17.1	7.7	7.8	9.5	9.5	90.2
ValueDICE	49.5	24.2	55.7	49.3	32.2	32.9	38.7	28.7	311.2
DemoDICE	83.2	31.5	81.6	28.5	0.9	-1.1	-1.7	-2.4	220.6
SMODICE	80.1	26.1	78.0	54.3	2.8	-1.0	1.0	-2.3	239.1
CEIL	87.4	74.3	81.2	82.4	44.0	30.4	25.0	17.1	441.9
C-off-LfO	ORIL (TD3+BC)	62.3	18.7	57.0	28.2	0.2	1.1	-0.3	-2.3	165.0
SMODICE	77.6	22.5	80.2	71.0	2.0	-0.9	0.8	-2.3	250.9
CEIL	56.4	58.6	56.7	65.2	5.5	36.5	5.0	5.0	288.7
Walker2d-v2	Ant-v2	sum
m	mr	me	e	m	mr	me	e
C-off-LfD	ORIL (TD3+BC)	22.0	24.5	23.9	33.1	16.0	18.6	2.5	0.4	141.0
SQIL (TD3+BC)	32.4	14.9	10.3	10.3	71.4	63.6	60.1	60.1	323.1
IQ-Learn	8.4	5.0	10.2	10.2	19.4	18.4	16.1	16.1	103.8
ValueDICE	31.7	21.9	22.9	27.7	70.5	68.5	69.3	68.5	380.9
DemoDICE	12.8	31.5	12.9	86.9	15.7	24.2	2.3	1.4	187.7
SMODICE	43.6	16.1	62.0	85.3	23.7	22.9	2.3	-5.9	249.9
CEIL	102.8	94.8	101.9	100.7	82.0	77.0	76.4	79.8	715.3
C-off-LfO	ORIL (TD3+BC)	22.4	15.2	17.8	12.6	13.6	20.7	5.5	-6.2	101.6
SMODICE	42.4	17.0	55.5	88.7	15.7	22.6	2.5	-6.3	238.1
CEIL	67.9	12.0	68.4	50.8	31.7	57.0	18.0	-1.9	304.0
Figure 5: Normalized performance improvement (left:  $C$ -off-LfD, right:  $C$ -off-LfO) when we ablate the cross-domain regularization (Equation 9 in the main paper) in cross-domain IL settings. We can observe the general trend (in 26 out of 32 tasks) that ablating the cross-domain regularization causes negative performance improvement. hop: Hopper-v2. hal: HalfCheetah-v2. wal: Walker2d-v2. ant: Ant-v2. m: medium. me: medium-expert. mr: medium-replay. e: expert.
Figure 6: Aggregate median, IQM, mean, and optimality gap over 16 offline IL tasks. Higher median, higher IQM, and higher mean and lower optimality gap are better. The shaded bar shows  $95\%$  stratified bootstrap confidence intervals. We can see that CEIL achieves consistently better performance across a wide range of offline IL settings.
ablate the cross-domain regularization in two cross-domain offline IL tasks ( $C$ -off- $LfD$  and  $C$ -off- $LfO$ ). We can find that in 26 out of 32 cross-domain tasks, ablating the regularization can cause performance to decrease (negative performance improvement), thus verifying the benefits of encouraging task-relevant embeddings.
Figure 7: Return curves in Walker2d-v2 (from left to right: S-on-LfD, C-on-LfD, S-on-LfO, and C-on-LfO), where the shaded area represents a  $95\%$  confidence interval over 30 trails. We can see that CEIL consistently achieves expert-level performance in LfD (S-on-LfD and C-on-LfD) tasks. Due to the lack of explicit exploration in online LfO settings, CEIL exhibits drastic performance degradation (in S-on-LfO and C-on-LfO) under the same environmental interaction steps.
8.4 Aggregate Results
According to Agarwal et al. [1], we report the aggregate statistics (for 16 offline IL tasks) in Figure 6. We can find that CEIL provides competitive performance consistently across a range of offline IL settings ( $S$ -off-LfD,  $S$ -off-LfO,  $C$ -off-LfD, and  $C$ -off-LfO) and outperforms prior offline baselines.
Table 8: Normalized scores (averaged over 30 trails for each task) when we vary the number of the expert demonstrations (#5, #10, #15, and #20). Scores within two points of the maximum score are highlighted
Offline IL settings	Hopper-v2	Halfcheetah-v2	Walker2d-v2	Ant-v2	sum
m	mr	me	m	mr	me	m	mr	me	m	mr	me
S-off-Lfd #5	ORIL (TD3+BC)	42.1	26.7	51.2	45.1	2.7	79.6	44.1	22.9	38.3	25.6	24.5	6.0	408.8
SQIL (TD3+BC)	45.2	27.4	5.9	14.5	15.7	11.8	12.2	7.2	13.6	20.6	23.6	-5.7	192.0
IQ-Learn	17.2	15.4	21.7	6.4	4.8	6.2	13.1	10.6	5.1	22.8	27.2	18.7	169.2
ValueDICE	59.8	80.1	72.6	2.0	0.9	1.2	2.8	0.0	7.4	27.3	32.7	30.2	316.9
DemoDICE	50.2	26.5	63.7	41.9	38.7	59.5	66.3	38.8	101.6	82.8	68.8	112.4	751.2
SMODICE	54.1	34.9	64.7	42.6	38.4	63.8	62.2	40.6	55.4	86.0	69.7	112.4	724.7
CEIL	94.5	45.1	80.8	45.1	43.3	33.9	103.1	81.1	99.4	99.8	101.4	85.0	912.5
S-off-Lfd #10	ORIL (TD3+BC)	42.0	21.6	53.4	45.0	2.1	82.1	44.1	27.4	80.4	47.3	24.0	44.9	514.1
SQIL (TD3+BC)	50.0	34.2	7.4	8.8	10.9	8.2	20.0	15.2	9.7	35.3	36.2	11.9	247.6
IQ-Learn	11.3	18.6	20.1	4.1	6.5	6.6	18.3	12.8	12.2	30.7	53.9	23.7	218.7
ValueDICE	56.0	64.1	54.2	-0.2	2.6	2.4	4.7	4.0	0.9	31.4	72.3	49.5	341.8
DemoDICE	53.6	25.8	64.9	42.1	36.9	60.6	64.7	36.1	100.2	87.4	67.1	114.3	753.5
SMODICE	55.6	30.3	66.6	42.6	38.0	66.0	64.5	44.6	53.8	86.9	69.5	113.4	731.8
CEIL	113.2	53.0	96.3	64.0	43.6	44.0	120.4	82.3	104.2	119.3	70.0	90.1	1000.4
S-off-Lfd #15	ORIL (TD3+BC)	38.9	22.3	46.8	44.7	1.9	83.8	37.9	4.2	69.9	59.4	22.3	12.4	444.6
SQIL (TD3+BC)	42.8	44.4	5.2	6.8	17.1	9.1	16.9	13.5	6.9	21.2	17.2	12.6	213.6
IQ-Learn	14.6	8.2	29.3	4.0	3.4	5.1	7.3	14.5	11.4	54.2	15.2	61.6	228.6
ValueDICE	66.3	58.3	53.6	2.3	2.3	1.2	5.2	-0.1	17.0	45.2	72.0	74.3	397.8
DemoDICE	52.2	29.6	67.3	41.9	37.6	58.1	66.4	42.9	103.5	86.6	68.3	114.3	768.7
SMODICE	55.9	25.7	72.7	42.5	37.6	66.4	67.0	43.2	55.1	86.7	69.7	118.2	740.6
CEIL	116.4	56.7	103.7	80.4	43.0	43.8	120.3	84.8	103.8	126.8	87.0	90.6	1057.3
S-off-Lfd #20	ORIL (TD3+BC)	50.9	22.1	72.7	44.7	30.2	87.5	47.1	26.7	102.6	46.5	31.4	61.9	624.3
SQIL (TD3+BC)	32.6	60.6	25.5	13.2	25.3	14.4	25.6	15.6	8.0	63.6	58.4	44.3	387.1
IQ-Learn	21.3	19.9	24.9	5.0	7.5	7.5	22.3	19.6	18.5	38.4	24.3	55.3	264.5
ValueDICE	73.8	83.6	50.8	1.9	2.4	3.2	24.6	26.4	44.1	79.1	82.4	75.2	547.5
DemoDICE	54.8	32.7	65.4	42.8	37.0	55.6	68.1	39.7	95.0	85.6	69.0	108.8	754.6
SMODICE	56.1	28.7	68.0	42.7	37.7	66.9	66.2	40.7	58.2	87.4	69.9	113.4	735.9
CEIL (ours)	110.4	103.0	106.8	40.0	30.3	63.9	118.6	110.8	117.0	126.3	122.0	114.3	1163.5
8.5 Varying the Number of Expert Trajectories
As a complement to the experimental results in the main paper, we continue to compare the performance of CEIL and baselines on more tasks when we vary the number of expert trajectories. Considering offline IL settings, we provide the results in Table 8 for the number of expert trajectories of 5, 10, 15, and 20 respectively. We can find that when varying the number of expert behaviors, CEIL can still obtain higher scores compared to baselines, which is consistent with the findings in Figure 3 in the main paper.
8.6 Limitation (Failure Modes in Online LfO Setting)
Meanwhile, we find that in the online LfO settings, CEIL's performance deteriorates severely on a few tasks, as shown in Figure 7 (Walker2d). In LfD (either on single-domain or on cross-domain IL) settings, CEIL can consistently achieve expert-level performance, but when migrating to LfO settings, CEIL suffers collapsing performance under the same number of environmental interactions. We believe that this is due to the lack of expert actions in LfO settings, which causes the agent to stay in the collapsed state region and therefore deteriorates performance. Thus, we believe a rich direction for future research is to explore the online exploration ability.
9 Implementation Details
9.1 Imitation Learning Tasks
In our paper, we conduct experiments across a variety of IL problem domains: single/cross-domain IL, online/offline IL, and LfD/LfO IL settings. By arranging and combining these IL domains, we obtain 8 IL tasks in all:  $S$ -on-LfD,  $S$ -on-LfO,  $S$ -off-LfD,  $S$ -off-LfO,  $C$ -on-LfD,  $C$ -on-LfO,  $C$ -off-LfD, and  $C$ -off-LfO, where S/C denotes single/cross-domain IL, on/off denotes online/offline IL, and LfD/LfO denote learning from demonstrations/observations respectively.
$S$ -on-LfD. We have access to a limited number of expert demonstrations and an online interactive training environment. The goal of  $S$ -on-LfD is to learn an optimal policy that mimics the provided demonstrations in the training environment.
$S$ -on-LfO. We have access to a limited number of expert observations (state-only demonstrations) and an online interactive training environment. The goal of  $S$ -on-LfO is to learn an optimal policy that mimics the provided observations in the training environment.
$S$ -off-LfD. We have access to a limited number of expert demonstrations and a large amount of pre-collected offline (reward-free) data. The goal of  $S$ -off-LfD is to learn an optimal policy that mimics the provided demonstrations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert demonstrations and the environment that was used to collect the offline data are the same environment.
$S$ -off-LfO. We have access to a limited number of expert observations and a large amount of pre-collected offline (reward-free) data. The goal of  $S$ -off-LfO is to learn an optimal policy that mimics the provided observations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert observations and the environment that was used to collect the offline data are the same environment.
$C$ -on-LfD. We have access to a limited number of expert demonstrations and an online interactive training environment. The goal of  $C$ -on-LfD is to learn an optimal policy that mimics the provided demonstrations in the training environment. Note that here the environment that was used to collect the expert demonstrations and the online training environment are not the same environment.
$C$ -on-LfO. We have access to a limited number of expert observations (state-only demonstrations) and an online interactive training environment. The goal of  $C$ -on-LfO is to learn an optimal policy that mimics the provided observations in the training environment. Note that here the environment that was used to collect the expert observations and the online training environment are not the same environment.
$C$ -off-LfD. We have access to a limited number of expert demonstrations and a large amount of pre-collected offline (reward-free) data. The goal of  $C$ -off-LfD is to learn an optimal policy that mimics the provided demonstrations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert demonstrations and the environment that was used to collect the offline data are not the same environment.
$C$ -off-LfO. We have access to a limited number of expert observations and a large amount of pre-collected offline (reward-free) data. The goal of  $C$ -off-LfO is to learn an optimal policy that mimics the provided observations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert observations and the environment that was used to collect the offline data are not the same environment.
Figure 8: MuJoCo environments and our modified versions. From left to right: Ant-v2, HalfCheetah-v2, Hopper-v2, Walker2d-v2, our modified Ant-v2, our modified HalfCheetah-v2, our modified Hopper-v2, and our modified Walker2d-v2.
9.2 Online IL Environments, Offline IL Datasets, and One-shot tasks
Our experiments are conducted in four popular MuJoCo environments (Figure 8): Hopper-v2, HalfCheetah-v2, Walker2d-v2, and Ant-v2. For offline IL tasks, we take the standard (reward-free) D4RL dataset [22] (medium, medium-replay, medium-expert, and expert domains) as the offline dataset. For cross-domain (online/offline) IL tasks, we collect the expert behaviors (demonstrations or observations) on a modified MuJoCo environment. Specifically, we change the height of the agent's torso (as shown in Figure 8). We refer the reader to our code submission, which includes our modified MuJoCo assets. For one-shot IL tasks, we train the policy only in the single-domain IL settings ( $S$ -on-LfD,  $S$ -on-LfO,  $S$ -off-LfD, and  $S$ -off-LfO). Then we collect only one expert trajectory in the modified MuJoCo environment, and roll out the fine-tuned/inferred policy in the modified environment to test the one-shot performance.
Collecting expert behaviors. In our implementation, we use the publicly available rlkit<sup>7</sup> implementation of SAC to learn an expert policy and use the learned policy to collect expert behaviors (demonstrations in LfD or observations in LfO).
9.3 CEIL Implementation Details
Trajectory self-consistency loss. To learn the embedding function  $f_{\phi}$  and a corresponding contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\mathbf{z})$ , we minimize the following trajectory self-consistency loss:
$$
\pi_ {\theta}, f _ {\phi} = \min  _ {\pi_ {\theta}, f _ {\phi}} - \mathbb {E} _ {\boldsymbol {\tau} _ {1: T} \sim \mathcal {D} (\boldsymbol {\tau} _ {1: T})} \mathbb {E} _ {(\mathbf {s}, \mathbf {a}) \sim \boldsymbol {\tau} _ {1: T}} \left[ \log \pi_ {\theta} \big (\mathbf {a} | \mathbf {s}, f _ {\phi} (\boldsymbol {\tau} _ {1: T}) \big) \right],
$$
where  $\tau_{1:T}$  denotes a trajectory segment with window size of  $T$ . In the online setting, we sample trajectory  $\pmb{\tau}$  from the experience replay buffer  $\mathcal{D}(\pmb{\tau})$ ; in the offline setting, we sample trajectory  $\pmb{\tau}$  directly from the given offline data  $\mathcal{D}(\pmb{\tau})$ . Meanwhile, if we can access the expert actions (i.e., LfD settings), we also incorporate the expert demonstrations into the empirical expectation (i.e., storing the expert demonstrations into the online/offline experience  $\mathcal{D}(\pmb{\tau})$ ).
In our implementation, we use a 4-layer MLP (with ReLU activation) to encode the trajectory  $\pmb{\tau}_{1:T}$  and a 4-layer MLP (with ReLU activation) to predict the action respectively. To regularize the learning of the encoder function  $f_{\phi}$ , we additionally introduce a decoder network (4-layer MLP with ReLU activation)  $\pi_{\theta}^{\prime}(\mathbf{s}^{\prime}|\mathbf{s},f_{\phi}(\pmb{\tau}_{1:T}))$  to predict the next states:  $\min_{\pi_{\theta}^{\prime},f_{\phi}} - \mathbb{E}_{\pmb{\tau}_{1:T}\sim \mathcal{D}(\pmb{\tau}_{1:T})}\mathbb{E}_{(\mathbf{s},\mathbf{a},\mathbf{s}^{\prime})\sim \pmb{\tau}_{1:T}}[\log \pi_{\theta}^{\prime}(\mathbf{s}^{\prime}|\mathbf{s},f_{\phi}(\pmb{\tau}_{1:T}))]$ . Further, to circumvent issues of "posterior collapse" [67], we encourage learning quantized latent embeddings. In a similar spirit to VQ-VAE [67], we incorporate ideas from vector quantization (VQ) and introduce the following regularization:  $\min_{f_{\phi}}||\mathrm{sg}[z_e(\pmb{\tau}_{1:T})] - e||^2 +||z_e(\pmb{\tau}_{1:T}) - \mathrm{sg}[e]||^2$ , where  $e$  is a dictionary of vector quantization embeddings (we set the size of this embedding dictionary to be 4096),  $z_{e}(\pmb{\tau}_{1:T})$  is defined as the nearest dictionary embedding to  $f_{\phi}(\pmb{\tau}_{1:T})$ , and  $\mathrm{sg}[\cdot ]$  denotes the stop-gradient operator.
Out-level embedding inference. In Section 4.2 (main paper), we approximate  $\mathcal{I}_{\mathrm{MI}}$  with  $\mathcal{I}_{\mathrm{MI}(f_{\phi})} \triangleq \mathbb{E}_{p(\mathbf{z}^{*})\pi_E(\tau_E)\pi_\theta(\tau_\theta|\mathbf{z}^*)}\left[-\|\mathbf{z}^* - f_\phi(\tau_E)\|^2 + \|\mathbf{z}^* - f_\phi(\tau_\theta)\|^2\right]$ , where we replace the mutual information with  $-\|\mathbf{z}^* - f_\phi(\tau)\|^2$  by leveraging the learned embedding function  $f_{\phi}$ . Empirically, we find that we can ignore the second loss  $\|\mathbf{z}^* - f_\phi(\tau_\theta)\|^2$ , and directly conduct outer-level embedding inference with  $\max_{\mathbf{z}^*, f_{\phi}} \mathbb{E}_{p(\mathbf{z}^*)\pi_E(\tau_E)}\left[-\|\mathbf{z}^* - f_{\phi}(\tau_E)\|^2\right]$ . Meanwhile, this simplification makes the support constraints  $(\mathcal{R}(\mathbf{z}^*))$  in Equation 7 in the main paper) for the offline OOD issues naturally satisfied, since  $\max_{\mathbf{z}^*} \mathbb{E}_{p(\mathbf{z}^*)\pi_E(\tau_E)}\left[-\|\mathbf{z}^* - f_{\phi}(\tau_E)\|^2\right]$  and  $\min_{\mathbf{z}^*} \mathcal{R}(\mathbf{z}^*)$  are equivalent.
Cross-domain IL regularization. To encourage  $f_{\phi}$  to capture the task-relevant embeddings and ignore the domain-specific factors, we set the regularization  $\mathcal{R}(f_{\phi})$  in Equation 5 to be:
<sup>7</sup>https://github.com/rail-berkeley/rlkit.
$\mathcal{R}(f_{\phi}) = \mathcal{I}(f_{\phi}(\pmb{\tau}); \mathbf{n})$ , where we couple each trajectory  $\pmb{\tau}$  in  $\{\pmb{\tau}_{E}\} \cup \{\pmb{\tau}_{E^{\prime}}\}$  with a label  $\mathbf{n} \in \{\mathbf{0}, \mathbf{1}\}$ , indicating whether it is noised. In our implementation, we apply MINE [6] to estimate the mutual information and conduct encoder regularization. Specifically, we estimate  $\mathcal{I}(\mathbf{z}; \mathbf{n})$  with  $\hat{\mathcal{I}}(\mathbf{z}; \mathbf{n}) := \sup_{\delta} \mathbb{E}_{p(\mathbf{z}, \mathbf{n})}[f_{\delta}(\mathbf{z}, \mathbf{n})] - \log \mathbb{E}_{p(\mathbf{z})p(\mathbf{n})}[\exp(f_{\delta}(\mathbf{z}, \mathbf{n}))]$  and regularize the encoder  $f_{\phi}$  with  $\max_{f_{\phi}} \hat{\mathcal{I}}(f_{\phi}(\pmb{\tau}); \mathbf{n})$ , where we model  $f_{\delta}$  with a 4-layer MLP (using ReLU activations).
Hyper-parameters. In Table 9, we list the hyper-parameters used in the experiments. For the size of the embedding dictionary, we selected it from a range of [512, 1024, 2048, 4096]. We found 4096 to almost uniformly attain good performance across IL tasks, thus selecting it as the default. For the size of the embedding dimension, we tried four values [4, 8, 16, 32] and selected 16 as the default. For the trajectory window size, we tried five values [2, 4, 8, 16, 32] but we did not observe a significant difference in performance across these values. Thus we selected 2 as the default value. For the learning rate scheduler, we tried the default Pytorch scheduler and CosineAnnealingWarmRestarts, and found CosineAnnealingWarmRestarts enables better results (thus we selected it). For other hyperparameters, they are consistent with the default values of most RL implementations, e.g. learning rate 3e-4 and the MLP policy.