Beyond-Expert Performance with Limited Demonstrations: Efficient Imitation Learning with Double Exploration
Authors: Heyang Zhao
Xingrui Yu
David Mark Bossens
Ivor Tsang
Quanquan Gu
Venue: ICLR 2025 Poster (2025)
Keywords: Optimistic exploration
Curiosity
Imitation learning
Inverse reinforcement learning
Reinforcement learning
Primary area: reinforcement learning

Abstract: Imitation learning is a central problem in reinforcement learning where the goal is to learn a policy that mimics the expert's behavior. In practice, it is often challenging to learn the expert policy from a limited number of demonstrations accurately due to the complexity of the state space. Moreover, it is essential to explore the environment and collect data to achieve beyond-expert performance. To overcome these challenges, we propose a novel imitation learning algorithm called Imitation Learning with Double Exploration (ILDE), which implements exploration in two aspects: (1) optimistic policy optimization via an exploration bonus that rewards state-action pairs with high uncertainty to potentially improve the convergence to the expert policy, and (2) curiosity-driven exploration of the states that deviate from the demonstration trajectories to potentially yield beyond-expert performance. Empirically, we demonstrate that ILDE outperforms the state-of-the-art imitation learning algorithms in terms of sample efficiency and achieves beyond-expert performance on Atari and MuJoCo tasks with fewer demonstrations than in previous work. We also provide a theoretical justification of ILDE as an uncertainty-regularized policy optimization method with optimistic exploration, leading to a regret growing sublinearly in the number of episodes.

Published as a conference paper at ICLR 2025
BEYOND-EXPERT PERFORMANCE WITH LIMITED DEMONSTRATIONS: EFFICIENT IMITATION LEARNING
WITH DOUBLE EXPLORATION
Heyang Zhao' Xingrui Yu David M. Bossens Ivor W. Tsang 23t, Quanquan Gu
"Department of Computer Science. University of California, Los Angeles 2THPC, Agency for Science, Technology and Research, Singapore 8CFAR, Agency for Science, Technology and Research, Singapore (hyzhao, @gu)ecs ucla.edi (yuxxingrui, david.bossens, wao.:sssss..o...s.ooooco..oo. edur.eduu.. sg
ABSTRACT
Imitation learning is central problem in reinforcement learning where the goal is to learn a policy that mimics the expert's ' behavior. practice, is often challenging learn the expert policy from limited number of demonstrations accurately due to the complexity of the state space. Moreover, it is essential to explore the environment and collect data achieve beyond-expert performance. To overcome these challenges, we propose a novel imitation learning algorithm called Imitation Learning with Double Exploration (ILDE), which implements exploration in two aspects: (1) optimistic policy optimization via exploration bonus that rewards state-action pairs with high uncertainty potentially improve the convergence to the expert policy, and (2) curiosity-driver exploration of the states that deviate from the demonstration trajectories to potentially yield beyond- expert performance. Empirically, we demonstrate that ILDE outperforms the state- of-the-art imitation learning algorithms in terms of sample efficiency and achieves beyond-expert performance Atari and MuJoCo tasks with fewer demonstrations than in previous work. We also provide a theoretical justification of IL.DE as an eeeertainty-rgullriize policy optimization method with optimistic exploration, leading to regret growing sublinearly in the number of episodes.
INTRODUCTION
Imitation learning (II. is an important subfield of reinforcement learning (RL.), in which ground truth rewards are not available and the goal is to learn a policy based expert demonstrations. Often such demonstrations are limited, making it challenging to achieve expert-level performance. In practice, imitation learning is widely used various applications such as autonomous vehicle navigation (Codevilla et al. 2018), robotic control (Finn et al. 2016 Zhang et al. 2018) and surgical assistance (Osa al. 2014 et al. 2022).
As one of the simplest approaches of imitation learning, behavior cloning (Pomerleau) 1988) lears a policy directly from the atate-action pairs in the demonstration dataset. More recent approaches to imitation learning, such as generative adversarial imitation learning (GAIL) (Ho & Ermon 2016).
discriminator to guide policy learning rather than directly learning reward function based demonstration trajectories. Due to matching occupancy based on a limited set of trajectories, it is still difficult for such algorithms achieve better-than-exper performance and solve high-dimensional problems with limited demonstration data. Moreover, these techniques are typically unstable since the reward function is continually changing as the reward function is updated during the RL training. Alternative approaches, based on the curiosity (Yu et al., [020 Pathak et al., 2017 Burda et al., 2018). a form of self-supervised exploration, encourage the agent to explore transitions
are distinct from the expert thereby potentially yielding beyond-expert performance. Another benefit of the approach is related to the GIRIL algorithm specifically; that is, it uses a pre-trained auto auto-encoder, which makes the technique more stable (Yu et al. 2020).
Equal technical contribution. 'Corresponding author.

---

Published as a conference paper at ICLR 2025
As illustrated in Table 3 in Appendix A this work seeks formulate a sample-efficient imitation learning algorithm that combines the best of both worlds. That is, we integrate the stability and self- supervised exploration properties of GIRIL into discriminator-base imitation learning approaches. Additionally improving the sample-efficiency with an exploration-bonus, we formulate the framework of Imitation Learning with Double Exploration (IIDE).
In detail, ILDE learns a policy based on an aucertainty-regulariin discrepancy. The uncertainty- regularized discrepancy combines the distance to the expert policy, the cumulative uncertainty of the policy during exploration, and the cumulative uncertainty of the policy with respect to the demonstration dataset. To optimize the uncertainty-regulariie based imitation reward while integrating two distinct forms of exploration. Firstly, through optimistic policy optimization, augmented by exploration bonus, ILDE incentivizes the exploration of state- action pairs characterized by high uncertainty, thereby facilitating the training of the discriminator network, which functions as the reward model. Secondly, leveraging curiosity-driven exploration, ILDE targets transitions deviating from demonstration trajectories, paving the way for nuanced policy improvements.
discrepancy, ILDE utilizes GAIL-
The key contributions of this paper are summarized as follows:
In Section 3, we theoretically formulate the problem of solving the aforementioned policy which minimizes the reeruaanty-reullrizzedd discrepancy between the learner trajectories and demonstra- tion trajectories.
In Section 4 we present our proposed algorithm, ILDE with natural policy gradient (ILDE-NPG).
We are then able to demonstrate the regret guarantee of ILDE-NPG in Theorem 4.77 which is the first theoretical guarantee for imitation learning in MDPs with nonlinear function approximation.
Although ILDE-NPG is computationally expensive, we offer more flexible version for the practical implementation of ILDE in Section 5 focusing on two key exploration strategies: (1) encouraging the learner to explore state-action pairs that are 'far' from the current rollout dataset, and (2) rewarding benign deviations from the demonstration dataset using curiosity-based intrinsic rewards.
6 we present experimental results for a practical implementation of ILDE. Our findings demonstrate that ILDE surpasses existing baselines in terms of both average return and sample complexity. This underscores the practical advantages of our proposed objective for imitation learning.
Section
2 RELATED WORK
Imitation learning with limited demonstrations for beyond-expert performance. Learning from limited demonstrations in high-dimensional environment is challenging, and the effectiveness of imitation learning methods is hampered (Li et al.) (023). Inverse reinforcement learning (IRL) methods (Ziebart et al. 2008: Ziebart, Brian D. and Bagnell, J. Andrew and Dey, Anind K. 2010 Boularias et al. 2011)) seek reward function that best explains the demonstration data, which makes hard for them to achieve thater-than-expertt performance when the data are extremely limited. GAIL achieves impressive performance in low-dimensional environments via adversarial learning- based distribution matching. Variational adversarial imitation learning (VAIL) improves GAIL, by compressing the information flow with variational discriminator bottleneck (Peng et al. 2019). Recent adversarial imitation learning methods include IQ-Learn (Garg et al. 2021), which directly infers the Q- (bypassing the reward function), and HyPE (Ren et al. 2024). which trains the policy on a mixture of expert and on-policy data.
Unfortunately, GAIL and VAIL do not scale well in high-dimensional environments (Brown et al. 2019), and still require many episodes of demonstrations. Recent work has focused on reducing the number of required demonstration episodes to between 5 and 20 episodes using techniques such as MCTS-based RL (Yin al. 2022), }, patch rewards (Liu et al. 2023a), and applying demonstration augmentation and a prior policy baseline Li et al.] (2022). Techniques based on curiosity-driven exploration, such as CDIL (Pathak et al., 2017) and GIRIL (Yu et al., 2020 potentially provide more scalable approach to imitation learning. In particular, Yu et al. (2020) propose the GIRIL algorithm within a setting with only a so-called one-life demonstration, which is only a partial episode of Atari game. The system demonstrates favorable performance on Atari games when compared to GAIL, VAIL, and CDIL. Such techniques with intrinsic motivation have the advantage of allowing beyond-expert performance but the disadvantage of potentially optimizing objective that is not
2

---

assume that the initial state $1 is fixed across Published as a conference paper at ICLR 2025
implied by the demonstrations. Our proposed ILDE system seeks to use curiosity- driven exploration as one of two sources of exploration to supplement a traditional imitation learning objective. This is based on the fact that the curiosity reward has high empirical success and enables agents to learn non-trivial skills by exploring hard-to-visit states (Rajaraman et al.) 2020).
Beyond reducing the sheer amount of demonstration data, the problem of varying optimality scores and the reliability of demonstrations has also garnered more attention recently. Confidence-Awan Imitation Learning (CAIL) extends adversarial inverse reinforcement learning with a bi-level op- timization technique to concurrently learn the policy along with confidence scores reflecting the optimality of the demonstrations (Zhang et al. 2021). The setting of imbalanced demonstrations has been studied from semi-supervised learning perspective (Fu et al., 2023) as well as within multi- modal imitation learning (Gu & Zhu 2024). We focus on an alternative approach to the problem of demonstration reliability which combines curiosity with a traditional imitation reward.
Reinforcement learning from demonstrations (RLID.. An alternative framework, related to im- itation learning, reinforcement learning from demonstrations (RLiD) (Hester al.) Christiano et al. 2017] Zhu et al.) (022). While potentially offering accelerated policy learning, such techniques require ground-trute truth rewards, which are not available in pure imitation learning settings.
Exploration bonuses and optimistic RL. Optimism a long-standing principle RL, where typically refers to overestimating the value of particular state state-action pair in order to make the state or state-action pair known more rapidly through increased visitations essentially, this implies that unknown state-action pairs given large exploration bonus while known acteinction pairs are given little to no exploration bonus. Analogous to upper confidence bound (UCB) bandit algorithms (Auer et al.] 2002), various techniques have been implemented in RL which use confidence bounds to upper bound the value in a probabilistic sense. Such techniques traditionally use concentration inequalities based on visitation counts for each state-action pair, which gives such techniques firm
statistical grounding but which limits them to discrete state-action spaces (Shani et al. 2020 Fruit
ct al. 2018 Jaksch et al., 2010).
Recent techniques have aimed at making exploration bonuses more scalable. E3B provides an approximate approach to visitation counts by maintaining covariance matrix which implicitly captures the relations between visited states (Henaff et al.) 2023). State Entropy directly aims to maximze the entropy over the states, rewarding states in batch according the log of their distance to its k-nearest neighbor state (Seo et al. 2021).
Theoretical guarantees for imitation learning. There is a large body of literature on imitation 22019/Chene learning for tabular MDPs (Cai et al. 2019 Chen et al. 2019 /Zhang al. 2020 Xu et al. 2020 Shani
et al. 2022 Chang et al. 2021). Subsequently, [iuu et al. (2021) studied imitation learning for linear
kernel MDPs and provided an D(VHd&3T regret bound, where d is the dimension of the feature space, H is the horizon of the MDPs, T the number of episodes. Viano et al. (2022) considered linear MDPs and proposed PPIL., which achieved sample complexity of O(d7((T - 7)*²² in
discounted linear MDPs. Our work considers MDPs with general function approximation, utilizing the generalized Eluder dimension (Agarwal 2023 complexity of the function class.
Zhao et al. 2023) to characterize the
PRELIMINARIES
We consider an episodic MDP H,P,P) r). where S and A are the state and action spaces, respectively, H is the length of each episode, TP, is the Markov transition kernel of the h-th step of each episode for any h E [H], and r: S x A [-1, 1] is the reward function. We assume without loss of generality that the reward function is deterministic.
While performing imitation learning episodic MDPs, the agent interacts with the environment as follows. At the beginning of each episode E [7]. the agent chooses a policy n' = heja E ((ASS, H). Then the agent takes an action af. ~ Tf.. (-1ss. at the h-th step of the f-th episode and transits to the next state si.41 ~ [P a. The agent does not receive the true reward r" (sf of)) but instead it receives a surrogate reward r(sh,- af)). The episode terminates when the agent reaches the state HH+1- Without loss of generality, different episodes. We remark that our algorithms and corresponding analyses readily generalize to the setting where the initial state 81 is sampled from a fixed distribution.
3

---

Published as a conference paper at ICLR 2025
We now define the value function in episodic MDPs. For any policy т ( hee and reward function S x A [1,1], the state value function V and action-value function Q are defined for any (s, a, h) E x [H] as follows,
Σ r(s;, 8h =s]: (s, = ETTA Σ r(s. 0, =
Vi, Ex
(3.1)
where the expectation E, is taken with respect to the action a; ~ [T(([ss)) and the state S.+1 ~ Pi(-18;; a;) for any i {h, h 1,..,, H.. With slight abuse of notation, also denote by Ps. the operator form of the transition kernel such that (xx a) = EN (+le,))) (s') for any f + IR.
By the definitions of the value functions 3. for any (s, E x A x [H]. any policy æ, and any reward function r, we have
vr (s) (Qi, -), (T((s))) (s,a) = (s a) 1
Vi (s)
(3.2)
where (;,)) A denotes the inner product over the action space A. We further define the expected cumulative reward as follows,
J (TT = (x).
We assume that there an unknown expert policy T { ME(M E (AASS, H) that achieves a high expected cumulative reward J J r*) under the unknown underlying reward function r*.
Given demonstration trajectories 7 {(s ( ) for j E [nl]. the goal of imitation learning to learn a policy æ that achieves a potentially high expected cumulative reward J(T,, (T,r") under the unknown reward function r* based the expert demonstration. As introduced in Chen et al. (2019). we characterize the discrepancy between the expert policy π and the learner policy æ by the tollowing Integral Probability Metric (IPM) over the stationary distributions of the MDP.
Definition 3.1 (Definition 2. Chen et al. 2019). Let R denote a class of symmetric reward functions r A + [-1, 1], i.e., if R, then -r R. Given two policies म, n E A(A H), the R-distance is defined as
dR dR(
( J(x,
rER
Remark 3.2. The IPM distance is a versatile tool for evaluating GAN models. In the context of imitation learning, can choose diffferent classes of reward functions R to measure the discrepancy between the expert policy and the learner policy. For instance, we can choose R to be the class of symmetric reward functions that are 1-Lipschitz continuous with respect to the state-action pair (s, a), which corresponds to the Wasserstein distance. Or we can choose R. to be the unit ball in an RKHS, which yields kernel maximum mean discrepancy (MMD).
As proposed in Yu et al. (2020), to encourage the agent to explore the environment and learn a beyond- expert policy, it is essential to incorporate uncerlainty-driven kitrinsic-reward-dae into the imitation learning framework. In contrast GIRIL (Yu al., 2020) which purely relies on the intrinsic reward to explore transitions that distinct from the expert demonstration, we consider learning a policy which mimics the expert policy subject to an intrinsic reward regularisation term To this end, we defined the following objective function:
exploration
min max (J( .) - (TT,)) Int(T;;
(3.3)
(E)(A))H
PER
where Int (nt(T;)) F) is the expected cumulative intrinsic reward the policy x under the expert demon- stration 7 More specifically,
H C 1
Int (T; B) = n. P(llan.a)))
h=1
where PP is the estimated transition probability that is estimated follows:
H
= argmin
E
((5 22
Pe pmadel
4

---

Published as a conference paper at ICLR 2025
is trained transition model over the demonstration trajectories T and L is distance metric (e.g- [(i++
$1+1) + state proposed in Yu et al. (2020) and Pathak et al. (2017).) For simplicity, we denote the intrinsic reward as the following form
|| 2 113) (Empirically, is the trained VAE model which samples the next
(, Rig * (('5",''))
(3.4)
As shorthand, we define the uncertaintt--rgullariedd loss function for a policy и and a reward function follows,
((T,r) = J(x _ -
In)(
(3.5)
Imitation Learning. As shown in 33. and 3..5 our goal to learn an optimal policy which minimizes the following worst-case loss:
argmin max ((T,) (EA(A rER
(3.6)
Conjecture. To achieve the optimal policy as defined in (366) we hypothesize that more principled exploration technique needed to adapt to the dynamic rewards during the training process. Addi- tionally, conjecture that such uncertainty regularization will help the learner attain performance beyond the expert level.
2019 Liu et al. 2021 In the previous theoretical formulation of online imitation learning (Chen et al.
the objective is learn policy which is close the expert policy with respect the R-distance. Thus, the optimal policy under their formulation is æ and the loss at the saddle point is zero. By contrast, our optimal policy x² is non-trine con-trivially different from the expert policy =
In this work, we aim find x² by proposing novel theory-guided policy optimization algorithm with uncertainty- exploration.
Within the online learning setting, we define the uncertainty regularized regret as follows:
Regret(T) max
Σ (((tt)) ((x",r).
rER
To achieve the objective (333 in the online learning regime, we aim to design an imitation learning algorithm to achieve (egret(T)) growing sublinear as a function of T. Then we can show that E B[e(Tooee r)] converges {( ;r) as T - oc for any r E IR, where Tout refers the hybrid policy
sampled uniformly from (x1 he(T)
IMITATION LEARNING WITH DOUBLE EXPLORATION (ILDE)
We now turn introducing the ILDE framework theoretically before its implementation in Section We analyze ILDE using mirror descent as the policy optimization technique which is more amenable to theoretical analysis than scalable techniques such as PPO. The theoretical analysis considers a general data collection subroutine of which on-policy data collection used in the practical implementation a special case. The thcoretical analysis further considers an IPM-base GAN and a bonus based on state-action uncertainty proposed by (Agarwal al. 2023)) In practice,
apply GAIL and state entropy (Seo et al.) (021). Since PPO already applies action entropy in the policy optimization, together the state entropy and action entropy help explore unvisited state action pairs in a similar manner.
4.1 ALGORITHM
This subsection instantiates a theoretical version of ILDE, where we apply Optimistic Natural Policy Gradient (Liu al., (0236) policy optimization module. Since Algorithm is phasic policy optimization algorithm, the number of episodes where the learner interacts with the environment is T K N/m.
Imitation reward Module. Line 4 of Algorithm [ the learner samples a data set with N trajectories. With carefully chosen N. the expected uncertainty of each state-action pair (s, a)
5

---

Published as a conference paper at ICLR 2025
Algorithm ILDE with Natural Policy Gradient
input: number of iterations K, period of collecting fresh data m, batch size N, learning rate 7, demonstration trajectories _ reward function r loss function L,E. 2: initialize: for all (h, s) [H] set ( s) Uniform(AA)
3: for k = 1...., K do
4: if l: mod m = then
i.d. 5: D - {N fresh trajectories ~ π J. where is chosen uniformly at random from {max [Max(1,K) k
1)....,k 1,
6:
else
7:
D t- D-D--1 p.k t- ph1-
8:
9:
end if
10:
Update n by projected gradient descent with estimated - L rd) [QhhhejH OPE(E* Db: . n²: LrE).
12: for all (h, 8) E [H] x S update Th (- (-18) x Tx( 8) exp(1 Qh(s,
13: end for
14: output: Tout that is sampled uniformly at random from ]
Algorithm 2 OPE(T, D,E.
1: Split D evenly into H disjoint sets Dí for h [M].
2: Set ( + for all s E S
3: for h = H, ... do
4:
Least-squares regression: Jn argmin E(mnen (Ja( - Vn+1 (84+11111))
5: for (s,a) E S x do
6: Compute exploration bonus bn a) as in Hli (fr an( bu,(,)))
7:
+ (s(
where (s, )) as) + (C)
8: (s) + BewTa(1)s) (s,
9: end for
10:
end for
11:
output {
S is upper bounded by Ö(L/VN), where omit the dependency of H and generalized Eluder dimension.
After updating the on-policy dataset, in Line 10 the agent can optimize the reward function according the following loss function:
L(m,r) =(T,r)= J( (T,
(4.1)
While L is unknown for the learner, in Line of Algorithm I we compute an empirical estimate of L:
M (an
sw)
H ΣΣ1 TED*
7(
II
(7)).
(4.2)
j=1A=1
N
Th
(ax
sx)
where t; the index of the last policy which is used to collect fresh data at iteration k.
Assumption 4.1 (Convexity and Lipschitz Continuity of the reward function). We assume that the reward function is parameterized by set of parameters 0 θ c R and the estimated loss function Σ is convex with respect to 0. For any 0,, 02 E θ, e - 0zl12 = O(1), and for any O, O(1).
8, a E S A, | Voro(s, a)
As a result, in Line 10 we essentially update the reward function , by the following projected gradient descent:
g++1 projo (0* B70 Vol-Ī(m> ,)))).
(4.3)

---

Published as a conference paper at ICLR 2025
To ensure that L can potentially approximate L accurately, also require the following assumption on the quality of the demonstration trajectories. Assumption 4.2 (Quality of the demonstration trajectories). The demonstration data 7 satisfies the following property for any reward function r E R:
H
(1) r(h
( J S EE
for some ég >
Note that the assumption is weak since the error can be bounded by Azuma-Hoeftding inequality. Policy optimization module. The algorithm is a phasic policy optimization algorithm built on Liu et al. (2023b). In Line I2 the policy *h -++ is obtained by taking mirror descent step from T maximizing the following Kl.-regularized return:
k+1 arg max (QN (s, )), (-(+s)) - AEDOKI((TA (
h
where ÀED 2 0 is regularization parameter, dki. is the Kullbach-Leiblr divergence, and Q is computed by an optimistic policy evaluation (OPE) subroutine to guide the direction of policy updates.
Double Exploration in Algorithm While algorithms widely used in empirical studies, such as PPO, directly use empirical returns to guide the policy updating procedure, we employ Algorithm 2 estimate the expected return from all state-action pairs.
In OPE (Algorithm)
2) we set an exploration bonus according to Lemma
by (s. a) v 8H2 log H (N 4c + Dy,
(4.4)
where Dy, is defined in Definition 43 exploration bonus to encourage the policy explore the state-action pairs whose expected return values are hard to predict given the on-policy dataset D
Additionally, utilize pretrained uncertainty evaluation function L, explore states beyond the demonstration trajectories. Then, in lines 3 10 we use value iteration to compute the value function considering the sum of the imitation reward the exploration bonus b/, and the uncertainty
4.2 THFORETICAL. ANALYSIS
With the ILDE framework introduced in Section we provide a thorough regret analysis for Algorithm!1 [
Our analysis addresses broad category of MDPs, specifically those whose value functions can be effectively approximated and generalized across various state-action pairs. These are referred to as MDPs with bounded generalized Eluder dimension (Agarwal et al. 2023).
Definition 4.3 (Generalized Eluder dimension, (Agarwal et al. 2023). Let AED 2 0, and let (= (S A)&T be sequence of state-action pairs Then the generalized Eluder dimension of a value function class F S x A + [-H, H. H) with respect to XED is defined by dimr(F) - dim Z).
dim(F, Z) = Σ min (1, ()))
(2()))
where (3 ;= sup
Ju.EEF Exett- f2(2z))2 AED
We write (F) -- - Me dimr(FF) for short when F is collection of function classes
F= (F the context.
Remark 4.4. The concept of generalized Eluder dimension was first introduced in Agarwal et al. (002),, where a value- algorithm was proposed to achieve a nearly optimal regret bound in the online RL setting. Notably, our definition of the function class is somewhat broader, an unweighted definition of the D, distance. Consequently, the definition of generalized Eluder dimension does require taking the supremum over weights, allowing for wider class of MDPs.
apply

---

Published as a conference paper at ICLR 2025
The analyzed policy optimization method needs policy evaluation subroutine introduced in Algorithm 2] assumption.
To make such a policy evaluation tractable, we make the following realizability
Assumption 4.5 (Strong realizability of state-action value function class). For all h E H and Vh+1 : + [-H, H). there exists function fi. Fh such that for all S, E A, have (s = PA as)
Definition 4.6 (Covering numbers of function classes), For each E [H] there exists an EF cover C(F S Fi with size (C(F,. (N,,,, EF), such that for any f E F, there exists f' E C(F,., such that Iff-f'll S EJ.. For any EF > O, we define the uniform covering number of FF with respect to Ey (N [M (F),
Theorem 4.7 (Regret bound for Algorithm [] Suppose Assumptions 4.5
4. and 4.2 hold. if we
set Y 11, Ex = 1/N, Vlog |A/ H KK m = vK/H log [A,, 100
O
////////// (HK), and
2/3
N = KH (imr((F) log Nex (F)/ Vlvlooggg log |A,, where K =
then with
H1th dimyr (F) log N.
probability least 1
5, Algorithm yields a regret of
δ 8// (Fimy(F) log Ne, (F)) (2)) (F)) T8/3 + +CET)
(4.5)
Remark 4.8. This is the first theoretical guarantee for imitation learning in MDPs with nonlinear function approximation. As shown in Zhao et al. (2023). the considered setting captures MDPs with bounded Eluder dimension (Russo & Van Roy, 2013) and thus also captures linear MDPs (Jin et al.) ky K
2020) as special cases. If we select policy from { Jk=1 uniformly at random, the expected loss for that policy is O( + eg) when T' @(1* dimr (F) which provides sample complexity bound for Algorthm
5 PRACTICAL IMPLEMENTATION OF ILDE
Algorithm 3 Imitation Learning with Double Exploration (ILDE)
1: input: number of episodes T, period of collecting fresh data m, batch size N, learning rate f). demonstration trajectories 7 reward function r', pretrained loss function Log.-
2: for t ,,,.., T' do
3: Sample trajectories D' from T'.
4: Update the imitation reward r² by minimizing the loss of the discriminator trained to distinguish between the demonstration trajectories and D' (e.g. (5.I).
5: Compute the aggregated reward (*ss a) = r (s, a) + A (s,a) for all state-action pairs
(s, a) in D'.
6: Update the current policy n' via policy optimization method (e.g., PPO) with reward 7* a) b(s, a) where b(s, a) is the exploration bonus and obtain n"
7: end for
While Algorithm|] provides solid sample-complexity guarantee, lacks computational efficiency and is challenging to implement at scale. In this section, we focus on the crucial elements of Algorithm I and introduce the proposed Imitation Learning with Double Exploration (ILDE) framework. The agorittmm summarized in pseudocode in Algorithm 3 includes the following key modules:
Imitation reward module: This module minimizes the loss of a discriminator which distinguishes
between demonstration data F and data obtained during policy optimization D'. Theoretically, the imitation reward function is progressively optimized according to the gap hetween the average return of demonstration trajectories and that of current trajectories. In practice, to implement the module, maintain discriminator Dp which continue minimizing the following loss:
- log (s a)] + E( [log(1 Ds a))] + BE [dxa. (E<a.(""( (=]s)]r") - 1.]. (5.1)
where encoder E' is introduced to incorporate variational discriminator bottleneck (Peng et al. 2019) to improve GAIL, 3 is the scaling weight, and I, is the information constraint. Further
details can be found in Appendix [.2.
Pretrained uncertainty evaluation module: To implement this module, we make use of GIRIL. That is, we pretrain VAE model estimate the transition kernel To over the demonstration trajectories
8

---

Published as a conference paper at ICLR 2025
T and compute the intrinsic reward L, L, (3.4) for all (s,a) in D'. Further details can be found in Appendix B.2.2
Exploration bonus b: To analyze the ILDE framework theoretically in Section 4 we choose the Dy-uncertainty of state-action pairs as the exploration bonus, which reflects the uncertainty of state-action pair after collecting enough data using the current policy. To efficiently implement the exploration bonus in practice, we utilize state entropy (Seo et al. 2021) in a representation space of a feature extractor f. We utilize a k-nearest neighbor entropy estimator to calculate the exploration bonus as b(s, a) = log(Il y y + 1), where y = f()) is a state representation from f and y*-NN is the k-nearest neighbor of within a set of N representations {1/1 {//1,//2 1/2," .UN J. where . is the batch size. The state- dependeeedd rather than state-action dependency of the bonus is motivated by the technique of entropy regularisation in the policy optimization already encouraging spread over the action space. Further details can be found in Appendix [32.3
Policy optimization module: While the above modules provide the design of the reward, the policy optimization module optimises the policy for the objective. For policy optimization, we use an actor-critic variant of proximal policy optimization (PPO) (Schulman et al. 2017) with generalized advantage updating.
6 EXPERIMENTS
We now conduct experiments based on four key hypotheses. First, hypothesize that pure imitation learning methods such as VAIL will fail with limited demonstrations. Second, we hypothesize that a pre-trained model-based exploration reward such as GIRIL alone provides stable improvement but may not reach a wide set of states due to the curiosity bias preferring certain states. Third, the state entropy reward bonus provides an exploration reward which provides a high reachability across a wide variety of states, thereby improving sample efficiency and overall performance. Fourth, VAIL's imitation reward can provide additional performance gains compared to pure exploration (i.e. exploration bonus and/or curiosity reward) when the imitation reward is reliable while there is no effect when the imitation reward is unreliable.
6.1
ATARI GAMES
To test these hypotheses, we compare our proposed ILDE to VAIL, GIRIL, and ablations on six Atari games within OpenAl Gym (Brockman al. 2016) in a challenging setting where the agent receives even less demonstration data than in so-called one-life demonstration data (Yu et al., (020)) The imitation learning agents are trained according to the design of the imitation reward and/or any potential bonuses but are evaluated based on their actual scores the Atari games. We summarize more experimental details in the Appendix B
Demonstrations. Our setting is very challenging compared to existing work because the demon- stration data are very sparse, since they are based on partial one-life demonstrations. A one-life demonstration only contains the states and actions performed by an expert player until they die for the first time in game. We generate the one-life demonstrations using a PPO agent trained with ground-truth rewards for 10 million simulation steps. Table 4 in the Appendix B. shows that one-life demonstration contains much less data than a full-episode demonstration. To make the problem even more challenging, the imitation learning agents receive demonstrations which comprise only 10% of a full one-life demonstration.
Baselines. We compare our ILDE with the following baselines: (1) VAIL (Peng et al. 2019) an
improved GAIL variant using variation discriminator bottleneck, which updates the policy based on p only; (2) GIRIL. (Yu et al. 2020), the generative intrinsic reward-driven imitation learning method: (3) ILDE w/o b, an ablation of ILDE without the exploration bonus b; and (4) ILDE w/o
an ablation of IDLE without the imitation reward p*
Results. Table D compares the performance of our ILDE and other baselines. ILDE outperforms the other baselines in terms of achieving higher average returns and outperforming experts in more games, Notably, ILDE achieves an average performance that is 5.33 times that of the expert demonstrator across the 6 Atari games. More specifically, ILDE outperforms the expert in all of the 6 Atari games and outperforms GIRIL in 4 games. The ablation study in Table Tibiicatss reward component. ILDE achieves much higher sample efficiency improvements than GIRIL, as shown in Table 10 and 11 of Appendix B.2.5 Further, we also show in Appendix B.2.7 that ILDE
indicates the significance of each
9

---

Published as a conference paper at ICLR 2025
Table I: Average return of Expert demonstrator, VAIL, GIRIL, ILDE wo ALTg. ILDE w/o b, ILDE w/o r* and ILDE with one-life demonstration data on six Atari games. The results reported in the format of mean = std over 10 random seeds with Sniee-tthn-Giiee performance in bold. "#Games> outperforms the expert. "Improve vs Expert" denotes the average performance improvements of IL methods versus the expert across Atari games.
denotes the number and ratio of games that an imitation learning method
Atari Games
Expert
VAIL.**)
GIRIL
ILDE w/o AC
ILDE w/o
ILDE w/o
ILDE
BeamRider
1.9181645
911116
1,245
3,23333,340
8,521
11,453:22416
11,453-2.319
DemonAttack
8,820+-2.891
979+863
12.381-19.851
78.513+11,5 11,583
64,598
BattleZone
22,410= 839
5,08144,971
25.203::22.866 868
1,412
18, 28,179:21.181
43,708
59.109-25,577
Obert
6,24612,964
3,7541-4460
16.491881.265
10,84915,0188
26,9001.52.146
54,169-68.85
70_566_:75.310
Krull
6.520:1911
7,93748.188
20,935::2,796
16,715:23.153
27,495-288.77
23,716=16,790
StrCunner
11,495
219
5424:188
764::455
478::239
16,526
25,361
Kiames>Expeet
0/0%
1/.6.7%
83.3%
33.3%
46.667%
5/83,3%
6/ 100%
Improve vs Expert
0.37
4.71
robust to noise in the demonstrations, Last, we note that the limited demonstration setting is challenging for IQ-Learn (see Table 12 in Appendix B.2.6).
6.2 CONTINUOUS CONTROL TASKS
Additionally, we also conduct experiments on the MuloCo environments. Table 2 compares the performance of ILDE and other baselines. ILDE consistently outperforms other baselines in the four continuous control tasks by achieving higher average returns over 10 random seeds. We summarized more experimental details for MuJoCo experiments in Appendix [.2.8] Based on our current experiments, observed that ILDE performs exceptionally well in games, while its improvement on MuJoCo is less significant. This may suggest that the current design of the intrinsic reward is better suited for taskswhere the extrinsic reward is more aligned with the objective of "seeking novelty". Thhis may also indicate that our algorithm does not (only) extrapolate the intention in the partial demonstration. Additional experiments (see Table 13 in Appendix B.2.6) show the potential of ILDE to generalize to other adversarial IL methods.
Table 2: Average return of VAIL, GIRIL, and ILDE the MuloCo tasks. The average returns calculated on the last 4 evaluation points. The results are reported in the format of mcan + std over 10 random seeds with the best performance in bold.
Tasks
VAIL (7**)
GIRII.
ILDE
Reacher
189
-60+:24
-51::15
Hopper
1,264-1783
1,447=127
1,878:1690
Walker2d
8791:825
997:5575
998:1694
HumanoidStanduup
52,63516,3770
76,1341110,26
77,887-16.123
CONCLUSION
Achieving beyond-expert performance and improving sample complexity in complicated tasks are two crucial challenges in imitation learning. To address these difficulties, we propose Imitation Learning with Double Exploration (ILDE), which optimizes an uncertainty-regularizee discrepancy that combines the distance to the expert policy with cumulative uncertainty during exploration. Theo- retically, we propose ILDE-NPG with a theoretical regret guarantee, a first for imitation learning in MDPs with nonlinear function approximation. Empirically, we introduce a flexible ILDE framework with efficient exploration strategies and demonstrate ILDE's outstanding performance over existing baselines in average return and sample complexity. It is also worth noting that our analysis cannot be directly applied to the practical version implemented in experiments.
Limitation of our work. Our method assumes that the target task be imitated is aligned with the concept of seeking novelty or curiosity. Its ability to achieve performance beyond the demonstration not only due the algorithm's 's capacity understand and extrapolate the intention behind the partial demonstration but also due to an alignment between the extrinsic reward and seeking novelty.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank the anonymous reviewers and area chair for their helpful comments. HZ and QG are supported in part by the NSF grants CPS-2312094, IIS-2403400 and Sloan Research Fellowship. HZ is also supported by UCLA-Amazon Science Hub Fellowship. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies. XY, DB, and IWT are supported by CFAR, Agency for Science, Technology and Research, Singapore. This project is supported by the National Research Foundation, Singapore and Infocomm Media Development Authority under its Trust Tech Funding Initiative. Any opinions, findings and conclusions or recommendations expressed in this material are those of the authors and do not reflect the views of National Research Foundation, Singapore and Infocomm Media Development Authority,
REFERENCES
Alekh Agarwal, Yujia Jin, and Tong Zhang. VO Q L: Towards Optimal Regret Model-free RL with Nonlinear Function Approximation. In Annual Conference Learning Theory (COLT 2023), pp. 987-1063, 2023.
Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-tima analysis of the multiarmed bandit problem. Machine learming, 47:235-256, 2002.
Abdeslam Boularias, Jens Kober, and Jan Peters. Relative entropy inverse reinforcement learning. In International Conference on Artificial Intelligence and Statistics (AISTATS 2011), pp. 182-189, 2011.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. OpenAl Gym. arXiv preprint arXiv:
2016.
Daniel Brown, Wonjoon Goo, Prabhat Nagarajan, and Scott Niekum. Extrapolating beyond sub- optimal demonstrations via inverse reinforcement learning from observations. In International Conference Machine Learning (ICML 2019), pp. 783-792, 2019.
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A Efros. Large-scale study of curiosity-driven learning. arXiv preprint arXiv: /808. 04355, 2018.
Qi Cai, Mingyi Hong, Yongxin Chen, and Zhaoran Wang. On the global convergence of imitation learning: A case for linear quadratic regulator. arXiv preprint arXiv:1901.03674. 2019.
Jonathan D Chang, Masatoshi Uchara, Dhruv Sreenivas, Rahul Kidambi. and Wen Sun. Mitigating covariate shift in imitation learning via offline data without great coverage. arXiv preprint arXiv:2/06.03207 2021.
Minshuo Chen, Yizhou Wang, Tianyi Liu, Zhuoran Yang. Xingguo Li, Zhaoran Wang, and Tuo Zhao. On Computation and Generalization of Generative Adversarial Imitation Learning. In International Conference on Learning Representations (ICLR 2019), 2019.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Advances in neural information processing systems (NeurIPS 2017), pp. 4302-4310, 2017.
Felipe Codevilla, Matthias Müller, Antonio López, Vladlen Koltun, and Alexey Dosovitskiy. End-to- end driving via conditional imitation learning. In IEEE International Conference on Roboties and Automation (ICRA 2018), pp. 4693-4700 2018.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International Conference on Machine Learning (ICML 2016), pp. 49-58. PMLR, 2016.
Ronan Fruit, Matteo Pirotta, Alessandro Lazaric, and Ronald Ortner. Efficient Bias Stas-Span-Cor ConntCoont Exploration-Exxllotttio in Reinforcement Learning. In International Conference on Machine Learning (ICMI. 2018), pp. 1578-1586, 2018.

---

Published as a conference paper at ICLR 2025
Huiqiao Fu, Kaiqiang Tang, Yuanyang Lu, Yiming Qi, Guizhou Deng, and Chunlin Chen. Ess- InfoGAIL: semi-supervised Imitation Learning from Imbalanced Demonstrations. In Advances in Neural Information Processing Systems (NeurIPS 2023), pp. 60048 00448,6009, 2023.
Divyansh Garg. Shuvam Chakraborty, Chris Cundy, Jiaming Song. and Stefano Ermon. IQ-Learn: Inverse soft-Q Learning for Imitation. In Advances in Neural Information Processing Systems
(NeurlPS 202 l), pp. 4028-4039, 2021.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11))139-14)) 2020.
Sijia Gu and Fei Zhu. BAGAIL: Multi-modal imitation learning from imbalanced demonstrations. Neural Networks, 174: 106:11,,,, 2024.
Mikael Henaff, Roberta Raileanu, Minqi Jiang, and Tim Rocktäschel. Exploration via Elliptical Episodic Bonuses. arXiv preprint arXiv:22/0.05800, 2023.
Todd Hester, Matej Vecerik, Olivier Pietquin, Marc Lanctot, Tom Schaul, Bilal Piot, Dan Horgan, John Quan, Andrew Sendonaris, Ian Osband, al. Deep Q-learning from Demonstrations. In Proceedings of the AAAI conference on artificial intelligence (AAAI 2018), pp. 3223-3230.
Jonathan Ho and Stefano Ermon. Generative ad adversarial imiittiin learning. In Advances in Neural Information Processing Systems (NeurIPS 2016), pp. 4572-4580, 2016.
Thomas Jaksch, Ronald Orner. and Peter Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research, 11:1563-1600, 2010.
Chi Jin, Zhuoran Yang. Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on learning theory (COLT 2020), pp. 2137-2143. PMLR, 2020.
Ilya Kostrikov. PyTorch Implementations of Reinforcement Learning Algorithms. https:/ // 2018.
gi github.ccomm httppp/.ww...oaaaaeoo//////eeeea
Bin Li, Ruofeng Wei, Jiaqi Xu, Bo Lu, Chi Hang Yee, Chi Fai Ng. Pheng Ann Heng, Qi Dou, and Yun Hui Liu. 3D Perception based Imitation Learning under Limited Demonstration for Laparoscope Control in Robotic Surgery. In IEEE International Conference on Roboties and Automation (ICRA 2022), pp- 7664-7670, 2022.
Ziniu Li, Tian Xu, Zeyu Qin, Yang Yu, and Zhi-Quan Luo. Imitation Learning from Imperfection: Theoretical Justifications and Algorithms. Advances in Neural Information Processing Systems (NeurIPS 2023), pp. 18404-18443, 2023.
Minghuan Liu, Tairan He, Weinan Zhang, Shuicheng Yan, and Zhongwen Xu. Visual imitation learning with patch rewards. In International Conference Learning Representations (ICLR 2023), 2023a.
Qinghua Liu, Gellért Weisz, András György, Chi Jin, and Csaba Szepesvári. Optimistic natural policy gradient: a simple efficient policy optimization framework for online RL. Advances Neural Information Processing Systems (NeurIPS 2023), pp. 3577,, 2023b.
Zhihan Liu, Yufeng Zhang, Zuyue Fu, Zhuoran Yang, and Zhaoran Wang. Learning from Demon- stration: Provably Efficient Adversarial Policy Imitation with Linear Function Approximation. In International Conference on Machine Learning (ICML 2021), pp. 14094-14138, 2021.
Francesco Orabona. A modern introduction to online learning, arXiv preprint arXiv: 2019.
Takayuki Osa, Naohiko Sugita, and Mamoru Mitsuishi. Online Trajectory Planning in Dynamic Environments for Surgical Task Automation. In Robotics: Science and Systems (RSS 2014), pp. 1-9, 2014.
12

---

Published as a conference paper at ICLR 2025
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity driven exploration by self-supervised prediction. In Interationall Conference on Machine Learning (ICML 2017), pP- 2778-2787,
Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine. Variational discriminator bottleneck: Improving imitation learning, inverse RL, and GANs by constraining information flow. In International Conference on Learning Representations (ICLR 2019), pp. 1-27, 2019.
Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Advances in Neural Information Processing Systems (NeurIPS 1998), pp. 305-313, 1988.
Nived Rajaraman, Lin Yang, Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. In Advances in Neural Information Processing Systems (NeurIPS 2020), pp- 2914-2924, 2020.
Juntao Ren, Gokul Swamy, Zhiwei Steven Wu, J Andrew Bagnell, and Sanjiban Choudhury. Hybrid inverse reinforcement learning. arXiv preprint arXiv:2402.08885 2024.
Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic ex- ploration. In Advances in Neural Information Processing Systems (NeurIPS 2013), pp. 2256-2264, 2013.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint (rXiv:1707.06347 2017.
Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. State entropy maximization with random encoders for efficient exploration. International Conference on 2021.
Machine Learning (ICML 2021), pp.
Lior Shani, Yonathan Efroni, Aviv Rosenberg, and Shie Mannor. Optimistic policy optimization with bandit feedback. In International Conference on Machine Learning (ICML 2020), pp. 8604-8613 2020.
Lior Shani, Tom Zahavy, and Shie Mannor. Online apprenticeship learning. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI 2022), pp. 8240 02400 2022.
Harshinder Singh, Neeraj Misra, Vladimir Hnizdo, Adam Fedorowicz, and Eugene Demchuk. Nearest neighbor estimates of entropy. American Journal of Mathematical and Management Sciences, 23 (3-4):301-321, 2003.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in Neural Information Processing Systems (NeurIPS 2015), pp. 3483-3491, 2015.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems (IROS 2012), pp. 5026-5033. IEEE, 2012.
Luca Viano, Angeliki Kamoutsi, Gergely Neu, Igor Krawezuk, and Volkan Cevher. Proximal point imitation learning. In Advances in Neural Information Processing Systems (NeurIPS 2022), pp. 23309-24326. 2022.
Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. In Advances in Neural Information Processing Systems (NeurIPS 2020), pp. 15737-15749, 2020.
Zhao-Heng
Zhao- Heng Yin, Weirui Ye, Qifeng Chen, and Yang Gao. Planning for sample efficient imitation learning. In Advances in Neural Information Processing Systems (NeurIPS 2022), pp. 2577-2589 2022.
Xingrui Yu, Yueming Lyu, and Ivor Tsang. Intrinsic reward driven imitation learning via generative model. In International Conference on Machine Learning (ICMI 2020), pp. 10925-10935, 2020.
13

---

Published as a conference paper at ICLR 2025
Songyuan Zhang, Zhangjie Cao, Dorsa Sadigh, and Yanan Sui. Confidence-Aware Imitation Learning from Demonstrations with Varying Optimality. In Advances in Neural Information Processing Systems (NeurIPS 2021), pp. 12340-12350, 2021.
Tianhao Zhang, Zoe McCarthy, Owen Jow, Dennis Lee, Xi Chen, Ken Goldberg, and Pieter Abbeel. Deep imitation learning for complex manipulation tasks from virtual reality teleoperation. In IEEE International Conference on Robotics and Automation (ICRA 2018). pp.
2018.
Yufeng Zhang, Qi Cai, Zhuoran Yang. and Zhaoran Wang. Generative adversarial imitation learning with neural network parameterization: Global optimality and convergence rate. In International Conference on Machine Learning (ICML 2020), pp. 11044-11054, 2020.
Heyang Zhao, Jiafan He, and Quanquan Gu. Nearly Optimal and Low-Switching Algorithm for Reinforcement Learning with General Function Approximation. arXiv preprint arXiv:2311.15238. 2023.
Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Self-adaptive imitation learning: Learning tasks with delayed rewards from sub-optimal demonstrations. Proceedings of the AAAI conference on artificial intelligence (AAAI 2022), pp.
2022.
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. AAAI Conference on Artificial Intelligence (AAAI 2008), pp. 1433-1438, 2008.
Ziebart, Brian D. and Bagnell, J. Andrew and Dey, Anind K. Modeling interaction via the principle of maximum causal entropy. In International Conference on Machine Learning (ICMI. 2010), pp. 1255-1262, 2010.
A MOTIVATION TABLE
Table 3 illustrates the strengths of each component in ILDE, which helps explain the motivation of ILDE.
Table 3: Comparison of three reward components with three attributes.
Stability
Exploration
Imitation
GIRIL reward
++
+
Bonus b
+
+
VAIL y.k
+
B AUXILIARY EXPERIMENTAL DETAILS
B. DEMONSTRATION LENGTHS IN THE ATARI ENVIRONMENT
Table compares the lengths of one-life demonstrations and full-episode demonstrations Atari games. In the experiments, we only use 10% of the one-life demonstrations for imitation learning. which more challenging than previous work in learning with limited demonstrations (Yu al. 2020).
B.2
EXPERIMENTAL SETUP AND ADDITIONAL RESULTS
This subsection presents details of VAIL, GIRIL, state entropy bonus, experimental setup, and additional results.
B.2.1
DETAILS OF THE VARIATIONAL ADVERSARIAL IMITATION LEARNING (VAIL)
Generative adversarial imitation learning (GAIL) (oo & Ermon 2016) regards imitation learning as a
distribution matching problem, and updates policy using adversarial learning (Goodfellow al. 2020).
14

---

Published as a conference paper at ICLR 2025
Table 4: Demonstration lengths in the Atari environment.
One-life Demo.
Full-episode Demo.
Atari Games
Length
# Life
Length
# Life
Beam Rider
702
1
1,412
3
Demon Attack
2,004
1
11,335
Battle Zone
260
1,738
5
bert
787
1,881
4
Krull
662
1,105
3
Star Gunner
118
1
2,117
5
Previous work has demonstrated that GAIL does not work well in high-dimensional environments such as Atari games (Brown et al. 2019). Variational adversarial imitation learning (VAIL) (Peng al. 2019) improves GAIL by compressing the information via a variational information bottleneck (VDB). B. VDB constrains the information flow in the discriminator using an information bottleneck. By enforcing constraint on the mutual information between the observations and the discriminator's internal representation, VAIL significantly outperforms GAIL by optimizing the following objective:
min maxB ( BI- (lg( log(-Do(z))]] + (s,a))n (E [- log(1 1 Do(2))]]
E. $20
E ddx E3' (E((2()) Te].
(B.1)
where म 1 F. + R represents mixture of the expert policy and the agent's policy, E' is the encoder for VDB, 8 is the scaling weight, and Ie is the information constraint. The reward for a is then specified by the discriminator T, log (1 (s))).
B.2.2 DETAILS OF THE GENERATIVE INTRINSIC REWARD DRIVEN IMITATION LEARNING (GIRIL)
Previous inverse reinforcement learning (IRL. methods usually fail to achieve expert-level perfor- mance when learning with limited demonstrations in high-dimensional environments. To address this challenge, Yu al. (2020) proposed generative intrinsic reward-driven imitation learning (GIRIL) to empower the agent with the demonstrator's intrinsic intention and better exploration ability. This was achieved by training a novel reward model to generate intrinsic reward signals via generative model. Specifically, GIRIL leverages a conditional VAE (Sohn et al. 2015) to combine backward action encoding model and a forward dynamics model into single generative model. The module is composed of several neural networks, including recognition network (=(2 8t,, 8e+1) 1), generative network pe (#141/2, 8r), and prior network (2(8))) GIRIL refers to the recognition network (i.e. the probabilistic encoder)
backward action encoding model, and the generative network (i.c. the probabilistic decoder) as a forward dynamics model. Maximizing the following objective to optimize the module:
J(po 4o) (gg llgg
dKl. (41))ooll
(B.2)
(qs (4) (411) s:)]
where is the latent variable, ( is the expert policy distribution, a Softmax(2) is the transformed latent variable, a is positive scaling weight. The reward model will be pre-trained on the demonstration data and used for inferring intrinsic rewards for the policy data. The intrinsic reward is calculated as the reconstruction error between $141 and 8111:
Ft S st41 11
(B.3)
where || |l2 denotes the L2 norm, and $ decoder st).
B.2.3 DETAILS OF STATE ENTROPY BONUS
State entropy maximization has been demonstrated be simple and compute-efficient method for exploration (Seo et al. 202T). The key idea for this method to work a high-dimensional environment is to utilize k-neerest neighbor state entropy estimator in the state representation space.
15

---

Published as a conference paper at ICLR 2025
k-nearest neighbor entropy estimator, Let X be random variable with probability density function whose support is a set X' c IRd. Then its differential entropy is given by
H(X) =0 p(x) log p())d.
Since estimating p is not tractable for high- nigh-dimensional data, a particle-based k-nearest neighbors (k-NN) entropy estimator (Singh et al., 2003) can be used:
N
N.
|
2; k NN 11 Ā
H(X)=
Rd
log
+
Ck
N
b: 7.1 + 1)
N
log ||0
k NN
2;
(B.4)
N
where = - is the k-nearest neighbor of x; within a set of N representations {21, 22. xN I' is the gamma function, ग refers to an estimate of the number (as opposed to the policy), and log(k= 1)) where & the digamma function.
State entropy estimate as bonus, Following Seo et al. (021), we define the bonus to proportional to the state entropy estimate in (B.4).
b(s;) = log( (1lv: 1/ y/ -NNN 112 + 1),
(B.5)
where yi = (s;) a fixed representation from a state feature extractor f and y - is the k-nearest neighbor of y/ within set of N representations {v1, 1/....................... yN). The use of fixed representation space produces a more stable intrinsic reward since the distance between a given pair of states does not change during training (Seo et al. 2021). In our implementation, we use the identity function as the feature extractor f.
B.2.4 EXPERIMENTAL SETUP
For the GIRIL and ILDEs, our first step was to train a reward model for each game using the 10% of one-life demonstration. The training was conducted with the Adam optimizer a learning rate of 3e-4 and a mini-batch size 32 for 1,000 epochs. In each training epoch, we sampled a mini-batch of data every four states. To evaluate the quality of our learned reward, used the trained reward learning module to produce intrinsic rewards for policy data and trained a policy to maximize the inferred intrinsic rewards via PPO. For VAIL. we trained the discriminator using the Adam optimizer with learning rate of 3e-4. The discriminator was updated at every policy step. We trained the PPO policy for all imitation learning methods for 50 million steps. We compare ILDE with baselines in the measurement of average return and sample efficiency. The average return was measured using the true rewards in the environment. We measured the sample efficiency based on the step after which an imitation learning method can continuously outperform the expert until the training ends. All the experiments are executed in an NVIDIA A100 GPU with GB memory.
Network architectures. We implement our ILDE and all baselines in the feature space of a state feature extractor f with 3 convolutional layers. The dimension of the state feature is 5184. In the state feature space, we use 3-layer MLPs to implement the discriminator for VAIL and the encoder and decoder for GIRIL. For fair comparison, we used an identical policy network for all methods. We used the actor-critic approach for training PPO policy for all imitation learning methods. Table 5 shows the architecture details of the state feature extractor and the actor-critic network. Table 6 shows the MLP architectures, i.c., GIRIL's encoder and decoder and VAIL's discriminator.
Hyperparameter settings. We implement all imitation learning methods and experiments based on the GitHub repository by Kostrikov (2018). Additional parameter settings are explained in Table 7 Note that X, which controls the curiosity bonus was set without tuning to À = 10 in our experiments. We use the same hyperparameter setting for the different experiments within a domain (Atari vs MuJoCo), apart from multiplicative constant based on the range of the curiosity. However, to further explore its effect, we perform a parametric study with À ranging in (20, 10, 5, 2}. The results in Table & show that ILDE's performance slightly increases with decreasing X.
16

---

Published as a conference paper at ICLR 2025
Table 5: Architectures of state feature extractor and actor-critic network for Atari games.
State feature extractor
Actor-Critic network
4 x 84 x 84 States
x 84 x 84 States
3 conv, 32 LeakyReLU
8 conv, 32, stride ReLU
x 3 conv, 32 LeakyReLU RU
4 x 4 conv, 64, stride 2, ReLU
3 conv, 64 LeakyReLU
3 conv, 32, stride ReLU
flatten 64 9 + 5184
dense 32
x 7 512 Categorical Distribution
dense 32 777-1
State features
Actions
Values
Table Architectures of MLP-based GIRIL's encoder and decoder, and VAIL's discriminator.
GIRIL
VAIL
encoder
decoder
discriminator
5184 State features and next-state features
Actions and 5184 State features dense Actions 1024, LeakyRel.U dense 1024 + 1024. LeakyReLU dense 1024 - 5184
Actions and 5184 State features
dense 5 x 2 1024, LeakyRel.U dense 1024 1024, LeakyReLU
dense 5182 1024, LeakyReI.U
dense 1024 - 1024, LeakyReLU split 1024 into 512, 512 dense 512 - fks. dense 512 + dr
dense 1024 - μ4, dense 1024 +
reparameterization + # Actions
reparamcterization -
Latent variables
5184 Predicted next state features
Mg,, og, discrimination
Table 7: Hyperparameter settings for Atari games.
Parameter
Setting
Actor-critic learning rate
2.5e-4,
linear decay
PPO clipping
0.1
Reward model learning rate
3e-4
Discount factor
0.99
GAE Lambda
0.95
Critic loss coefficient
0.5
Entropy regularisation
0.01
Epochs per batch
4
Batch size
32
128 (parallel workers x time steps)
Mini batch size
32
Evaluation frequency
every 200 policy updates
GIRIL objective
100.0
Mini batch size for training the reward model
32
Total training epoch of reward model
0.000
VAIL's bottleneck loss weight B Information constraint Ic
0.2
X in Algorithm
10.0
Table 8: The effect of À ILDE performance in the BeamRider task.
20
10
8.973 + 1,315
11,453 + 2,319
11,947 + 1,228 12,600 + 1,351
B.2.5 ADDITIONAL RESULTS
To better illustrate the significance of the improvement of ILDE against the expert demonstrator, normalize the results by stating that the expert was 1.0. The normalized results are summarized in Table On average, VAIL is worse than the expert, only outperforming the expert in Krull. GIRIL
17

---

Published as a conference paper at ICLR 2025
archives performance that is 4.87 times the expert, performing the best in Q*bert. ILDE is the best, achieving a performance that is 5.33 times the expert. Figure Iillustrates the learning curves of imitation learning methods in the 6 Atari games over 10 random seeds.
Table Average performance improvements of imitation learning methods versus the expert demon- strator. The expert performance is regarded as 1.0. "Improve vs Expert" denotes the average improvements of methods versus the expert across Atari games.
Atari Games
Expert
VAIL (r*)
GIRIL
ILDE w/o b
ILDE w/o rk
ILDE
BeamRider
1.0
0.05
4.44
4.44
5.97
5.97
DemonAttack
1.0
0.11
8.21
8.90
6,91
7.31
BattleZone
1.0
0.23
1.12
0.81
1.95
2.64
Qebert
1.0
0.60
12.25
4.31
8.67
11.30
Krull
1.0
1.22
3.21
2.56
4.22
3.64
StarGunner
1.0
0.01
0.02
0.02
0.73
1.13
Improve vs Expert
1.0
0.37
4.87
3.50
4.71
5.33
Expert
VAIL (+*)
GIRIL
ILDE w/o b
ILDE w/o
ILDE
100000
120000
12500
Average Return
Averagge Resurn
80000
Average Return
100000
10000
60000
80000
7500
5000
40000
40000
20000
20000
Time steps
Time steps
Time steps
(a) BeamRider.
(b) DemonAttack.
(c) BattleZone.
50000
Average Retarn
80000
Average Retar
40000
Average Retur
30000
60000
40000
20000
10000
20000
10000
0
3 Time steps
Time steps
107
Time steps
107
107
(d) Q"bert.
(c) Krull.
(1) StarGunner.
Figure 1: Average return vs. number of simulation steps Atari games. The solid lines show the mean performance over 10 random seeds. The shaded area represents the standard deviation from the mean. The blue dotted line denotes the average return of the expert demonstrator. The area above the blue dotted line means performance beyond the expert.
Additionally, we summarize the sample efficiency improvements of VAIL and ILDE methods GIRIL in Table TI The sample efficiency of VAIL in outperforming experts is much worse than GIRIL and ILDE variants. ILDE shows impressive sample efficiency improvements across several games, and especially so in BattleZone, where the improvement over GIRIL is 90.06%.
B.2.6 EXPERIMENTS WITH IQ-LEARN AND HYPE
As additional baseline comparisons, we include IQ-Leamn (Garg et al. 2021 and HyPE (Ren et al. 2024
First, we compare IQ-Learn on our Atari task setting (with 10% of a one-life demonstration). The results in Table 12 show that IQ-Learn does perform well Atari tasks, likely because the
18

---

Published as a conference paper at ICLR 2025
Table 10: The sample efficiency of imitation learning methods. The metric for sample efficiency is based the ratio t/T, where is the first time step such that the agent achieves least expert performance continuously within the interval [t, T) In Atari, is 50 million time steps.
Atari Games
VAIL (r²))
GIRIL
ILDE w/o
ILDE w/o ndc
ILDE
BeamRider
9.84%
8.20%
9.84%
11.48%
DemonAttack
13.12%
9.84%
11.48 %
8.20%
BattleZone
98.31%
11.48%
9.84%
Qbert
9.84%
21.31%
6.56% %
8.20%
Krull
96.67%
85.20%
91.76%
72.10%
83.57 %
StarGunner
93.40%
Table 11: Sample efficiency improvements VAIL and ILDE methods versus the GIRIL, defined fa=tx where Lo the time for GIRIL reach expert-level and is the time for the method to of comparison reach expert-level (continuously until time 7). "Improve GIRIL" denotes the average improvements of IL methods versus the GIRIL across Atari games. To make the scores feasible, we limit the maximum of 1x to
Atari Games
VAIL (1*)
GIRIL
ILDE w/o b
ILDE w/o x
ILDE
BeamRider
-916.26%
0%
16.67%
0%
-6.67%
Demon
-662.21% %
0%
25%
12.5%
37.5%
BattleZone
-1.71%
0%
-1.72%
88.32%
90.06%
Q*bert
-916.26%
0%
-116.57%
33.33%
16.67%
Krull -13.46%
0%
-7.69%
15.37%
1.91%
StarGunner 0%
0%
0%
0%
6.6%
Average -418.32%
0%
-14.05%
24.92%
22.68%
demonstration data are extremely limited. We do not include HyPE in the Atari tasks because the HyPE implementation does not support Atari environments.
Table 12: Performance of IQ-Learn on Atari benchmarks.
Atari Game
IQ-Leern performance
BeamRider
0.0 + 0.0
DtmokAttack
88.5 + 92
BattleZone
3,800 870
Qbert
225 125
Krull
4,500 2,030
StarGunner
200 100
Second. we compare HyPE and ILDE,, a variant of our framework implemented with HyPE, on two MuJoCo tasks based on only a single demonstration (which is significantly more challenging than the 64 demonstrations as in |Ren al. (204 The results Table 13 show that HyPF-II.DE clearly outperforms HyPE.
Table 13: Comparison of HyPE and HyPE-ILDE MuJoCo tasks.
MuloCo Tasks
HyPE
HyPE-II.DE
Ant-v3 Hopper-v3
785.8 + 528.6
813.8 + 833.0
2.995.0# 745.4
3,372.6= 359.4
19

---

Published as a conference paper at ICLR 2025
B.2.7 ILDE FOR NOISY DEMONSTRATIONS
To test the robustness of ILDE, we generate noisy demonstrations by injecting noise into the environ- ment in the following manner: with probability Puemble random action executed instead of the one chosen by the policy (Ren et al.) 2024). Due to the limited time and computational resources, we test the robustness of ILDE for noisy demonstrations The results in Table 14 show that ILDE consistently outperforms the expert (1,918,645) and GIRIL (8.5241, .085) on the noisy demonstrations with noise rate of Premble E {0.01, 0.05, 0.1, 0.3, 0.5}. providing insights into ILDE's s robustness in real-world settings.
BeamRider without tuning any hyperparameters.
Table 14: The effect of noise rate (Phemble) on ILDE performance for noisy demonstrations on the BeamRider task.
0.0
0.01
0.05
0.1
0.3
0.5
11,453 + 2,319
12,353
2,753 8,893 + 951
9,935 + 1,952
9,728 + 1,847
10,735 + 1,198
B.2.8 CONTINUOUS CONTROL TASKS
Except for the above evaluation on the Atari games with nigh-dimensional state space and discrete action space, we also evaluated our method on continuous control tasks where the state space is low-dimensional and the action space continuous. The continuous control tasks were from MuloCo (Todorov et al. 2012).
Demonstrations. We use the demonstrations from the open-source repository "pytorch-a2c--po- -po-kkr---- gail" (Kostrikov] 2018). We compare ILDE with VAIL and GIRIL on the continuous control tasks (i.e., Reacher, Hopper, Walker2d, and HumanoidStanduup) with 4 trajectories in the demonstration for each task.
Experimental setups. Similar to the setups in Atari games, we need to pretrain a reward model for GIRIL and ILDE using the demonstrations. The training is conducted with the Adam optimizer a learning rate of 3e-4 and a mini-batch size of 32 for 10,000 epochs. In each training epoch, we sample mini-batch of data every 20 states. To evaluate the quality of our learned reward, we use the trained reward learning module to produce intrinsic rewards for policy data and trained a policy to maximize the inferred intrinsic rewards via PPO. For VAIL, we train the discriminator using the Adam optimizer with learning rate of 3e-4. The discriminator is updated at every policy step. We train the PPO policy for all imitation learning methods for 10 million steps. We compare ILDE with baselines in the measurement of average return. The average return is measured using the true rewards in the environment. All the experiments are executed in an NVIDIA A100 GPU with 40 GB memory.
Network architectures. We use 3-layer MLPs implement the discriminator for VAI, and the encoder and the decoder for GIRIL and ILDE. For fair comparison, we used an identical policy network for all methods. We used the actor-critic approach for training PPO policy for all the imitation learning methods. The number of hidden layers in each MLP is 100. Table 15 shows the architecture details of the actor-critic network for MuJoCo tasks. Table shows the MLP architectures, i.e.. GIRIL's encoder and decoder and VAIL's discriminator.
Table 15: Architectures of the actor-critic network for MuJoCo tasks, dim(S) and dim(.A represent the state dimension and action dimension, respectively.
Actor-Criti network
1 dim States
dense dim S) + 100, Tanh
dense 100 - 100, Tanh
dense 100 + dim (AA dense 100
Gaussian Distribution
Actions
Values
Hyperparameter settings. Table 17 summarized the parameter settings for MuloCo tasks.
20

---

Published as a conference paper at ICLR 2025
Table 16: Architectures of ML.P-based GIRIL's encoder and decoder, and VAIL's discriminator.
GIRIL
VAIL
encoder
decoder
discriminator
1 x dSm(S) States und Next State
x dim((A Actions and x dim(S) States
1 x dimn(A) Actions and x dim(S) States
dense dimí(S) 2 + 100, Tanh
dense diai dim(S))dim((A )dlim(A)
100, Tanh dense 100 100, Tanh
dense dim( - 100. Tanh
dense 100 100, Tanh
dense 100 100, Tanh
dense 100 - dense 100 +
dense 100 - dim(S)
split 100 into 50, 50 dease 50 + per dense - 0'
reparameterization + diau(A) A)
reparumeterization
Latent variables
dim(S) Predicted next states
pa. discrimination
Table 17: Hyperparameter settings for MuJoCo tasks.
Parameter
Setting
Actor-critic learning rate
3e-4,
linear decay
PPO clipping
0.1
Reward model learning rate
3e-4
Discount factor
0.99
GAE Lambda
0.95
Critic loss coefficient
0.5
Entropy regularisation
0.0
Epochs per batch
10
Batch size
32 x 2048 (parallel workers x time steps)
Mini batch size
32
Evaluation frequency
every 6 policy updates
GIRIL objective a
0.01
Mini batch size for training the reward model
32
Total training epoch of reward model
10, 000
VAIL's bottleneck loss weight
1.0
Information constraint I
0.2
X in Algorithm 3
10.0
Results. Figure 2 iillustrates the curves of VAIL, GIRIL, and ILDE in the four continuous control tasks. The ablation study results for MuJoCo are illustrated in Figure
C PROOF OF THEOREM 4.7
Our proof structure for Theorem largely adapts from the proof's Liu (2023b). However, our analysis different from theirs in three aspects: (1) we consider imitation learning, where the reward is time-varys varying, (2) we consider MDPs with generalized Eluder dimension, which is slightly more general than Eluder dimension (Russo & Van Roy] 2013) as shown in Zhao et al. (2023). and (3) the analysis leads sublinear batch-regret, which is stronger than sample complexity bound.
First, we have the following extended value difference lemma.
Lemma C.1 (Value difference lemma). For any policy æ, the value functions {V/ H returned by OPB(T, D, satisfy the following equation,
Jk- and (Q) Fh
Vi (81) Vi = Σ Bag ~ (Qh :), TA (-1sn) ++((|s>)))]
H
E 1 an) - F(sh an) -PP V. (Sh, an))
h=1
21

---

Published as a conference paper at ICLR 2025
2000
Average Return
-100
Average Return
200
1500
-300
1000
-400
500
VAIL (r*)
500
GIRIL
-600
0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Time steps
1e7
Time steps
(a) Reacher.
(b) Hopper.
80000
1250
Average Return
Average Return
1000
70000
750
60000
250
50000
0-
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Time steps
1e7
Time steps
(c) Walker2d.
(d) HumanoidStandup.
Figure 2: Average return vs. the number simulation steps on MuJoCo tasks. The solid lines show the mean performance over 10 random seeds. The shaded area represents the standard deviation from the mean.
Proof. For any [H], S, we have
V (s) - (s) (R,, mn(+ss) (s.
(s, - (s,,s,)) Tn (-1s) - (0% m)(ss) (
(PPAVit1, (s, PA. (s.
(s, (f(ss)))))))))-))))))))))))))))) - (-ls)
-QQi(s,-) - F(s.-)
P, 1 (x(s)))
(C.1)
where the first equality follows from the definition of the value functions V and Q+ the last equality follows from the Bellman equation for the state-action value function Qt.= in (3.2) From C.. we further obtain that for any [H], policy IT,
E. - [VT, (sN) - (sw) - E., ~a [Vă+ V (8)+1) - vol (
E.g [1 :), TA (-1sk) + (-)))>)))] E., ~u (sk.)) F(sh.)) PA V ( :), m-
Taking the sum over h E [H] yields the following result,
Vil((11) (s1)
Bxyaa (, nf((sss))]]]
h=1
#1
Th Eapes ((2. F(sn..) -) PA Vhl (ah, (-18a))]
h=l
Exy (sn.) (-1(m)))) -f((118)))]]
Σ
~ (Kh -((sn, an) PA (h an)].
h=1
h=1
22

---

Published as a conference paper at ICLR 2025
GIRIL
2000
ILDE w/o
Average Return
40
ILDE w/o
Average Return
ILDE
1500
50
1000
-60
500
0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Time steps
1e7
Time steps
(a) Reacher.
(b) Hopper.
1250
80000
Average Return
Average Return
1000
75000
750
70000
500
65000
250
60000
0
55000
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Time steps
1e7
Time steps
(c) Walker2d.
(d) HumanoidStandup.
Figure 3: Ablations on MuJoCo tasks. Average return vs. the number of simulation steps on MuloCo tasks. The solid lines show the mean performance over 10 random seeds. The shaded area represents the standard deviation from the mean.
Lemma C.2 (Lemma 17, Shani et 2020). For any policy æ, S, h E [H], we have
[s 1), mn(sss) [ < log |A/n K/2
Lemma C.3. Suppose Assumption 4.5 holds. Then with probability least 1 5, for all h [H],
E S, E A, we have
(an((s PPAV < 8/2 log(H -N((e))88) + 4e + Dr, Dr,((s, a() Dn).
Proof. We have
Σ (Ran(snaan) PA (en+s((
Σ -P (Ven)1(5u411) (acin)-PPPP...e (1
Σ (nnu(sn,au)
- (on4))))
Σ
(Veat1(5u41) - PA. Va. 1 (sh, 0,
(SN.en)DD
where the last inequality follows from the definition of PA.A in Algorithm 2.
23

---

Published as a conference paper at ICLR 2025
Then it follows that for all h [H].
Σ
(Fea(sh, an) Pa Pi an))
2
Σ (Fan(sn,
Pa (F (Px (P V
(C.2)
Note that VK, is function only depends on data in DR+1, "h+ Dr. Hence the right-hand side of the above inequality is martingale difference sequence.
Consider a function f E C(Fr.cF) & Fr. By Lemma D.2 we have with probability at least 1 1- 8/ -N((FF)).
Σ
(f(sh: PAVK,n- an)) (Pk Vr.nt1(sh Vinn+((5h+1))
(4h JEDD
Σ
(f(s))a)) Py (s), 2H 2HH 1og(HooN(HHH)l -Nz(c=)88).
(C.3)
(Sh.as)EDk
Combining C.2 and C.3 and from the definition ex- EF), we have with probability at
least 1 8/H.
Σ
(Fr,h(Shaa)) an) PP (1
log(H -Nr(cr)/8) H N/H.
(Nh (EDD
By the definition of D',,, we have with probability at least - 8/HH
(s) PA V2 P /&H² log(H N=((=)/8) + 4EFN 7 Dy, a(s,a);
from which we can complete the proof by applying the union bound over h E [H].
Lemma C.4. For any policy æ, suppose we sample trajectories D (0) " (e( Then,
with probability at least -- 8. we have
H² + B n DA, (2n: =O
dimN (Fr) (F)) (F log(1/8)
Proof. It follows from the definition of dimy (Fi.) that
dim (Fr) 2 Σ min (3, (2 :0:
(C.4)
From Definition
we further have the following upper bound for cumulative uncertainty,
D
(2h
2
1)
2
Σ
(2h
0 1-- 1 >
Σ min (1, DJ (=1²;
i=1
H
+ 1)
dim (Fr).
(C.5)
7
Thus,
Σ
Dr (2 : [t-1]] I
N ( 7
H2 +1) dimy
(C.6)
1m1
by AM-GM inequality.
24

---

Published as a conference paper at ICLR 2025
On the other hand, from Azuma-Hoeffding inequality (Lemma D.I. have with probability at least 1 6,
Drn (2h 2h ) 2y [D5, (2h:
H Mlog(1/8)
2//
2 N E Beg DD (2h: -0
H
v
/Noog(1/0)
(C.7)
Vy
Substituting
(C.6)
into C.7 yields:
Eza~
Drs (2n: D)]=
(v V
H + N
dim (Fm)(FF))onn log(1/8) (ogg1)))))
(C.8)
Lemma C.5 (Lemma 3, Liu et al. 2023b). Let ty be the index of the last policy which is used to collect fresh data at iteration k. Suppose we choose yi and m such that 1pm I 1/ H22 then for any k E N and any function g S
R+:
+ [s( ax) θ (E (g(s). an)))
Lemma C.6 (Bounding batched regret (1 + I3)). Let 7% be defined in 7 and update the policy as in Algorithm!] Suppose that Assumption 4.5 holds. Then with probability at least 1 26.
[i ph) - F(T H log 4
R=1
(V
H 1 +
+ 2HK- v logHlog(F N=((=)/0) + +cF
dim (Fr) log(1/6)
N
Proof.
Suppose that the high- probability events in Lemma
C.3 and Lemma C.5 hold.
K
Σ ( - J(x,
J(=, x*) =
Vi (81) v(()s + >
Vi (81) - v (s1)
En (sn (sh. Fu (-1sn) Thi(|19)3]]]
K
(en wutt (hh( F(sm. Pr (1+
H log Al/nn + K/2+ V8H log(H -Nse))/)(er)))) 4 +
H
H + (Fm log(1/8)
N
where the first inequality is obtained by applying the value difference lemma C.1 to both terms and then applying Lemma C.3
Lemma C.7 (Upper bound for 12). If we choose ne = O(L/VH²K) in Algorithm ] and update the reward function as shown (4.3). then with probability 1 5,
K
[L
- , z (VH2K+CB)
Proof. From Lemma D.3 we have for any R.
K
[1 .) - -(aa + < O(1/no + no 122112 K/2).
(C.9)
25

---

Published as a conference paper at ICLR 2025
From Lemmas D.1 and
D.4
with probability at least 6,
L( r) - L(T ,r']]
1=1
I (1/ne + 118 k/2 + OOV H log(1/6)) + O(EgK),
(C.10)
from which we can complete the proof by substituting the value of into C.10)
Theorem C.8 (Restatement of Theorem 4.7 Suppose Assumptions
hold. If we set = eF 1/N, 77 Vlog A H HvK, K, m IvK/H H log Al. Tl O(1/v
4.5 and
H² ), and
2/3
where K = 172 dimy (F) log N.x with probability at least 6, Algorithm yields a regret
N = [KH dimy [F) log Ne (F)) log
then
δ 8/3 (1 (dimr(F) log Ne (F) (22) +
Proof for Theorem 4.7
Throughout the proof, we suppose the events Lemma C.6 and Lemma C.7
hold simultaneously.
From the definition of Regret,
K
Regret(T) S (N/m) max
El((T" r)] ((x="r))
k=1
S (N/m).
max rER
Σ - [(Tt :1))
(C.11)
where the last inequality follows from the definition of lg. Then consider the following decomposition for the regret term,
(( (( r) E(TT, <
[J ;r) -[x ,(**tr*))) [
mAx
1=1
1=1
+ sup rR
L(T r) L(x >>]] Int(#";; (77) :n=((==;="> r])] Σ
(C.12)
To
Ts
where 7, + T3 can be controlled by the classic analysis for optimistic policy optimization, Tn is the regret of an online learning algorithm for the reward function. From Lemma C.6
H HPK
I + I3 < Õ(H log Al/n K/2) +Ō
N
②ō 11 K Kog log A
(C.13)
where the last inequality holds due to the value of n, we are choosing.
From Lemma C.7
1 < Sō @ ([IIK+EKK))) +
(C.14)
Substituting
CC133 and C.14) into C.11
(( r) ((===( .) 5 112 K log AA +EEK).
1=1
26

---

Published as a conference paper at ICLR 2025
After further substitute the value of hyperparameters into (C.11) obtain
(T) S Ō (H (d1)a (dmm((F) log (o) 2/3 + EgT)
D AUXILIARY LEMMAS
Lemma D.1
inequality). Let {r:]== be a martingale difference sequence with respect to filtration {9.) satisfying |ary] S M for some constant M, æ; meas--easur,ble, || = 0. Then for any 0 < 8 < 1, with probability at least 1 - 6, we have
Σ x4 S 2o
Self-normalized bound for scaler martingales). Consider random variables (Un, N) adapted to the filtration (Hz 0, 1, Let {na ] be sequence of real-valued random variables which is -eaa-meauurable and is conditionally g-sub-Gaussian. Then for an arbitrarily chosen À 0, for any ô O, with probability at least 6, it holds that
Lemma D.2
2a 2
Σ E. S
+ log(1/8)/A
Vn E N.
Lemma D.3 (Regret bound of online gradient descent, Orabona 2019). Let V S Rd non-empty closed convex set with diameter D. i.c., maX. || yll < D. Let l,,-.. by an arbitrary sequence of convex functions V differentiable open sets containing V. Pick any æn V and assume 7/++1 I Te, t = 1,..., Т. Then, Vu E V, the following regret bound holds:
D
T
Σ (4(xz) (((u)) -
9b
21TT
[=1
Moreover, if is constant, i.e., 11t Vt 1,... T, we have
T
(((()))) - (2) ||
21
11
19
2m
2 t=1
Lemma D.4 (I (Lemma 9, Liu et al. 2023b). There exists an absolute constant e 0 such that for any
pair of policies म, से satisfying (a 8) x an(a 8) x exp( ( (s, a)), where (Ta(s, ()) is set of functions from S A to [-1 H, 1/ H.. it holds that for any тH (s1, 01........ 10. SH, aH) e (S A)":
p (TT) S eex IPT (Tn).
27


Reviews:
[{"review_id": "odKNVeZLau", "reviewer": "Reviewer_4joU", "summary": "This paper addresses the challenge of surpassing expert performance through imitation learning. The authors propose ILDE which modifies the Generative Adversarial Imitation Learning (GAIL) framework by incorporating an intrinsic reward term to encourage exploration. The method is supported by theoretical analysis quantifying the algorithm's regret bounds. Empirical evaluation on 6 Atari games and 4 MuJoCo tasks demonstrates ILDE's improvements over baseline methods.", "questions": "Majors:\n- As mentioned in the weakness part, the biggest concern from me is the formulation. Could you please elaborate why, intuitively, adding an intrinsic reward term to the GAIL objective (Eq. 3.3) can help the policy to achieve a better performance than the expert? I think to achieve a better performance than the expert, the agent must understand the intention of the expert and extrapolate it. I cannot see how an intrinsic reward guarantee to help for this. Maybe there are some hidden assumptions about the task or the demonstrations?\n- In the introduction, you list the formulation of the problem itself as one of the contribution, but you write the formulation in the Preliminary section. Could you elaborate what is the difference between the formulation in this paper with previous works, for example [1]?\n- From the results in Table 1, ILDE w/o $r^k$ is much better than ILDE w/o $b$. Is this a desired behaviour that using only the intrinsic reward allows achieving a better performance than the imitation reward? Does this not imply that the selected 6 Atari tasks can be solved by only doing exploration?\n\nMinors:\n- In line 41, your citation of GAIL is wrong.\n- No number is bold in the Qbert line of Table 1.\n- In the result section, you only include Table 1 in the main text but spend the majority of the text discussing the results on Tables 9 and 10 in the appendix. Can you reorganize the paper to make the text and table better align with each other?\n\n[1] Yu *et al.*, Intrinsic Reward Driven Imitation Learning via Generative Model, ICML 2020", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- This paper studies a seemingly interesting problem of learning a policy with imitation learning to be much better than the demonstration.\n- The paper offer some theoretical analysis of the proposed method.", "weaknesses": "- The problem formulation of optimizing the GAIL loss + intrinsic reward to achieve a better performance than the expert is not convincing.\n- The proposed method is only evaluated on limited tasks and the results lack explanation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u5Dq4IBklx", "reviewer": "Reviewer_HCFX", "summary": "This paper focuses on imitation learning from limited demonstrations, where the goal is to learn a policy that either imitates the expert directly or learns a policy that is better than the expert. Prior techniques focused on exploration-based approaches for this problem, with the goal of designing useful intrinsic motivation bonuses to incentivize the agent to explore its (large) state space to ultimately learn a better-than-expert policy. This paper proposes the use of two bonuses: an exploration bonus that encourages the agent to visit novel state-action pairs in the environment, and another bonus that incentivizes the agent to deviate from the expert trajectories themselves. The authors propose theoretical justification of their method, achieving a poly(H, T) regret w.r.t. the true expert.", "questions": "I don't have any additional questions.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I am not the best theoretician by any means, but this paper is very theoretically sound. I think the regret analysis was done well, which is the main theoretical contribution of this work.\n\nExperimental results were quite expansive and good. The authors seemed to ablate the right variables (e.g. what happens if we turn our state-action based exploration bonus off, what happens if we turn our imitation reward off) to analyze exactly what is different about their method compared to a standard baseline (e.g. GIRIL). In terms of domains of interest, the authors also covered a variety of benchmarks across a variety of RL benchmarks, including Atari and MuJoCo tasks.", "weaknesses": "It would be interesting to see how just the state-action exploration bonus does as well (e.g. the state entropy bonus stays, but the GIRIL-like intrinsic reward is not used). To me, both of these methods seem to be doing the same thing (one on the state-action space as a whole, and one focusing on the expert trajectories), so it would be good to see an experiment that keeps the state entropy bonus, doesn't use the demonstration-based intrinsic reward, and keeps the standard imitation learning reward, which I understand is different from GIRIL (where their intrinsic reward bakes in both the true demonstration actions and the demonstration-based intrinsic reward). This is pretty minor, but would love to see this myself.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WcI5JUZkN1", "reviewer": "Reviewer_FEJp", "summary": "This paper introduces a novel imitation learning algorithm, ILDE (Imitation Learning with Double Exploration). The goal of ILDE is to achieve beyond-expert performance by improving sample efficiency and addressing the challenges posed by limited demonstrations. ILDE implements two exploration strategies: (1) optimistic policy optimization with an exploration bonus that encourages the exploration of high-uncertainty state-action pairs, and (2) curiosity-driven exploration that rewards states deviating from the expert’s demonstration trajectory. Empirical results show that ILDE outperforms state-of-the-art algorithms in terms of sample efficiency and return on Atari and MuJoCo tasks. The authors also provide a theoretical framework, showing that ILDE’s regret grows sublinearly with the number of episodes.", "questions": "Typos:\n\n1. In line 41, the reference of GAIL is wrong.\n2. In line 256, “desccent” should be descent.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "1. This paper proposes a novel ILDE framework, which combines optimistic exploration with curiosity-driven exploration for imitation learning.", "weaknesses": "1. **Theoretical Limitations**\n    1. First, the assumptions made in this paper are overly strong. In Assumption 4.2, the authors directly assume that the statistical error resulting from using a finite set of expert demonstrations is small. This assumption is both strong and unreasonable, as a primary goal in imitation learning theory is to examine the relationship between the number of expert trajectories and the error bound [1]. Additionally, Assumption 4.1 presumes that the loss function is convex with respect to the reward parameters, making the reward optimization a straightforward online convex optimization problem. In practice, however, the reward optimization problem may be highly non-convex due to the use of neural networks.\n    2. Second, the proposed method in Algorithm 1 requires knowledge of the entire MDP transition function. In line 7 of Algorithm 2, computing the intrinsic reward $\\mathcal{L}_{\\tau^{E}, h}$ depends on the true transition function. However, if the true transition function were known, online exploration would no longer be necessary.\n2. **Experimental Limitations**\n    \n    In the experiments, this paper does not compare its approach with current state-of-the-art (SOTA) imitation learning (IL) methods. In Atari games, the SOTA method is IQ-Learn [2]. In MuJoCo environments, the SOTA methods include IQ-Learn and HyPE [3]. The absence of these methods in the experiments makes it unclear whether the proposed approach outperforms current methods.\n    \n3. **Unsubstantiated Conjecture**\n    \n    The goal of this paper is to develop an IL approach capable of achieving performance beyond the expert level. To this end, the paper is based on the conjecture that optimizing the uncertainty-regularized loss shown in Eq. (3.6) enables performance surpassing the expert’s. However, no evidence or justification is provided to support this conjecture.\n    \n\nReferences:\n\n[1] Nived Rajaraman et al., Toward the fundamental limits of imitation learning.\n\n[2] Divyansh Garg et al., Iq-learn: Inverse soft-q learning for imitation.\n\n[3] Juntao Ren et al., Hybrid inverse reinforcement learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MtfYszAdMX", "reviewer": "Reviewer_4bF6", "summary": "This paper proposes a new framework for imitation learning, which succinctly, leverage informed interaction to improve policy performance in the absence of an abundance of expert demonstrations. In particular, two exploration strategies are considered: optimistic policy optimization with an exploration bonus, and curiosity-driven exploration of states beyond the demonstration trajectories. The authors provide extensive theoretical justification and prove performance bounds to support their algorithm, and also provide empirical results to support their method, testing metrics in both sample efficiency and final performance.", "questions": "1. While the authors provide a practical instantiation, I am under the impression that proposed ILDE framework should be flexible, and can be generalized to other empirical instantiations. If this is true, have the authors explored applying these exploration bonuses on top of existing methods and seeing how they fair?\n2. How sensitive is ILDE to the choice of hyperparameters, such as the exploration bonus weight and the learning rate for curiosity-driven exploration? What is the amount of hyperparameter tuning that one had to search over in order to reach the optimal performance for these exploration bonus trainings? \n3. How does ILDE perform in environments where the expert demonstrations are noisy or imperfect? What is the consequence of data quality on the two exploration bonuses, and how much additional tuning may be needed to stabilize performance? This would provide insight into ILDE’s robustness in real-world settings.\n4. I'm wondering whether ILDE could be combined with techniques in model-based (inverse) reinforcement learning to further improve sample efficiency?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The proposed method seems novel and to be well-grounded in theoretical derivations. While learning from limited expert demonstrations has been explored in prior work, this work introduces a unique combination of two exploration strategies to tackle this regime. \n* The paper provides solid theoretical foundations for ILDR, and is thorough in stating their assumptions. As noted, Theorem 4.7 is the first theoretical guarantee for imitation learning in MDPs with nonlinear\nfunction approximation.\n* The problem is well motivated, and algorithms are well presented. Particularly, the integration of optimistic policy optimization and curiosity-driven exploration are clearly labeled. The empirical experiments also do a good job ablating on these two exploration models, highlighting both are needed for best performance.\n* Relatedly, the empirical results are promising, demonstrating that ILDE achieves better performance over baselines on metrics such as final reward, # games better than expert, and average performance against experts over the considered environments.", "weaknesses": "* There are more modern baselines that come to mind when looking at the few expert-demonstration regime, such as IQLearn [1], which also demonstrate reasonable performance in the low expert-data regime, or the model-free and/or model-based inverse RL algorithms in Hybrid Inverse Reinforcement Learning [2], which also prove competitive sample complexity bounds. The experiments section could benefit from slightly more baselines.\n* Although the theoretical framework is well-explained, the practical implementation details could use a bit more clarity. In particular, how are exploration bonuses computed in practice, and how much work is involved in hyperparameter search (i.e. the original ranges considered, type of search used, and how the authors recommend searching over these hyperparameters when faced with a new environment/dataset)? (Relatedly, see Question 2 below.)\n* Relatedly, I personally felt the tradeoff between exploration and imitation could have been explored a bit more thoroughly in the paper. For example, I'm curious how algorithms may perform in situations where the demonstrations may not be fully optimal. (Relatedly, see Question 3 below.)\n* The ILDR-NPG variant as I understood it seems susceptible to incur large computational costs due to using natural policy gradients for policy optimization. The paper doesn't seem to fully address this issue/potential ways to mitigate it.\n\n\n[1] Garg, Divyansh, et al. \"Iq-learn: Inverse soft-q learning for imitation.\" Advances in Neural Information Processing Systems 34 (2021): 4028-4039.\n\n[2] Ren, Juntao, et al. \"Hybrid inverse reinforcement learning.\" International Conference on Machine Learning (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)