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 assumes no access to the reward values and attempts to learn a policy by starting just with one demonstration to define the reward. For obtaining the reward, the authors rely on the ideas from Value Iteration Networks (VIN) method and they add the modules that help to deal with cases with complex transition dynamics. The resulting method is tested on atari domain and on continuous control tasks.