Confidence-Aware Imitation Learning from Demonstrations with Varying Optimality
Year: 2022

Abstract:
Most existing imitation learning approaches assume the demonstrations are drawn from experts who are optimal, but relaxing this assumption enables us to use a wider range of data. Standard imitation learning may learn a suboptimal policy from demonstrations with varying optimality. Prior works use confidence scores or rankings to capture beneficial information from demonstrations with varying optimality, but they suffer from many limitations, e.g., manually annotated confidence scores or high average optimality of demonstrations. In this paper, we propose a general framework to learn from demonstrations with varying optimality that jointly learns the confidence score and a well-performing policy. Our approach, Confidence-Aware Imitation Learning (CAIL) learns a well-performing policy from confidence-reweighted demonstrations, while using an outer loss to track the performance of our model and to learn the confidence. We provide theoretical guarantees on the convergence of CAIL and evaluate its performance in both simulated and real robot experiments. Our results show that CAIL significantly outperforms other imitation learning methods from demonstrations with varying optimality. We further show that even without access to any optimal demonstrations, CAIL can still learn a successful policy, and outperforms prior work.

heading: 1 Introduction
text: We consider an imitation learning setting that learns a well-performing policy from a mixture of demonstrations with varying optimality ranging from random trajectories to optimal demonstrations. As opposed to standard imitation learning, where the demonstrations come from experts and thus are optimal, this benefits from a larger and more diverse source of data. Note that different from setting that the demonstrations are optimal but lack some causal factors [31], in our setting, the demonstrations can be suboptimal. However, this introduces a new set of challenges. First, one needs to select useful demonstrations beyond the optimal ones. We are interested in settings where we do not have sufficient expert demonstrations in the mixture so we have to rely on learning from sub-optimal demonstrations that can still be successful at parts of the task. Second, we need to be able to filter the negative effects of useless or even malicious demonstrations, e.g., demonstrations that implicitly fail the tasks.
To address the above challenges, we propose to use a measure of confidence to indicate the likelihood that a demonstration is optimal. A confidence score can provide a fine-grained characterization of each demonstration’s optimality. For example, it can differentiate between near-optimal demonstrations or adversarial ones. By reweighting demonstrations with a confidence score, we can simultaneously learn from useful but sub-optimal demonstrations while avoiding the negative effects of malicious ones. So
∗Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
our problem reduces to learning an accurate confidence measure for demonstrations. Previous work learns the confidence from manually annotated demonstrations [30], which are difficult to obtain and might contain bias—For example, a conservative and careful demonstrator may assign lower confidence compared to an optimistic demonstrator to the same demonstration. In this paper, we remove restrictive assumptions on the confidence, and propose an approach that automatically learns the confidence score for each demonstration based on evaluation of the outcome of imitation learning. This evaluation often requires access to limited evaluation data.
We propose a new algorithm, Confidence-Aware Imitation Learning (CAIL), to jointly learn a wellperforming policy and the confidence for every state-action pair in the demonstrations. Specifically, our method adopts a standard imitation learning algorithm and evaluates its performance to update the confidence scores with an evaluation loss, which we refer to as the outer loss. In our implementation, we use a limited amount of ranked demonstrations as our evaluation data for the outer loss. We then update the policy parameters using the loss of the imitation learning algorithm over the demonstrations reweighted by the confidence, which we refer to as the inner loss. Our framework can accommodate any imitation learning algorithm accompanied with an evaluation loss to assess the learned policy.
We optimize for the inner and outer loss using a bi-level optimization [5], and prove that our algorithm converges to the optimal confidence assignments under mild assumptions. We further implement the framework using Adversarial Inverse Reinforcement Learning (AIRL) [14] as the underlying imitation learning algorithm along with its corresponding learning loss as our inner loss. We design a ranking loss as the outer loss, which is compatible with the AIRL model and only requires easy-to-access ranking annotations rather than the exact confidence values.
The main contributions of the paper can be summarized as:
• We propose a novel framework, Confidence-Aware Imitation Learning (CAIL), that jointly learns confidence scores and a well-performing policy from demonstrations with varying optimality.
• We formulate our problem as a modified bi-level optimization with a pseudo-update step and prove that the confidence learned by CAIL converges to the optimal confidence in O(1/ √ T ) (T
is the number of steps) under some mild assumptions. • We conduct experiments on several simulation and robot environments. Our results suggest that
the learned confidence can accurately characterize the optimality of demonstrations, and that the learned policy achieves higher expected return compared to other imitation learning approaches.

heading: 2 Related Work
text: Imitation Learning. The most common approaches for imitation learning are Behavioral Cloning (BC) [20, 4, 23, 22, 3], which treats the problem as a supervised learning problem, and Inverse Reinforcement Learning (IRL), which recovers the reward function from expert demonstrations and finds the optimal policy through reinforcement learning over the learned reward [1, 21, 32]. More recently, Generative Adversarial Imitation Learning (GAIL) [18] learns the policy by matching the occupancy measure between demonstrations and the policy in an adversarial manner [15]. Adversarial Inverse Reinforcement Learning (AIRL) [14] and some other approaches [13, 17] improve upon GAIL by simultaneously learning the reward function, and the optimal policy. However, these approaches assume that all the demonstrations are expert demonstrations, and cannot learn a well-performing policy when learning from demonstrations with varying optimality.
Learning from Demonstrations with Varying Optimality: Ranking-based. Ranking-based methods learn a policy from a sequence of demonstrations annotated with rankings [2, 25, 29, 10]. T-REX learns a reward from the ranking of the demonstrations and learns a policy using reinforcement learning [8]. In our work, we assume access to rankings of a small subset of the demonstrations. The reward function learned from such a small number of rankings by T-REX may have low generalization ability to out of distribution states. D-REX improves T-REX by automatically generating the rankings of demonstrations [9], and SSRR further finds the structure of the reward function [12]. These techniques automatically generate rankings under the assumption that a perturbed demonstration will have a lower reward than the original demonstration, which is not necessarily true for random or malicious demonstrations that can be present in our mixture. DPS utilizes partial orders and pairwise comparisons over trajectories to learn and generate new policies [19]. However, it requires interactively collecting feedback, which is not feasible in our offline learning setting.
Learning from Demonstrations with Varying Optimality: Confidence-based. Confidence-based methods assume each demonstration or demonstrator holds a confidence value indicating their optimality and then reweight the demonstrations based on this value for imitation learning. To learn the confidence, 2IWIL requires access to ground-truth confidence values for the demonstrations to accurately learn a confidence predictor [30]. Tangkaratt et al. require that all the actions for a demonstration are drawn from the same noisy distribution with sufficiently small variance [27]. ICGAIL implicitly learns the confidence score by aligning the occupancy measure of the learned policy with the expert policy, but requires a set of ground-truth labels to estimate the average confidence [30]. Following works relax the assumption of access to the ground-truth confidence, but still require more optimal demonstrations than non-optimal ones in the dataset [26]. Other works require access to the reward of each demonstration [11]. All of these methods either rely on a specific imitation learning algorithm or require strong assumptions on the confidence. To move forward, we propose a general framework to jointly learn the confidence and the policy. Our framework is flexible as it can use any imitation learning algorithm as long as there exists a compatible outer loss, i.e., the outer loss can evaluate the quality of the imitation learning model.

heading: 3 Problem Setting
text: We formulate the problem of learning from demonstrations with varying optimality as a Markov decision process (MDP): M = 〈S,A, T ,R, ρ0, γ〉, where S is the state space, A is the action space, T : S × A × S → [0, 1] is the transition probability, ρ0 is the distribution of initial states, R : S × A → R is the reward function, and γ is the discount factor. A policy π : S × A → [0, 1] defines a probability distribution over the action space in a given state. The expected return, which evaluates the quality of a policy, can be defined as ηπ = Es0∼ρ0,π [ ∑∞ t=0 γ
tR(st, at)], where t indicates the time step.
We aim to learn a policy that imitates the behavior of a demonstrator d following policy πd who provides a set of demonstrations Ξ = {ξ1, . . . , ξD} and ξi ∼ πd. Each trajectory is a sequence of stateaction pairs ξ = {s0, a0, . . . , sN}, and the expected return of a trajectory is ηξ = ∑N−1 t=0 γ
tR(st, at). A common assumption in classical imitation learning work is that the demonstrations are drawn from the expert policy πd = π∗, i.e., , the policy that optimizes the expected return of the MDPM [18, 14]. Here, we relax this assumption so that the demonstrations may contain non-expert demonstrations or even failures—drawn from policies other than π∗. Given the demonstration set D, we need to assess our confidence in each demonstration. To achieve learning confidence over this mixture of demonstrations, we rely on the ability to evaluate the performance of imitation learning. This can be achieved by using an evaluation loss trained on evaluation data, DE (as shown in Fig. 1). In our implementation, we rely on a small amount of rankings between trajectories as our evaluation data: DE = ηξ1 ≥ · · · ≥ ηξm . To summarize, our framework takes a set of demonstrations with varying optimality D as well as a limited amount of evaluation data DE along with an evaluation loss to find a well-performing policy. Note that unlike prior work [30], we do not assume that optimal demonstrations always exist in the demonstration set, and CAIL can still extract useful information from D while avoiding negative effects of non-optimal demonstrations.

heading: 4 Confidence-Aware Imitation Learning
text: In our framework, we adopt an imitation learning algorithm with a model Fθ parameterized by θ and a corresponding imitation learning loss Lin, which we refer to as inner loss (as shown in Figure 1). We assign each state-action pair a confidence value indicating the likelihood of the state-action pair appearing in the well-performing policy. The confidence can be defined as a function mapping from a state-action pair to a scalar value β : S ×A → R. We aim to find the optimal confidence assignments β∗ to reweight state-action pairs within the demonstrations. We then conduct imitation learning from the reweighted demonstrations using the inner imitation loss Lin to learn a well-performing policy. Here, we first define the optimal confidence β∗ and describe how to learn it automatically.
Defining the Optimal Confidence. We define the distribution of state-action pairs visited by a policy π based on the occupancy measure ρπ : S×A → R: ρπ(s, a) = π(a|s) ∑∞ t=0 γ
tP (st = s|π), which can be explained as the un-normalized distribution of state transitions that an agent encounters when navigating the environment with the policy π. We can normalize the occupancy measure to form the
state-action distribution: pπ(s, a) = ρπ(s,a)∑ s,a ρπ(s,a)
. Recall that πd is the policy that the demonstrations are derived from, which can potentially be a mixture of different expert, suboptimal, or even malicious policies. We reweight the state-action distribution of the demonstrations to derive a new state-action distribution, which corresponds to another policy πnew: pπnew(s, a) = β(s, a)pπd(s, a). Our goal is to find the optimal confidence β∗ that ensures the derived policy πnew maximizes the expected return:
β∗(s, a) = arg max β ηπnew . (1)
With such β∗(s, a), we can conduct imitation learning from the reweighted demonstrations to maximize the expected return with the provided demonstrations.
Learning the Confidence. We will learn an estimate of the confidence score β without access to any annotations of the ground-truth values based on optimizing two loss functions: The inner loss and the outer loss. The inner loss Lin is accompanied with the imitation learning algorithm encouraging imitation, while the outer loss Lout captures the quality of imitation learning, and thus optimizing it finds the confidence value that maximizes the performance of the imitation learning algorithm.
Specifically, we first learn the imitation learning model parameters θ∗ that minimize the inner loss:
θ∗(β) = arg min θ E(s,a)∼β(s,a)p πd(s,a) Lin(s, a; θ, β) (2)
We note that the inner loss Lin(s, a; θ, β) refers to settings where (s, a) is sampled from the distribution β(s, a)pπd(s,a), and hence implicitly depends on β. Thus we need to find the optimal β∗, which can be estimated by minimizing an outer loss Lout:
β∗out = arg min β Lout(θ∗(β)). (3)
This evaluates the performance of the underlying imitation learning algorithm with respect to the reward with limited evaluation dataDE (e.g. limited rankings if we select a ranking loss as our choice of Lout; which we will discuss in detail in Sec. 4.3).

heading: 4.1 Optimization of Outer and Inner Loss
text: We design a bi-level optimization process consisting of an inner-level optimization and an outerlevel optimization to simultaneously update the confidence β and the model parameters θ. Within the outer-level optimization, we first pseudo-update the imitation learning parameters to build a connection between β and the optimized parameters θ′ with the current β. We then update β to make the induced θ′ minimize the outer loss Lout in Eqn. (3). The inner-level optimization is to find the imitation learning model parameters that minimize inner loss Lin with respect to the confidence β. We introduce the details of the optimization below. We use τ to denote the number of iterations. Note that the losses in this section are all computed based on the expectation over states and actions.
Outer-Level Optimization: Updating β. Let βτ be the confidence at time τ . Using βτ , we first pseudo-update the imitation learning parameters θ using gradient descent. Here pseudo-update means that the update aims to compute the gradients of β but does not really change the value of θ. Let θ′0 = θτ be the current imitation learning model parameters, and we update θ ′ as:
θ′t+1 = θ ′ t − µ∇θ′Lin(s, a; θ′t, βτ ), (4)
where µ is the learning rate, t is the pseudo-updating time step for θ′. We will update θ′ with respect to the fixed βτ after convergence of Eqn. (4). After updating θ′, we now update β using gradient descent with the outer loss Lout from Eqn. (3):
βτ+1 = βτ − α∇βLout(θ′), (5)
where α is the learning rate for updating β. Intuitively, updating β as in Eqn. (5) aims to find the fastest update direction of θ′ for decreasing the outer loss Lout. Though we compute gradients of gradients for β here, β is only a one-dimension scalar for each state-action pair and within each iteration of training, we only sample a mini-batch of thousands of state-pairs for update. Thus, within each iteration, the total dimension of β is small and computing the gradient of gradient is not costly.
Inner-Level Optimization: Updating θ. With the updated βτ+1, we now will update θ using gradient descent, where we denote the initialization as θ0 = θτ .
θt+1 = θt − µ∇θLin(s, a; θt, βτ+1). (6)
After convergence, we set θτ+1 = θ. With the two updates introduced above (outer and inner optimization), we finish one update iteration with setting βτ to βτ+1 using the converged value from Eqn. (5) and θτ to θτ+1 using the converged value from Eqn. (6).
In each iteration of the above optimization—in the steps of pseudo-updating and the steps of updating the imitation learning model—multiple gradient steps are required for convergence, meaning that there is a nested loop of gradient descent algorithms. The nested loop costs quadratic time and is inefficient especially for deep networks. To further accelerate the optimization, we propose an approximation, which only updates θ once in the pseudo-updating and the updating steps. Therefore, the new updating rule can be formalized as follows:
θ′τ+1 = θτ − µ∇θLin(s, a; θτ , βτ ), βτ+1 = βτ − α∇βLout(θ′τ+1), θτ+1 = θτ − µ∇θLin(s, a; θτ , βτ+1).
(7)

heading: 4.2 Theoretical Results
text: We analyze the convergence of the proposed bi-level optimization algorithm for the CAIL framework and derive the following theorems. We provide the detailed proofs of these theorems in Appendix. Theorem 1. (Convergence) Suppose the outer loss Lout is Lipschitz-smooth with constant L, the inequality ∇θLout(θτ+1)>∇θLin(θτ , βτ+1) ≥ C||∇θLin(θτ , βτ+1)||2 (8) holds for a constant C ≥ 0 in every step τ ,2 and the learning rate satisfies µ ≤ 2CL , then the outer loss decreases along with each iteration: Lout(θτ+1) ≤ Lout(θτ ), and the equality holds if ∇βLout(θτ ) = 0 or θτ+1 = θτ . Remark 1. The inequality in the assumption of Theorem 1 (Eqn. 8) indicates that the directions of the gradients of Lout and Lin with respect to θ should be close. Intuitively only when the two gradient directions align, we can decrease the evaluation loss Lout by updating θ with Lin.
Theorem 1 ensures that the confidence and the imitation learning parameters monotonically decrease the outer loss. When the gradient of the outer loss with respect to β is zero, β converges to the optimal confidence that minimizes the outer loss, i.e., , β∗ in Eqn. (1). With the optimal confidence, we can learn a well-performing policy from more useful demonstrations by reweighting them. Thus, the learned imitation model induces lower outer loss (has higher-quality) than the imitation learning model learned from the original demonstrations in the dataset without reweighting.
2We remove (s, a) in Lin for notation simplicity.
Theorem 2. (Convergence Rate) Under the assumptions in Theorem 1, let g(θ, β) = θ − µ∇θLin(s, a; θ, β) (9)
We assume that Lout(g(θ, β)) is Lipschitz-smooth w.r.t. β with constant L1, Lin and Lout have σbounded gradients, and the norm of ∇β∇θLin(θ;β) is bounded by σ1. L is the Lipschitz-smooth constant for Lout w.r.t. g(θ, β) as shown in Theorem 1. Consider the total training steps as T , we set α = C1√
T , for some constant C1 where 0 < C1 ≤ 2L1 and µ = C2 T for some constant C2. Then:
min 1≤τ≤T
E[||∇βLout(θτ )||2] ≤ O (
1√ T
) . (10)
Remark 2. The assumptions of Theorem 2 are Lipschiz-smoothness and bounded first-order and second-order gradients of Lin and Lout, which are satisfied for typical Lin and Lout such as the cross-entropy loss of AIRL and the ranking loss in our implementation of CAIL in Section 4.3.
With the bound on the convergence rate, the gradient of the outer loss with respect to β is gradually getting close to 0, which means that β gradually converges to the optimal confidence β∗ that minimizes the outer loss if Lout is convex with respect to β.

heading: 4.3 An Implementation of CAIL
text: To implement CAIL, we need to adopt an imitation learning algorithm whose imitation loss will be the inner loss. We also need to design an outer loss on the imitation learning algorithm to evaluate the quality of imitation given some evaluation data DE (e.g. partial ranking annotations). Based on the above considerations, as an instance of the implementation of CAIL, we use Adversarial Inverse Reinforcement Learning (AIRL) [14] as our imitation learning model. We use the imitation loss of AIRL as the inner loss, and a ranking loss (based on a partial ranking of trajectories) as the outer loss. AIRL and the ranking loss are compatible since AIRL can induce the reward function from the discriminator within the model, and the ranking loss can penalize the mismatches of the trajectory rankings computed by the induced reward function and the ground-truth rankings from the evaluation data DE . Furthermore, the implementation only requires the ranking of a subset of demonstrations {ξi}mi=1 ⊂ Ξ, i.e., , DE = ηξ1 ≥ ηξ2 ≥ · · · ≥ ηξm , which is much easier to access than the exact confidence value annotations [6, 19] since confidence not only reflects the rankings of different demonstrations but also how much one demonstration is better than the other.
AIRL consists of a generator G parameterized by θG as the policy, and a discriminator parameterized by θD. The generator and the discriminator are trained in an adversarial manner as in [15] to match the occupancy measures of the policy and the demonstrations. We write the loss Lin as: LDin (s, a; θD, β) = E(s,a)∼β(s,a)pπd(s,a) [− logD(s, a)] + E(s,a)∼πθG [− log(1−D(s, a))], (11)
LGin (s, a; θG) = E(s,a)∼πθG [logD(s, a)− log(1−D(s, a))], (12) where LDin is the inner loss for the discriminator, LGin is the inner loss for the generator and πθG is the policy derived from the generator. The discriminator D is learned by minimizing the loss LDin , which aims to discriminates the state-action pair (s, a) drawn from πθG and the state-action pair (s, a) drawn from πd. The generator parameter θG is trained to minimize the loss LGin , which enables the generator to generate state-action pairs that are similar to the state transitions in the demonstrations.
For the outer loss, AIRL approximates the reward function by the discriminator parameters, i.e.,R′θD . We compute η′ξi = ∑N t=0 γ tR′θD (st, at) as the expected return of a trajectory using the rewardR ′ θD . Then we penalize the mismatches of the rankings derived by η′ξi and the ground-truth rankings:
Lout(θD) = ∑ i ∑ j>i RK [ η′ξi , η ′ ξj ; I[ηξi > ηξj ] ] , (13)
where I[ηξi > ηξj ] is 1 if ηξi > ηξj and otherwise is −1. RK is defined as a revised version of the widely-used margin ranking loss with margin as 0:
RK [ η′ξi ; η ′ ξj , ηξi , ηξj ] =  max(0,−I[ηξi > ηξj ](η′ξi − η ′ ξj )), |η ′ ξi − η ′ ξj | >
max(0, 1
4 (I[ηξi > ηξj ](η′ξi − η ′ ξj )− ) 2), |η′ξi − η ′ ξj | ≤
(14)
We revised the original margin ranking loss within a range around the point of (η′(ξi)−η′(ξj)) = 0 to make it Lipschitz smooth. If we adopt small enough , the functionality of the revised marginal ranking loss is close to the original one. In all the experiments, we use = 10−5.

heading: 5 Experiments
text: In this section, we conduct experiments on the implementation of CAIL in Sec. 4.3. We verify the efficacy of the CAIL in simulated and real-world environments. We report the results on various compositions of demonstrations with varying optimality. The code is available on our website3
We conduct experiments in four environments including two MuJoCo environments (Reacher and Ant) [28] in OpenAI Gym [7], one Franka Panda Arm4 simulation environment, and one real robot environment with a UR5e robot arm5. For each environment, we collect a mixture of optimal and non-optimal demonstrations with different optimality to show the efficacy of CAIL. We investigate the performance with respect to the optimality of demonstrations ranging from failures to near-optimal or optimal demonstrations. We provide the implementation details and more results on the sensitivity of the parameters, and visualize the learned confidence in the supplementary materials.
Source of Demonstrations. For MuJoCo environments, following the demonstration collecting method in [30], we train a reinforcement learning algorithm and select four intermediate policies as policies with varying optimality and the converged policy as the optimal policy, so that the demonstrations range from worse-than-random ones to near-optimal ones. We draw 20% of demonstrations from each policy. For the RL algorithm, we use SAC [16] for the Reacher environment and PPO [24] for the Ant environment. For the Franka Panda Arm simulation and the real robot environment with UR5e, we hand-craft demonstrations with optimality varying continuously from near-optimal ones to unsuccessful ones to approximate the demonstration collecting process from demonstrators with different levels of expertise. We label only 5% of the demonstrated trajectories with rankings since we target realistic settings where only a small number of rankings are available for the demonstrations.
Baselines. We compare CAIL with the most relevant works in our problem setting including: the state-of-the-art standard imitation learning algorithms: GAIL [18], AIRL [14], imitation learning from suboptimal demonstration methods including two confidence-based methods, 2IWIL and ICGAIL [30], and three ranking-based methods, T-REX [8], D-REX [9], and SSRR [12]. GAIL and AIRL learn directly from the mixture of optimal and non-optimal demonstrations. T-REX needs demonstrations paired with rankings, so we provide the same number of rankings as our approach. For D-REX and SSRR, we further generate rankings by disturbing demonstrations as done in their papers. For 2IWIL and IC-GAIL—that need a subset of demonstrations labeled with confidence—we label the subset of ranked demonstrations with evenly-spaced confidence, i.e., , the highest expected return as confidence 1, and the lowest expected return as 0. This is a reasonable approximation of the confidence score with no prior knowledge available. For a fair comparison, we re-implement 2IWIL with AIRL as its backbone imitation learning method. For the RL algorithm in T-REX, D-REX, and SSRR, we also use PPO. DPS [19] requires interactively collecting demonstrations and the approach in Cao et al. [11] requires the ground truth reward of demonstrations, which are both not implementable under the assumptions in our setting, so we do not include them.

heading: 5.1 Results
text: Reacher and Ant. In the Reacher, the end effector of the arm is supposed to reach a final location. Figure 2(a) shows the optimal trajectories of the joint and the end effector in green, which illustrates the policy reaching the location with the minimum energy cost, and the trajectories with lower optimality in red and orange, where the agent just spins around the center and wastes energy without reaching the target. We collect 200 trajectories in total, where each trajectory has 50 interaction steps.
In Ant, the agent has four legs, each with two links and two joints. Its goal is to move in the x-axis direction as fast as possible. Figure 2(b) illustrates the demonstrated trajectories, where green shows the optimal one, and red shows suboptimal trajectories (darker colors show lower optimality). In
3https://sites.google.com/view/cail 4https://www.franka.de 5https://www.universal-robots.com/products/ur5-robot
optimal demonstrations, the agent moves quickly along the x-axis, while in suboptimal ones, it moves slowly to other directions. We collect trajectories with 200,000 interaction steps in total.
As shown in Figure 2(e) and 2(f), CAIL achieves the highest expected return compared to the other methods and experiences fast convergence. For Reacher, the p-value6 between CAIL and the closest baseline method, T-REX, is 5.4054× 10−6 (statistically significant). For Ant, the p-value between CAIL and the closest baseline method, 2IWIL, is 0.1405. CAIL outperforms standard imitation learning methods, GAIL and AIRL, because CAIL selects more useful demonstrations, and avoids the negative influence of harmful demonstrations. We observe that 2IWIL and IC-GAIL do not perform well because neighboring demonstrations in a given ranking are not guaranteed to have the same distance in terms of confidence score and thus the evenly-spaced confidence values derived from rankings are likely not accurate. All the ranking-based methods do not perform well. For T-REX, the potential reason can be that the rankings of a subset of demonstrations are not enough to learn a generalizable reward function covering states. For D-REX and SSRR, the automatically generated rankings can be incorrect since we also have unsuccessful demonstrations— which can at times be worse than random actions—and perturbing such demonstrations is not guaranteed to produce demonstrations that imply rankings.
Robot Arm. We further conduct experiments in more realistic environments: a simulated Franka Panda Arm and a real UR5e robot arm. As shown in Figure 2(c) and 2(d), we design a task to let the robot arm pick up a bottle, avoid the obstacle, and put the bottle on a target platform. In the optimal demonstrations in green, the arm takes the shortest path to avoid the obstacle, and puts the bottle on the target, while in suboptimal ones in red (where, similar to before, the brightness of the trajectories indicates their optimality), the arm detours, does not reach the target, and even at times collides with the obstacle. The suboptimal demonstrations represent a wide range of optimality from near-optimal ones (small detour) to adversarial ones (colliding). We vary the initial position of the robot end-effector and the goal position within an initial area and goal area respectively. For both simulated and real robot environments, we collect trajectories with 200,000 interaction steps in total.
As shown in Figure 2(g) and 2(h), CAIL outperforms other methods in expected return in both the simulated and real robot environments. For the simulated robot arm environment, the p-value between CAIL and the closest baseline, 2IWIL, is 0.0974. For the real robot environment, the p-value between CAIL and the closest baseline, AIRL, is 0.0209 (statistically significant). In particular, in the real robot environment, CAIL achieves a low standard deviation while other methods especially AIRL, IC-GAIL, D-REX and GAIL suffer from an unstable performance. The results demonstrate that
6All p-values are computed by the student’s t-test and the null hypothesis is the performance of CAIL is equal to or smaller than the baseline methods.
CAIL can work stably in the real robot environment. We report the success rate—rate that the robot successfully reaches the target within the time limit without colliding with the obstacle—and videos of sample policy rollouts in the supplementary materials.
Demonstrations with Different Optimality. We show the performance of different methods with demonstrations at different levels of optimality in the Reacher environment. We fix 20% of the total demonstrations to be optimal and make the remaining 80% demonstration drawn from the same suboptimal policy. We vary the optimality of this policy to investigate the performance change with respect to different optimalities. Another way to obtain different optimality is to vary the ratio of optimal demonstrations. We show the results of both varying optimality in Figure 3(a) and 3(b) respectively. We observe that CAIL consistently outperforms or performs comparably to other methods with demonstrations at different optimality. Also, CAIL performs more stably while the baselines suffer from a performance drop at specific optimality levels.
Learning from Only Non-optimal Demonstrations. We verify that CAIL can also learn from solely non-optimal demonstrations without relying on any optimal demonstrations. We remove the optimal demonstrations in the Ant environment and use the remaining demonstrations to conduct imitation learning. As shown in Figure 3(c), CAIL still achieves the best performance among all the methods, which demonstrates that even with all demonstrations being non-optimal, CAIL still can learn useful knowledge from those demonstrations with higher expected return and induce a better policy. The highest p-value between CAIL and the closest baseline (2IWIL) is 0.0067, which indicates the performance gain is statistically significant. We observe that the performance of AIRL first increases and then decreases. This is because even though the demonstrations are suboptimal, there are potentially optimal state-action transitions leading to the initial high performance. However, at this early training stage, the AIRL model does not converge yet and the model parameters can still change rapidly. After training a sufficient number of steps, the AIRL model observes both useful and less useful transitions and converges to the average return of all the demonstrations.

heading: 6 Conclusion
text: Summary. We propose a general learning framework, Confidence-Aware Imitation Learning, for imitation learning from demonstrations with varying optimality. We adopt standard imitation learning algorithms with their corresponding imitation loss (inner loss), and leverage an outer loss to evaluate the quality of the imitation learning model. We simultaneously learn a confidence score over the demonstrations using the outer loss and learn the policy through optimizing the inner loss over the confidence-reweighted distribution of demonstrations. Our framework is applicable to any imitation learning model with compatible choices of inner and outer losses. We provide theoretical guarantees on the convergence of CAIL and show that the learned policy outperforms baselines on various simulated and real-world environments under demonstrations with varying optimality.
Limitations and Future Work. Although we propose a flexible framework to address the problem of imitation learning from demonstrations with varying optimality, our work is limited in a few ways: To learn a well-performing policy, we still require that the dataset consists of demonstrations that encode useful knowledge for policy learning. We also require that the demonstrations and the imitation agent have the same dynamics. In the future, we plan to learn from demonstrations with more failures and relax the assumptions of the demonstrator and imitator having the same dynamics.

heading: 7 Acknowledgement
text: We would like to thank Tong Xiao for inspiring discussions about the backbone imitation learning algorithms and her help on the experiments. This work is funded by Tsinghua GuoQiang Research Institute, FLI grant RFP2-000, NSF Awards 1941722 and 1849952, and DARPA HiCoN project.

Review summaries:
* This paper proposes a novel imitation learning algorithm, Confidence-Aware Imitation Learning (CAIL), to learn policy from suboptimal data.   * In particular, this proposed algorithm iteratively does two gradient updates:     1. learns a function, \beta, that maps (state, action) to a confidence score using an evaluation dataset (e.g., a limited amount of partial rankings) with an outer loss.     2. learn the imitation learning model parameters \theta by reweighting the full IL dataset (could be suboptimal, without labeled rankings) using \beta, and optimizing for an inner loss.   * The authors proposed a convergence rate of the proposed method as a bi-level optimization problem (Theorem 2).  * The authors conduct a comprehensive evaluation between the proposed algorithm with a set of baseline algorithms.   * Domains: MuJoCo reacher, MuJoCo ant, Franka Panda simulation, UR5e real robot.   * Baselines:     * Standard imitation learning algorithms: GAIL, AIRL.     * Confidence-based methods: 2IWIL, ICGAIL.     * Ranking-based methods: T-REX, D-REX, SSRR.   * Results:     * Reacher: CAIL outperformed other methods in expected return.       * All comparisons (between CAIL and another method) are statistically significant.     * Ant: CAIL outperformed other methods in expected return.       * The p-value between CAIL and the closest baseline method, 2IWIL, is 0.1405.     * Franka Panda simulation: CAIL outperformed other methods in expected return.       * The p-value between CAIL and the closest baseline, 2IWIL, is 0.0974.     * UR5e real robot: CAIL outperformed other methods in expected return.       * All comparisons (between CAIL and another method) are statistically significant.     * Demonstrations with Different Optimality       * CAIL consistently outperforms or performs comparably to other methods with demonstrations at different optimality.       * CAIL performs more stably while the baselines suffer from a performance drop at speciﬁc optimality levels.     * Learning from Only Non-optimal Demonstrations       * Ant: CAIL outperformed other methods in expected return.         * All comparisons (between CAIL and another method) are statistically significant. 

This paper considers the problem of offline imitation learning from demonstrations when these demonstrations may have been generated from a mixture of optimal, suboptimal, and even adversarial (i.e. worse than random) agents. To address this problem, the authors introduce Confidence-Aware Imitation Learning or CAIL, an imitation learning methodology where the policy is optimized jointly with a weight function that dynamically reweighs the "expert" demonstrations so that policy more closely imitates the high quality demonstrations. To enable the training of this weight function, CAIL requires, similarly to prior work, some additional supervision in the form of demonstration ranks, namely the quality of a small collection of demonstrations are ranked relative one another. Two theorems show that CAIL converges at a given rate under intuitive assumptions. Finally a collection of experiments show that, for both simulated environments and real-world robotic settings, CAIL outperforms competitive baselines. 

This paper investigates imitation learning under the setting where there are non-interactive state-action demonstrations with varying optimality, and among them a small percentage has trajectory-level ranking annotations. The authors propose an imitation learning framework that can leverage all demonstrations, although some of them are suboptimal or even malicious, and potentially outperform the demonstrators. The idea is to formulate learning a state-action confidence function (which weights demonstrations) as a bi-level optimization problem. This closely ties IRL and ranking matching. The authors prove the convergence rate of the resulting algorithm -- CAIL -- and show empirically both in simulation and in real-word that CAIL outperforms recent algorithms designed for similar settings.  

The paper proposes Confidence-Aware Imitation Learning (CAIL) that learns from sub-optimal demonstrations collected by different policies. The main idea is to train an imitating policy from demonstrations weighted by an estimated confidence scores, while the confidence scores are simultaneously estimated by using the imitating policy and ranked demonstrations. The paper also presents the convergence proofs of the algorithm. Experiments on simulated and real-world environments indicate that CAIL is more effective at learning from sub-optimal demonstrations when compared to existing algorithms.