Title: CEIL: Generalized Contextual Imitation Learning
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

In this paper, we present C ont E xtual I mitation L earning (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.

1 Introduction

Imitation learning (IL) allows agents to learn from expert demonstrations. Initially developed with a supervised learning paradigm [58, 63], IL can be extended and reformulated with a general expert matching objective, which aims to generate policies that produce trajectories with low distributional distances to expert demonstrations [30]. This formulation allows IL to be extended to 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 a 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 is observed for a new IL task.

Modern IL algorithms introduce various designs or mathematical principles to cater to the expert matching objective in a specific scenario. For example, the LfO setting requires particular considerations regarding the absent expert actions, e.g. , learning an inverse dynamics function [5, 65]. Besides, out-of-distribution issues in offline IL require specialized modifications to the learning objective, such as introducing additional policy/value regularization [32, 72]. However, such a methodology, designing an individual formulation for each IL setting, makes it difficult to scale up a specific IL 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 a particular IL setting but consist of a mixture of them. For example, we may have access

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-deploy 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 a 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 to retain its generality in IL tasks, how to bias the hindsight towards expert behaviors remains a significant barrier as the extrinsic rewards are missing.

To design a general IL formulation and tackle the above problems, we propose C ont E xtual I mitation L earning (CEIL), which readily incorporates the hindsight information matching principle within a bi-level expert matching objective. In the inner-level optimization, we explicitly learn a hindsight embedding function to deal with the challenges of unknown rewards. In the outer-level optimization, we perform IL expert matching via inferring an optimal embedding ( i.e. , hindsight embedding biasing), replacing the naive reward biasing in hindsight. Intuitively, we find that such a bi-level objective results in a 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 matching 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 a 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 (LfD) to learning from observations (LfO) [7, 9, 35, 54, 62, 66], from online IL to offline IL [11, 15, 33, 53, 73], and from single-domain IL 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 IL method. Indeed, a number of prior works have studied part of the above IL settings, such as offline LfO [78], cross-domain LfO [48, 60], and cross-domain offline IL [56]. While such works demonstrate the feasibility of tackling multiple IL 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, as we strive to 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-matching methods [2, 8, 24, 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 RL [4, 42, 64] or a handcrafted target return in offline RL [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 transformer-based RL/IL algorithms that explicitly model a long sequence of transitions [10, 12, 31, 36, 43, 71]. We find that simple fully-connected networks can also elicit useful embeddings and guide expert behaviors when conditioned on a wellcalibrated 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-tuning, with the main motivation that we tune the prompt (the optimal contextual variable) with an expert matching objective in IL settings.

3 Background

Before discussing our method, we briefly introduce the background for IL, including learning from demonstrations (LfD), learning from observations (LfO), online IL, offline IL, and 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 a discrete-time Markov decision process (MDP) 2 M = {S , A , T , r, γ, p 0 } , where S is the state (observation) space, A is the action space, T : S×A×S → R is the transition dynamics function, r : S×A → R is the reward function, γ is the discount factor, and p 0 is the distribution of initial states. The goal in a reinforcement learning (RL) control task is to learn a policy π θ ( a | s ) maximizing the expected sum of discounted rewards E π θ ( τ ) [ ∑ T -1 t =0 γ t r ( s t , a t ) ] , where τ := { s 0 , a 0 , · · · , s T -1 , a T -1 } denotes the trajectory and the generated trajectory distribution π θ ( τ ) = p 0 ( s 0 ) π θ ( a 0 | s 0 ) ∏ T -1 t =1 π θ ( a t | s t ) T ( s t | s t -1 , a t -1 ) .

In IL, the ground truth reward function ( i.e. , r in M ) is not observed. Instead, we have access to a set of demonstrations (or observations) { τ | τ ∼ π E ( τ ) } that are collected by an unknown expert policy π E ( a | s ) . The goal of IL tasks is to recover a 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:

where D is a distance measure. Meanwhile, we emphasize that the given expert data { τ | τ ∼ π E ( τ ) } 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 { ( s t , a t , s t +1 ) } (learning from demonstrations, LfD), as well as a set of state-only transitions { ( s t , s t +1 ) } (learning from observations, LfO). When it is clear from context, we abuse notation π E ( τ ) to denote both demonstrations in LfD and observations in LfO for simplicity.

Besides, we can also divide IL settings into two orthogonal categories: online IL and offline IL. In online IL, the learning policy π θ can interact with the environment and generate online trajectories τ ∼ π θ ( τ ) . In offline IL, the agent cannot interact with the environment but has access to an offline

static dataset { τ | τ ∼ π β ( τ ) } , collected by some unknown (sub-optimal) behavior policies π β . By leveraging the offline data { π β ( τ ) } ∪ { π E ( τ ) } without any interactions with the environment, the goal of offline IL is to learn a policy recovering the expert behaviors (demonstrations or observations) generated by π E . Note that, in contrast to the typical offline RL problem [46], the offline data { π β ( τ ) } in offline IL does not contains any reward signal.

Cross-domain IL. Beyond the above two IL branches (online/offline and LfD/LfO), 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 in which the learning policy is to be learned, and 2) the cross-domain 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:

where the learner Alg ( π θ ; · , · ) could be any RL methods, τ g ∼ π θ ( τ g | g ) denotes the trajectory generated by a goal-conditioned policy π θ ( a t | s t , g ) , and f HER denotes a pre-defined (hindsight information extraction) function, e.g. , returning the last state in trajectory τ g .

HER can also be applied to the (single-goal) reward-driven online/offline RL tasks, setting the return (sum of the discounted rewards) of a trajectory as an implicit goal for the corresponding trajectory. Thus, we can reformulate the (single-goal) reward-driven RL task , learning policy π θ ( a t | s t ) that maximize the return, as a multi-goal RL task , learning a return-conditioned policy π θ ( a t | s t , · ) that maximize the following log-likelihood:

where f R ( τ ) denotes the return of trajectory τ . At test, we can then condition the contextual policy π θ ( a | s , · ) on a desired target return. In offline RL, the empirical distribution D ( τ ) in Equation 2 can be naturally set as the offline data distribution; in online RL, D ( τ ) can be set as the replay/experience buffer, and will be updated and biased towards trajectories that have high expected returns.

Intuitively, biasing the sampling distribution ( D ( τ ) towards higher returns) leads to an implicit policy improvement operation . However, such an operator is non-trivial to obtain in the IL problem, where we do not have access to a pre-defined function f R ( τ ) to bias the learning policy towards recovering the given expert data { π E ( τ ) } (demonstrations or observations).

4 Method

In this section, we will formulate IL as a bi-level optimization problem, which will allow us to derive our method, contextual imitation learning (CEIL). Instead of attempting to train the learning policy π θ ( a | s ) with the plain expert matching objective (Equation 1), our approach introduces an additional contextual variable z for a contextual IL policy π θ ( a | s , · ) . The main idea of CEIL is to learn a contextual policy π θ ( a | s , z ) and an 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 π θ ( a | s , z ∗ ) . 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 in online/offline RL (Section 3.2), we propose to learn a general hindsight embedding function f ϕ , which encodes trajectory τ (with window size T ) into a latent variable z ∈ Z , |Z| ≪ T ∗ |S| . Formally, we learn the embedding function f ϕ and a corresponding contextual policy π θ ( a | s , z ) by minimizing the trajectory self-consistency loss:

where in the online setting, we sample trajectory τ from buffer D ( τ ) , known as the experience replay buffer in online RL; in the offline setting, we sample trajectory τ directly from the given offline data.

If we can ensure that the learned contextual policy π θ and the embedding function f ϕ are accurate on the empirical data D ( τ ) , then we can convert the IL policy optimization objective (in Equation 1) into a bi-level expert matching objective:

where R ( f ϕ ) is an added regularization over the embedding function (we will elaborate on it later), and supp ( D ) denotes the support of the trajectory distribution { τ |D ( τ ) > 0 } . Here f ϕ is employed to map the trajectory space to the latent variable space ( Z ). Intuitively, by optimizing Equation 4, we expect the induced trajectory distribution of the learned π θ ( 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-matching loss but does not preserve the trajectory self-consistency (Equation 3). Therefore, we formalize CEIL into a bi-level optimization problem, where, in Equation 5, we explicitly constrain the inferred z ∗ lies in the ( f ϕ -mapped) support of the training trajectory distribution.

At a high level, CEIL decouples the learning policy into two parts: an expressive contextual policy π θ ( a | s , · ) and an optimal contextual variable z ∗ . By comparing CEIL with the plain expert matching objective, min π θ D ( π θ ( τ ) , π E ( τ )) , in Equation 1, we highlight two merits: 1) CEIL's expert matching loss (Equation 4) does not account for updating π θ 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 [74], and 2) we learn π θ by simply performing supervised regression (Equation 5), which is more stable compared to vanilla inverse-RL/adversarial-IL 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 3 , i.e. , D ( q, p ) = D KL ( q ∥ p ) + D KL ( p ∥ q ) , and derive an alternative form for Equation 4:

where I ( x ; y ) denotes the mutual information (MI) between x and y , which measures the predictive power of y on x (or vice-versa), the latent variables are defined as τ E := τ ∼ π E ( τ ) , τ θ := τ ∼ p ( z ∗ ) π θ ( τ | z ∗ ) , and π θ ( τ ) = E z ∗ [ π θ ( τ | z ∗ )] .

Intuitively, the second term J D on RHS of Equation 6 is similar to the plain expert matching objective in Equation 1, except that here we optimize a latent variable z ∗ over this objective. Regarding the MI terms J MI, we can interpret the maximization over J MI as an implicit policy improvement, which incentivizes the optimal latent variable z ∗ for having high predictive power of the expert data τ E and having low predictive power of the non-expert data τ θ .

Further, we can rewrite the MI term ( J MI in Equation 6) in terms of the learned embedding function f ϕ , yielding an approximate embedding inference objective J MI ( f ϕ ) :

where we approximate the logarithmic predictive power of z ∗ on τ with -∥ z ∗ -f ϕ ( τ ) ∥ 2 , by taking advantage of the learned embedding function f ϕ in Equation 5.

Algorithm 1 Training CEIL: Online or Offline IL Setting

Require: Expert demonstrations { π E ( τ ) } , 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 π θ ( a | s , · ) , embedding function f ϕ ( z | τ ) , and latent variable z ∗ .

2: for k = 1 , · · · , K do

3: (Online only) Run policy π θ ( a | s , z ∗ ) in environment and store experience into buffer D .

5: Learn π θ and f ϕ over sampled { τ } n 1 using the trajectory self-consistency loss.

4: Sample a batch of data { τ } n 1 from D for online IL or D for offline IL.

6: Update z ∗ and f ϕ over sampled { τ } n 1 by maximizing J MI ( f ϕ ) -α J D .

7: (Offline only) Update z ∗ by minimizing R ( z ∗ ) . # eliminating the offline OOD issues.

8: end for

Return: the learned contextual policy π θ ( a | s , · ) and the optimal latent variable z ∗ .

By maximizing J MI ( f ϕ ) , 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, J MI ( f ϕ ) can also be thought of as an instantiation of contrastive loss, which manifests two facets we consider significant in IL:

1) the "anchor" variable 4 z ∗ is 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 5.

In summary, by comparing J MI ( f ϕ ) and J D , we can observe that J MI ( f ϕ ) actually encourages expert matching in the embedding space, while J D encourages expert matching in the original trajectory space. In the next section, we will see that such an embedding-level expert matching objective naturally lends itself to cross-domain IL settings.

4.3 Practical Implementation

In this section, we describe how we can convert the bi-level IL problem above (Equations 4 and 5) into a feasible online/offline IL objective and discuss some practical implementation details in LfO, offline IL, cross-domain IL, and one-shot IL settings (see more details 5 in Appendix 9.3).

As shown in Algorithm 1 (best viewed in colors), CEIL alternates between solving the bi-level problem with respect to the support constraint (Line 3 for online IL or Line 7 for offline IL), the trajectory self-consistency loss (Line 5), and the optimal embedding inference (Line 6).

To satisfy the support constraint in Equation 5, for online IL (Line 3), we directly roll out the z ∗ -conditioned policy π θ ( a | s , z ∗ ) in the environment; for offline IL (Line 7), we minimize a simple regularization 6 over z ∗ , bearing a close resemblance to the one used in TD3+BC [23]:

where we apply a stop-gradient operation to f ¯ ϕ . To ensure the optimal embedding inference ( max z ∗ J MI ( f ϕ ) - J D ) retaining the flexibility of seeking z ∗ across different instances of f ϕ , we jointly update the optimal embedding z ∗ and the embedding function f ϕ with

where we use α to control the weight on J D .

LfO. In the LfO setting, as expert actions are missing, we apply our expert matching objective only over the observations. Note that even though expert data contains no actions in LfO, we can still

0

100

200

Rollout steps (k)

10

1

10 0

L2 distance

Hopper-v2

0

1k

2k

3k

Return

0

100

200

Rollout steps (k)

10

1

10 0

L2 distance

0

2k

4k

6k

Return

Return

||

z

*

f

(

E

)||

2

||

z

*

f

(

)||

2

leverage a large number of suboptimal actions presented in online/offline D ( τ ) . Thus, we can learn the contextual policy π θ ( a | s , 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 IL. Cross-domain IL considers the case in which the expert's and learning agent's MDPs are different. Due to the domain shift, the plain idea of min J D may not be a sufficient proxy for the expert matching objective, as there may never exist a trajectory (in the learning MDP) that matches the given expert data. Thus, we can set (the weight of J D ) α to 0 .

Further, to make embedding function f ϕ useful for guiding the expert matching in latent space ( i.e. , max J MI ( f ϕ ) ), we encourage f ϕ to capture the task-relevant embeddings and ignore the domainspecific factors. To do so, we generate a set of pseudo-random transitions { τ E ′ } by independently sampling trajectories from expert data { π E ( τ E ) } and adding random noise over these sampled trajectories, i.e. , τ E ′ = τ E + noise. Then, we couple each trajectory τ in { τ E } ∪ { τ E ′ } with a label n ∈ { 0 , 1 } , indicating whether it is noised, and then generate a new set of { ( τ , n ) } , where τ ∈ { τ E } ∪ { τ E ′ } and n ∈ { 0 , 1 } . Thus, we can set the regularization R ( f ϕ ) in Equation 5 to be:

Intuitively, maximizing R ( f ϕ ) encourages embeddings to be domain-agnostic and task-relevant: f ϕ ( τ E ) 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 τ new, we can use the learned embedding function f ϕ to generate a corresponding latent embedding z new. When conditioning on such an embedding, we can directly roll out π θ ( a | s , z new ) to recover the one-shot expert behavior.

5 Experiments

In this section, we conduct experiments across a variety of IL problem domains: single/cross-domain IL, online/offline IL, and LfD/LfO IL settings. By arranging and combining these IL domains, we obtain 8 IL tasks in all: S-on-LfD , S-on-LfO , S-off-LfD , S-off-LfO , C-on-LfD , C-on-LfO , C-off-LfD , and C-off-LfO , where S/C denotes single/cross-domain IL, on/off denotes online/offline IL, and LfD/LfO denote learning from demonstrations/observations respectively. Moreover, we also verify the scalability of CEIL on the challenging one-shot IL setting.

Our experiments are conducted in four popular MuJoCo environments: Hopper-v2 (Hop.), HalfCheetah-v2 (Hal.), Walker2d-v2 (Wal.), and Ant-v2 (Ant.). In the single-domain IL setting, we train a SAC policy in each environment and use the learned expert policy to collect expert trajectories (demonstrations/observations). To investigate the cross-domain IL setting, we assume the two domains (learning MDP and the expert-data collecting MDP) have the same state space and action space, while they have different transition dynamics. To achieve this, we modify the torso length of the MuJoCo agents (see details in Appendix 9.2). Then, for each modified agent, we train a separate expert policy and collect expert trajectories. For the offline IL setting, we directly take the reward-free D4RL [22] as the offline dataset, replacing the online rollout experience in the online IL setting.

5.1 Evaluation Results

To demonstrate the versatility of the CEIL idea, we collect 20 expert trajectories (demonstrations in LfD or state-only observations in LfO) for each environment and compare CEIL to GAIL [30], AIRL [21], SQIL [61], IQ-Learn [28], ValueDICE [41], 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/C-on/off-LfD/LfO ), thus we only provide experimental comparisons with compatible baselines in each IL setting.

Online IL. In Figure 2, we provide the return (cumulative rewards) curves of our method and baselines on 4 online IL settings: S-on-LfD ( top-left ), S-on-LfO ( top-right ), C-on-LfD ( bottom-left ), and Con-LfO ( bottom-right ) settings. As can be seen, CEIL quickly achieves expert-level performance in S-on-LfD . When extended to S-on-LfO , CEIL also yields better sample efficiency compared to baselines. Further, considering the complex cross-domain setting, we can see those baselines SQIL

0

100

200

300

400

500

Rollout steps (k)

0

1.5k

3k

Return

(a)

Hopper-v2

0

100

200

300

400

500

Rollout steps (k)

0

2k

4k

Single-domain online LfD (S-on-LfD)

0

100

200

300

400

500

Rollout steps (k)

0

1.5k

3k

Return

(b)

Hopper-v2

0

100

200

300

400

500

Rollout steps (k)

0

2k

4k

Single-domain online LfO (S-on-LfO)

0

100

200

300

400

500

Rollout steps (k)

0

0.5k

1k

1.5k

Return

(c)

Hopper-v2

0

100

200

300

400

500

Rollout steps (k)

0

2k

4k

Cross-domain online LfD (C-on-LfD)

0

100 200 300 400 500

Rollout steps (k)

0

0.4k

0.8k

1.2k

Return

(d)

Hopper-v2

0

100 200 300 400 500

Rollout steps (k)

0

2k

Cross-domain online LfO (C-on-LfO)

CEIL

AIRL

SQIL

IQ-Learn

ValueDICE

(legend for S/C-on-LfD)

CEIL

AIRL (state only)

GAIfO

IQ-Learn

(legend for S/C-on-LfO)

and IQ-Learn (in C-on-LfD and C-on-LfO ) suffer from the domain mismatch, leading to performance degradation at late stages of training, while CEIL can still achieve robust performance.

Offline IL. Next, we evaluate CEIL on the other 4 offline IL settings: S-off-LfD , S-off-LfO , C-off-LfD , and C-off-LfO . In Table 2, we provide the normalized return of our method and baseline methods on reward-free D4RL [22] medium (m), medium-replay (mr), and medium-expert (me) datasets. We can

0

100

(online)

(a)

1

2

3

4

5

20 15 10 8

0

100

(offline)

Demonstration Num.

Normalized Return

Hopper-v2

0

100

(online)

(b)

1

2

3

4

5

20 15 10 8

0

100

(offline)

Demonstration Num.

Normalized Return

2

4

8

16

32

0

100

Normalized Return

(c)

Hopper-v2

S-on-LfD (CEIL)

S-off-LfD (CEIL)

C-on-LfD (CEIL)

C-off-LfD (CEIL)

2

4

8

16

32

0

100

Normalized Return

(d)

S-on-LfD (CEIL)

S-off-LfD (CEIL)

C-on-LfD (CEIL)

C-off-LfD (CEIL)

CEIL

IQ-Learn

SQIL

ORIL

ValueDICE

SMODICE   (legend for varying Demo Num.)

observe that CEIL achieves a significant improvement over the baseline methods in both S-off-LfD and S-off-LfO settings. Compared to the state-of-the-art offline baselines, CEIL also shows competitive results on the challenging cross-domain offline IL settings ( C-off-LfD and C-off-LfO ).

One-shot IL. Then, we explore CEIL on the one-shot IL tasks, where we expect CEIL can adapt its behavior to new IL tasks given only one trajectory for each task (mismatched MDP, see Appendix 9.2).

We first pre-train an embedding function and a contextual policy in the training domain (online/offline IL), then infer a new contextual variable and evaluate it on the new task. To facilitate comparison to baselines, we similarly pre-train a policy network (using baselines) and run BC on top of the pre-trained policy by using the provided demonstration. Consequently, such a baseline+BC procedure cannot be applied to the (one-shot) LfO tasks. The results in Table 3 show that baseline+BC struggles to transfer their expertise to new tasks. Benefiting from the hindsight framework, CEIL shows better one-shot transfer learning performance on 7 out of 8 one-shot LfD tasks and retains higher scalability and generality for both one-shot LfD and LfO IL tasks.

5.2 Analysis of CEIL

Hybrid IL settings. In real-world, many IL tasks do not correspond to one specific IL setting, and instead consist of a hybrid of several IL settings, each of which passes a portion of task-relevant information to the IL agent. For example, we can provide the agent with both demonstrations and state-only observations and, in some cases, cross-domain demonstrations (S-LfD+S-LfO+C-LfD).

To examine the versatility of CEIL, we collect a separate expert trajectory for each of the four offline IL settings, and study CEIL's performance under hybrid IL settings. As shown in Table 4, we can see that by adding new expert behaviors on top of LfD, even when carrying relatively less supervision ( e.g. , actions are absent in LfO), CEIL can still improve the performance.

Varying the number of demonstrations. In Figure 3 (a, b), we study the effect of the number of expert demonstrations on CEIL's performance. Empirically, we reduce the number of training demonstrations from 20 to 1, and report the normalized returns at 1 M training steps. We can observe that across both online and offline (D4RL *-medium) IL settings, CEIL shows more robust performance with respect to different numbers of demonstrations compared to baseline methods.

Varying the window size of trajectory. Next we assess the effect of the trajectory window size ( i.e. , the length of trajectory τ used for the embedding function f ϕ in Equation 3). In Figure 3 (b, c), we ablate the number of the window size in 4 LfD IL instantiations. We can see that across a range of window sizes, CEIL remains stable and achieves expert-level performance.

0

100

200

300

400

500

Rollout steps (k)

0

1.5k

3k

Return

(a)

Hopper-v2

CEIL

CEIL(ablating

f

)

CEIL(ablating

J

MI

)

0

100

200

300

400

500

Rollout steps (k)

0

2k

4k

CEIL

CEIL(ablating

f

)

CEIL(ablating

J

MI

)

Single-domain online LfD (S-on-LfD)

0

100

200

300

400

500

Rollout steps (k)

0

0.5k

1k

1.5k

Return

(b)

Hopper-v2

CEIL

CEIL(ablating

f

)

CEIL(ablating

J

MI

)

0

100

200

300

400

500

Rollout steps (k)

0

2k

4k

CEIL

CEIL(ablating

f

)

CEIL(ablating

J

MI

)

Cross-domain online LfD (C-on-LfD)

Ablation studies on the optimization of f ϕ and the objective of J MI. In Figure 4 and Table 5, we carried out ablation experiments on the loss of f ϕ and J MI in both online IL and offline IL settings. We can see that ablating the f ϕ loss (optimizing with Equation 5) does degrade the performance in both online and offline IL tasks, demonstrating the effectiveness of optimizing with Equation 8. Intuitively, Equation 8 encourages the embedding function to be task-relevant, and thus we use the expert matching loss to update f ϕ . We can also see that ablating J MI does lead to degraded performance, further verifying the effectiveness of our expert matching objective in the latent space.

6 Conclusion

In this paper, we present CEIL, a novel and general Imitation Learning framework applicable to a wide range of IL settings, including C/S-on/off-LfD/LfO and few-shot IL settings. This is achieved by explicitly decoupling the imitation policy into 1) a contextual policy, learned with the self-supervised hindsight information matching objective, and 2) a latent variable, inferred by performing the IL expert matching objective. Compared to prior baselines, our results show that CEIL is more sample-efficient in most of the online IL tasks and achieves better or competitive performances in offline tasks.

Limitations and future work. Our primary aim behind this work is to develop a simple and scalable IL method. We believe that CEIL makes an important step in that direction. Admittedly, we also find some limitations of CEIL: 1) Offline results generally outperform online results, especially in the LfO setting. The main reason is that CEIL lacks explicit exploration bounds, thus future work could explore the exploration ability of online CEIL. 2) The trajectory self-consistency cannot be applied to cross-embodiment agents once the two embodiments/domains have different state spaces or action spaces. Considering such a cross-embodiment setting, a typical approach is to serialize state/action from different modalities into a flat sequence of tokens. We also remark that CEIL is compatible with such a tokenization approach, and thus suitable for IL tasks with different action/state spaces. Thus, we encourage the future exploration of generalized IL methods across different embodiments.

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. 2022ZD0208800), and NSFC General Program (Grant No. 62176215).

References

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems , 34:29304-29320, 2021.

Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. Advances in neural information processing systems , 30, 2017.

Kai Arulkumaran, Dylan R Ashley, Jürgen Schmidhuber, and Rupesh K Srivastava. All you need is supervised learning: From imitation learning to meta-rl with upside down rl. arXiv preprint arXiv:2202.11960 , 2022.

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.

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:1801.04062 , 2018.

Damian Boborzi, Christoph-Nikolas Straehle, Jens S Buchner, and Lars Mikelsons. Imitation learning by state-only distribution matching. arXiv preprint arXiv:2202.04332 , 2022.

David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? arXiv preprint arXiv:2206.01079 , 2022.

Daniel S Brown, Wonjoon Goo, and Scott Niekum. Better-than-demonstrator imitation learning via automatically-ranked demonstrations. In Conference on robot learning , pages 330-359. PMLR, 2020.

Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu, Mingfei Sun, David Bignell, Stephanie Milani, Katja Hofmann, Matthew Hausknecht, Anca Dragan, et al. Unimask: Unified inference in sequential decision problems. arXiv preprint arXiv:2211.10869 , 2022.

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.

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:15084-15097, 2021.

Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning. arXiv preprint arXiv:2006.04678 , 2020.

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.

Branton DeMoss, Paul Duckworth, Nick Hawes, and Ingmar Posner. Ditto: Offline imitation learning with world models. arXiv preprint arXiv:2302.03086 , 2023.

Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAI Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. Advances in neural information processing systems , 30, 2017.

Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.10751 , 2021.

Arnaud Fickinger, Samuel Cohen, Stuart Russell, and Brandon Amos. Cross-domain imitation learning via optimal transport. arXiv preprint arXiv:2110.03684 , 2021.

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. PMLR, 2022.

Tim Franzmeyer, Philip HS Torr, and João F Henriques. Learn what matters: cross-domain imitation learning with task-relevant embeddings. arXiv preprint arXiv:2209.12093 , 2022.

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. arXiv preprint arXiv:1710.11248 , 2017.

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.

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems , 34:20132-20145, 2021.

Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10364 , 2021.

Tanmay Gangwani and Jian Peng. State-only imitation with transition dynamics mismatch. arXiv preprint arXiv:2002.11879 , 2020.

Tanmay Gangwani, Yuan Zhou, and Jian Peng. Imitation learning from observations under transition model disparity. arXiv preprint arXiv:2204.11446 , 2022.

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 arXiv:2110.04544 , 2021.

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.

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.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems , 29, 2016.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:12731286, 2021.

Firas Jarboui and Vianney Perchet. Offline inverse reinforcement learning. arXiv preprint arXiv:2106.05068 , 2021.

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.

Shengyi Jiang, Jingcheng Pang, and Yang Yu. Offline imitation learning with a misspecified simulator. Advances in neural information processing systems , 33:8510-8520, 2020.

Kshitij Judah, Alan Fern, Prasad Tadepalli, and Robby Goetschalckx. Imitation learning with demonstrations and shaping rewards. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 28, 2014.

Yachen Kang, Diyuan Shi, Jinxin Liu, Li He, and Donglin Wang. Beyond reward: Offline preference-guided policy optimization. arXiv preprint arXiv:2305.16217 , 2023.

Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srinivasa. Imitation learning as f-divergence minimization. In International Workshop on the Algorithmic Foundations of Robotics , pages 313-329. Springer, 2020.

Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srinivasa. Imitation learning as f-divergence minimization. In Algorithmic Foundations of Robotics XIV: Proceedings of the Fourteenth Workshop on the Algorithmic Foundations of Robotics 14 , pages 313-329. Springer International Publishing, 2021.

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 on Learning Representations , 2022.

Kuno Kim, Yihong Gu, Jiaming Song, Shengjia Zhao, and Stefano Ermon. Domain adaptive imitation learning. In International Conference on Machine Learning , pages 5286-5295. PMLR, 2020.

Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. arXiv preprint arXiv:1912.05032 , 2019.

Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-conditioned policies. arXiv preprint arXiv:1912.13465 , 2019.

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: //openreview.net/pdf?id=j0miEWtw87 .

Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J Lim. Generalizable imitation learning from observation via inferring goal proximity. Advances in neural information processing systems , 34:16118-16130, 2021.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.08691 , 2021.

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.

Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. Advances in Neural Information Processing Systems , 30, 2017.

Fangchen Liu, Zhan Ling, Tongzhou Mu, and Hao Su. State alignment-based imitation learning. arXiv preprint arXiv:1911.10947 , 2019.

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.

Jinxin Liu, Hongyin Zhang, and Donglin Wang. Dara: Dynamics-aware reward augmentation in offline reinforcement learning. arXiv preprint arXiv:2203.06662 , 2022.

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:2306.12755 , 2023.

Minghuan Liu, Tairan He, Minkai Xu, and Weinan Zhang. Energy-based imitation learning. arXiv preprint arXiv:2004.09395 , 2020.

Minghuan Liu, Hanye Zhao, Zhengyu Yang, Jian Shen, Weinan Zhang, Li Zhao, and Tie-Yan Liu. Curriculum offline imitating learning. Advances in Neural Information Processing Systems , 34:6266-6277, 2021.

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.

Yicheng Luo, Zhengyao Jiang, Samuel Cohen, Edward Grefenstette, and Marc Peter Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303.13971 , 2023.

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.

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.

Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation , 3(1):88-97, 1991.

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.

Dripta S Raychaudhuri, Sujoy Paul, Jeroen Vanbaar, and Amit K Roy-Chowdhury. Crossdomain imitation from observations. In International Conference on Machine Learning , pages 8902-8912. PMLR, 2021.

Siddharth Reddy, Anca D Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. arXiv preprint arXiv:1905.11108 , 2019.

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 international conference on artificial intelligence and statistics , pages 627-635. JMLR Workshop and Conference Proceedings, 2011.

Stefan Schaal. Learning from demonstration. Advances in neural information processing systems , 9, 1996.

Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Ja´ skowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv:1912.02877 , 2019.

Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. arXiv preprint arXiv:1805.01954 , 2018.

Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation. arXiv preprint arXiv:1807.06158 , 2018.

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems , 30, 2017.

Luca Viano, Yu-Ting Huang, Parameswaran Kamalaruban, Craig Innes, Subramanian Ramamoorthy, and Adrian Weller. Robust learning from observation with model misspecification. arXiv preprint arXiv:2202.06003 , 2022.

Tianyu Wang, Nikhil Karnwal, and Nikolay Atanasov. Latent policies for adversarial imitation learning. arXiv preprint arXiv:2206.11299 , 2022.

Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning , pages 24725-24742. PMLR, 2022.

Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In International Conference on Machine Learning , pages 24631-24645. PMLR, 2022.

Sheng Yue, Guanbo Wang, Wei Shao, Zhaofeng Zhang, Sen Lin, Ju Ren, and Junshan Zhang. Clare: Conservative model-based reward learning for offline inverse reinforcement learning. arXiv preprint arXiv:2302.04782 , 2023.

Wenjia Zhang, Haoran Xu, Haoyi Niu, Peng Cheng, Ming Li, Heming Zhang, Guyue Zhou, and Xianyuan Zhan. Discriminator-guided model-based offline imitation learning. In Conference on Robot Learning , pages 1266-1276. PMLR, 2023.

Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu. Learning to prompt for vision-language models. International Journal of Computer Vision , 130(9):2337-2348, 2022.

Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off-policy imitation learning from observations. Advances in Neural Information Processing Systems , 33:12402-12413, 2020.

Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo. Behavior proximal policy optimization. arXiv preprint arXiv:2302.11312 , 2023.

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.

Konrad Zolna, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcehre, Ziyu Wang, Yusuf Aytar, Misha Denil, Nando de Freitas, and Scott Reed. Offline learning from demonstrations and unlabeled experience. arXiv preprint arXiv:2011.13885 , 2020.

Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. In Conference on Robot Learning , pages 247-263. PMLR, 2021.

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 D ( · , · ) as the sum of reverse KL and forward KL divergence, i.e. , D ( q, p ) = D KL ( q ∥ p ) + D KL ( p ∥ q ) , and derive an alternative form for Equation 4:

where I ( x ; y ) denotes the mutual information (MI) between x and y , which measures the predictive power of y on x (or vice-versa), the latent variables are defined as τ E := τ ∼ π E ( τ ) , τ θ := τ ∼ p ( z ∗ ) π θ ( τ | z ∗ ) , and π θ ( τ ) = E z ∗ [ π θ ( τ | z ∗ )] .

Below is our derivation:

8 More Comparisons and Ablation Studies

8.1 Offline Comparison on D4RL Expert Domain Dataset

In Table 6, we provide the normalized return of our method and baseline methods on the reward-free D4RL [22] expert dataset. Consistently, we can observe that CEIL achieves a significant improvement over the baseline methods in both S-off-LfD and S-off-LfO settings. Compared to the state-of-the-art offline IL baselines, CEIL also shows competitive results on the challenging cross-domain offline IL settings ( C-off-LfD and C-off-LfO ).

8.2 Generalizability on Cross-domain Offline IL Settings

In the standard cross-domain IL setting, the goal is to extract expert-relevant information from the mismatched expert demonstrations/observations (expert domain) and to mimic such expert behaviors in the training environment (training domain). Thus, we validate the performance of the learned policy in the training environment ( i.e. , the environment where the offline data was collected). Here, we also study the generalizability of the learned policy by evaluating the learned policy in the expert environment ( i.e. , the environment where the mismatched expert data was collected). We provide the normalized scores ( evaluated in the expert domain ) in Table 7. We can find that across a range of cross-domain offline IL tasks, CEIL consistently demonstrates better (zero-shot) generalizability compared to baselines.

8.3 Ablating the Cross-domain Regularization

We now conduct ablation studies to evaluate the importance of cross-domain regularization in Equation 9 (in the main paper). In Figure 5, we provide the performance improvement when we

hop-m

hop-mr

hop-me

hop-e

hal-m

hal-mr

hal-me

hal-e

wal-m

wal-mr

wal-me

wal-e

ant-m

ant-mr

ant-me

ant-e

6

4

2

0

2

Improvement

C-off-LfD

hop-m

hop-mr

hop-me

hop-e

hal-m

hal-mr

hal-me

hal-e

wal-m

wal-mr

wal-me

wal-e

ant-m

ant-mr

ant-me

ant-e

4

3

2

1

0

1

Improvement

C-off-LfO

30

60

90

normalized return

CEIL

SMODICE

DemoDICE

ValueDICE

ORIL(TD3+BC)

SQIL(TD3+BC)

IQ-Learn

Median

25

50

75

100

normalized return

IQM

25

50

75

100

normalized return

Mean

20

40

60

80

normalized return

(a) S-off-LfD

Optimality Gap

45

60

75

normalized return

CEIL

SMODICE

ORIL(TD3+BC)

Median

45

60

75

normalized return

IQM

50

60

70

normalized return

Mean

30

40

50

60

normalized return

(b) S-off-LfO

Optimality Gap

25

50

75

normalized return

CEIL

SMODICE

DemoDICE

ValueDICE

ORIL(TD3+BC)

SQIL(TD3+BC)

IQ-Learn

Median

20

40

60

80

normalized return

IQM

20

40

60

80

normalized return

Mean

40

60

80

normalized return

(c) C-off-LfD

Optimality Gap

20

40

60

80

normalized return

CEIL

SMODICE

ORIL(TD3+BC)

Median

30

45

60

normalized return

IQM

30

45

60

normalized return

Mean

45

60

normalized return

(d) C-off-LfO

Optimality Gap

ablate the cross-domain regularization in two cross-domain offline IL tasks ( C-off-LfD and C-off-LfO ). We can find that in 26 out of 32 cross-domain tasks, ablating the regularization can cause performance to decrease (negative performance improvement), thus verifying the benefits of encouraging taskrelevant embeddings.

0

100 200 300 400 500

Rollout steps (k)

0

2000

4000

Return

S-on-LfD

0

100 200 300 400 500

Rollout steps (k)

0

2000

4000

Return

C-on-LfD

0

100 200 300 400 500

Rollout steps (k)

0

250

500

750

1000

Return

S-on-LfO

0

100 200 300 400 500

Rollout steps (k)

0

250

500

750

1000

Return

C-on-LfO

CEIL

AIRL

SQIL

IQ-Learn

ValueDICE

(legend for S/C-on-LfD)

CEIL

AIRL (state only)

GAIfO

IQ-Learn

(legend for S/C-on-LfO)

8.4 Aggregate Results

According to Agarwal et al. [1], we report the aggregate statistics (for 16 offline IL tasks) in Figure 6. We can find that CEIL provides competitive performance consistently across a range of offline IL settings ( S-off-LfD , S-off-LfO , C-off-LfD , and C-off-LfO ) and outperforms prior offline baselines.

8.5 Varying the Number of Expert Trajectories

As a complement to the experimental results in the main paper, we continue to compare the performance of CEIL and baselines on more tasks when we vary the number of expert trajectories. Considering offline IL settings, we provide the results in Table 8 for the number of expert trajectories of 5, 10, 15, and 20 respectively. We can find that when varying the number of expert behaviors, CEIL can still obtain higher scores compared to baselines, which is consistent with the findings in Figure 3 in the main paper.

8.6 Limitation (Failure Modes in Online LfO Setting)

Meanwhile, we find that in the online LfO settings, CEIL's performance deteriorates severely on a few tasks, as shown in Figure 7 (Walker2d). In LfD (either on single-domain or on cross-domain IL) settings, CEIL can consistently achieve expert-level performance, but when migrating to LfO settings, CEIL suffers collapsing performance under the same number of environmental interactions. We believe that this is due to the lack of expert actions in LfO settings, which causes the agent to stay in the collapsed state region and therefore deteriorates performance. Thus, we believe a rich direction for future research is to explore the online exploration ability.

9 Implementation Details

9.1 Imitation Learning Tasks

In our paper, we conduct experiments across a variety of IL problem domains: single/cross-domain IL, online/offline IL, and LfD/LfO IL settings. By arranging and combining these IL domains, we obtain 8 IL tasks in all: S-on-LfD , S-on-LfO , S-off-LfD , S-off-LfO , C-on-LfD , C-on-LfO , C-off-LfD , and C-off-LfO , where S/C denotes single/cross-domain IL, on/off denotes online/offline IL, and LfD/LfO denote learning from demonstrations/observations respectively.

S-on-LfD . We have access to a limited number of expert demonstrations and an online interactive training environment. The goal of S-on-LfD is to learn an optimal policy that mimics the provided demonstrations in the training environment.

S-on-LfO . We have access to a limited number of expert observations (state-only demonstrations) and an online interactive training environment. The goal of S-on-LfO is to learn an optimal policy that mimics the provided observations in the training environment.

S-off-LfD . We have access to a limited number of expert demonstrations and a large amount of pre-collected offline (reward-free) data. The goal of S-off-LfD is to learn an optimal policy that mimics the provided demonstrations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert demonstrations and the environment that was used to collect the offline data are the same environment.

S-off-LfO . We have access to a limited number of expert observations and a large amount of precollected offline (reward-free) data. The goal of S-off-LfO is to learn an optimal policy that mimics the provided observations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert observations and the environment that was used to collect the offline data are the same environment.

C-on-LfD . We have access to a limited number of expert demonstrations and an online interactive training environment. The goal of C-on-LfD is to learn an optimal policy that mimics the provided demonstrations in the training environment. Note that here the environment that was used to collect the expert demonstrations and the online training environment are not the same environment .

C-on-LfO . We have access to a limited number of expert observations (state-only demonstrations) and an online interactive training environment. The goal of C-on-LfO is to learn an optimal policy that mimics the provided observations in the training environment. Note that here the environment that was used to collect the expert observations and the online training environment are not the same environment .

C-off-LfD . We have access to a limited number of expert demonstrations and a large amount of pre-collected offline (reward-free) data. The goal of C-off-LfD is to learn an optimal policy that mimics the provided demonstrations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert demonstrations and the environment that was used to collect the offline data are not the same environment .

C-off-LfO . We have access to a limited number of expert observations and a large amount of precollected offline (reward-free) data. The goal of C-off-LfO is to learn an optimal policy that mimics the provided observations in the environment in which the offline data was collected. Note that here the environment that was used to collect the expert observations and the environment that was used to collect the offline data are not the same environment .

9.2 Online IL Environments, Offline IL Datasets, and One-shot tasks

Our experiments are conducted in four popular MuJoCo environments (Figure 8): Hopper-v2, HalfCheetah-v2, Walker2d-v2, and Ant-v2. For offline IL tasks, we take the standard (reward-free) D4RL dataset [22] (medium, medium-replay, medium-expert, and expert domains) as the offline dataset. For cross-domain (online/offline) IL tasks, we collect the expert behaviors (demonstrations or observations) on a modified MuJoCo environment. Specifically, we change the height of the agent's torso (as shown in Figure 8). We refer the reader to our code submission, which includes our modified MuJoCo assets. For one-shot IL tasks, we train the policy only in the single-domain IL settings ( S-on-LfD , S-on-LfO , S-off-LfD , and S-off-LfO ). Then we collect only one expert trajectory in the modified MuJoCo environment, and roll out the fine-tuned/inferred policy in the modified environment to test the one-shot performance.

Collecting expert behaviors. In our implementation, we use the publicly available rlkit 7 implementation of SAC to learn an expert policy and use the learned policy to collect expert behaviors (demonstrations in LfD or observations in LfO).

9.3 CEIL Implementation Details

Trajectory self-consistency loss. To learn the embedding function f ϕ and a corresponding contextual policy π θ ( a | s , z ) , we minimize the following trajectory self-consistency loss:

where τ 1: T denotes a trajectory segment with window size of T . In the online setting, we sample trajectory τ from the experience replay buffer D ( τ ) ; in the offline setting, we sample trajectory τ directly from the given offline data D ( τ ) . Meanwhile, if we can access the expert actions ( i.e. , LfD settings), we also incorporate the expert demonstrations into the empirical expectation ( i.e. , storing the expert demonstrations into the online/offline experience D ( τ ) ).

In our implementation, we use a 4-layer MLP (with ReLU activation) to encode the trajectory τ 1: T and a 4-layer MLP (with ReLU activation ) to predict the action respectively. To regularize the learning of the encoder function f ϕ , we additionally introduce a decoder network (4-layer MLP with ReLU activation) π ′ θ ( s ′ | s , f ϕ ( τ 1: T )) to predict the next states: min π ′ θ ,f ϕ -E τ 1: T ∼D ( τ 1: T ) E ( s , a , s ′ ) ∼ τ 1: T [log π ′ θ ( s ′ | s , f ϕ ( τ 1: T ))] . Further, to circumvent issues of "posterior collapse" [67], we encourage learning quantized latent embeddings. In a similar spirit to VQ-VAE [67], we incorporate ideas from vector quantization (VQ) and introduce the following regularization: min f ϕ || sg [ z e ( τ 1: T )] -e || 2 + || z e ( τ 1: T ) -sg [ e ] || 2 , where e is a dictionary of vector quantization embeddings (we set the size of this embedding dictionary to be 4096), z e ( τ 1: T ) is defined as the nearest dictionary embedding to f ϕ ( τ 1: T ) , and sg [ · ] denotes the stop-gradient operator.

Out-level embedding inference. In Section 4.2 (main paper), we approximate J MI with J MI ( f ϕ ) ≜ E p ( z ∗ ) π E ( τ E ) π θ ( τ θ | z ∗ ) [ -∥ z ∗ -f ϕ ( τ E ) ∥ 2 + ∥ z ∗ -f ϕ ( τ θ ) ∥ 2 ] , where we replace the mutual information with -∥ z ∗ -f ϕ ( τ ) ∥ 2 by leveraging the learned embedding function f ϕ . Empirically, we find that we can ignore the second loss ∥ z ∗ -f ϕ ( τ θ ) ∥ 2 , and directly conduct outer-level embedding inference with max z ∗ ,f ϕ E p ( z ∗ ) π E ( τ E ) [ -∥ z ∗ -f ϕ ( τ E ) ∥ 2 ] . Meanwhile, this simplification makes the support constraints ( R ( z ∗ ) in Equation 7 in the main paper) for the offline OOD issues naturally satisfied, since max z ∗ E p ( z ∗ ) π E ( τ E ) [ -∥ z ∗ -f ϕ ( τ E ) ∥ 2 ] and min z ∗ R ( z ∗ ) are equivalent.

Cross-domain IL regularization. To encourage f ϕ to capture the task-relevant embeddings and ignore the domain-specific factors, we set the regularization R ( f ϕ ) in Equation 5 to be:

R ( f ϕ ) = I ( f ϕ ( τ ); n ) , where we couple each trajectory τ in { τ E } ∪ { τ E ′ } with a label n ∈ { 0 , 1 } , indicating whether it is noised. In our implementation, we apply MINE [6] to estimate the mutual information and conduct encoder regularization. Specifically, we estimate I ( z ; n ) with ˆ I ( z ; n ) := sup δ E p ( z , n ) [ f δ ( z , n )] -log E p ( z ) p ( n ) [exp ( f δ ( z , n ))] and regularize the encoder f ϕ with max f ϕ ˆ I ( f ϕ ( τ ); n ) , where we model f δ with a 4-layer MLP (using ReLU activations).

Hyper-parameters. In Table 9, we list the hyper-parameters used in the experiments. For the size of the embedding dictionary, we selected it from a range of [512, 1024, 2048, 4096]. We found 4096 to almost uniformly attain good performance across IL tasks, thus selecting it as the default. For the size of the embedding dimension, we tried four values [4, 8, 16, 32] and selected 16 as the default. For the trajectory window size, we tried five values [2, 4, 8, 16, 32] but we did not observe a significant difference in performance across these values. Thus we selected 2 as the default value. For the learning rate scheduler, we tried the default Pytorch scheduler and CosineAnnealingWarmRestarts, and found CosineAnnealingWarmRestarts enables better results (thus we selected it). For other hyperparameters, they are consistent with the default values of most RL implementations, e.g. learning rate 3e-4 and the MLP policy.

9.4 Baselines Implementation Details

We summarize our 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 a 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

objective below, the policy is optimized via an RL algorithm to match the expert occupancy measure (minimize the objective below):

We used the implementation by Gleave et al. [29] on the GitHub page 8 , where there are two modifications introduced with respect to the original paper: 1) a higher output of the discriminator represents better, 2) PPO is used to optimize the policy instead of TRPO.

Generative Adversarial Imitation from Observations (GAIfO). GAIfO [66] is an online LfO 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:

Based on the implementation of GAIL, we implement GAIfO by changing the input of the discriminator to state transitions.

Adversarial Inverse Reinforcement Learning (AIRL). AIRL [21] is an online LfD/LfO method using an adversarial learning framework similar to GAIL. It modifies the form of the discriminator to explicitly disentangle the task-relevant information from the transition dynamics. To make the policy more generalized and less sensitive to dynamics, AIRL proposes to learn a parameterized reward function using the output of the discriminator:

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] is 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 a state planner and an inverse dynamics model, 2) globally, a Wasserstein divergence of state occupancy between the policy and the expert. The policy is optimized using:

where D is a state-based discriminator trained via J ( D ) = E π E [ D ( s )] -E π [ D ( s )] , f is the pretrained VAE-based state planner, and g inv is the inverse dynamics model trained by supervised regression.

In the online setting, we use the official implementation published by the authors 9 , where SAIL is optimized using PPO with the reward definition: r ( s t , s t +1 ) = 1 T [ D ( s t +1 ) -E π E ( s ) D ( s ) ] . 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-LfD ), our early experimental results find that SAIL(TD3+BC) with heavy hyperparameter tuning failed on the offline setting. This indicates that SAIL is rather sensitive to 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 LfD IL algorithm that is easy to implement with both online and offline Q-learning algorithms. The main

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:

The overall objective of the Q-function is to maximize the following objective:

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 an 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 as unlabelled data that could potentially include both successful and failed trajectories. More specifically, ORIL aims to train a reward function to distinguish between the expert and the suboptimal data without explicitly knowing the negative labels. By incorporating Positive-unlabeled learning (PU-learning), the objective of the reward model can be written as follows (for the LfD setting):

where η 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 is 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 a state-only reward function:

Inverse soft-Q learning (IQ-Learn). IQ-Learn [28] is an IRL-based method that can solve IL tasks in the online/offline and LfD/LfO settings. It proposes to directly learn a 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:

where V π ( s ) ≜ E a ∼ π ( ·| s ) [ Q ( s, a ) -log π ( a | s )] , ψ ( r ) is a 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-LfD 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:

In the offline setting, we find that using the above objective easily leads to an overfitting issue, causing collapsed performance. Thus, we follow the instruction provided in the paper and only penalize the expert samples:

Imitation Learning via Off-Policy Distribution Matching (ValueDICE). ValueDICE [41] is a DICE-based 11 LfD algorithm which minimizes the divergence of state-action distributions between the policy and the expert. In contrast to the state-conditional distribution of actions π ( ·| s ) used in the

above methods, the state-action distribution, d π ( s, a ) : S × A → [0 , 1] , can uniquely characterize a one-to-one correspondence,

Thus, the plain expert matching objective can be reformulated and expressed in the Donsker-Varadhan representation:

The objective above can be expanded further by defining x ( s, a ) = v ( s, a ) -B π v ( s, a ) and using a zero-reward Bellman operator B π to derive the following (adversarial) objective:

We use the official Tensorflow implementation 12 in our experiments. In the online setting, the rollouts collected are used as an additional replay regularization. The overall objective in the online setting is:

where d mix ≜ (1 -α ) d π E + αd RB and α is a non-negative regularization coefficient (we set α as 0.1 following the specification of the paper).

In the offline setting, ValueDICE only differs in the source of sampling data. We change the online replay buffer to the offline pre-collected dataset.

Offline Imitation Learning with Supplementary Imperfect Demonstrations (DemoDICE). DemoDICE [39] is a DICE-based offline LfD method that assumes to have access to an offline dataset collected by a behavior policy π β . Using this supplementary dataset, the expert matching objective of DemoDICE is instantiated over ValueDICE:

where α is a 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 a binary discriminator D : S × A → [0 , 1] :

2) a value function optimization objective:

and 3) a policy optimization step:

We report the offline results using the official Tensorflow implementation 13 .

State Matching Offline DIstribution Correction Estimation (SMODICE). SMODICE [56] proposes to solve offline IL tasks in LfO and cross-domain settings and it optimizes the following state occupancy objective:

To incorporate the offline dataset, SMODICE derives an f-divergence regularized state-occupancy objective:

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 stateaction occupancy. Similarly, we can divide the objective into three steps: 1) deriving a state-based reward by learning a state-based discriminator:

2) learning a value function using the learned reward:

and 3) training the policy via weighted regression:

where f ∗ is the Fenchel conjugate of f-divergence (please refer to Ma et al. [56] for more details).

We conduct experiments using the official Pytorch implementation 14 , where the f-divergence used is X 2 -divergence. On the LfD tasks, we change the input of the discriminator to state-action pairs.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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).\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThis works does not seem to have any negative societal impact. Authors discussed its limitations and planned to leave them for future work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n''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?\n\nLIMITATIONS:\nSee weakness and questions.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\n1. 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nMinor 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", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nOne 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.\n\nSTRENGTHS:\n1. 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)\n\nWEAKNESSES:\n1. 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.)\n\nQUESTIONS:\nHow 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?\n\nLIMITATIONS:\nThe 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).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}