DemoDICE: Offline Imitation Learning with Supplementary Imperfect Demonstrations

1 INTRODUCTION . Reinforcement learning ( RL ) ( Sutton et al. , 1998 ) aims to learn a intelligent behavioral strategy based on reward feedback . Although RL has achieved remarkable success in many challenging domains , its practicality and applicability are still limited in two respects . First , we need to specify the reward function , which may be non-trivial to do so in many real-world problems that require complex decision making . Second , the standard RL setting assumes online interaction with the environment during the intermediate stages of learning , which is infeasible for mission-critical tasks . Imitation learning ( IL ) ( Pomerleau , 1991 ; Ng & Russell , 2000 ) addresses the first limitation of RL , where the agent is trained to mimic the expert from demonstration instead of specifying the reward function . It is well known that adopting supervised learning for training the imitating agent , commonly referred to as behavioral cloning ( BC ) , is vulnerable to the distribution drift ( Ross et al. , 2011 ) . Thus , most of the successful IL algorithms rely on online experiences collected from the environment by executing intermediate policies during training . Recent progress made by adversarial imitation learning ( AIL ) ( Ho & Ermon , 2016 ; Ke et al. , 2019 ; Kostrikov et al. , 2020 ) achieving stateof-the-art results on challenging imitation tasks still relies on such an online training paradigm . Unfortunately , in many realistic tasks such as robotic manipulation and autonomous driving , online interactions are either costly or dangerous . Offline RL ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; 2020 ; Levine et al. , 2020 ; Wang et al. , 2020 ; Lee et al. , 2021 ; Kostrikov et al. , 2021 ) aims to address these concerns by training the agent from the pre-collected set of experiences without online interactions . To prevent an issue caused by the distributional shift , offline RL algorithms mitigate the phenomenon by constraining the shift or making a conservative evaluation of the policy being learned . In this paper , we are concerned with offline IL problems . Finding an effective algorithm for these problems is tricky . For instance , naively extending the offline RL algorithms ( which assume a reward function ) to the offline IL setting does not work . In practice , expert demonstrations are scarce due to the high cost of obtaining them . Thus , they typically cover only a small fraction of the state and action spaces , which in turn makes the distribution drift issue even more stand out compared with the standard offline RL setting with a reward function . We mitigate this issue by assuming a large number of supplementary imperfect demonstrations , without requiring any level of optimality for these imperfect demonstrations ; they may contain expert or near-expert trajectories ( Wu et al. , 2019 ; Wang et al. , 2021 ) as well as non-expert ones all together . This generality covers the situations from real-world applications , but at the same time , it poses a significant challenge for the design of a successful offline IL algorithm . In this paper , we propose DemoDICE , a novel model-free algorithm for offline IL from expert and imperfect demonstrations . We formulate an offline IL objective which not only mitigates distribution shift from the demonstration-data distribution but also naturally utilizes imperfect demonstrations . Our new formulation allows us to compute a closed form solution , which learns a policy in the space of stationary distributions but suffers from the instability issue in practice . We tackle the issue by proposing an alternative objective , which leads to a stable algorithm in practice while keeping the optimal stationary distribution . Finally , we introduce a method to extract the expert policy from a learned stationary distribution in a simple yet effective way . Our extensive evaluations show that DemoDICE achieves performance competitive to or better than a state-of-the-art off-policy IL algorithm in the offline-IL tasks with expert and imperfect demonstrations . 2 PRELIMINARIES . 2.1 MARKOV DECISION PROCESS ( MDP ) . We assume an environment modeled as a Markov Decision Process ( MDP ) , defined by tuple M = 〈S , A , T , R , p0 , γ〉 , where S is the set of states , A is the set of actions , T : S × A → ∆ ( S ) is the probability p ( st+1|st , at ) of making transition from state st to state st+1 by executing action at at timestep t , R : S × A → R is the reward function , p0 ∈ ∆ ( S ) is the distribution of the initial state s0 , and γ ∈ [ 0 , 1 ] is the discount factor . A policy π : S → ∆ ( A ) of MDP M is a mapping from states of M to distributions over actions . For the given policy π , the stationary distribution dπ is defined as follows : dπ ( s , a ) = ( 1− γ ) ∞∑ t=0 γtp ( st = s , at = a ∣∣s0 ∼ p0 ( · ) , st ∼ T ( ·|st−1 , at−1 ) , at ∼ π ( ·|st ) ) We assume a precollected dataset DE of ( s , a , s′ ) tuples generated by the expert , and a precollected imperfect dataset DI ( generated by unknown degrees of optimality ) . More precisely , for the ( underlying ) expert policy ’ s stationary distribution dE ( s , a ) , we assume that ( s , a , s′ ) ∈ DE is sampled as ( s , a ) ∼ dE , s′ ∼ T ( ·|s , a ) . We define DU = DE ∪ DI , the union of two datasets , and denote the corresponding stationary distribution of the dataset DU as dU . We denote the trajectories generated by expert policy as expert trajectories , and trajectories generated by non-expert policies as non-expert trajectories . The expert demonstrations consist only of expert trajectories and imperfect demonstrations consist of a mixture of expert and non-expert trajectories . In this paper , we assume that the quality of imperfect demonstrations is unknown . 2.2 IMITATION LEARNING . Behavior cloning ( BC ) is a classical IL approach , which attempts to find a function that maps s to a via supervised learning . The standard BC finds a policy π by minimizing the negative log-likelihood : min π JBC ( π ) : = min π − 1 |D| ∑ ( s , a ) ∈D log π ( a|s ) . ( 1 ) However , it is known to be brittle ( Ross et al. , 2011 ) when the interaction with the environment deviates from the scarce trajectories in DE . In such cases , BC fails to recover expert policies . One of the notable approaches for IL is to formulate the problem as distribution matching ( Ho & Ermon , 2016 ; Ke et al. , 2019 ; Kostrikov et al. , 2020 ) . When instantiated with the KL divergence widely used in previous IL works ( Ke et al. , 2019 ; Kostrikov et al. , 2020 ) , the approach amounts to finding a policy π by optimizing the following objective : max π −DKL ( dπ‖dE ) = E ( s , a ) ∼dπ [ log dE ( s , a ) dπ ( s , a ) ] . ( 2 ) Since we can not directly access the exact value of dE ( s , a ) and dπ ( s , a ) , rwe estimate their ratio using samples from dE ( s , a ) and dπ ( s , a ) , given as follows : max c : S×A→ ( 0,1 ) E ( s , a ) ∼dE [ log c ( s , a ) ] + E ( s , a ) ∼dπ [ log ( 1− c ( s , a ) ) ] . ( 3 ) Here , the optimal discriminator c∗ recovers log c∗ ( s , a ) − log ( 1 − c∗ ( s , a ) ) = log d E ( s , a ) dπ ( s , a ) . Based on this connection between generative adversarial networks ( GANs ) and IL , AIL algorithms focus on recovering the expert policy ( Ho & Ermon , 2016 ; Kostrikov et al. , 2019 ) . However , the agent obtain samples from dπ through interaction with environment , which is impossible in offline setting . Therefore , to tackle the offline IL problems , we should derive an alternative estimation without on-policy samples . 3 DEMODICE . In this section , we present a novel model-free offline IL algorithm named offline imitation learning using additional imperfect Demonstrations via stationary DIstribution Correction Estimation ( DemoDICE ) . Starting from a regularized offline IL objective which accords with offline RL algorithms , we present a formulation that does not require on-policy samples . Such formulation allows us to construct a nested optimization for offline IL from expert and imperfect demonstrations ( Section 3.1 ) . Then , we derive the closed-form solution to the sub-problem of the aforementioned optimization and obtain a simple convex optimization objective ( Section 3.2 ) . Since the objective is unstable in practice , we transform the objective to an alternative yet still convex objective ( Section 3.3 ) . Finally , we show how to extract the policy from the learned correction term ( Section 3.4 ) . 3.1 TRANSFORM CONSTRAINED OPTIMIZATION INTO NESTED OPTIMIZATION . In the context of offline RL , most works use expected return maximization with some regularization to overcome the extrapolation error in offline settings ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Nachum et al. , 2019b ; Kumar et al. , 2020 ; Lee et al. , 2021 ) . In this work , we use KL divergence minimization between dπ and dE with KL-regularization : π∗ : = arg max π −DKL ( dπ‖dE ) − αDKL ( dπ‖dU ) , ( 4 ) where α ≥ 0 is a hyperparameter that controls the balance between minimizing KL divergence with dE and preventing deviation of dπ from dU . Many online AIL algorithms estimate divergence between expert and current policy using on-policy samples , which is not available in offline scenario . In contrast , to construct a tractable optimization problem in the offline setting , we consider a problem equivalent to Equation 4 in terms of stationary distribution d : max d −DKL ( d‖dE ) − αDKL ( d‖dU ) ( 5 ) s.t ∑ a d ( s , a ) = ( 1− γ ) p0 ( s ) + γ ∑ s̄ , ā T ( s|s̄ , ā ) d ( s̄ , ā ) ∀s , ( 6 ) d ( s , a ) ≥ 0 ∀s , a . ( 7 ) The constraints ( 6-7 ) are called the Bellman flow constraints . The dual problem for the above constrained optimization problem is max d≥0 min ν −DKL ( d‖dE ) − αDKL ( d‖dU ) + ∑ s ν ( s ) ( ( 1− γ ) p0 ( s ) + γ ( T∗d ) ( s ) − ( B∗d ) ( s ) ) , ( 8 ) where ν ( s ) are the Lagrange multipliers , ( B∗d ) ( s ) : = ∑ a d ( s , a ) is the marginalization operator , and ( T∗d ) ( s ) : = ∑ s̄ , ā T ( s|s̄ , ā ) d ( s̄ , ā ) is the transposed Bellman operator . We introduce following derivations for the optimization ( 8 ) to obtain tractable optimization in the offline setting : −DKL ( d‖dE ) − αDKL ( d‖dU ) + ∑ s ν ( s ) ( ( 1− γ ) p0 ( s ) + γ ( T∗d ) ( s ) − ( B∗d ) ( s ) ) = ( 1− γ ) Es∼p0 [ ν ( s ) ] + E ( s , a ) ∼d [ γ ( T ν ) ( s , a ) − ν ( s ) − log d ( s , a ) dE ( s , a ) − α log d ( s , a ) dU ( s , a ) ] ( 9 ) = ( 1− γ ) Ep0 [ ν ( s ) ] + Ed [ γ ( T ν ) ( s , a ) − ν ( s ) + log d E ( s , a ) dU ( s , a ) ︸ ︷︷ ︸ : =r ( s , a ) − ( 1 + α ) log d ( s , a ) dU ( s , a ) ︸ ︷︷ ︸ : =w ( s , a ) ] , ( 10 ) the equality in Equation 9 holds from the following properties of transpose operators : ∑ s ν ( s ) ( B∗d ) ( s ) = ∑ s , a d ( s , a ) ( Bν ) ( s , a ) and ∑ s ν ( s ) ( T∗d ) ( s ) = ∑ s , a d ( s , a ) ( T ν ) ( s , a ) , where ( Bν ) ( s , a ) = ν ( s ) , ( T ν ) ( s , a ) = ∑ s′ T ( s ′|s , a ) ν ( s′ ) , with assumption dE ( s , a ) > 0 when d ( s , a ) > 0 ( Nachum et al. , 2019a ) . We introduce another log ratio , denoted by r ( s , a ) in Equation 10 to avoid using log d E ( s , a ) d ( s , a ) , which requires on-policy samples to estimate . Unlike log dE ( s , a ) d ( s , a ) , we can estimate r ( s , a ) in the offline setting using dE and dU , as we will discuss in the next section in detail . We change the distribution used in the expectation of Equation 10 from d to dU by following the standard trick of importance sampling as follows : ( 1− γ ) Es∼p0 [ ν ( s ) ] + E ( s , a ) ∼d [ r ( s , a ) + γ ( T ν ) ( s , a ) − ( Bν ) ( s , a ) ︸ ︷︷ ︸ : =Aν ( s , a ) ( advantage using ν ) − ( 1 + α ) logw ( s , a ) ] = ( 1− γ ) Es∼p0 [ ν ( s ) ] + E ( s , a ) ∼dU [ w ( s , a ) ( Aν ( s , a ) − ( 1 + α ) logw ( s , a ) ) ] = : L ( w , ν ; r ) . ( 11 ) As an alternative , one can convert expectation of d to dE instead of dU in Equation 10 by the similar application of the trick of importance sampling . ( 1− γ ) Es∼p0 [ ν ( s ) ] + E ( s , a ) ∼dE [ exp ( −r ( s , a ) ) w ( s , a ) ( Aν ( s , a ) − ( 1 + α ) logw ( s , a ) ) ] . In practice , due to the small number of demonstrations in dE , there may be a lack of diversity . To release this practical issue , we prefer to use dU instead of dE . In summary , DemoDICE solves the following maximin optimization : max w≥0 min ν L ( w , ν ; r ) , ( 12 ) where r is trained by using precollected datasets . Note that the solution w∗ of Equation 12 with the ground-truth ratio r is the ratio of two distributions , the stationary distribution dπ ∗ of the expert policy π∗ and the stationary distribution dU of union of expert and imperfect demonstrations .

Summary: The paper considers the problem of imitation learning, and specifically the setting where a small number of expert demonstrations is paired with some amount of suboptimal data. To tackle this setting, the paper proposes to optimize a convex constrained optimization, with linear variables denoting d^pi, linear constraints expressing the MDP transitions, and an objective composed of a linear combination of KL from the expert and KL from the suboptimal data. The paper presents derivations transforming this convex constrained optimization to a practical, unconstrained objective. The final algorithm is composed of three stages: (1) train a discriminator to get a density ratio estimator; (2) use the density ratios as a reward in a separate log-expected-exp objective (similar to REPS or ValueDICE objectives) to learn a different density ratio w; (3) extract a policy from w using weighted max likelihood training on the suboptimal data. The algorithm is evaluated on a variety of simulated robotics (mujoco) environments, and shows favorable results compared to baselines.