Confidence-Aware Imitation Learning from Demonstrations with Varying Optimality

1 Introduction . 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 . 2 Related Work . 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 can not 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 . 3 Problem Setting . 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 . 4 Confidence-Aware Imitation Learning . 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 ) .

Summary: * 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.