Learning Efficient Planning-based Rewards for Imitation Learning

1 INTRODUCTION . Imitation learning ( IL ) offers an alternative to reinforcement learning ( RL ) for training an agent , which mimics the demonstrations of an expert and avoids manually designed reward functions . Behavioral cloning ( BC ) ( Pomerleau , 1991 ) is the simplest form of imitation learning , which learns a policy using supervised learning . More advanced methods , inverse reinforcement learning ( IRL ) ( Ng & Russell , 2000 ; Abbeel & Ng , 2004 ) seeks to recover a reward function from the demonstrations and train an RL agent on the recovered reward function . In the maximum entropy variant of IRL , the aim is to find a reward function that makes the demonstrations appear near-optimal on the principle of maximum entropy ( Ziebart et al. , 2008 ; 2010 ; Boularias et al. , 2011 ; Finn et al. , 2016 ) . However , most state-of-the-art IRL methods fail to meet the performance of demonstrations in highdimensional environments with limited demonstration data , e.g. , a one-life demonstration in Atari domain ( Yu et al. , 2020 ) . This is due to the main goal of these IRL approaches is to recover a reward function that justifies the demonstrations only . The rewards recovered from limited demonstration data would be vulnerable to the overfitting problem . Optimizing these rewards from an arbitrary initial policy results in inferior performance . Recently , Yu et al . ( 2020 ) proposed generative intrinsic reward learning for imitation learning with limited demonstration data . This method outperforms expert and IRL methods in several Atari games . Although GIRIL uses the prediction error as curiosity to design the surrogate reward that encourages ( pushes ) states away from the demonstration and avoids overfitting , the curiosity also results in ambiguous quality of the rewards in the environment . In this paper , we focus on addressing the two key issues of previous methods when learning with limited demonstration data , i.e. , 1 ) overfitting problem , and 2 ) ambiguous quality of the reward function . To address these issues , we propose to learn a straightforward surrogate reward function by learning to plan from the demonstration data , which is more reasonable than the previous intrinsic reward function ( i.e. , the prediction error between states ) . Differential planning modules ( DPM ) is potentially useful to achieve this goal , since it learns to map observation to a planning computation for a task , and generates action predictions based on the resulting plan ( Tamar et al. , 2016 ; Nardelli et al. , 2019 ; Zhang et al. , 2020 ) . Value iteration networks ( VIN ) ( Tamar et al. , 2016 ) is the representative one , which represents value iteration as a convolutional neural network ( CNN ) . Meaningful reward and value maps have been learned along with the useful planning computation , which leads to policies that generalize well to new tasks . However , due to the inefficiency of summarizing complicated transition dynamics , VIN fails to scale up to the Atari domain . To address this challenge , we propose a novel method called variational planning-embedded reward learning ( vPERL ) , which is composed of two submodules : a planning-embedded action back-tracing module and the transition dynamics module . We leverage a variational objective based on the conditional variational autoencoder ( VAE ) ( Sohn et al. , 2015 ) to jointly optimize the two submodules , which greatly improves the generalization ability . This is critical for the success of achieving a straightforward and smooth reward function and value function with limited demonstration data . As shown in Figure 1 , vPERL learns meaningful reward and value maps that attends to the resulting region of the agent executing an action , which indicates meaningful planning computation . However , directly applying VIN in Atari domain in the way of supervised learning ( Tamar et al. , 2016 ) only learns reward and value maps that attend no specific region , which usually results in no avail . Empirical results show that our method outperforms state-of-the-art IRL methods on multiple Atari games and continuous control tasks . Remarkably , our methods achieve performance that is up to 58 times of the demonstration . Moreover , the average performance improvement of our method is 1,139.1 % of the demonstration over eight Atari games . 2 BACKGROUND AND RELATED LITERATURE . Markov Decision Process ( MDP ) ( Bellman , 1966 ) is a standard model for sequential decision making and planning . An MDP M is defined by a tuple ( S , A , T , R , γ ) , where S is the set of states , A is the set of actions , T : S ×A× S → R+ is the environment transition distribution , R : S → R is the reward function , and γ ∈ ( 0 , 1 ) is the discount factor ( Puterman , 2014 ) . The expected discounted return or value of the policy π is given by V π ( s ) = Eτ [ ∑ t=0 γ tR ( st , at ) |s0 = s ] , where τ = ( s0 , a0 , s1 , a1 , · · · ) denotes the trajectory , in which the actions are selected according to π , s0 ∼ T0 ( s0 ) , at ∼ π ( at|st ) , and st+1 ∼ T ( st+1|st , at ) . The goal in an MDP is to find the optimal policy π∗ that enables the agent to obtain high long-term rewards . Generative Adversarial Imitation Learning ( GAIL ) ( Ho & Ermon , 2016 ) extends IRL by integrating adversarial training technique for distribution matching ( Goodfellow et al. , 2014 ) . GAIL performs well in low-dimensional applications , e.g. , MuJoCo . However , it does not scale well to high-dimensional scenarios , such as Atari games ( Brown et al. , 2019a ) . Variational adversarial imitation learning ( VAIL ) ( Peng et al. , 2019 ) improves on GAIL by compressing the information via variational information bottleneck . GAIL and VAIL inherit problems of adversarial training , such as instability in training process , and are vulnerable to overfitting problem when learning with limited demonstration data . We have included both methods as comparisons to vPERL in our experiments . Generative Intrinsic Reward driven Imitation Learning ( GIRIL ) ( Yu et al. , 2020 ) leverage generative model to learn generative intrinsic rewards for better exploration . Though GIRIL outperforms previous IRL methods on several Atari games , the reward map of GIRIL is ambiguous and less informative , which results in inconsistent performance improvements in different environments . In contrast , our vPERL learns efficient planning-based reward that is more straightforward and informative . We have included GIRIL as a competitive baseline in our experiments . Differentiable planning modules perform end-to-end learning of planning computation , which leads to policies that generalize to new tasks . Value iteration ( VI ) ( Bellman , 1957 ) is a well-known method for calculating the optimal value V ∗ and optimal policy π∗ : Vn+1 ( s ) = maxaQn ( s , a ) , where Qn ( s , a ) = R ( s , a ) + γ ∑ s′ T ( s ′|s , a ) Vn ( s′ ) denotes the Q value in the nth iteration . The value function Vn in VI converges as n→∞ to V ∗ , from which the optimal policy may be derived as π∗ ( s ) = argmaxaQ∞ ( s , a ) . Value iteration networks ( VIN ) ( Tamar et al. , 2016 ) proposes to embed value iteration ( VI ) ( Bellman , 1957 ) process with a recurrent convolutional network , and generalizes well in conventional navigation domains . VIN assumes there is some unknown embedded MDP M where the optimal plan in M contains useful information about the optimal plan in the original MDP M . VIN connects the two MDPs with a parametric reward function R = fR ( s ) . Nardelli et al . ( 2019 ) proposes value propagation networks ( VPN ) generalize VIN for better sample complexity by employing value propagation ( VProp ) . Recently , universal value iteration networks ( UVIN ) extends VIN to spatially variant MDPs ( Zhang et al. , 2020 ) . Although VIN can be extended to irregular spatial graphs by applying graph convolutional operator ( Niu et al. , 2018 ) , most of the VIN variants still focus on solving the conventional navigation problems ( Zhang et al. , 2020 ) . In this paper , we extend differentiable planning module to learn an efficient reward function for imitation learning on limited demonstration data . We dig more on leveraging the learned reward function for imitation learning ; while previous related work of VIN focuses more on the value function . Therefore , our work is complementary to the research of VIN and its variants . Note that any differentiable planning module can be embedded in our method . As a simple example , we utilize the basic VIN as a backbone to build our reward learning module . 3 VARIATIONAL PLANNING-EMBEDDED REWARD LEARNING . In this section , we introduce our solution , variational planning-embedded reward learning ( vPERL ) . As illustrated in Figure 2 , our reward learning module is composed of two submodules to accomplish planning-embedded action back-tracing and explicit forward transition dynamics modeling . 3.1 ACTION BACK-TRACING AND FORWARD DYNAMICS MODELLING IN VPERL . Planning-embedded action back-tracing . Instead of directly applying VIN for policy learning ( Tamar et al. , 2016 ) , we build our first submodule qφ ( at|st , st+1 ) for action back-tracing . As illustrated in the top section of Figure 2 , we first obtain the reward map R = fR ( st , st+1 ) on an embedded MDP M , where fR is a convolutional layer . A VI module fV I takes in the reward map R , and effectively performs K times of VI by recurrently applying a convolutional layer Q for K times ( Tamar et al. , 2016 ) . The Q layer is then max-pooled to obtain the next-iteration value V . The right-directed circular arrow in a light-blue color denotes the direction of convolutions . Then , we simply obtain the action from the intermediate optimal value V ∗ by an action mapping function : at = fa ( V ∗ ) . On these terms , we build our planning-embedded action back-tracing submodule , which is formally represented as qφ ( at|st , st+1 ) = fa ( fV I ( fR ( st , st+1 ) ) ) . Since the convolutional kernel is incapable of summarizing the transition dynamics in a complex environment , directly training this submodule is still insufficient for learning efficient reward function and planning computation in an environment like Atari domain . Explicit transition dynamics modeling via inverse VI . To address this , we further build upon another submodule pθ ( st+1|at , st ) for explicit transition dynamics modeling . We build the submodule based on the inverse VI module , which is a NN architecture that mimics the process of the inverse version of VI . The implementation of the inverse VI module is straightforward . We first map the action for the intermediate optima value in another embedded MDP M ′ by a value mapping function : V ′ ∗ = fV ′ ( st , at ) . Then , we apply the inverse VI module to obtain the reward map R′ . The inverse VI module f ′V I takes in the intermediate value V ′ and recurrently apply a deconvolutional layer Q′ for K times on the value to obtain the reward map R′ . The left-directed circular arrow in a purple color denotes the direction of deconvolutions . To accomplish the transition , we map the obtained R′ to the future state by : st+1 = fs′ ( R′ ) . The transition modeling is therefore presented as pθ ( st+1|at , st ) = fs′ ( f ′V I ( fV ′ ( st , at ) ) ) , which is a differentiable submodule , and can be trained simultaneously with the action back-tracing submodule . Variational solution to vPERL . A variational autoencoder ( VAE ) ( Kingma & Welling , 2013 ) can be defined as being an autoencoder whose training is regularised to avoid overfitting and ensure that the latent space has good properties that enable generative process . To avoid the learned planningbased reward overfitting to the demonstration , we optimize both submodules in a unified variational solution , which follows the formulation of conditional VAE ( Sohn et al. , 2015 ) . Conditional VAE is a conditional generative model for structured output prediction using Gaussian latent variables , which is composed of a conditional encoder , decoder and prior . Accordingly , we regard the action back-tracing module qφ ( z|st , st+1 ) as the encoder , pθ ( st+1|z , st ) as the decoder , and pθ ( z|st ) as the prior . Our vPERL module is maximized with the following objective : L ( st , st+1 ; θ , φ ) = Eqφ ( z|st , st+1 ) [ log pθ ( st+1|z , st ) ] −KL ( qφ ( z|st , st+1 ) ‖pθ ( z|st ) ) − αKL ( qφ ( ât|st , st+1 ) ‖πE ( at|st ) ) ] ( 1 ) where z is the latent variable , πE ( at|st ) is the expert policy distribution , ât = Softmax ( z ) is the transformed latent variable , α is a positive scaling weight . The first two terms on the RHS of Eq . ( 1 ) in the first line denote the evidence lower bound ( ELBO ) of the conditional VAE ( Sohn et al. , 2015 ) . These two terms are critical for our reward learning module to perform planning-based action backtracing and transition modeling . Additionally , we integrate the third term on the RHS of Eq . ( 1 ) in the second line to further boost the action back-tracing . The third term minimizes the KL divergence between the expert policy distribution πE ( at|st ) and the action distribution qφ ( ât|st , st+1 ) , where ât = Softmax ( z ) is transformed from the latent variable z . In this way , we train the forward state transition and action back-tracing simultaneously . Algorithm 1 Imitation learning via variational planning-embedded reward learning ( vPERL ) . 1 : Input : Expert demonstration data D = { ( si , ai ) } Ni=1 . 2 : Initialize policy π , and the dual planning networks . 3 : for e = 1 , · · · , E do 4 : Sample a batch of demonstration D̃ ∼ D. 5 : Train vPERL module on D̃ to converge . 6 : end for 7 : for i = 1 , · · · , MAXITER do 8 : Update policy via any policy gradient method , e.g. , PPO on the learned surrogate reward rt . 9 : end for 10 : Output : Policy π . Note that the full objective in Eq . ( 1 ) is still a variational lower bound of the marginal likelihood log ( pθ ( st+1|st ) ) . Accordingly , it is reasonable to maximize this as an objective of our reward learning module . By optimizing the objective , we improve the forward state transition and action back-tracing . As a result , our reward learning module efficiently models the transition dynamics of the environment . During training , we use the latent variable z as the intermediate action . After training , we will calculate the surrogate rewards from the learned reward map . As shown in Figure 1 , our method learns meaningful reward map , which highlights the resulting region of an agent executing one action . To leverage such meaningful information , we calculate two types of rewards that both correspond to the highlighted informative region , i.e. , rt = RMax = maxR and rt = RMean = meanR , which uses the maximum and mean value of the reward map R , respectively . Algorithm 1 summarizes the full training procedure of imitation learning via vPERL . The process begins by training a vPERL module for E epochs ( steps 3-6 ) . In each training epoch , we sample a mini-batch demonstration data D̃ with a mini-batch size of B and maximize the objective in Eq . ( 1 ) . Then in steps 7-9 , we update the policy π via any policy gradient method , e.g. , PPO ( Schulman et al. , 2017 ) , so as to optimize the policy π with the learned surrogate reward function rt .

Summary: This paper proposes a method for inverse reinforcement learning that incorporates a differential planning module. Explicit transition dynamics modeling with inverse value iteration is added to promote meaningful reward learning. Empirical evaluations on several high-dimensional Atari environments and 2 continuous control environments are provided which show improvements over existing inverse reinforcement learning baselines when given only one-life demonstrations. Some visuals are also presented to show that the proposed method is able to learn more meaningful reward maps than previous methods.