SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathcal {J} _ {\mathrm {M I}} = \mathbb {E} _ {\pi_ {E} (\mathbf {z} ^ {*}, \boldsymbol {\tau} _ {E})} \log p (\mathbf {z} ^ {*} | \boldsymbol {\tau} _ {E}) - \mathbb {E} _ {\pi_ {\theta} (\mathbf {z} ^ {*}, \boldsymbol {\tau} _ {\theta})} \log p (\mathbf {z} ^ {*} | \boldsymbol {\tau} _ {\theta}) \\ \approx \mathbb {E} _ {p (\mathbf {z} ^ {*}) \pi_ {E} (\boldsymbol {\tau} _ {E}) \pi_ {\theta} (\boldsymbol {\tau} _ {\theta} | \mathbf {z} ^ {*})} \left[ - \| \mathbf {z} ^ {*} - f _ {\phi} (\boldsymbol {\tau} _ {E}) \| ^ {2} + \| \mathbf {z} ^ {*} - f _ {\phi} (\boldsymbol {\tau} _ {\theta}) \| ^ {2} \right] \triangleq \mathcal {J} _ {\mathrm {M I} (f _ {\phi})}, \\ \end{array}
$$
where we approximate the logarithmic predictive power of  $\mathbf{z}^*$  on  $\pmb{\tau}$  with  $-\| \mathbf{z}^{*} - f_{\phi}(\pmb {\tau})\|^{2}$ , by taking advantage of the learned embedding function  $f_{\phi}$  in Equation 5.
$^3 D_{\mathrm{KL}}(p\| q)\coloneqq \mathbb{E}_{p(\mathbf{x})}\left[\log \frac{p(\mathbf{x})}{q(\mathbf{x})}\right]$  denotes the (forward) KL divergences. It is well known that reverse KL ensures that the learned distribution is mode-seeking and forward KL exhibits a mode-covering behavior [37]. For analysis purposes, here we define  $D(\cdot ,\cdot)$  as the sum of reverse KL and forward KL, and set the weights of both reverse KL and forward KL to 1.
Algorithm 1 Training CEIL: Online or Offline IL Setting
Require: Expert demonstrations  $\{\pi_E(\tau)\}$ , empty buffer  $\mathcal{D}$  for online IL or reward-free offline data  $\mathcal{D}$  for offline IL, training iteration  $K$ , and batch size  $N$ .
1: Initialize contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\cdot)$ , embedding function  $f_{\phi}(\mathbf{z}|\pmb{\tau})$ , and latent variable  $\mathbf{z}^*$ .
2: for  $k = 1, \dots, K$  do
3: (Online only) Run policy  $\pi_{\theta}(\mathbf{a}|s,\mathbf{z}^{*})$  in environment and store experience into buffer  $\mathcal{D}$ .
4: Sample a batch of data  $\{\tau\}_{1}^{n}$  from  $\mathcal{D}$  for online IL or  $\mathcal{D}$  for offline IL.
5: Learn  $\pi_{\theta}$  and  $f_{\phi}$  over sampled  $\{\tau\}_{1}^{n}$  using the trajectory self-consistency loss.
6: Update  $\mathbf{z}^*$  and  $f_{\phi}$  over sampled  $\{\tau\}_{1}^{n}$  by maximizing  $\mathcal{J}_{\mathrm{MI}(f_{\phi})} - \alpha \mathcal{J}_D$
7: (Offline only) Update  $\mathbf{z}^*$  by minimizing  $\mathcal{R}(\mathbf{z}^*)$ . # eliminating the offline OOD issues.
8: end for
Return: the learned contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\cdot)$  and the optimal latent variable  $\mathbf{z}^*$ .
By maximizing  $\mathcal{J}_{\mathrm{MI}(f_{\phi})}$ , the learned optimal  $\mathbf{z}^*$  will be induced to converge towards the embeddings of expert data and avoid trivial solutions (as shown in Figure 1). Intuitively,  $\mathcal{J}_{\mathrm{MI}(f_{\phi})}$  can also be thought of as an instantiation of contrastive loss, which manifests two facets we consider significant in IL:
1) the "anchor" variable<sup>4</sup>  $\mathbf{z}^*$  is unknown and must be estimated, and 2) it is necessary to ensure that the estimated  $\mathbf{z}^*$  lies in the support set of training distribution, as specified by the support constraints in Equation 5.
In summary, by comparing  $\mathcal{I}_{\mathrm{MI}(f_{\phi})}$  and  $\mathcal{I}_D$  we can observe that  $\mathcal{I}_{\mathrm{MI}(f_{\phi})}$  actually encourages expert matching in the embedding space, while  $\mathcal{I}_D$  encourages expert matching in the original trajectory space. In the next section, we will see that such an embedding-level expert matching objective naturally lends itself to cross-domain IL settings.
Figure 1: During learning, the distance between  $\mathbf{z}^*$  and  $f_{\phi}(\pmb{\tau}_{E})$  decreases rapidly (green lines). Meanwhile, as policy  $\pi_{\theta}(\cdot|\cdot,\mathbf{z}^{*})$  gets better (blue lines),  $f_{\phi}(\pmb{\tau}_{\theta})$  gradually approaches  $\mathbf{z}^{*}$  (red lines).
4.3 Practical Implementation
In this section, we describe how we can convert the bi-level IL problem above (Equations 4 and 5) into a feasible online/offline IL objective and discuss some practical implementation details in LfO, offline IL, cross-domain IL, and one-shot IL settings (see more details<sup>5</sup> in Appendix 9.3).
As shown in Algorithm 1 (best viewed in colors), CEIL alternates between solving the bi-level problem with respect to the support constraint (Line 3 for online IL or Line 7 for offline IL), the trajectory self-consistency loss (Line 5), and the optimal embedding inference (Line 6).
To satisfy the support constraint in Equation 5, for online IL (Line 3), we directly roll out the  $z^{*}$ -conditioned policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\mathbf{z}^{*})$  in the environment; for offline IL (Line 7), we minimize a simple regularization over  $\mathbf{z}^{*}$ , bearing a close resemblance to the one used in TD3+BC [23]:
$$
\mathcal {R} (\mathbf {z} ^ {*}) = \min  \left(\| \mathbf {z} ^ {*} - f _ {\bar {\phi}} (\boldsymbol {\tau} _ {E}) \| ^ {2}, \| \mathbf {z} ^ {*} - f _ {\bar {\phi}} (\boldsymbol {\tau} _ {\mathcal {D}}) \| ^ {2}\right), \quad \boldsymbol {\tau} _ {E} := \boldsymbol {\tau} \sim \pi_ {E} (\boldsymbol {\tau}), \quad \boldsymbol {\tau} _ {\mathcal {D}} := \boldsymbol {\tau} \sim \mathcal {D} (\boldsymbol {\tau}), \tag {7}
$$
where we apply a stop-gradient operation to  $f_{\bar{\phi}}$ . To ensure the optimal embedding inference  $(\max_{\mathbf{z}^*}\mathcal{J}_{\mathrm{MI}(f_\phi)} - \mathcal{J}_D)$  retaining the flexibility of seeking  $\mathbf{z}^*$  across different instances of  $f_{\phi}$ , we jointly update the optimal embedding  $\mathbf{z}^*$  and the embedding function  $f_{\phi}$  with
$$
\max  _ {\mathbf {z} ^ {*}, f _ {\phi}} \mathcal {J} _ {\mathrm {M I} \left(f _ {\phi}\right)} - \alpha \mathcal {J} _ {D}, \tag {8}
$$
where we use  $\alpha$  to control the weight on  $\mathcal{J}_D$ .
LfO. In the LfO setting, as expert actions are missing, we apply our expert matching objective only over the observations. Note that even though expert data contains no actions in LfO, we can still
4The triplet contrastive loss enforces the distance between the anchor and the positive to be smaller than that between the anchor and the negative. Thus, we can view  $\mathbf{z}^*$  in  $\mathcal{I}_{\mathrm{MI}(f_{\phi})}$  as an instance of the anchor.
Our code will be released at https://github.com/wechtto/GeneralizedCEIL.
In other words, the offline support constraint in Equation 5 is achieved through minimizing  $\mathcal{R}(\mathbf{z}^*)$ .
leverage a large number of suboptimal actions presented in online/offline  $\mathcal{D}(\tau)$ . Thus, we can learn the contextual policy  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\mathbf{z})$  using the buffer data in online IL or the offline data in offline IL, much owing to the fact that we do not directly use the plain expert matching objective to update  $\pi_{\theta}$ .
Cross-domain IL. Cross-domain IL considers the case in which the expert's and learning agent's MDPs are different. Due to the domain shift, the plain idea of  $\min \mathcal{J}_D$  may not be a sufficient proxy for the expert matching objective, as there may never exist a trajectory (in the learning MDP) that matches the given expert data. Thus, we can set (the weight of  $\mathcal{J}_D$ )  $\alpha$  to 0.
Further, to make embedding function  $f_{\phi}$  useful for guiding the expert matching in latent space (i.e.,  $\max \mathcal{I}_{\mathrm{MI}(f_{\phi})}$ ), we encourage  $f_{\phi}$  to capture the task-relevant embeddings and ignore the domain-specific factors. To do so, we generate a set of pseudo-random transitions  $\{\tau_{E'}\}$  by independently sampling trajectories from expert data  $\{\pi_E(\tau_E)\}$  and adding random noise over these sampled trajectories, i.e.,  $\tau_{E'} = \tau_E + \text{noise}$ . Then, we couple each trajectory  $\tau$  in  $\{\tau_E\} \cup \{\tau_{E'}\}$  with a label  $\mathbf{n} \in \{\mathbf{0}, \mathbf{1}\}$ , indicating whether it is noised, and then generate a new set of  $\{(\tau, \mathbf{n})\}$ , where  $\tau \in \{\tau_E\} \cup \{\tau_{E'}\}$  and  $\mathbf{n} \in \{\mathbf{0}, \mathbf{1}\}$ . Thus, we can set the regularization  $\mathcal{R}(f_{\phi})$  in Equation 5 to be:
$$
\mathcal {R} \left(f _ {\phi}\right) = \mathcal {I} \left(f _ {\phi} (\boldsymbol {\tau}); \mathbf {n}\right). \tag {9}
$$
Intuitively, maximizing  $\mathcal{R}(f_{\phi})$  encourages embeddings to be domain-agnostic and task-relevant:  $f_{\phi}(\pmb{\tau}_{E})$  has high predictive power over expert data  $(\mathbf{n} = 0)$  and low that over noised data  $(\mathbf{n} = \mathbf{1})$ .
One-shot IL. Benefiting from the separate design of the contextual policy learning and the optimal embedding inference, CEIL also enjoys another advantage — one-shot generalization to new IL tasks. For new IL tasks, given the corresponding expert data  $\pmb{\tau}_{\mathrm{new}}$ , we can use the learned embedding function  $f_{\phi}$  to generate a corresponding latent embedding  $\mathbf{z}_{\mathrm{new}}$ . When conditioning on such an embedding, we can directly roll out  $\pi_{\theta}(\mathbf{a}|\mathbf{s},\mathbf{z}_{\mathrm{new}})$  to recover the one-shot expert behavior.
In this section, 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. Moreover, we also verify the scalability of CEIL on the challenging one-shot IL setting.
Our experiments are conducted in four popular MuJoCo environments: Hopper-v2 (Hop.), HalfCheetah-v2 (Hal.), Walker2d-v2 (Wal.), and Ant-v2 (Ant.). In the single-domain IL setting, we train a SAC policy in each environment and use the learned expert policy to collect expert trajectories (demonstrations/observations). To investigate the cross-domain IL setting, we assume the two domains (learning MDP and the expert-data collecting MDP) have the same state space and action space, while they have different transition dynamics. To achieve this, we modify the torso length of the MuJoCo agents (see details in Appendix 9.2). Then, for each modified agent, we train a separate expert policy and collect expert trajectories. For the offline IL setting, we directly take the reward-free D4RL [22] as the offline dataset, replacing the online rollout experience in the online IL setting.
5.1 Evaluation Results
To demonstrate the versatility of the CEIL idea, we collect 20 expert trajectories (demonstrations in LfD or state-only observations in LfO) for each environment and compare CEIL to GAIL [30], AIRL [21], SQIL [61], IQ-Learn [28], ValueDICE [41], GALfO [66], ORIL [78], DemoDICE [39], and SMODICE [56] (see their implementation details in Appendix 9.4). Note that these baseline methods cannot be applied to all the IL task settings ( $S/C$ -on/off-LfD/LfO), thus we only provide experimental comparisons with compatible baselines in each IL setting.
Online IL. In Figure 2, we provide the return (cumulative rewards) curves of our method and baselines on 4 online IL settings:  $S$ -on-LfD (top-left),  $S$ -on-LfO (top-right),  $C$ -on-LfD (bottom-left), and  $C$ -on-LfO (bottom-right) settings. As can be seen, CEIL quickly achieves expert-level performance in  $S$ -on-LfD. When extended to  $S$ -on-LfO, CEIL also yields better sample efficiency compared to baselines. Further, considering the complex cross-domain setting, we can see those baselines SQIL
(a) Single-domain online LfD (S-on-LfD)
(c) Cross-domain online LfD (C-on-LfD)
Figure 2: Return curves on 4 online IL settings: (a)  $S$ -on-LfD, (b)  $S$ -on-LfO, (c)  $C$ -on-LfD, and (d)  $C$ -on-LfO, where the shaded area represents a  $95\%$  confidence interval over 30 trails. Note that baselines cannot be applied to all the IL task settings, thus we only provide comparisons with compatible baselines (two separate legends).
(b) Single-domain online LfO (S-on-LfO)
(d) Cross-domain online LfO (C-on-LfO)
Table 2: Normalized scores (averaged over 30 trails for each task) on 4 offline IL settings:  $S$ -off-LfD,  $S$ -off-LfO,  $C$ -off-LfD, and  $C$ -off-LfO. 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-Lfd	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
S-Lfo	ORIL (TD3+BC)	43.4	25.7	73.0	44.9	2.4	81.8	58.9	16.8	78.2	33.7	29.6	67.1	555.4
SMODICE	54.5	26.4	73.7	42.7	37.9	66.2	60.6	38.5	70.9	85.7	68.3	116.3	741.7
CEIL (ours)	54.2	51.4	90.4	43.5	40.1	47.7	78.5	20.5	110.0	97.0	67.8	120.5	821.7
C-Lfd	ORIL (TD3+BC)	52.8	27.6	46.5	38.3	8.0	74.0	25.3	28.4	26.3	26.0	17.6	11.9	382.6
SQIL (TD3+BC)	34.4	19.1	11.4	19.2	25.1	19.9	15.8	16.5	8.8	21.8	23.2	21.2	236.2
IQ-Lean	37.3	35.4	25.9	27.4	27.1	31.2	27.7	22.2	31.7	63.7	63.3	55.8	448.8
ValueDICE	22.0	18.3	18.9	14.0	11.7	8.7	11.5	10.0	8.6	24.1	21.4	19.2	188.4
DemoDICE	52.9	15.2	77.2	42.8	38.9	53.8	58.4	26.4	77.8	87.8	69.3	114.9	715.6
SMODICE	55.4	21.4	71.2	42.7	38.0	64.6	68.4	34.2	80.4	87.4	70.4	115.7	749.7
CEIL (ours)	58.4	39.8	81.6	42.6	38.3	46.6	76.5	21.1	81.1	91.6	88.0	115.3	780.9
C-Lfo	ORIL (TD3+BC)	55.5	18.2	55.5	40.6	2.9	73.0	26.9	19.4	22.7	11.2	21.3	10.8	358.0
SMODICE	53.7	18.3	64.2	42.6	38.0	63.0	68.9	37.5	60.7	87.5	75.1	115.0	724.4
CEIL (ours)	44.7	44.2	48.2	42.4	36.5	46.9	76.2	31.7	77.0	95.9	71.0	112.7	727.3
and IQ-Learn (in  $C$ -on- $LfD$  and  $C$ -on- $LfO$ ) suffer from the domain mismatch, leading to performance degradation at late stages of training, while CEIL can still achieve robust performance.
Offline IL. Next, we evaluate CEIL on the other 4 offline IL settings:  $S$ -off-LfD,  $S$ -off-LfO,  $C$ -off-LfD, and  $C$ -off-LfO. In Table 2, we provide the normalized return of our method and baseline methods on reward-free D4RL [22] medium (m), medium-replay (mr), and medium-expert (me) datasets. We can
Figure 3: Ablating (a, b) the number of expert demonstrations and (c, d) the trajectory window size.
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 baselines, CEIL also shows competitive results on the challenging cross-domain offline IL settings ( $C$ -off-LfD and  $C$ -off-LfO).
One-shot IL. Then, we explore CEIL on the one-shot IL tasks, where we expect CEIL can adapt its behavior to new IL tasks given only one trajectory for each task (mismatched MDP, see Appendix 9.2).
We first pre-train an embedding function and a contextual policy in the training domain (online/offline IL), then infer a new contextual variable and evaluate it on the new task. To facilitate comparison to baselines, we similarly pre-train a policy network (using baselines) and run BC on top of the pre-trained policy by using the provided demonstration. Consequently, such a baseline+BC procedure cannot be applied to the (one-shot) LfO tasks. The results in Table 3 show that baseline+BC struggles to transfer their expertise to new tasks. Benefiting from the hindsight framework, CEIL shows better one-shot transfer learning performance on 7 out of 8 one-shot LfD tasks and retains higher scalability and generality for both one-shot LfD and LfO IL tasks.
One-shot IL	Hop.	Hal.	Wal.	Ant.
Online	SQIL	16.8	1.1	3.5	4.2
IQ-Learn	4.6	0.2	1.7	7.5
CEIL (LfD)	29.9	2.5	31.7	20.5
CEIL (LfO)	17.8	3.2	5.6	29.7
Offine	ORIL	14.7	0.2	6.9	17.4
SQIL	7.4	0.8	4.6	12.5
IQ-Learn	18.8	1.2	4.0	19.3
DemoDICE	76.5	-0.5	-0.1	19.5
SMODICE	78.0	1.1	8.1	24.6
CEIL (LfD)	85.6	5.6	67.1	24.3
CEIL (LfO)	72.2	5.1	70.0	19.4
5.2 Analysis of CEIL
Hybrid IL settings. In real-world, many IL tasks do not correspond to one specific IL setting, and instead consist of a hybrid of several IL settings, each of which passes a portion of task-relevant information to the IL agent. For example, we can provide the agent with both demonstrations and state-only observations and, in some cases, cross-domain demonstrations (S-LfD+S-LfO+C-LfD).
To examine the versatility of CEIL, we collect a separate expert trajectory for each of the four offline IL settings, and study CEIL's performance under hybrid IL settings. As shown in Table 4, we can see that by adding new expert behaviors on top of LfD, even when carrying relatively less supervision (e.g., actions are absent in LfO), CEIL can still improve the performance.
Table 3: Normalized results on one-shot IL, where CEIL shows prominent transferability.
Hybrid offline IL settings	Hop.	Hal.	Wal.	Ant.
S-LfD	29.4	69.9	42.8	84.9
S-LfD + S-LfO	30.4	68.6	42.3	91.6
S-LfD + S-LfO + C-LfD	30.7	71.7	42.9	89.2
S-LfD + S-LfO + C-LfD + C-LfO	58.6	79.6	43.7	98.0
Table 4: The normalized results of CEIL, showing that CEIL can consistently digest useful (task-relevant) information and boost its performance, even under a hybrid of offline IL settings.
Varying the number of demonstrations. In Figure 3 (a, b), we study the effect of the number of expert demonstrations on CEIL's performance. Empirically, we reduce the number of training demonstrations from 20 to 1, and report the normalized returns at 1M training steps. We can observe that across both online and offline (D4RL  $*$ -medium) IL settings, CEIL shows more robust performance with respect to different numbers of demonstrations compared to baseline methods.
Varying the window size of trajectory. Next we assess the effect of the trajectory window size (i.e., the length of trajectory  $\pmb{\tau}$  used for the embedding function  $f_{\phi}$  in Equation 3). In Figure 3 (b, c), we ablate the number of the window size in 4 LfD IL instantiations. We can see that across a range of window sizes, CEIL remains stable and achieves expert-level performance.
(a) Single-domain online LfD (S-on-LfD)
Figure 4: Ablation studies on the optimization of  $f_{\phi}$  (ablating  $f_{\phi}$ ) and the objective of  $\mathcal{I}_{\mathrm{MI}}$  (ablating  $\mathcal{I}_{\mathrm{MI}}$ ), where the shaded area represents  $95\%$  CIs over 5 trails. See ablation results for offline IL tasks in Table 5.
(b) Cross-domain online LfD (C-on-LfD)
Table 5: Ablation studies on the optimization of  $f_{\phi}$  (ablating  $f_{\phi}$ ) and the objective of  $\mathcal{I}_{\mathrm{MI}}$  (ablating  $\mathcal{I}_{\mathrm{MI}}$ ), where scores (averaged over 5 trails for each task) 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-LfD	CEIL (abating fφ)	97.9	92.5	99.3	41.3	30.3	66.7	103.6	88.1	114.4	97.6	98.4	100.7	1030.8
CEIL (abating JMI)	83.2	89.0	98.7	27.1	28.3	53.5	107.4	68.0	75.6	116.9	97.8	105.9	951.4
CEIL	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
S-LfO	CEIL (abating fφ)	51.5	41.1	83.3	43.8	40.1	63.7	76.3	20.3	103.0	78.0	52.5	105.5	759.2
CEIL (abating JMI)	54.3	44.9	84.7	42.2	39.9	51.6	77.4	22.7	94.0	92.1	67.9	118.4	792.0
CEIL	54.2	51.4	90.4	43.5	40.1	47.7	78.5	20.5	110.0	97.0	67.8	120.5	821.7
Ablation studies on the optimization of  $f_{\phi}$  and the objective of  $\mathcal{I}_{\mathrm{MI}}$ . In Figure 4 and Table 5, we carried out ablation experiments on the loss of  $f_{\phi}$  and  $\mathcal{I}_{\mathrm{MI}}$  in both online IL and offline IL settings. We can see that ablating the  $f_{\phi}$  loss (optimizing with Equation 5) does degrade the performance in both online and offline IL tasks, demonstrating the effectiveness of optimizing with Equation 8. Intuitively, Equation 8 encourages the embedding function to be task-relevant, and thus we use the expert matching loss to update  $f_{\phi}$ . We can also see that ablating  $\mathcal{I}_{\mathrm{MI}}$  does lead to degraded performance, further verifying the effectiveness of our expert matching objective in the latent space.
In this paper, we present CEIL, a novel and general Imitation Learning framework applicable to a wide range of IL settings, including  $C/S$ -on/off- $LfD/LfO$  and few-shot IL settings. This is achieved by explicitly decoupling the imitation policy into 1) a contextual policy, learned with the self-supervised hindsight information matching objective, and 2) a latent variable, inferred by performing the IL expert matching objective. Compared to prior baselines, our results show that CEIL is more sample-efficient in most of the online IL tasks and achieves better or competitive performances in offline tasks.
Limitations and future work. Our primary aim behind this work is to develop a simple and scalable IL method. We believe that CEIL makes an important step in that direction. Admittedly, we also find some limitations of CEIL: 1) Offline results generally outperform online results, especially in the LfO setting. The main reason is that CEIL lacks explicit exploration bounds, thus future work could explore the exploration ability of online CEIL. 2) The trajectory self-consistency cannot be applied to cross-embodiment agents once the two embodiments/domains have different state spaces or action spaces. Considering such a cross-embodiment setting, a typical approach is to serialize state/action from different modalities into a flat sequence of tokens. We also remark that CEIL is compatible with such a tokenization approach, and thus suitable for IL tasks with different action/state spaces. Thus, we encourage the future exploration of generalized IL methods across different embodiments.
We sincerely thank the anonymous reviewers for their insightful suggestions. This work was supported by the National Science and Technology Innovation 2030 - Major Project (Grant No. 2022ZD0208800), and NSFC General Program (Grant No. 62176215).
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-29320, 2021.
[2] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657, 2022.
[3] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. Advances in neural information processing systems, 30, 2017.
[75] Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off-policy imitation learning from observations. Advances in Neural Information Processing Systems, 33:12402-12413, 2020.
[76] Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo. Behavior proximal policy optimization. arXiv preprint arXiv:2302.11312, 2023.
[77] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pages 1433-1438. Chicago, IL, USA, 2008.
[78] 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.
[79] 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.
7 Additional Derivation