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 an offline imitation learning (IL) problem with an addition of supplementary imperfect demonstrations. To solve this problem, the paper proposes DemoDICE which regularizes a distribution-matching objective of IL by a KL divergence between the agent distribution and a mixed of expert and imperfect distributions. DemoDICE finds an optimal state-action distribution of this regularized objective by using a dual-program technique similar to that of OptiDICE (Lee et al., 2021) for offline RL with an improvement in terms of stability. Given the optimal state-action distribution, DemoDICE extracts the expert policy by performing weighted behavioral cloning.  Empirical evaluation on Mujoco tasks with D4RL datasets show that DemoDICE can efficiently and effectively solve the offline IL problem.   ### Contributions - A new learning problem combining offline IL and IL with imperfect demonstrations.  - A new model-free offline IL method based on dual-program optimization.