CEIL: Generalized Contextual Imitation Learning
Authors: Jinxin Liu
Li He
Yachen Kang
Zifeng Zhuang
Donglin Wang
Huazhe Xu
Venue: NeurIPS 2023 poster (2023)
Keywords: imitation learning
reinforcement learning
offline imitation learning

Abstract: In this paper, we present ContExtual Imitation Learning (CEIL), a general and broadly applicable algorithm for imitation learning (IL). Inspired by the formulation of hindsight information matching, we derive CEIL by explicitly learning a hindsight embedding function together with a contextual policy using the hindsight embeddings. To achieve the expert matching objective for IL, we advocate for optimizing a contextual variable such that it biases the contextual policy towards mimicking expert behaviors. Beyond the typical learning from demonstrations (LfD) setting, CEIL is a generalist that can be effectively applied to multiple settings including: 1) learning from observations (LfO), 2) offline IL, 3) cross-domain IL (mismatched experts), and 4) one-shot IL settings. Empirically, we evaluate CEIL on the popular MuJoCo tasks (online) and the D4RL dataset (offline). Compared to prior state-of-the-art baselines, we show that CEIL is more sample-efficient in most online IL tasks and achieves better or competitive performances in offline tasks.

CEIL: Generalized Contextual Imitation Learning
Jinxin Liu
Li He!
Yachen Kang 1.2
Zifeng Zhuang 1.2
Donglin Wang
Huazhe Xu
WWestlkee University
ZZeeiang University Tsinghua University
Weettaaee Institute for Advanced Study *Shanghai Qi Zhi Institute "Shangaai Al Lab
Abstract
In this paper, we present ContExtual Imitation Learning (CEIL), a general and broadly applicable algorithm for imitation learning (IL.). Inspired by the formula- tion of hindsight information matching, we derive CEIL by explicitly learning a hindsight embedding function together with a contextual policy using the hindsight embeddings. To achieve the expert matching objective for IL, advocate for optimizing a contextual variable such that it biases the contextual policy towards mimicking expert behaviors. Beyond the typical learning from demonstrations (LiD) setting, CEIL a generalist that can be effectively applied to multiple settings including: 1) learning from observations (LfO), 2) offline IL. 3) cross cross-domain IL (mismatched experts), and 4) one-shot settings. Empirically, we evaluate CEIL the popular MuJoCo tasks (online) and the D4RL dataset (offline). Compared to prior state-of-tte of-the-art baselines, we show that CEIL more sample-efficient in most online IL. tasks and achieves better competitive performances in offline tasks.
Introduction
Imitation learning (IL) allows agents to learn from expert demonstrations. Initially developed with supervised learning paradigm [58, 63], IL can be extended and reformulated with general expert matching objective, which aims to generate policies that produce trajectories with low distributional distances to expert demonstrations [30]. This formulation allows IL be extended various new settings: 1) online IL where interactions with the environment are allowed, 2) learning from observations (LfO) where expert actions are absent, 3) offline IL where agents learn from limited expert data and fixed dataset of sub-optimal and reward-free experience, 4) cross-domain IL where the expert demonstrations come from another domain (i.e. environment) that has different transition dynamics, and 5) one-shot IL which expects to recover the expert behaviors when only one expert trajectory observed for a new IL task.
Modern IL algorithms introduce various designs or mathematical principles to cater to the expert matching objective in specific scenario. For example, the L.fO setting requires particular considera- tions regarding the absent expert actions, e.g., learning an inverse dynamics function [5, 65]. Besides, -uitt--distributionn issues offline IL require specialized modifications the learning objective, such as introducing additional policy/value regularization [32, 72]. However, such a methodology, designing an individual formulation for each setting, makes difficult scale up a specific II. algorithm to more complex tasks beyond its original IL setting, e.g., online IL methods often suffer severe performance degradation in offline IL settings. Furthermore, realistic IL tasks are often not subject to particular IL setting but consist of a mixture of them. For example, we may have access
Equul contributions. Corresponding author: Donglin Wang <wangdonglin @ westlake.edu.con>
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Table 1: A coarse summary of IL methods demonstrating 1) different expert data modalities they can handle (learning from demonstrations or observations), 2) disparate task settings they consider (learning from online environment interactions or pre-collected offline static dataset). 3) the specific cross-domain setting they assume (the transition dynamics between the learning environment and that of the expert behaviors are different), and 4) the unique one-shol merit they desire (the learned policy is capable of one-shot transfer to new imitation tasks). We highlight that our contextual imitation learning (CEIL) method can naturally be applied to all the above III settings.
Expert data
Task setting
Cross-domain
One-shot
LiD
L1O
Online
Offline
S-on-I,fD [9, 13, 21, 30, 38, 57, 771
S-on-LfO 17, 54, 65, 66, 75]
S-off-LfD [19, 32, 33, 39, 72,
X
X
x
S-off-LfO [78]
V
X
x
C-on-LID [18, 69,
X
V
X
x
C-on-L.tO [20, 25, 26, 48, 59, 60]
C-off-LffD [34]
C-off-LfO [56, 68]
S-on/off-LffO [28]
Online one-shot [14, 16, 40]
Offline one-shot [24, 71]
CEIL (ours)
to both demonstrations and observation-only data in offline robot tasks; however, it could require significant effort to adapt several specialized methods to leverage such mixed/hybrid data. Hence, a problem naturally arises: How can we accommodate various design requirements of different IL settings with a general and practically ready-to-deppe IL formulation?
Hindsight information matching, a task-relabeling paradigm in reinforcement learning (RL), views control tasks as analogous to a general sequence modeling problem, with the goal to produce sequence of actions that induces high returns [12]. Its generality and simplicity enable it to be extended to both online and offline settings [17, 42]. In its original RL context, an agent directly uses known extrinsic rewards to bias the hindsight information towards task-related behaviors. However, when we attempt retain its generality in IL tasks, how bias the hindsight towards expert behaviors remains a significant barrier as the extrinsic rewards are missing.
To design general IL formulation and tackle the above problems, propose ContExtual Imitation Learning (CEIL), which readily incorporates the hindsight information matching principle within a bi-level expert matching objective. In the inner-level optimization, we explicitly learn hindsight embedding function to deal with the challenges of unknown rewards. In the outer-level optimization, perform IL expert matching via inferring biasing), replacing the naive reward biasing in hindsight. Intuitively, we find that such a bi-level objective results
optimal embedding (i.e., hindsight embedding
spectrum of expert matching objectives from the embedding space to the trajectory space. To shed light on the applicability and generality of CEIL, we instantiate CEIL to various IL settings, including online/offline IL, LfD/LfO, cross-domain IL, and one-shot IL settings.
In summary, this paper makes the following contributions: 1) We propose a bi-level expert match- ing objective ContExtual Imitation Learning (CEIL), inheriting the spirit of hindsight information matching, which decouples the learning policy into a contextual policy and an optimal embedding. 2) CEIL. exhibits high generality and adaptability and can be instantiated over range of IL. tasks. 3) Empirically, we conduct extensive empirical analyses showing that CEIL is more sample-efficient in online IL and achieves better or competitive results in offline IL tasks.
2 Related Work
Recent advances in decision-making have led to rapid progress in IL settings (Table 1), from typical learning from demonstrations (LffD to learning from observations (LfO) [7, 9, 35, 54, 62, 66), from online II. to offline II. [11, 15, 33, 53, 73], and from single-domain II. to cross-domain IL [34, 48, 56, 68]. Targeting a specific IL setting, individual works have shown their impressive

---

ability to solve the exact IL setting, However, it is hard to retrain their performance in new unprepared IL settings. In light of this, it is tempting to consider how we can design a general and broadly applicable II. method. Indeed, a number of prior works have studied part of the above II. settings, such as offline L.FO [78], cross-domain L.f [48, 60], and cross-domain offline II. [56]. While such works demonstrate the feasibility of tackling multiple II. settings, they still rely on standard online/offline RL algorithmic advances to improve performance [25, 32, 44, 47, 50, 51, 55, 72, 76]. Our objective diverges from these works, we strive minimize the reliance on the RL pipeline by replacing it with a simple supervision objective, thus avoiding the dependence on the choice of RL algorithms.
Our approach to IL is most closely related to prior hindsight information-mater 49], both learning a contextual policy and using a contextual variable to guide policy improvement. However, these prior methods typically require additional mechanisms to work well, such as extrinsic rewards in online RI. [4, 42, 64] or a handcrafted target return in offline RI [12, 17]. Our method does not require explicit handling of these components. By explicitly learning an embedding space for both expert and suboptimal behaviors, we can bias the contextual policy with an inferred optimal embedding (contextual variable), thus avoiding the need for explicit reward biasing in prior works. Our method also differs from most prior offline crunsoomerr-baee RL/IL algorithms that explicitly model long sequence of transitions [10, 12, 31, 36, 43, 71J. We find that simple fully-connected networks can also elicit useful embeddings and guide expert behaviors when conditioned on a well- calibrated embedding. In the context of the recently proposed prompt-tuning paradigm in large language tasks or multi-modal tasks [27, 45, 74], our method can be interpreted as a combination of IL and prompting tuuiing.t with the main motivation that we tune the prompt (the optimal contextual variable) with an expert matching objective in II. settings.
methods [2, 8, 24,
3 Background
Before discussing our method, we briefly introduce the background for II., including learning from demonstrations (Lff),, learning from observations (LfO), online II., offline II., and cross cross-domain settings in Section 3.1, and introduce the hindsight information matching in Section 3.2.
3.1 Imitation Learning
Consider a control task formulated as discrete-time Markov decision process (MDP) M (S, A, T, r, 7, P.Po. where S is the state (observation) space, A is the action space, T AxxxS + IR is the transition dynamics function, r SAAA + I is the reward function, 7 is the discount factor, and Po is the distribution of initial states. The goal in reinforcement learning (RL) control task is to learn a policy (als) maximizing the expected sum of discounted rewards ET((T) Lt=0 7'n(s,, ag) where T := {su. an. ... ST-1.ar-1 denotes the trajectory and the generated trajectory distribution (T( (So) ma (ao (so) [ (a
In IL, the ground truth reward function (i.e., in M) is not observed. Instead, we have access set of demonstrations (or observations) (7|7 TT] that are collected by an unknown expert policy Ta)ss)) The goal of IL tasks is to recover policy that matches the corresponding expert policy. From the mathematical perspective, IL achieves the plain expert matching objective by minimizing the divergence of trajectory distributions between the learner and the expert:
min ((TO(T),TEE(
where D is distance measure. Meanwhile, we emphasize that the given expert data {r|T ~
[T) may not contain the corresponding expert actions. Thus, in this work, we consider two IL cases where the given expert data consists of a set of state-action demonstrations ((St. ar, 8 }}
(learning from demonstrations, LfD), as well as a set of state-only transitions [(S), $141)} (learning from observations, LfO). When it is clear from context, demonstrations in LfD and observations in Lf0 for simplicity.
abuse notation (T( to denote both
Besides, we can also divide IL settings into two orthogonal categories: online IL and offline IL. In online II, the learning policy To can interact with the environment and generate online trajectories To (T). In offline IL, the agent cannot interact with the environment but has access to an offline
Pn this paper, we use environment and MDP interchangeably, and state and observation interchangeably.
3

---

static dataset {T|T ~ TT(T)), collected by some unknown (sub-optimal) behavior policies Tg. By leveraging the offline data (T)) [Tx(T)} without any interactions with the environment, the goal of offline II is to learn policy recovering the expert behaviors (demonstrations or observations) generated by mg- Note that, in contrast to the typical offline RI. problem [46], the offline data (T)(T)) in offline II. does not contains any reward signal.
Cross-domain IL. Beyond the above two IL branches (online/offline and LID/LIO), we can also divide IL into: 1) single-domain IL and 2) cross-domain IL, where 1) the single-domain IL assumes that the expert behaviors are collected in the same MDP which the learning policy is to be learned, and 2) the cross domainn IL studies how to imitate expert behaviors when discrepancies exist between the expert and the learning MDPs (e.g., differing in their transition dynamics or morphologies).
3.2 Hindsight Information Matching
In typical goal-conditioned RL problems, hindsight experience replay (HER) [3] proposes to leverage the rich repository of the failed experiences by replacing the desired (true) goals of training trajectories with the achieved goals of the failed experiences:
Alg(Tei 9. Tq) + Alg(TTe; SHER (Tg) Tg);
where the learner Alg (g((=o:->) could be any RL methods, Ty ~ = (g) denotes the trajectory generated by goal-conditioned policy (s: and fHER denotes a pre-defined (hindsight information extraction) function, e.g., returning the last state in trajectory Tg.
HER can also be applied to the (single-goal) reward-driven online/offline RL tasks, setting the return (sum of the discounted rewards) of trajectory an implicit goal for the corresponding trajectory. Thus, we can reformulate the (single-goul reward-driven RI. task, learning policy na Ta(a, (S:) that maximize the return, as a multi-goal RL task, learning a return-conditioned policy my (ar(s))).-) that maximize the following log- -ogliikeihood:
max ED(T) log TO (als,, (aals, (T)), Wg
(2)
where fR( denotes the return of trajectory т. At test, we can then condition the contextual policy (a(a)s on a desired target return. In offline RL, the empirical distribution D(T) in Equation 2 can be naturally set as the offline data distribution; in online RL., D(T) can set as the replaylexperience buffer, and will be updated and biased towards trajectories that have high expected returns.
Intuitively, biasing the sampling distribution (D(T) towards higher returns) leads an implicit policy improvement operation. However, such an operator is non-trivial to obtain in the IL problem, where we do not have access pre-defined function fR(T) (7) bias the learning policy towards recovering the given expert data {Tg(T)} (demonstrations or observations).
4 Method
In this section, will formulate IL a bi-level optimization problem, which will allow to derive our method, contextual imitation learning (CEIL). Instead of attempting to train the learning policy To(a(s) with the plain expert matching objective (Equation 1), our approach introduces an additional contextual variable z for a contextual II. policy To (a)s...)). The main idea CEIL. is to learn contextual policy wo To (a|s, z) and optimal contextual variable z" such that the given expert data (demonstrations in LfD or observations in LfO) can be recovered by the learned z -conditioned policy To ( ). We begin by describing the overall framework of CEIL in Section 4.1, and make a connection between CEIL and the plain expert matching objective in Section 4.2, which leads to a practical implementation under various IL settings in Section 4.3.
4.1 Contextual Imitation Learning (CEIL)
Motivated by the hindsight information matching online/offline RI. (Section 3.2), we propose to learn general hindsight embedding function fe. which encodes trajectory (with window size T') into a latent variable z E Z, |z| < T * [S]. Formally, we learn the embedding function fg and a corresponding contextual policy no Talss, Is, z) by minimizing the trajectory self-consistency loss:
no, J6 = min - -E [log To (T((f((T)) fo
min -R,~D(r) B B(aa)~ [log (alss, fo(T)),, (3)
No.fo

---

where in the online setting, we sample trajectory т from buffer D(T), known as the experience replay buffer in online RL: in the offline setting, we sample trajectory T directly from the given offline data.
we can ensure that the learned contextual policy To and the embedding function fo are accurate on the empirical data (T)), then we can convert the II. policy optimization objective (in Equation 1) into a bi-level expert matching objective:
min D(To (T())) x=
(4)
s.L. ne. So
min -Ep(r) [log (F( R(fo); and z' E fo o supp(D), #o.Je
(5)
where R(fa) an added regularization over the embedding function (we will elaborate on later), and (P) denotes the support of the trajectory distribution {T( > 0}. Here fo is employed to map the trajectory space to the latent variable space (2). Intuitively, by optimizing Equation 4, we expect the induced trajectory distribution of the learned (T(a)s, z²) will match that of the expert. However, in the offline IL setting, the contextual policy can not interact with the environment. If we directly optimize the expert matching objective (Equation 4), such an objective can easily exploit generalization errors in the contextual policy model to infer a mistakenly overestimated z" that achieves low expert-matchinn Therefore, we formalize CHIL into a bi-level optimization problem, where, in Equation 5, we explicitly constrain the inferred z" lies in the (fa-mapped) support of the training trajectory distribution.
loss but does not preserve the trajectory self-consistency (Equation 3).
At a high level, CEIL decouples the learning policy into two parts: an expressive contextual policy (a(a)s,-) and optimal contextual variable z'. By comparing CEIL with the plain expert match- ing objective, min.,, (T (T), (T))), in Equation 1, we highlight two merits: 1) CEIL's expert matching loss (Equation 4) does not account for updating Te and is only incentivized to update the dimeennnsi latent variable z", which enjoys efficient parameter learning similar to the prompt tuning in large language models [74], and 2) we learn me by simply performing supervised regression (Equation 5), which is more stable compared to vanilla inverse- inverse-RL/adverssin RL/adversarial- methods.
4.2 Connection to the Plain Expert Matching Objective
To gain more insight into Equation 4 that captures the quality of IL (the degree of similarity to the expert data), we define D(,,)) as the sum of reverse KL and forward KL divergence", i.e. D(q.r) p) (gl(pplp)) + DRI. (pll)) and derive an alternative form for Equation 4:
arg min ((Fa(T|Z*). (T)) = arg max I(z"; TE)
TT) ((TT(T),TT(T))
(6)
x
Tsn
where I [(x; y) denotes the mutual information (MI) between x and y, which measures the predictive power of y x (or vice-versa), the latent variables are defined as Tg == (T) To := (Te")=(((("") ), and (T(T)) )-= = Ez [T=(T[z")]
Intuitively, the second term To on RHS of Equation 6 is similar to the plain expert matching objective Equation 1, except that here we optimize a latent variable z" over this objective. Regarding the MI terms TMI. we can interpret the maximization over JMI as an implicit policy improvement, which incentivizes the optimal latent variable z" for having high predictive power of the expert data TE and having low predictive power of the non-expert data Tp-
Further, we can rewrite the MI term (JMI in Equation 6) in terms of the learned embedding function fö. yielding an approximate embedding inference objective Tmu(sa):
TMil E Ex (2) log p(z" (TE) (x (T)) logpplz"") (Te)
(ppr))) [- Iz" (TE) ||z" [T( & Fuicson
where we approximate the logarithmic predictive power of z on T with - fo())( by taking advantage of the learned embedding function f in Equation 5.
denotes the (forward) KL divergences. It 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 -)--)) as the sum of reverse KL and forward KL, and set the weights of both reverse KL and forward KL to I.
3 DDi (pl14) == Ep(x) log p(x)
q(x)
5

---

Algorithm Training CEIL: Online or Offline IL Setting
Require: Expert demonstrations (T)) empty buffer D for online IL or reward-free offline data D for offline IL, training iteration K, and batch size N.
1: Initialize contextual policy Ta(als, -), embedding function fz(z/)), and latent variable z".
2: for k = 1, K do
3: (Online only) Run policy naa (alss, z=" in environment and store experience into buffer D.
4: Sample a batch of data {T} from for online or 'D for offline IL.
5: Learn To and fa over sampled { } using the trajectory self-consistency loss.
6: Update z" and fe over sampled {T} by maximizing FMI(ss) aJD.
7: (Offline only) Update z' by minimizing R(z"). # climinating the offline OOD issues.
8: end for
Return: the learned contextual policy (a(al,, and the optimal latent variable zz.
By maximizing (M the learned optimal z" will be induced to converge towards the embeddings of expert data and avoid trivial solutions (as shown in Figure 1). Intuitively, JMI(f.)) fu)) can also be thought of as an instantiation of contrastive loss, which manifests two facets we consider significant in IL: 1) the "anchor" variable* z² unknown and must be estimated, and 2) it is necessary to ensure that the estimated z lies in the support set of training distribution, as specified by the support constraints in Equation
Happer-v2
Ant-v2
10e
12 distance
1.00
Retur
distance
Retwrn
00
100
In summary. by comparing IMI(Sa) and Jn. can observe that JMI(fa) actually encour- ages expert matching in the embedding space, while TD encourages expert matching in the original trajectory space. In the next section, we will see that such an embedding level ex- pert matching objective naturally lends itself to cross-domain settings.
Rollout steps (k)
Rollout steps (K)
Retum
1|z
-ylTelll
Figure 1: During learning, the distance between z and (T() decreases rapidly (green lines). Mean- while, policy no(-),,2" 22) gets better (blue lines), fo(TT) (T)) gradually approaches 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 II., cross-domain II., and one-shot II. settings (see more details in Appendix 9.3).
(best viewed in colors), CEIL alternates between solving the bi-level problem with respect to the support constraint (Line for online or Line for offline IL), the trajectory self-consistency loss (Line 5), and the optimal embedding inference (Line 6).
As shown in Algorithm 1
To satisfy the support constraint in Equation 5, for online IL (Line 3), we directly roll out the 2" --conditioned policy To (a(ss, z²) ) in the environment; for offline IL (Line 7), we minimize simple regularization over z bearing a close resemblance to the one used in TD3+BC [23]:
R(z") min (|1z - fa(T))l ||2" fa(TT) TE := T (E(T),
TD (-
D(T), (7)
where apply stop- -raddiett operation to fa. To ensure the optimal embedding inference (max² JMILA)) - JD) retaining the flexibilityy of seeking z" across different instances of fa. we jointly update the optimal embedding z and the embedding function fs with
max JMULsO) aJDp 26
(8)
where we use to control the weight on Jp-
LfO. In the Lf0 setting, expert actions 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
"The 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 z" in (Milsa)) as an instance of the anchor.
5 "Our code will be released at httpp::// //t/t//ttttbb// com/wecht.. Steecht//Generlll izedCEII. In other words, the offfine support constraint in Equation 5 is achieved through minimizing R(z").

---

leverage a large number of suboptimal actions presented in online/offline D(T). Thus, we can learn the contextual policy ( is 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
Cross-domain II.. Cross-domain II. 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 JD may not be a sufficient proxy for the expert matching objective, as there may never exist trajectory (in the learning MDP) that matches the given expert data. Thus, we can set (the weight of JD) a to O.
Further, to make embedding function fo useful for guiding the expert matching in latent space (i.e., max JMI.f()), we encourage fo to capture the task-relevant embeddings and ignore the domain- specific factors. To do so, we generate a set of pseudo-random transitions (T } by independently sampling trajectories from expert data {TE(TE)) and adding random noise over these sampled trajectories, i.e., TE = TE + noise. Then, we couple each trajectory t in {TE} U (TE') } with a label n E {0, 1), indicating whether it is noised, and then generate T E {Te} U (T and n E {0,1}. Thus, we can set the regularization R(faa) in Equation 5 to be:
set of [(T,)),, where
R(fs) I()((T);n)
(9)
Intuitively, maximizing R()a) encourages embeddings to be domain-agnostic and task-relevant: fT(TE) has high predictive power over expert data (n 0) and low that over noised data (n = 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 Tnew, function fo to generate a corresponding latent embedding Zaew- When conditioning on such an embedding, can directly roll out a (s, Znew) to recover the one-shot expert behavior.
the learned embedding
5 Experiments
In this section, we conduct experiments across a variety of IL problem domains: single/cross-domain IL, online/offliee IL, and LID/LIO IL settings. By arranging and combining these IL domains, we obtain 8 IL tasks in all: S-on-LjD, S-on-LJO, S-off-L/D, S-off-LfO, C-on-LJD, C-on-LfO, C-off-LfD, and C-off-LfO, where S/C denotes single/cross-domale LiD/LfO denote learning from ceemonstrations/oseeraaa the scalability of CEIL on the challenging one-shot IL setting.
IL, on/off denotes online/offine IL, and respectively. Moreover, we also verify
Our experiments are conducted in four popular MuJoCo environments: Hopper-v2 (Hop.), HalfCheetah-v2 (Hal.), Walker2d-v2 (Wal.), and Ant-v (Ant.). In the single-domain II setting, train a SAC policy in each environment and use the learned expert policy to collect expert trajecto- To investigate the cross-domain IL setting,
ries idemonstaations/obervvaons 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 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.
assume the two
5.1 Evaluation Results
To demonstrate the versatility of the CEIL idea, we collect 20 expert trajectories (demonstrations in LID or state-only observations in LfO) for each environment and compare CEIL to GAIL [30], AIRL [21], SQIL [61], IQ-Learn [28]. ValueDICE [41], GAIfO [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//Coon/of-Lff//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 online IL settings: S-on-LfD (top-lefi), S-on-I,fO (top-right), C-on-L.D (botiom-left), and C- on-I-O (bottom-right) settings. As can be seen, CEIL quickly achieves expert-level performance in S-on-I.DD. When extended S-on-IfO, CEII. also yields better sample efficiency compared to baselines. Further, considering the complex cross- cros--ommain setting, we can see those baselines SQIL
7

---

(a) Single-domain online LfD (S-on-LfD)
(b) Single-domain online LfO (S-on-LfO)
Hopperv2
Ant-v2
Hopper-v2
Ant-v2
Return
3K-
4K-
Return
3K
4k:
1.5k
2k
1.5k
2k
0-
100
500
200
500
200
300 400
Rollout steps (k)
Rollout steps (K)
Rollout steps
Rollout steps (K)
(c)
Cross-domaile online LfD (C-on-LfD)
(d)
Cross moos--dmain online LfO (C-on-LfO)
Hopperv22
Ant-v2
Hopper-v2
Ant-v2
4k
1.2k
Return
Return
0.8k
2k
0.5k
100 200 300
400 500
0 100 200
300 400 500
100 200 300 400 500
0 100 200 300 400 500
Roilout sceps
Rol lout (K)
Rollout steps
Rollout steps
CEIL
IQ-Leam
ValueDICE
CEIL
GAIO
IQ-Learn
AIRL
SQIL
(legend for S/C on LTD)
AIRL (state only)
(legend for S/C on LFO)
Figure 2: Return curves online settings: (a) S-on- LJD. (b) S-on-LjO, (c) C-on-LfD, (d) C-on-LfO, where the shaded area represents 95% confidence interval over 30 trails. Note that baselines cannot be applied all the task settings, thus we only provide comparisons with compatible baselines (two separate legends).
Table 2: Normalized scores (averaged over 30 trails for each task) 4 offline IL settings: S-off-LID, S-off-LjO. C-off-IfDD, and C-off-LfO LfO. Scores within two points of the maximum score are highlighted
Hopper-v2
Halfchectab-v2
Walker2d-v2
Ant-v2
Offiine settings
mr
me
me
me
ORIL (TD3+BCC
50.9
22.1 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
25.5
13.2
25.3
14.4
25.6
15.6
63.6
387.1
SALDD
1Q-Leam
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
117.0
126.3
122.0
114.3
1163.5
$4.00
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)
78.5
20.5
120.5
821.7
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
CLD
IQ-Lean
35.4
25.9
27.4
27.1
31.2
27.7
22.2
31.7
63.7
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
77.2
38.9
26.4
77.8
114.9
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
CEH (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
C1180
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
CELL (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-IfO) 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 we provide the normalized return of our method and baseline methods on reward-free D4RL [22] medium (m), medium-replay (mr), and medium- (me) datasets. We can
Normalized Return
Happer-v2
Normalized Return
Ant-v2
Normalized Return
(c)
Hopper-v2
Normalized Return
(d)
Ant-v2
(online)
(online)
LFD (CEIL)
S-on-LfD (CEIL)
100-
300
5-off-LID (CEIL)
Gon-LID (CEIL)
Coon-LID (CEIL)
loffline)
Coff-LID (CEIL)
C-off-LID (CEIL)
10
20 à
16
32
16
32
Demonstration Num.
Demonstration Num.
Window Size
Window Size
10-izam
SQIL
ORIL
ValueDICE
SMODICE (legend varying Num.
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 Lffff-LfD and Loff-Lfo settings. Compared to the state-of-hh-ard offline baselines, CEIL also shows competitive results on the challenging cross-domain offlin II. settings (C-off-I.fD and C-off-1fOO.
One-shot II. Then, we explore CEIL the one-shot II. tasks, where expect CEIL 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 contex- tual policy in the training domain (online/offline IL), then infer new contextual variable and evaluate it on
One-shot IL
Hop.
Hal.
Wal.
Ant.
Online
SQUL
16.8
1.1
3.5
4.2
IQ-Learn
4.6
0.2
1.7
7.5
the new task. To facilitate comparison to baselines, we similarly pre-train policy network (using baselines) and run BC on top of the pre-trained policy by using the provided demonstration. Consequently, such base- line+BC procedure cannot be applied to the (one-shot) LIO tasks. The results in Table 3 show that baseline+BC struggles to transfer their expertise to new tasks. Bene- fiting 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.
CEIL (LiD)
29.9
2.5
31.7
20.5
CEIL (LIO(
17.8
3.2
5.6
29.7
ORIL
14.7
0.2
6,9
17.4
SQIL
7.4
0.8
4.6
12.5
Offine
1Q-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 (LiD)
85.6
24.3
CELL (LiO)
72.2
5.1
70.0
19.4
Table 3: Normalized results one-shot IL, where CEIL shows prominent transferability.
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 several IL settings, each of which passes a portion of task-relevant information to the agent. For exam- ple, we can provide the agent with both demonstrations and state-only observa- tions and, in some cases, cross-domain demonstrations (S-LID+S-LFO+C-LID)
Hybrid offline IL settings
Hop.
Hal.
Wal.
Ant.
S-LID
29.4
69.9
42.8
84.9
S-LID + S-LfO
30.4
68.6
42.3
91.6
S-LID S-LfO C-LID S-LID + C-LIO C-LIO
30.7
71.7
42.9
89.2
58.6
79.6
43.7
98.0
Table 4: The normalized results of CEIL, showing that CEIL consistently digest useful (task-relevant) information and boost its performance, under hybrid of offline IL settings.
To examine the versatility of CEIL, we collect 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.
Varying the number of demonstrations. In Figure 3 (a. b), 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 trajectory. Next we assess the effect of the trajectory window size (i.e., the length of trajectory 7 used for the embedding function Sa in Equation 3). In Figure 3 (b, c), we ablate the number of the window size in 4 LiD 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-Lfi) Hopper-v2
(b) Cross-domain online LfD (C-on-LfD) Hopper-v2
Ant-v2
Ant-v2
1.5k
4k
Return
3k
4k-
Retum
1k
CEIL
2k
CEIL
1.5k
CElL/ablating (J)
2k
CEIL(ablsting ( CEIL/ablating.e)
CEILiabiating fg)
CEILiableting fy)
CEIL/ablating.Jo) Joe)
CEILiabisting Jwi
CEILiablating jwl
0
200 400 Rol lout steps (k)
300 500 Rollout steps (k)
300 500 100
200
Rollout steps (k)
Rollout steps (k)
Figure 4: Ablation studies on the optimization of fo (ablating fo) and the objective of Jsu (ablating JMn), where the shaded represents 95% Cls over trails. See ablation results for offline IL tasks in Table 5.

---

Table 5: Ablation studies on the optimization of fo (ablating fa) and the objective of Jmi (ablating JMi),
where scores (averaged over trails each task) within two points the maximum score are highlighted.
Hopper-v2
HalfCheetah-v2
Walker2d-v2
Ant-v2
Offline settings
sum
me
S-LID
CEIL (ablating fe)
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 (ablating Fun)
83,2
89.0
53.5
107.4
68.0
75.6
97.8
951.4
CEII.
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
CEIL (ablating fe)
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. (ablating Jsu)
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ø and the objective of JMI- In Figure and Table 5, we carried out ablation experiments on the loss of fe and JMi in both online IL and offline IL. settings. We can that ablating the fo 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 encourages the embedding function be task-relevant, and thus we use the expert matching loss to update fs. We can also see that ablating JMI does lead to degraded performance, further verifying the effectiveness of our expert matching objective in the latent space.
6 Conclusion
In this paper, present CEIL, a novel and general Imitation Learning framework applicable to wide range of IL settings, including C/S-on/off-LjD//..0 and few-shot IL settings. This achieved by explicitly decoupling the imitation policy into contextual policy, learned with the self- sef-suupe hindsight information matching objective, and 2) a latent variable, inferred by performing the IL expert matching objective. Compared prior baselines, in most of the online II. tasks and achieves better or competitive performances in offline tasks.
results show that CEIL is more sample-efficiet
Limitations and future work. Our primary aim behind this work is to develop 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 LIO 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/domain have different state spaces or action spaces. Considering such cross 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 tokenization approach, and thus suitable for II. tasks with different action/state spaces. Thus, we encourage the future exploration of generalized IL methods across different embodiments.
Acknowledgments and Disclosure of Funding
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. 20222D0208800), and NSFC General Program (Grant No. 62176215).
References
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-290 293220, 2021.
[2] Anurag Ajay, Yilun Du Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision 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, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. Advances in neural information processing systems, 30, 2017.
10

---

[4] Kai Arulkumaran, Dylan R Ashley, Jürgen Schmidhuber, and Rupesh K Srivastava. All you need supervised learning: From imitation learning meta-rl with upside down rl. arXiv preprini arXiv:2202. / 2022.
[5] Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639-24654, 2022.
[6] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint arXiv: 040020000..0...... 2018.
[7] Damian Boborzi, Prristoph-Nioolat Strachle, Jens S Buchner, and Lars Mikelsons. Imitation learning by state-only distribution matching, arXiv preprint arXiv:2202.04332, 2022.
[8] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? arXiv preprini arXivv2206.01077 2022.
[9] Daniel Brown, Wonjoon Goo, and Scott Niekum. Better-than-demonsttan imitation learning via automatically-rankes demonstrations. In Conference robot learning, pages 330-359. PMLR, 2020.
[10] Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu. Mingfei Sun, David Bignell, Stephanie Milani, Katja Hofmann, Matthew Hausknecht, Anea Dragan, al. Unimask: Unified inference in sequential decision problems. arXiv preprint wrXiv:221.10080 2022.
[11] Jonathan Chang, Masatoshi Uehara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift in imitation learning via offline data with partial coverage. Advances in Neural Information Processing Systems, 34:965-979, 2021.
[12] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 500441 2021.
[13] Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning. arXiv preprint arXiv:2006.04457 2020.
[14] Christopher R Dance, Julien Perez, and Théo Cachet. Conditioned reinforcement learning for few-shot imitation. In International Conference on Machine Learning, pages 2376-2387 PMLR, 2021.
[15] Branton DeMoss, Paul Duckworth, Nick Hawes, and Ingmar Posner. Ditto: Offline imitation learning with world models. arXiv preprint arXiv:2302.0300 2023.
[16] Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAl Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. Advances in neural information processing systems, 30, 2017.
[17] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.1075:: 107511, 2021.
[18] Arnaud Fickinger, Samuel Cohen, Stuart Russell, and Brandon Amos. Cross-domain imitation learning via optimal transport. arXiv preprint arXiv:2 10.:22/00003844 2021.
[19] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168 168. PMLR, 2022.
[20] Tim Franzmeyer, Philip HS Torr, and João Henriques. Learn what matters: cross-domain imitation learning with task-relevant embeddings. arXiv preprint 220..222000...93,. 2022.
11

---

[21] Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. arXiv preprint arXiv:1710.1248,, 2017.
[22] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 2020.
[23] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances neural information processing systems, 34:20132-20145, 2021.
[24] Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10364. 2021.
[25] Tanmay Gangwani and Jian Peng. State-only imitation with transition dynamics mismatch. arXiv preprint arXiv:2002. 11879, 2020.
[26] Tanmay Gangwani, Yuan Zhou, and Jian Peng. Imitation learning from observations under transition model disparity. arXiv preprint arXiv:2204. 1446, 2022.
[27] Peng Gao, Shijie Geng, Renrui Zhang, Teli Ma, Rongyao Fang, Yongfeng Zhang, Hongsheng Li, and Yu Qiao. Clip-adapter: Better vision-language models with feature adapters. arXiv preprint
2021.
[28) Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34: 4028-4039, 2021.
[29] Adam Gleave, Mohammad Taufeeque, Juan Rocamonde, Erik Jenner, Steven H. Wang, Sam Toyer, Maximilian Ernestus, Nora Belrose, Scott Emmons, and Stuart Russell, imitation: Clean imitation learning implementations, 2022.
[30] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
[31] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273- 1286, 2021.
[32] Firas Jarboui and Vianney Perchet. Offline inverse reinforcement learning. arXiv preprint arXiv:2106.05068 2021.
[33] Daniel Jarrett, Ioana Bica, and Mihaela van der Schaar. Strictly batch imitation learning by energy-based distribution matching. Advances in Neural Information Processing Systems, 33: 7354-7365, 2020.
[34] Shengyi Jiang, Jingcheng Pang, and Yang Yu. Offline imitation learning with misspecified simulator. Advances in neural information processing systems,
2020.
[35] Kshitij Judah, Alan Fern, Prasad Tadepalli, and Robby Goetschalekx. Imitation learning with demonstrations and shaping rewards. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 28, 2014.
[36] Yachen Kang, Diyuan Shi, Jinxin Liu, Li He, and Donglin Wang. Beyond reward: Offline preference-uuider policy optimization. arXiv preprint arXiv:2305.16217, 2023.
[37] Liyiming Ke. Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srini- vasa. Imitation learning as f-divergence minimization. In International Workshop on the Algorithmic Foundations of Robotics, pages 313-329, Springer, 2020.
[38] Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srini- vasa. Imitation learning as f-divergence minimization. In Algorithmic Foundations of Roboties XIV: Proceedings of the Fourteenth Workshop on the Algorithmic Foundations of Robotics 14, pages 313-329. Springer International Publishing, 2021.
12

---

[39] Geon-Hyeong Kim, Seokin Seo, Jongmin Lee. Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. Demodice: Offline imitation learning with supplementary imperfect demonstrations. In International Conference Learning Representations, 2022.
[40] Kuno Kim, Yihong Gu, Jiaming Song, Shengjia Zhao, and Stefano Ermon. Domain adaptive imitation learning. In International Conference on Machine Learning, pages 1295. PMLR, 2020.
[41] Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off off-policy distribu- tion matching. arXiv preprint arXiv:1912.05032 2019.
[42] Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-conditioned policies. arXiv preprint arXiv:1912.13465. 2019.
[43] Yao Lai, Jinxin Liu, Zhentao Tang, Bin Wang. HAO Jianye, and Ping Luo. Chipformer: Transferable chip placement via offline decision transformer. ICML., 2023. URL. https:
/openrevieww hwwpeeeeewww..eeeeeeeeeeeee heetppdffiiddjmmlittttt
[44] Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J Lim. Generalizable imitation learn- ing from observation via inferring goal proximity, Advances in neural information processing systems, 18-16111- 2021.
[45] Brian Lester, Rami A1-Rfou, and Noah Constant. The power of scale for parameter-efffcient prompt tuning. arXiv preprint arXiv:2104.0880 18691, 2021.
[46] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 2020.
[47] Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. Advances in Neural Information Processing Systems, 30, 2017.
[48] Fangchen Liu, Zhan Ling, Tongzhou Mu, and Hao Su. State alignment- based imitation learning. arXiv preprint 1911 2019.
[49] Jinxin Liu, Donglin Wang, Qiangxing Tian, and Zhengyu Chen. Learn goal-conditioned policy with intrinsic motivation for deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7558-7566, 2022.
[50] Jinxin Liu, Hongyin Zhang, and Donglin Wang. Dara: Dynamics-aware reward augmentation in offline reinforcement learning. arXiv preprint wwXiv:2003.06662, 2022.
[51] Jinxin Liu, Ziqi Zhang, Zhenyu Wei, Zifeng Zhuang, Yachen Kang, Sibo Gai, and Donglin Wang. Beyond ood state actions: Supported cross-domain offline reinforcement learning. arXiv preprint arXiv: 12755, 2023.
[52] Minghuan Liu, Tairan He, Minkai Xu, and Weinan Zhang. Energy-based imitation learning. arXiv preprint 200.::20004.93355, 2020.
[53] Minghuan Liu, Hanye Zhao, Zhengyu Yang, Jian Shen, Weinan Zhang, Li Zhao, and Tie- Liu. Curriculum offline imitating learning. Advances in Neural Information Processing Systems, 34:6266 00066666772222 2021.
[54] YuXuan Liu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Imitation from observation: Learning to imitate behaviors from raw video via context translation. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 1118 - 1125. IEEE, 2018
[55] Yicheng Luo, Zhengyao Jiang, Samuel Cohen, Edward Grefenstette, and Mare Peter Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303. /3971, 2023.
[56] Yecheng Jason Ma, Andrew Shen, Dinesh Jayaraman, and Osbert Bastani. Smodice: Versatile offline imitation learning via state occupancy matching. arXiv e-prints, pages arXiv-2202, 2022.
13

---

[57] Tianwei Ni, Harshit Sikchi, Yufei Wang, Tejus Gupta, Lisa Lee, and Ben Eysenbach. f-irl: Inverse reinforcement learning via state marginal matching. In Conference on Robot Learning. pages 529-551. PMLR, 2021.
[58] Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1):88--8))) 97, 1991.
[59] Yiwen Qiu, Jialong Wu, Zhangjie Cao, and Mingsheng Long. Out-of-dynamics imitation learning from multimodal demonstrations. In Conference on Robot Learning, pages 1071-1080. PMLR, 2023.
[60] Dripta S Raychaudhuri, Sujoy Paul, Jeroen Vanbaar, and Amit K Roy-Chowdhury. Cross- domain imitation from observations. International Conference on Machine Learning, pages 8902 8912. PMLR, 2021.
[61] Siddharth Reddy, Anca D Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. arXiv preprint arXiv:1905.11108 2019.
[62] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth inferna- sional conference on artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
[63] Stefan Schaal. Learning from demonstration. Advances in neural information processing systems, 9, 1996.
[64] Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaškowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv: 1rXiv:1912.028777 2019.
[65] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. arXiv preprint arXiv:1805.01954 2018.
[66] Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observa- tion. arXiv preprint arXiv: 1807.088770 2018.
[67] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
[68] Luca Viano, Yu-Ting Huang, Parameswaran Kamalaruban, Craig Innes, Subramanian Ra- mamoorthy, and Adrian Weller, Robust learning from observation with model misspecification. arXiv preprint arXiv:2202.00 06000, 2022.
[69] Tianyu Wang, Nikhil Karnwal, and Nikolay Atanasov. Latent policies for adversarial imitation learning. arXiv preprini arXiv:2206. 1299, 2022.
[70] Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning, pages 24725-24742. PMLR, 2022.
[71] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. International Conference Machine Learning, pages 24631-24645, PMLR, 2022.
[72] Sheng Yue, Guanbo Wang, Wei Shao, Zhaofeng Zhang, Sen Lin, Ju Ren, and Junshan Zhang. Clare: Conservative model- reward learning for offline inverse reinforcement learning. arXiv preprint arXiv:2302.04782. 2023.
[73] Wenjia Zhang, Haoran Xu, Haoyi Niu, Peng Cheng. Ming Li, Heming Zhang, Guyue Zhou, and Xianyuan Zhan. Discriminator-gan -uided model-based offline imitation learning. In Conference on Robot Learning, pages
PMLR, 2023.
[74] Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu. Learning to prompt for lasion-language models. International Journal of Computer Vision, 23309))2337- 2022.
14

---

[75] Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off-policy imitation learning from observations. Advances in Neural Information Processing Systems, 33:12402-12444 12413, 2020.
[76] Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo. Behavior proximal policy optimization. arXiv preprint 2tt::::3322.111122 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 P.Xiv:2011.13888 2020.
[79] Konrad Zoina, 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.
15

---

Appendix
7 Additional Derivation
(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 -). as the sum of reverse KL and forward KL divergence, i.e., (g DxL (qll)) + DxL(p) q). and derive an alternative form for Equation 4: arg min D(To(T|Z") TE(T)) = arg max T(z", TT) I(z"; TD) DKL (To (T)
Jmn
ID
where [(x;y) denotes the mutual information (MI) between x and y, which measures the predictive power of y x (or vice-versa), the latent variables are defined as TE == (T(T), = p(z" (ee())TT("""") [2"", and (T(T)) = Ez [Tσ(T|z*) }}
Below is our derivation:
min D(TT(T)Z"), (T
min Ep(r)) [xL ( (T/2 (=)(((2")][TT(T) +
min pp(e [log no (T|z") log TE (T)
x"
+ E (T)
log T ( (T) log ne (T|z²))
log p(z" (TT
min (a
log ग (T()
p(z*)
log T (T) - -log p(z") p(z"|T) (TT(
+ P
p(z*)
log p(2*(T) + log TT(T)
log p(z²1T)
TT(T)
min Epiz") (m )
Expt" (m((F)
p(2²) + log
p(z*)
(T)
TH (T)
max I(x" ; Tg)
:T) D(Fo (T)))
where TE T TT(T), To T p(z" ) (T(2²))
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 the reward-free D4RL [22] expert dataset. Consistently, we can observe that CEIL achieves significant improvement over the baseline methods in both S-off-I.fD and S-off-I.O settings. Compared to the state-of-the-arr offline IL baselines, CEIL also shows competitive results on the challenging cross-domain offline IL settings LC-off-LiD and L-fff-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 demoonsrationskooserraiins (expert domain) and to mimic such expert behaviors in the training environment (training domain). Thus, 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 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 We can find that across 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
16

---

Table 6: Normalized scores (averaged over 30 trails for each task) on D4RL expert dataset. Scores within two points of the maximum score highlighted. hop: Hopper-v2. hal: HaliCheetahh-2. wal: Walker2d-v2. ant: Ant-v2.
hop expert
hal expert
wal expert
expert
sum
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
S-off-I.fD
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
115.6
117.8
435.4
ORIL (TD3+BC)
64.2
92.1
44.3
212.8
Lff
SMODICE
111.3
93.7
108.0
122.0
435.0
CEIL.
103.3
96.8
110.0
126.4
436.5
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-Leamn
25.9
31.2
31.7
55.8
144.6
C-off-L.fD
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
ORIL (TD3+BC)
22.5
76.6
11.2
28.2
138.6
C-off-IfO
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 experi datasel over 30 trails for each task) on 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- e: expert.
Hopper-v2
HallCheetah--22
mnr
me
m
me
me
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-Leam
11.8
9.7
17.1
17.1
7.7
7.8
9.5
9.5
90.2
C-off-LiD
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
54.3
-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
ORIL (TD3+BC)
62.3
18.7
57.0
28.2
0.2
1.1
-0.3
-2.3
165.0
C-off-IfOO
SMODICE
77.6
22.5
80.2
71.0
2.0
-0.9
0.8
-2.3
250.9
CEI..
56.4
58.6
56.7
65.2
5.5
36.5
5.0
5.0
288.7
Walker2d-v2
Ant-v2
sum
m
mn
me
m
mn
me
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-Leam
8.4
5.0
10.2
10.2
19.4
18.4
16.1
16.1
103.8
C-off-IffD
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
ORIL (TD3+BC)
22.4
15.2
17.8
12.6
13.6
20.7
5.5
-6.2
101.6
C-off-1f0
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
17

---

C-off-LFD
C-off-LfO
Improvement
Improvement
hop-mir
hop-me
hop-e
hal-m
hal-mir
hal-me
hale
walmm
wal-mr
wal-me
wal-e
ant m
ant-me
ant-e
hop-m
hop hop-mm
hop-e
hal-m.
hal-hr
hal-me
hal-e
wal-m
wa wal-mi mn
walee
ant-m
ant-mi
ant-me ant
hop
ant
hop
wal
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 doos-nn regularization causes negative performance improvement. hop: Hopper-v2. hal: HalfCheetah-v2 wal: Walker2d-v2 v2. ant: Ant-v2. m: medium. medium-expert. medium-
expert.
(a) S-off-LfD
Median
IQM
Mean
Optimality Gap
10-Learn
SQIL(TD3+ [CTDD+++CC
ORIL(TD3+ BC)
ValueDICE
DemoDICE
SMODICE
CEIL
30
60
90
25 normalized return
50
75
100
25
50
75
100
20
40
60
8D
normalized return
normalized return
normalized return
S-off-LfO
Median
IQM
Mean
pptimality Gap
ORIL(TD3+BC)
SMODICE
CEIL
45
60
75
45
60
75
50 norzalized return
60
70
30
40
50
60
normalized return
normalized return
normalized retum
(c) C-off-LID
Median
JQM
Mean
Optimality
10-Learn
SQIL(TD3
ORIL(TD3+ BC)
VilueDICE
DemoDICE
SMODICE
CEIL
25
50
75
20
40
60
80
20
40
60
80
40
60
80
normalized return
normalized return
normalized return
normalized return
C-off-LfO
Median
IQM
Mean
Optimality Gap
ORIL(TD3 BC) SMODICE
CEIL
20
40
60
80
30 45
60
30
45
60
45
60
normalized return
normalized return
nomaalzed return
normalized return
Figure 6: Aggregate median, IQM, mean, and optimality gap over 16 offline II tasks. Higher median, higher IQM, and higher mean and lower optimality gap are better. The shaded bar shows 95% stratified bootstrap confidence intervals. We see that CEIL achieves consistently better performance across a wide range of offline IL settings.
ablate the cross-domain regularization two cross-domain offline IL tasks (C-off-LyD and C-off-IjO). 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.
18

---

C-on-LID
S-on-LIO
1000
1000
Return
4000
Return
4000
750
750
500
500
2000
250
250
500
steps
steps
Rollout steps
steps Ok)
CEIL
IQ-Loam
ValueDICE (legend for S/C-on-LfD)
CEIL
IQ-Learn (legend for S/C-on-LFO)
AIRL
SQIL
AIRL (state only)
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 95% confidence interval over 30 trails. We can see that CEIL consistently achieves expert-level performance in LfD (S-on-IfD and C-on-IfD) tasks. Due to the lack of explicit exploration in online LfO settings, CEIL. exhibits drastic performance degradation (in f- and C-on-I.OO 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 range of offline IL settings Lf-of-LLDD LfO, Lffff-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 highlighted
Hopper-v2
Halfcheetah-v2
Walker2d-v2
Ant-v2
Offline settings
sum
m
mr
me
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
S M5
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
1Q-Learn
21.7
13.1
10.6
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
63.7
41.9
38.7
59.5
66.3
38.8
101.6
82.8
68.3
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
ORIL (TD3+BC)
53.4
45.0
2.1
27.4
80.4
47.3
514.1
S-ff--L #10
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
1Q-1eaam
11.3
18.6
20.1
4.1
6.5
6.6
18.3
12.8
12.2
30.7
$3.9
23.7
218.7
ValueDICE
56.0
54.2
-0.2
2.6
31.4
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
30.3
66.6
42.6
38.0
66.0
64.5
44.6
53.8
86.9
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
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
S-off-1 #15
16.9
13.5
IQ-Leam
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
58.3
53.6
2.3
2.3
1.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
114.3
768.7
SMODICE
55.9
25.7
72.7
42.5
37.6
66.4
67.0
43.2
$5.1
86.7
69.7
118.2
740.6
CEIL
116.4
103.7
80.4
43.0
120.3
84.8
126.8
1057.3
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
S.of L.D.DDD2 N20
SQIL (TD3+BC)
32.6
60.6
25.5
13.2
25.3
14.4
25.6
15.6
8.0
63.6
$8.4
44,3
387.1
IQ-Learn
19.9
24.9
5.0
7.5
7.5
22.3
19.6
18.5
38.4
264.5
ValueDICE
73.8
83.6
50.8
1.9
244
3.2
24.6
26.4
44.1
79.1
82.4
75.2
$47.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 per- formance of CEIL and baselines more tasks when we vary the number of expert trajectories. Considering offline IL settings, we provide the results 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 consistent with the findings in Figure in the main paper.
19

---

8.6 Limitation (Failure Modes in Online LfO Setting)
Meanwhile, we find that in the online LfO settings, CEIL's performance deteriorates severely few tasks, as shown in Figure 7 (Walker2d). In Lff (either on single-domain or on cross-domain IL) settings, CEIL can consistently achieve expert-level performance, but when migrating to LIO 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.
Implementation Details
9.1 Imitation Learning Tasks
In our paper, we conduct experiments across variety of IL problem domains: single/cross-dome -omain IL, online/offline IL, and LiD/LfO IL settings. By arranging and combining these IL domains, we obtain 8 IL tasks in all: S-on-LJD, S-on-LfO, S-off-L/D, S-off-LfO, C-on-LiD, C-on-LfO, C-aff-LfD, and C-off-LfO, where S/C denotes single/cross-domain II, on/off denotes online/offtine II, and respectively.
LiD/LfO denote learning from demonstratioon/oobrrvatioo
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 to learn an optimal policy that mimics the provided demonstrations in the training environment.
S-on-IfO. We have access 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-L.fD. We have access to a limited number of expert demonstrations and large amount of pre-collected offline (reward-free) data. The goal of S-off-I.fD is to learn 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 large amount of pre- collected offline (reward- free) data. The goal of S-off-L.OO 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-I.fD. We have access limited number of expert demonstrations and 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 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 large amount of pre-collected offline (reward-free) data. The goal of C-off-LfD is to learn 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 large amount of pre- collected offline reward-free) data. The goal of C-off-I-O 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 collect the expert observations and the environment that was used to collect the offfine data are not the same environment.
20

---

MuloCo environments and our modified versions. From left to right: Ant-v2, HalfCheetah- v2, Hopper-v2, Walker2d-v2, our modified Ant-v2, our modified HalfCheetah-v22 modified Hopper-v2, and our modified Walker2d-v2.
Figure 8:
9.2 Online IL Environments, Offline IL Datasets, and One-shot tasks
Our experiments are conducted in four popular MuloCo environments (Figure 8): Hopper-v2, HalfCheetah-v2 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, 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
code submission, which includes modified MuJoCo assets. For one-shot IL tasks, we train the policy only in the single-doma settings LS-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 test the one-shot performance.
IL
Collecting expert behaviors. In our implementation, we use the publicly available rikit? imple- mentation of SAC to learn expert policy and use the learned policy to collect expert behaviors (demonstrations in LfD or observations in LfO).
9.3 CEIL Implementation Details
Trajectory seff-consistency loss. To learn the embedding function fo and a corresponding contextual policy Fe a|ss, z). we minimize the following trajectory self-consistency loss:
mo. f min ne.fe
[log ( (T(T))))
where TI.T denotes trajectory segment with window size of T'. In the online setting, we sample trajectory from the experience replay buffer D(T); in the offline setting, we sample trajectory directly from the given offline data D(T). Meanwhile, if we can access the expert actions (i.e., LiD settings), we also incorporate the expert demonstrations into the empirical expectation (i.e., storing the expert demonstrations into the onlinc/offline experience D(T)).
In our implementation, jectory TEr and a 4-layer MLP (with ReLU activation ) to predict the action respectively. To regularize the learning of the encoder function fa. additionally introduce decoder network (4-layer MLP with ReLU activation) Th(ss is fo(())) to predict the next states: min -Br..~P(TLz) (s...'"'' [log (s' [s, (f(TT)) Further, eircumvent issues of "posterior collapse" [67]. we encourage learning quantized latent embeddings. In a similar spirit to VQ-VAE [67]. incorporate ideas from vector quantization (VQ) and introduce the following regularization: min/, (TI.T)) ell2 + | (TIT) -sg[e]] gg]e]l?? where is dictionary of vector quantization embeddings (we set the size of this embedding dictionary to be 4096), ze (Tur) is defined as the nearest dictionary embedding to (T.r) and sg(:) denotes the stop-gradient operator.
4-layer MLP (with ReLU activation) encode the tra-
Out-level embedding inference. In Section 4.2 (main paper), approximate JMI with JMI(Fs) & (T + ||z - fo ( where we replace the mutual infor- mation with -1|z" fo (T) by leveraging the learned embedding function fo. Empirically, we find that we can ignore the second loss ||" (1 and directly conduct outer-level embedding inference with max the support constraints (R(z") in Equation 7 in the main paper) for the offline OOD issues naturally satisfied, since max,. (p(r)))((T)) [-||z* fa (Tx) and ming. R(z") equivalent.
Ep(rr) (T [-||z*
(Tg(Tg) (x(r))
[-| (TE)
Meanwhile, this simplification makes
Cross-domain II regularization. To encourage fo to capture the task-relevant embeddings and ignore the domain -emafn- factors, we set the regularization R(fo) in Equation 5 to be:
21

---

R(fx) = I[(a((T); where we couple each trajectory T in {TE) } U {TE) with a label n E {0,1}. indicating whether it is noised. In our implementation, we apply MINE [6] estimate the mutual information and conduct encoder regularization. Specifically, we estimate I(z; n) with I(z;n) sup3 Bp(z,n) [Z, n)] log (p(z)r(n) (exp (z, m))l and regularize the encoder fo with max,, (T) where we model fa with a 4-layer MLP (using ReLU activations).
Hyper-parameters, In Table 9, we list the hyper-parameter used in the experiments. For the size of the embedding dictionary, we selected it from 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 the default value. For the learning rate scheduler, tried the default Pytorch scheduler and CosineAnnealing WarenmmcallngWaaeetttttttt and found Aoncallnn WarmRestarts enables better results (thus 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.
Table 9: CEIL papper-prramterss...
Parameter
Value
size of the embedding dictionary size of the embedding dimension
4096
16
trajectory window size
2
encoder: optimizer
Adam
encoder: learning rate
3e-4
encoder: learning rate scheduler
CostneAnnealingWWm//essssst 0 1000,T mult=1, min=le-5)
encoder: number of hidden layers
4
encoder: number of hidden units per layer
512
encoder: nonlinearity
ReLU
policy: optimizer
Adam
policy: learning rate
3e-4
policy: learning rate scheduler
CosineAnnealing WarmRestaarts(T_O [000,,_mul=1,,, ela_ min=le-5)
policy: number of hidden layers
policy: number of hidden units per layer
512
policy: nonlinearity
ReLU
decoder: optimizer
Adam
decoder: learning rate
3e-4
decoder: learning rate scheduler
DostrcAAunealingaaekkkssrrrr = 1000,T_mul==1, eta min=le-5)
decoder: number of hidden layers
4
decoder: number of hidden units per layer
512
decoder: nonlincarity
ReLU
Table 10: Baseline methods and their code-bases.
Baselines
Code-basee bases
GAII,, GAIFO, AIRL
SAII.
iu/SAA..
IQ-Learn, SQIL
https///ittubbooomD/ii/////LLLat
ValueDICE
http://iiillooooooeccccaoooeeeee dice
DemoDICE
SMODICE, ORIL
9.4 Baselines Implementation Details
We summarize code-bases of our baseline implementations in Table 10 and describe each baseline as follows:
Generative Adversarial Imitation Learning (GAIL). GAIL [30] is a GAN-based online LfD method that trains policy (generator) to confuse a discriminator trained to distinguish between generated transitions and expert transitions, While the goal of the discriminator is to maximize the
22

---

objective below, the policy is optimized via an RL algorithm to match the expert occupancy measure (minimize the objective below):
[mr ; D) E, (Dss a)) Eg 1 log((Dss -AH{))
We used the implementation by Gleave al. [29] on the GitHub page", where there modifica- tions introduced with respect to the original paper: 1) a higher output of the discriminator represents better, PPO used optimize the policy instead of TRPO.
Generative Adversarial Imitation from Observations (GAIfO). GAIIO [66] an online LIO method that applies the principle of GAIL and utilizes a state-only discriminator to judge whether the generated trajectory matches the expert trajectory in terms of states. We provide the objective of GAIfO as follows:
J (T, D) = Ex [log (ss, + E - (o 5 s') (T)...
Based on the implementation of GAIL, we implement GAIfO by changing the input the discrimi- nator to state transitions.
Adversarial Inverse Reinforvement Learning (AIRL). AIRL [21] is an online LID/LiO method using adversarial learning framework similar to GAIL. It modifies the form of the discriminator to explicitly disentangle the rask-relevant information from the transition dynamics. To make the policy more generalized and less sensitive to dynamics, AIRL proposes to learn parameterized reward function using the output of the discriminator:
So,s (0 >ha(s') - hp(s);
ex,(fa.p(,a,,s") as))
(o, expífa, a, s')) + (als)
Similarly to GAIL, we used the code provided by Gleave et al. [29]. and the RL algorithm is also PPO.
State Alignment-based Imitation Learning (SAIL). SAIL [48] an online LfO method capable of solving cross-domain tasks. SAIL aims to minimize the divergence between the policy rollout and the expert trajectory from both local and global perspectives: 1) locally, a KL divergence between the policy action and the action predicted by state planner and inverse dynamics model, 2) globally, Wasserstein divergence of state occupancy between the policy and the expert. The policy is optimized using:
( -DW(T(s)](TE(S)) AD
D(s+1)) - Eng(s) D(s) ADKL T(( (st) I1, (. [(s)))
where D is state-based discriminator trained via J(D) Exg [D(s)] E. [D(),, S is the pretrained VAE-based state planner, and Dlvv is the inverse dynamics model trained by supervised regression.
In the online setting, we use the official implementation published by the authors where SAIL is optimized using PPO with the reward definition: r(st, 81+1) = (D(44)) Emg(s) E Besides, we further implement SAIL in the offline setting by using TD3+BC [23] to maximize the reward defined above.
In our experiments, we empirically discover that SAIL is computationally expensive. While SAIL is able to learn tasks in the typical IL setting (S-on-I,DD), our early experimental results find that SAIL(TD++BCC with heavy hyperparameter tuning failed on the offline setting. This indicates that SAIL is rather sensitive the dataset composition, which also coincides with the results gathered in Ma et al. [56]. Thus, we do not include SAIL in our comparison results.
Soft-Q Imitation Learning (SQIL). SQIL [61] is a simple but effective single-domain LID IL algorithm that is easy to implement with both online and offline Q-learning algorithms. The main
23

---

idea of SQIL is to give sparse rewards (+1) only to those expert transitions and zero rewards (0) to those experiences in the replay buffer. The Q-function of SQIL is updated using the squared soft Bellman Error:
(D,
Σ
(Q(s.a)-(r+ (Q(s. (r + log Σ
exp(Q(s,d)))
D
(e,0,8")E
a'E,A
The overall objective of the Q-function maximize the following objective:
T(Q) 62 (Dr, ( 0).
(D(,:
In our experiments, the online imitation policy is optimized using SAC which is also used in the original paper. To make a fair comparison among the offline IL baselines, the offline policy is optimized via TD3+BC.
Offline Reinforced Imitation Learning (ORIL). ORIL [78] is offline single-domain IL method that solves both LfD and LfO tasks. To relax the hard-label assumption (like the sparse rewards made in SQIL), ORIL treats the experiences stored in the replay buffer unlabelled data that could potentially include both successful and failed trajectories. More specifically, ORIL aims to train reward function distinguish between the expert and the suboptimal data without explicitly knowing the negative labels. By incorporating Positive-unt undibbled learning (PU-learning). the objective of the reward model can be written as follows (for the LID setting):
J [R) ((.a)().a) [log( [[x, a))] + (s.a) [oo((1 R(s, a))] T)(BTT((..)) [log(1 R(s, a))]
where 1) is the relative proportion of the expert data and we set it as 0.5 throughout our experiments. In the original paper, the policy learning algorithm of ORIL Critic Regularized Regression (CRR), while in this paper, we implemented ORIL using TD3+BC for fair comparisons. Besides, we adapted ORIL to the LfO setting by learning state-only reward function:
(RR) (s..)) [ogg[((s,5 s' s (s.s'") [log(1 - 8')1 -TT [log(1 R(s,s'))
Inverse soft-Q learning (IQ-Learn). IQ-Learn [28] is an IRL-based method that can solve IL tasks in the online/offline and LiD/LfO settings. It proposes to directly learn Q-function from demonstrations and avoid the intermediate step of reward learning. Unlike GAIL optimizing a min-max objective defined in the reward-policy space, IQ-Learn solves the expert matching problem directly in the policy-Q space. The Q-function is trained to maximize the objective:
E (x, [Q(s, a) (s) Ea(saa))) [Q(s, a) (s)) ((r).
where (** A (1-)( [Q(s, a) log π (a]s)), (6rr is regularization term calculated over the expert distribution. Then, the policy is learned by SAC.
We use the code provided in the official IQ-learn repository 10 and reproduce the online- results reported in the original paper. For online tasks, we empirically find that penalizing the Q-value on the initial states gives the best and most stabilized performance. The learning objective of the Q-function for the online tasks is:
J(Q) (T.e(....") (Q(s) a) -(7)))))) (1 [VT(so)] (r))
In the offline setting, we find that using the above objective easily leads an overfitting issue, causing collapsed performance. Thus, we follow the instruction provided in the paper and only penalize the expert samples:
I(Q) = Exn(s.a.") [Q(s - >V)((') - Emnls.. (V*(s) - >V'(')) w(r)
= Erg()..,)) [Q(s.)) V*((s) v(r).
Imitation Learning via Off-Policy Distribution Matching (ValueDICE). ValueDICE [41] is "iCe-basedd L.fD algorithm which minimizes the divergence of state-action distributions between the policy and the expert. In contrast to the state-conditional distribution of actions T-|ss) used in the
hhhttp://githhbbco "https//githuh.com///99/9/ www.Diiv99//Q-L.cmmmm
" DICE refers stationary DIstribution Estimation Correction

---

above methods, the acte-nction distribution, d= S x A- + [0, 1], can uniquely characterize one-to-one correspondence,
(s. - 1)
= s, s,
|sp Po. T(s)), 8141 ~ P(st,
1=0
Thus, the plain expert matching objective can be reformulated and expressed in the Donsker- Varadhan Donsker-Varadhan representation:
J(x) = -Dx (s.add
min =:5xAABB
log (s.a)aar [exp(a(s, a))] (x) [x(ss a)].
The objective above can be expanded further by defining z(s, a) v(s, a) B² D(s a) and using zero-reward Bellman operator 8" to derive the following (adversarial) objective:
JDICElE, v) = log EE (s,a))=d= eex (u(s,a) a) 5 v(s
(1
[u(5o. an)]
We the official Tensorflow implementation" 12 in experiments. In the online setting, the rollouts collected are used as an additional replay regularization. The overall objective in the online setting is:
JDC (TT, uT)
-Dx a)d (s. + ad a)) (s, + ad (s.
= log 16 (d))))ndm [exp (v(s,a) - (8, (1 (1-a)(1 - 7) Eso~po (o-a(() [mr an)] ((s.a)ydR (v(8,a) a) rs a)]
where dii (1 add + Rd and is a non-negative regularization coefficient (we set a as 0.1 following the specification of the paper).
In the offline setting, ValueDICE only differs in the source sampling data. We change the online replay buffer to the offine pre-collected dataset.
Offiine Imitation Learning with Supplementary Imperfect Demonstrations (DemoDICE). DemoDICE [39] is DICE-based offline LID method that assumes have access to offline dataset collected by a behavior policy тT. Using this supplementary dataset, the expert matching objective of DemoDICE is instantiated over ValueDICE:
-D ( a)ll"""((("""x (s, - - aDxr (s
where a is positive weight for the constraint.
The above optimization objective can be transformed into three tractable components: 1) a reward function r(s, a) derived by pre-training binary discriminator D S A + [0,
r(s,a) = log(
(s) 1),
(*, = arg max Ers [log D(s a)) Exr [log(1 (ss, a))l
2) value function optimization objective:
Tw) - -11-7) [o (1 + a:) log E (a.)aar' eepp( (exp( r(s, + (oo-))
1+a
and 3) a policy optimization step:
J(x) E (.. [V" (s) lags()
(s, a) = arg max (v)).
We report the offline results using the official Tensorflow implementation" 13
dice
25

---

State Matching Offline DIstribution Correction Estimation (SMODICE). SMODICE [56] pro- poses solve offline IL tasks LfO and cross-domain settings and optimizes the following state occupancy objective:
D (s)(dd))
To incorporate the offline dataset, SMODICE derives an f-divergence regularized state-occupancy objective:
log( (5) do
(x)
+ -D a))ld (s,
Intuitively, the first term can be interpreted as matching the offline states towards the expert states, while the second regularization term constrains the policy close to the offline distribution of state- action occupancy. Similarly, we can divide the objective into three steps: 1) deriving a state-based reward by learning a state-based discriminator:
r(s, a) - = (*)s) 1),
D' ('(saa) = arg max = Exr (log D(s) + Egto [log(1 D(s))),
2) learning a value function using the learned reward:
((v) -1 pp [u(s)] log (a.a)aa''s +bx v(s))
and 3) training the policy via weighted regression:
J (T) (s,a)aa'' (r(s, +
- (s()))o00a)))
0² (s, a) arg max [())
where ʃa is the Fenchel conjugate of f-divergence (please refer to Ma al. [56] for more details). We conduct experiments using the official Pytorch implementation 14 where the f-divergence used X² diverregee On the LID tasks, we change the input of the discriminator state-action pairs.
26


Reviews:
[{"review_id": "Pp61uwC61w", "reviewer": "Reviewer_vVku", "summary": "One recent idea in reinforcement learning is to learn a sequential model that can predict state-action transitions and rewards, and then obtain a strong policy by inferring actions conditioned on high reward. This has the major benefit that even low-reward trajectories provide useful information for learning.\n\nThe challenge in adapting this idea to imitation learning is that in imitation learning, we do not have access to rewards, but rather a dataset of expert demonstrations, plus (potentially) datasets of other suboptimal behavior (e.g. rollouts from a policy in the online setting, static dataset from suboptimal policies in the offline setting).\n\nThe proposed solution is to simply infer / learn a contextual variable $z$ that performs a similar role as reward: that is, it determines the “type” of policy / actions that make up the trajectory. To avoid $z$ from encoding too much information, we make it very limited and regularize it. We can then train a $z$-conditional policy that can produce actions that mimic both the suboptimal behavior as well as the expert demonstrations (using a self-consistency loss). We also learn an optimal setting $z^*$ that selects the high-performing policy, by ensuring that the $z^*$-conditional policy produces trajectories with high similarity to the expert trajectories. There is a significant amount of math to flesh this out, which I won’t get into.\n\nThe authors test their method on four MuJoCo environments, using expert demonstrations from a SAC policy, in both online / offline settings, in both full demos / observation-only settings, and with / without a distributional shift (modifying the torso length) in the test environment, and show that CEIL tends to match or improve upon the results from a variety of baselines.", "questions": "How would you expect CEIL to work in settings other than MuJoCo?\n\nHave you run CEIL with access only to expert demonstrations? How does it perform in that setting?", "limitations": "The paper should mention that it does not tackle the setting in which we only have expert demonstrations (even in the offline setting, the paper assumes access to a static dataset of suboptimal behavior, specifically D4RL in its experiments).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The area of imitation learning is important and relevant, and the approach suggested covers a wide variety of use cases.\n2. The idea proposed is conceptually simple: train a network that models a variety of policies, and then select an appropriate high-performing policy through the use of context variable that controls the policy. Similar ideas have seen significant success in reinforcement learning.\n3. The empirical evaluations are quite extensive with an impressive number of baselines, and show strong performance of the author’s method in a wide variety of settings.\n4. I particularly appreciated Table 1, which provides an excellent overview of related work.\n5. While I found the paper hard to read in an absolute sense, I think that is mostly because the ideas are quite technical: relative to other imitation learning papers focused on expert matching, I found this paper easier to read and understand.\n\nOverall, I recommend accepting the paper. However, I should note that I am not very familiar with the imitation learning literature, and so cannot provide an evaluation of the following aspects:\n1. How novel / original these ideas are\n2. Whether the baselines chosen are appropriate (e.g. perhaps the paper has not compared to current SOTA)\n3. Whether the performance of baselines is in line with expected numbers (e.g. perhaps the authors did not tune hyperparameters of baselines well)", "weaknesses": "1. The empirical evaluations are entirely based on MuJoCo. It is unclear whether the strong performance will generalize to very different settings (e.g. Atari).\n2. Though the paper aims to provide a general and broadly applicable imitation learning algorithm, it doesn’t tackle the most typical imitation learning setting: where we have access only to a dataset of expert demonstrations. (In particular, even in the offline setting, CEIL assumes access to a dataset of suboptimal behavior.) In principle we could run the algorithm in such a setting, though my guess is that it will not perform as well as baseline methods like behavior cloning, since the major benefit of CEIL is in its ability to leverage suboptimal data. (However, the expert-demos-only setting tends to be very vulnerable to spurious correlations / overfitting, and so is not as significant as the other settings.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and thoughtful feedback. We will address your concerns one by one. \n\n**Q1: whether the strong performance will generalize to very different settings (e.g. Atari).**\n \n**A1:** As suggested by the reviewer, we have carried out new experiments on Atari and Adroit domains. We refer the reviewer to our new results (Figure 1) in the PDF file in the general response. We can see that CEIL consistently achieves better or comparable performance compared to baseline methods in both online (Atari) and offline (Adroit) IL tasks. \n\n**Q2: run CEIL with access only to expert demonstrations.**\n\n**A2:** Thank you for the suggestion. In single-domain IL tasks (with only expert demonstrations), CEIL and BC (behavior cloning) are essentially the same, since Equation 5 is all about fitting expert demonstrations. However, if we consider cross-domain IL settings (with expert source and target data), CEIL can effectively use the source domain data to learn the contextual behavioral policy, and then use the expert target data to fit $z*$. Intuitively, although the source data is expert in the source domain, it can also be viewed as sub-optimal data for the target environment, thus CEIL can effectively utilize the cross-domain data to improve the performance. However, simple BC cannot do this. \n\nHere, we also run CEIL with access only to expert-level behaviors in the offline cross-domain IL settings and compare it to three baselines: 1) *BC over target* denotes that we train BC agent only over the target data; 2) *BC over source+target* denotes that we train BC agent over the combined source and target data; 3) *BC Fine-tuning* denotes that we first train BC agent over the source data and then fine-tune the agent over the target data.  \n\n|\t|BC over target\t|BC over source+target\t|BC Fine-tuning\t|CEIL|\n|  ----  | ----  | ----  | ----  | ----  |\n|Hopper\t|46.5 \t|37.8 \t|69.9 \t|**93.7** |\n|Halfcheetah\t|15.4 \t|11.5 \t|44.3 \t|**48.1** |\n|Walker2d\t|97.6 \t|92.5 \t|104.2 \t|**111.8** |\n|Ant\t|72.6 \t|72.0 \t|83.1 \t|**95.0** |\n\nWe can see that both *BC over source+target* and *BC over target* perform poorly (with 5 expert demonstrations in target), and *BC Fine-tuning* can slightly improve the performance. Our CEIL can effectively achieve the best performance. \n\n\nThanks again for all of your constructive suggestions, which have helped us improve the quality and clarity of our paper. \n"}, {"review_id": "zXLTgZZt3U", "reviewer": "Reviewer_PiwE", "summary": "This paper presents a method that aims to address Imitation Learning (IL) tasks by simultaneously updating the embedding function of a contextual variable, an optimal contextual variable, and a policy conditioned on that variable. The proposed method learns the conditional policy by minimizing the trajectory self-consistency loss based on the concept of hindsight information matching. The optimal contextual variable is updated by minimizing the discrepancy between the learned conditional policy and the expert policy. The embedding function is optimized based on both the self-consistency loss and the discrepancy loss. The experimental results demonstrated improved performance in the following tasks: (1) learning from observations (LfO), (2) online/offline IL, (3) cross-domain IL, and (4) one-shot generalization IL.", "questions": "1. As mentioned in Weakness (1), how are the hyperparameters chosen in the experiments? Could the authors offer guidance on the process of selecting the hyperparameters (those in Appendix Table 4)?\n2. As stated in Weakness (2), could the authors provide ablation studies and further clarification on the optimization of $f_{\\phi}$ as described in Eq.(8)?\n3. As mentioned in Weakness (3), could the authors clarify that there exist two regularization losses in this work? If possible, it would be best to include both $\\mathcal{R}(f_{\\phi})$ and $\\mathcal{R}(\\mathbf{z}^*)$ in Eq.(5). Alternatively, the authors could at least mention (in Line 164-165) that $\\mathbf{z}^* \\in  f_{\\phi}\\circ\\mathrm{supp}(\\mathcal{D})$ is achieved through minimizing $\\mathcal{R}(\\mathbf{z}^*)$. It is acceptable to use the symbol $\\mathcal{R}$ to represent both losses, as long as it is straightforward for readers to distinguish between the two.", "limitations": "Minor suggestions:\n- Section 3.1, Line 93: The transition dynamics function should be $\\mathcal{T}$, not $\\mathcal{P}$.\n- Section 3.1, Line 98: The transition dynamics function should be $\\mathcal{T}$, not $T$.\n- Section 3.2, Line 124: Missing citation for hindsight experience replay (HER) [[1]].\n- Algorithm 1, Line 4: The two $\\mathcal{D}$ here can be combined into a single $\\mathcal{D}$ for simplicity.\n\n[1]: https://arxiv.org/abs/1707.01495", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The proposed method is novel, and enables solving a variety of IL tasks with minimal adjustments.\n2. The experiments were conducted on 8 different IL settings and outperformed previous baselines in most environments.\n3. The empirical analysis in Section 5.2 offers interesting insights into the practical application of the proposed method.", "weaknesses": "1. A number of hyperparameters are introduced in this work, such as the embedding dimension, the trajectory window size, the architecture of the encoder/decoder networks. However, the authors do not provide any guidance or recommendations regarding the selection of these hyperparameters.\n2. There appears to be a disparity between the theoretical objective (Eq.(4), Eq.(5)) and the practical objective (Eq.(8), and Line 6 in Algorithm 1) when it comes to optimizing the hindsight embedding function $f_{\\phi}$. Theoretically, the update of $f_{\\phi}$ should solely be based on the trajectory self-consistency loss, as mentioned in Eq.(4). In practice, however, $f_{\\phi}$ is also updated according to Eq.(8).\n3. The notation for the regularization losses is ambiguous. Specifically, there are two regularization losses used in this work, one for regularizing the embedding function $f_{\\phi}$ (Eq.(9)), and another for regularizing the optimal contextual variable $\\mathbf{z}^*$ (Eq.(7)). Unfortunately, both of these losses are represented by the symbol $\\mathcal{R}$, and their definitions are scattered throughout the manuscript, resulting in unnecessary confusion.\n4. The confidence intervals are not reported in Table 2-4.\n5. The expert trajectories generated by SAC used in the experiments have not been provided, which could potentially pose challenges for replicating the results in future studies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and thoughtful feedback. We will address your concerns one by one. \n\n**Q1: guidance on the process of selecting the hyperparameters in the Appendix.**\n\n**A1:** 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. \n\n**Q2: ablation studies on the optimization of $f_\\phi$.**\n\n**A2:** Thank you for such a valuable suggestion. We have carried out new ablation experiments on the loss of $f_\\phi$ in both online IL and offline IL settings (see results, Figure 2 and Table 1, in the PDF file in the \"global\" response). We can see that ablating the $f_\\phi$ loss (optimizing $f_\\phi$ with Equation 5) does degrade the performance in both online and offline IL tasks, demonstrating the effectiveness of optimizing $f_\\phi$ with Equation 8. Intuitively, Equation 8 encourages the embedding function $f_\\phi$ to be task-relevant, and thus we use the expert matching loss to update $f_\\phi$. \n\n**Q3: the notation for the regularization losses.** \n\n**A3:** We appreciate your suggestion. Actually, the support constraint is achieved through minimizing $R(z*)$. There is only one regularization (for cross-domain IL) in this work. We will clarify it in our paper revision. \n\n**Q4: confidence intervals, expert data, and other suggestions.**\n\n**A4:** Thank you for the valuable suggestions. We will incorporate/elaborate on them in our revision. \n\nThanks again for all of your constructive suggestions, which have helped us improve the quality and clarity of our paper. \n\n"}, {"review_id": "HMoMnVgs0w", "reviewer": "Reviewer_7vt9", "summary": "This work proposed ContExtual Imitation Learning (CEIL), a general method that can be applied to multiple settings, including learning from observations (LfO), offline IL, cross-domain IL, and one-shot IL. CEIL incorporates the hindsight information-matching principle within a bi-level expert matching objective, which decouples the learning policy into a contextual policy and an optimal embedding. Empirical analysis demonstrates that CEIL is more sample efficient in online IL and performs well in offline IL settings. ", "questions": "1. What is the relationship between the return $f_R(\\tau)$ in Equation (2) and the hindsight embedding function $f_\\phi(\\tau)$ in Equation (3)?\n2. In Equation (7), what is the motivation for applying stop gradient operation to $f_{\\bar{\\phi}}$? Is there any guarantee that this operation will satisfy the support in Equation (5)?\n3. Why does CEIL perform worse than AIRL in the cross-domain online LfO task (Hopper, Figure 2(d))?\n", "limitations": "1. CEIL lacks explicit exploration bounds; thus, the offline results generally outperform the online results, especially in the LfO setting.\n2. The trajectory self-consistency cannot be applied to cross-embodiment agents once the two embodiments/domains have different state spaces or action spaces. \n\n\n---------------------\nAfter rebuttal\n---------------------\nThanks to the authors for their rebuttal. I'd like to increase my score to 6.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. CEIL is closely related to hindsight information-matching methods. CEIL introduces an additional context variable $z$ to learn a contextual policy $\\pi_\\theta(a|s, z)$ and an optimal contextual variable $z^*$. The idea is to use the learned $z^*$-conditioned policy $\\pi_\\theta(a|s, z^*)$ to recover the expert data. CEIL is learned through a bi-level expert matching objective: explicitly learn a hindsight embedding function in the inner-level optimization, and perform expert matching via inferring an optimal embedding in the outer-level optimization. Such a decoupling procedure enables CEIL to generalize to diverse IL settings. \n2. Unlike the prior hindsight information-matching methods, CEIL does not require explicit handling components such as explicit rewards in online RL and handcrafted target return in offline RL.\n3. CEIL is a scalable method that can be applied to LfD, LfO, offline IL, cross-domain IL, and one-shot IL. CEIL was evaluated in diverse settings in the experiments. The results on four MuJoCo environments show that CEIL achieves better sample efficiency than other baselines in the online IL tasks. Extensive ablations demonstrated that CEIL is not sensitive to the number of demonstrations and window size of trajectory. \n", "weaknesses": "1. It seems that the pre-defined return $f_R(\\tau)$ in Equation (2) is closely related to the hindsight embedding function $f_\\phi(\\tau)$ in Equation (3). However, their relationship has not been discussed. \n2. In the cross-domain online LfO experiments (Hopper, Figure 2(d)), CEIL performs worse than AIRL (state only). There is no explanation for this result. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for the detailed and thoughtful feedback. We will address your concerns one by one. \n\n**Q1: the relationship between the pre-defined return $f_R(\\tau)$ in Equation 2 and $f_\\phi(\\tau)$ in Equation 3.**\n\n**A1:** Compared to the pre-defined return $f_R(\\tau)$, $f_\\phi(\\tau)$ can be seen as a more general formulation representing the hindsight information of a trajectory. $f_\\phi(\\tau)$ could be any function of a trajectory that captures some statistical properties in state-space or trajectory-space, such as sufficient statistics of a distribution, such as mean, variance or higher-order moments [1]. Empirically, we also find that trajectories with high $f_R(\\tau)$ are closer to the encodings of expert trajectories in the latent space (as implied by Figure 1 in the main paper). Thank you for the valuable suggestion, we will add new discussions in our paper revision. \n\n[1] Furuta H, Matsuo Y, Gu S S. Generalized decision transformer for offline hindsight information matching[J]. arXiv preprint arXiv:2111.10364, 2021. \n\n**Q2: a stop gradient operation to $f_{\\bar{\\phi}}$ in Equation 7. The guarantee for the support constraint.** \n\n**A2:** In Equation 7, we are optimizing $z*$, thus we add the stop gradient operator to $f_{\\bar{\\phi}}$. The guarantee for the support constraint comes from BCQ [2]. Intuitively, Equation 7 minimizes the distance of selected $z*$ to the embeddings of the batch data ($\\tau_E$ and $\\tau_D$), forcing $z*$ towards behaving close to a subset of the offline and expert behaviors. The main difference is that BCQ implements offline support constraint over the action space while we restrict it over the latent space. \n\n[2] Fujimoto S, Meger D, Precup D. Off-policy deep reinforcement learning without exploration[C]//International conference on machine learning. PMLR, 2019: 2052-2062. \n\n**Q3: Why does CEIL perform worse than AIRL in the cross-domain online LfO task (Hopper, Figure 2(d))?**\n\n**A3:** As implied by the reviewer, we believe that this is due to the lack of explicit exploration bonus, which causes the agent to stay in the collapsed state region and therefore deteriorates performance. To address this limitation, we have explicitly imposed a lower bound regularization on the policy entropy to encourage exploration, borrowing the idea from SAC and online Decision Transformer [3]. We implement such regularization with the Lagrangian relaxation method. We refer the reviewer to our new results (Figure 3) in the PDF file in the general response, where we find such a simple exploration strategy can effectively bridge the gap between online IL and offline IL, improving the online LfO performance and outperforming online LfO baselines. \n\n[3] Zheng, Qinqing, Amy Zhang, and Aditya Grover. \"Online decision transformer.\" international conference on machine learning. PMLR, 2022. \n\n**Q4: the two embodiments/domains have different state spaces or action spaces.**\n\n**A4:** We appreciate your concern about the state/action space. If cross-domain agents have different action/state spaces, a typical approach is to serialize state/action from different modalities into a flat sequence of tokens [4]. We remark that CEIL is also compatible with such a tokenization approach, thus suitable for IL tasks with different actions/state spaces, yet it is out of the scope of the current work. \n\n[4] Reed S, Zolna K, Parisotto E, et al. A generalist agent[J]. arXiv preprint arXiv:2205.06175, 2022. \n\nThanks again for all of your constructive suggestions, which have helped us improve the quality and clarity of our paper. \n"}, {"review_id": "doasvlCAhD", "reviewer": "Reviewer_2c9A", "summary": "The objective of this work is to create an imitation learning (IL) method that functions effectively across a variety of common IL settings, which include online, offline, Learning from Demonstrations (LFD), Learning from Observation (LFO), and cross-domain. To achieve this, a dual-level expert matching goal is proposed. This involves not only trajectory matching with expert demonstration but also matching the latent variable, z, with the expert in latent space.\n\nContrasting with traditional trajectory matching methods, matching expert trajectories in latent space appears to more accurately replicate expert behavior. Experimental results reinforce the efficacy of this methodology, demonstrating good performance across four different MuJoCo tasks examined in this study.", "questions": "''Offline results generally outperform online results, especially in the LfO setting. ''\n\nQ1: Can you provide more details on this particular issue? Under similar conditions with the same volume of expert demonstration data, the online setting is generally capable of accessing more information. Therefore, it's not quite clear why the proposed method would underperform in an online setting.\n\nQ2: Similar to bi-level learning methods such as GAIL and IRLs, they first learn a discriminator or reward function, followed by policy training based on the learned discriminator or reward function. Frequently, this type of structure tends to struggle with instability and prolonged training durations. Does the proposed bi-level training process encounter similar issues and why?", "limitations": "See weakness and questions.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed method can adjust to five different Interactive Learning (IL) settings with mere minor adjustments in the algorithm, lending the model practical applications in real-world scenarios. \n\n- The document provides an extensive review of existing methodologies across these five IL settings, thereby allowing for a broad understanding of the field. \n\n- Additionally, by comparing the suggested methods against multiple baselines, good results have been demonstrated. \nThe presentation of the paper communicates ideas clearly.\n", "weaknesses": "- The claim that all Online, Offline, LFD, LFO, and cro-domain issues are addressed appears to be overly ambitious. Given that only two MuJoCo tasks were evaluated in an online setting and four in an offline setting, the experimental evidence provided is inadequate to substantiate this claim.\n\n\n- Furthermore, the conducted ablation study does not provide sufficient insights into the effectiveness of the J_MI term, which is designed to align with the expert demonstration in the latent space. To magnify our understanding of J_MI's contribution, an additional ablation baseline CEIL without the J_MI term could be quite useful. If, subsequently, we find that the removal of J_MI leads to a performance reduction of 10%, 30%, or even more, we can then quantify the effectiveness of matching expert demonstrations in the latent space with the specific technique used in this work. Providing such a comparison can significantly enhance the reader's understanding of the extent of performance improvements achieved through the application of J_MI.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and thoughtful feedback. We will address your concerns one by one. \n\n**Q1: the experimental evidence.** \n\n**A1:** We have carried out new experiments in Atari and Adroit domains (see results, Figure 1, in the PDF file in the \"global\" response). We can see that CEIL consistently achieves better or comparable performance compared to baseline methods in both online and offline IL tasks. \n\n**Q2: an additional ablation baseline CEIL without the J_MI term.**\n\n**A2:** Thank you for such a valuable suggestion for improving this paper. We have carried out new ablation experiments on the J_MI term in both online IL and offline IL settings (see results, Figure 2 and Table 1, in the PDF file in the general response). We can see that ablating J_MI does lead to degraded performance, further verifying the effectiveness of our expert matching objective in the latent space. \n\n**Q3: offline results generally outperform online results, especially in the LfO setting.**\n\n**A3:** Yes. In Appendix 2.6 (supplementary material), we also find this limitation in the online LfO setting. We believe this is due to the lack of explicit exploration bonus, which causes the agent to stay in the collapsed state region and therefore deteriorates performance. To address this limitation, here we explicitly impose a lower bound regularization on the policy entropy to encourage exploration, borrowing the idea from SAC and online Decision Transformer [1]. We implement such a regularization with the Lagrangian relaxation method. We refer the reviewer to our new results (Figure 3) in the PDF file in the \"global\" response, where we find such a simple exploration strategy can effectively bridge the gap between online IL and offline IL, improving the online LfO performance and outperforming online IL baselines. \n\n[1] Zheng, Qinqing, Amy Zhang, and Aditya Grover. \"Online decision transformer.\" international conference on machine learning. PMLR, 2022. \n\n**Q4: Does the proposed bi-level training process encounter similar issues (in GAIL and IRLs) and why?**\n\n**A4:** No, it does not. Many existing (GAIL or IRLs) methods are difficult to train in practice due to an adversarial optimization process over reward and policy approximators (biased or high variance gradient estimators). Another challenge in the GAN-style objective is balancing the performance of the generator (policy) and discriminator. A discriminator that achieves very high accuracy can produce relatively uninformative gradients, but a weak discriminator can also hamper the policy’s ability to learn. However, CEIL learns (a contextual) policy by simply performing supervised regression, which is more stable compared to vanilla GAIL and IRLs. Further, to recover the expert behaviors, CEIL's expert matching objective is optimized over the representation space, which does not depend on the accuracy of the policy. This benefit is even more significant in cross-domain IL settings, since we only need a good representation to characterize trajectories.\n\nThanks again for all of your constructive suggestions, which have helped us improve the quality and clarity of our paper. \n"}, {"review_id": "KWV9yAjGXD", "reviewer": "Reviewer_STjS", "summary": "This paper aims to develop a simple and scalable IL method, which is applicable to a wide range of IL settings (e.g., offline/online, LfD/LfO, etc.). To achieve this,  the imitation policy are decoupled into a contextual policy and a latent variable. Experiments on a wide range of tasks illustrate effectiveness of the proposed methods.", "questions": "1. I read the derivation of Equation (6) in the appendix and get confused on the constant $C$. If I am not wrong,\n$$\nC = \\mathbb{E}_{p(z^*)\\pi_E(\\tau)}\\left[\\log \\frac{\\pi_\\theta(\\tau)}{\\tau_E(\\tau)}\\right] = \\mathbb{E}_{p(z^*)\\pi_E(\\tau)}\\left[\\log \\frac{\\mathbb{E}_{p(z^*)}[\\pi_\\theta(\\tau|z^*)]}{\\tau_E(\\tau)}\\right],\n$$\nhow can it be a constant to $z^*$? By the way, what do you mean by $p(z^*)$? $z^*$ is a single variable not a distribution that you plan to optimize, right?\n2. Considering the one-shot IL setting and your claim in line 246-248, I wonder how you guarantee that $f_\\phi$ will generate a proper latent embedding? Will it be trained in a meta-style when meeting one-shot IL tasks?\n\nI am willing to raise my score if you can answer my questions and solve my concerns. Thanks in advance.", "limitations": "This works does not seem to have any negative societal impact. Authors discussed its limitations and planned to leave them for future work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper is well written and easy to follow. I enjoy reading it.\n2. Authors conduct extensive experiments. It is a solid work.\n3. I notice that authors present a detailed implementation details on baselines (like GAIL, GAIfO, et.) in the appendix, which will serve as a good reference for imitation learning researchers.\n4. The proposed framework is simple and compatible with many different IL settings.", "weaknesses": "1. Results of online HalfCheetah seems to be lost.\n2. Adding comparisons with recently popular (decision-)transformer-style methods in their compatible settings (e.g., [1] for offline IL, [2] for offline multi-task IL, [3] for one-shot IL and LfO) will make this work more promising.\n3. The proposed method, CEIL, is not a plug-in method for every IL setting. We still need to make considerate modifications and designations when meeting new tasks. And a lot of works upon the generative adversarial imitation learning (GAIL) framework has been proposed to deal with different IL settings and they work well (e.g., [4] for cross-domain IL, [5] for one-shot IL). A further discussion on the superiority of this work over existing GAIL-style works will be appreciated. \n\n[1] Carroll, Micah, et al. Unimask: Unified inference in sequential decision problems. arXiv preprint arXiv:2211.10869 (2022).\n\n[2] Furuta, Hiroki, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10364. 2021.\n\n[3] Xu, Mengdi, et al. Hyper-decision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08487 (2023).\n\n[4] Franzmeyer, Tim, Philip Torr, and João F. Henriques. Learn what matters: cross-domain imitation learning with task-relevant embeddings. Advances in Neural Information Processing Systems 35 (2022): 26283-26294.\n\n[5] Yu, Lantao, et al. Meta-inverse reinforcement learning with probabilistic context variables. Advances in neural information processing systems 32 (2019).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and thoughtful feedback. We will address your concerns one by one. \n\n**Q1: results of online HalfCheetah.**\n\n**A1:** Thank you for the suggestion. We have provided online HalfCheetah results, Figure 3, and more comparison results in Atari and Adroit domain, Figure 1, in the PDF file in the general response and we will add it in our revision. \n\n**Q2: comparisons with recently popular (decision-)transformer-style methods.**\n\n**A2:** Thank you for the suggestion. We have conducted new comparison against Unimask and GDT (generalized decision transformer). Please refer to the PDF file (Table 2) in the \"global\" response. We can see that CEIL consistently performs better than or comparable to Unimask and GDT. \n\n\n**Q3: the superiority of this work over existing GAIL-style works.**\n\n**A3:** Thank you for raising this concern. By comparing CEIL with the plain expert matching objective $\\min_\\theta D(\\pi_\\theta(\\tau), \\pi_E(\\tau))$ in existing GAIL-style works, we highlight two merits: 1) CEIL’s expert matching loss does not account for updating $\\pi_\\theta$ and is only incentivized to update the low-dimensional latent variable $z*$, which enjoys efficient parameter learning similar to the prompt tuning in large language models, and 2) we learn $\\pi_\\theta$ by simply performing supervised regression, which is more stable compared to vanilla inverse-RL/adversarial-IL (GAIL-style) methods. \n\n**Q4: the constant C in the derivation in the appendix.**\n\n**A4:** Thank you for the careful review. It is a typo. In Equation 6 and its derivation in the appendix, 'D_{KL}' should be 'D', where we defined 'D' as the sum of ‘D_{KL}’ and inverse 'D_{KL}' (see Footnote 3 in Page 5, main paper). Therefore, there is no constant C in the corresponding derivation. We will correct it. Thank you. \n\n**Q5: $z\\*$ is a single variable in the implementation, right?**\n\n**A5:** Yes, it is.\n\n**Q6: Will it be trained in a meta-style when meeting one-shot IL tasks?**\n\n**A6:** In response to your concern, we are actually training in a meta-style. Intuitively, Equation 3 treats each trajectory individually as a separate task (a separate z for each trajectory/task), in the same spirit of meta-learning. As the reviewer points out, our approach can also be adapted to multi-task IL and one-shot IL tasks, benefiting from the separating optimization of $\\pi_\\theta$ and $z*$. \nAs for the generalization, how well the model generalizes depends on the (multi-task) data distribution $D(\\tau)$. \n\nThanks again for all of your constructive suggestions, which have helped us improve the quality and clarity of our paper. \n"}]

Author rebuttal:
Dear reviewers, 

Thank you for all of your constructive suggestions, which have helped us improve the quality of our paper. 

This general response provides a summary of the experimental requirements you suggested. Below, we list the content of each chart presented in the submitted PDF. 
+ Figure 1: New experimental results in Atari and Adroit domains. [Reviewers vVku and 2c9A.]
+ Figure 2 and Table 1: Ablation studies on the optimization of $f_\phi$ and the objective of J_MI. [Reviewers PiwE and 2c9A respectively.]
+ Figure 3: Adding exploration bonus for C-on-LfO tasks. [Reviewers 7vt9 and 2c9A.]
+ Table 2: Comparison with respect to transformer-style methods. [Reviewer STjS.]

Thanks again for reviewing our submission, we do not take it for granted. We hope we have resolved all of your concerns. We are always willing to answer any of your further concerns. 


Meta-review: Dear Authors,

After careful evaluation and rigorous discussions with the reviewers, we recognize that the submitted paper presents a method poised to make a significant contribution to the IL field. The methodology, by updating the embedding function of a contextual variable, an optimal contextual variable, and a policy conditioned on that variable, showcases a novel approach. Particularly praiseworthy are the experimental results, highlighting improved performance across various IL tasks. The paper's breadth, covering eight different IL settings, and its depth, outperforming previous baselines in most environments, reflects its robustness. The empirical analysis provides valuable insights, further strengthening the research's practical significance. Furthermore, the paper's quality in terms of clarity, extensive review, and detailed implementation details in the appendix distinguishes it from other submissions.

Nevertheless, there are aspects of the work that merit attention. One of the most pressing concerns is the introduction of several hyperparameters without clear guidance on their selection. In addition, ambiguities persist in the representation of regularization losses, making it difficult for readers to follow without unnecessary confusion. The disparity between the theoretical and practical objectives, especially concerning the optimization of the hindsight embedding function, raises questions about the methodology's consistency. While the paper claims to address all online, offline, LFD, LFO, and cross-domain issues, the provided experimental evidence seems limited to back this broad claim comprehensively. Insights into the effectiveness of the J_MI term and its alignment with the expert demonstration in the latent space would be instrumental. Similarly, the inclusion of comparisons with recent transformer-style methods and an extensive discussion on CEIL's superiority over existing GAIL-style works would bolster the paper's overall impact.

In light of the above, while the paper demonstrates substantial promise and significant advancements in IL, addressing the aforementioned issues will enhance its value in the community. We hope that the feedback aids in the paper's refinement, making it an essential read for all researchers in the IL domain.

Best regads,

Area Chair
Decision: Accept (poster)