What Would pi* Do?: Imitation Learning via Off-Policy Reinforcement Learning

1 INTRODUCTION . Many real-world sequential decision-making problems can be tackled by imitation learning , where an expert demonstrates near-optimal behavior to an agent that then attempts to replicate that behavior in novel situations ( Argall et al. , 2009 ) . This paper considers the problem of training an agent to imitate an expert policy , given expert action demonstrations and access to the environment . The agent does not get to observe a reward signal or query the expert , and does not know the state transition dynamics ex-ante . Standard approaches to this problem based on behavioral cloning seek to imitate the expert ’ s actions , but do not reason about the consequences of actions ( Pomerleau , 1991 ) . As a result , they suffer from state distribution shift , and fail to generalize to states that are very different from those seen in the demonstrations ( Ross & Bagnell , 2010 ; Ross et al. , 2011 ) . Approaches based on inverse reinforcement learning ( IRL ) deal with this issue by fitting a reward function that represents preferences over trajectories rather than individual actions ( Ng et al. , 2000 ; Ziebart et al. , 2008 ) , and using the learned reward function to train the imitation agent through RL ( Wulfmeier et al. , 2015 ; Finn et al. , 2016 ; Fu et al. , 2017 ) . This is the core idea behind generative adversarial imitation learning ( GAIL ) , which implicitly combines IRL and RL using generative adversarial networks ( GANs ) ( Ho & Ermon , 2016 ; Goodfellow et al. , 2014 ) . GAIL is state-of-the-art in environments with low-dimensional observations , but requires additional reward augmentation and feature engineering when applied to environments with high-dimensional image observations ( Li et al. , 2017 ) , due to the difficulty of training GANs ( Kurach et al. , 2018 ) . In short , imitation learning from raw pixel inputs without prior knowledge of the domain remains a challenge . The key insight in this paper is that instead of using a learned reward function or discriminator to provide a reward signal to the imitation policy , as is done in IRL-based imitation methods and GAIL , the imitation agent can simply be rewarded for matching demonstrated actions in demonstrated states using off-policy soft Q-learning . The agent ’ s experience replay buffer is initially filled with demonstrations where rewards are set to some positive constant , and gradually accumulates additional experiences with rewards set to zero . The agent has an incentive to imitate the expert in demonstrated states , and to take actions that lead it back to demonstrated states when it encounters new states that differ from the demonstrations . We call this algorithm soft Q imitation learning ( SQIL , pronounced “ skill ” ) . We motivate the SQIL algorithm by deriving it from first principles as a method for performing approximate inference under the MaxCausalEnt model of expert behavior ( Ziebart et al. , 2010 ) . Our analysis shows that SQIL maximizes a lower bound on the log-posterior of the soft Q values given the demonstrations – a bound that trades off between a pure behavioral cloning loss and a regularization term that incorporates state transition information via the soft Bellman error . The main contribution of this paper is SQIL : a practical and general imitation learning algorithm that is effective in MDPs with high-dimensional , continuous observations – including raw pixel inputs – and has a natural theoretical interpretation as approximate inference . We run experiments in the image-based Car Racing and low-dimensional Lunar Lander game environments from OpenAI Gym , the Humanoid and HalfCheetah continuous control tasks from MuJoCo , and image-based Atari Pong , to compare SQIL to two prior methods : behavioral cloning and GAIL . The results show that SQIL matches GAIL in low-dimensional environments , and significantly outperforms GAIL on the image-based Car Racing and Pong tasks . 2 PRELIMINARIES . This work builds on the maximum causal entropy ( MaxCausalEnt ) model of expert behavior ( Ziebart et al. , 2010 ; Levine , 2018 ) . In a finite-horizon Markov Decision Process ( MDP ) with a finite , discrete action space A,1 the demonstrator is assumed to follow a policy π that maximizes reward R ( s , a ) . To simplify exposition , we assume the state transition dynamics are deterministic : let s′ denote the state that follows from taking action a in state s. The policy π forms a Boltzmann distribution over actions , π ( a|s ) , exp ( Q ( s , a ) − V ( s ) ) , ( 1 ) where Q is the soft Q function , and V is the soft value function , V ( s ) , { 0 if s is a terminal state log ( ∑ a∈A exp ( Q ( s , a ) ) ) otherwise ( 2 ) In the standard MaxCausalEnt model , the soft Q values are a deterministic function of the rewards and dynamics , given by the soft Bellman equation , Q ( s , a ) , R ( s , a ) + γV ( s′ ) . ( 3 ) 2.1 APPROXIMATE MAXCAUSALENT MODEL . In the analysis in Section 3 , the standard MaxCausalEnt model is approximated by a probabilistic model that treats the soft Q values as random variables , Q ( s , a ) ∼ N ( R ( s , a ) + γV ( s′ ) , σ2 ) , ( 4 ) where σ ∈ R is a constant hyperparameter , and R ( s , a ) is a random variable with zero mean and unit variance . To simplify the analysis in Section 3 , we assume that the dynamics of the environment satisfy the property that if state si is reachable from state sj , then sj is not reachable from si – this ensures that the probabilistic graphical model that captures the joint distribution over soft Q values is acyclic . The experiments in Section 5 show that even when this assumption is violated , SQIL performs well empirically . 1Assuming a discrete action space simplifies exposition . The analysis can be extended to continuous action spaces , as long as the entropy of policy π ( ·|s ) is bounded above by some constant ( see Equation 13 ) , where π is any policy in the policy class being optimized over . The SQIL algorithm can be applied to continuous control tasks using existing soft Q-learning methods ( Haarnoja et al. , 2017 ; 2018 ) , as illustrated in Section 5.4 . 3 IMITATION LEARNING VIA APPROXIMATE MAXCAUSALENT INFERENCE . We aim for an imitation learning algorithm that generalizes to new states , without resorting to adversarial optimization procedures that would be difficult to use with image observations . Our approach is to derive a lower bound on the log-posterior of the soft Q values given the demonstrations ; in particular , a lower bound that is maximized by the SQIL algorithm . 3.1 DERIVING THE APPROXIMATE INFERENCE OBJECTIVE . Consider a demonstration rollout τ = ( s0 , a0 , s1 , a1 , ... , sT ) , where sT is a terminal state . Standard MaxCausalEnt IRL searches for the reward function that maximizes the posterior probability of the rewards given the demonstration . Our goal , however , is not to infer the expert ’ s reward function ; it is to infer the expert ’ s policy , which is represented by their soft Q function . We approach this problem by formalizing uncertainty about the soft Q values using the probabilistic model described in Section 2.1 . We fit a soft Q function to optimize the posterior probability of the soft Q values Q given the demonstration τ , where Q , { Q ( s , a ) | ( s , a ) ∈ X } , and X is a finite subset of S × A that includes the state-action pairs in τ . Since Q ( s , a ) may depend on Q ( s′ , a′ ) ( see Equations 4 and 2 ) , X is assumed to satisfy the closure property , ∀ ( s , a ) ∈ X . ( s′ not terminal =⇒ ( ∀a′ ∈ A . ( s′ , a′ ) ∈ X ) ) . ( 5 ) The goal is to derive a lower bound on the log-posterior , log p ( Q|τ ) , that is maximized by the SQIL algorithm . The first step is to marginalize out the rewards , R , { R ( s , a ) | ( s , a ) ∈ X } , as follows . p ( Q|τ ) ∝ p ( τ |Q ) p ( Q ) = ER [ p ( τ |Q , R ) p ( Q|R ) ] = ER [ p ( τ |Q ) p ( Q|R ) ] . ( 6 ) First bound . Optimizing the log of the expectation in Equation 6 is hard , for a number of reasons : we haven ’ t assumed a particular prior distribution over rewards ( only that it has zero mean and unit variance ) , and the soft Q values in Q are not conditionally independent given R. Optimizing the expectation of the log is easier . By Jensen ’ s inequality , log ( ER [ p ( τ |Q ) p ( Q|R ) ] ) ≥ ER [ log p ( τ |Q ) + log p ( Q|R ) ] . ( 7 ) The conditional likelihood of the demonstration τ given the soft Q valuesQ is given by Equation 1 : p ( τ |Q ) = exp ( T−1∑ t=0 Q ( st , at ) − V ( st ) ) . ( 8 ) The conditional likelihood of the soft Q valuesQ given the rewardsR is given by Equation 4 : p ( Q|R ) = ∏ ( s , a ) ∈X N ( Q ( s , a ) ; R ( s , a ) + γV ( s′ ) , σ2 ) , ( 9 ) where the joint distribution over Q can be factored due to the state re-entry assumption about the dynamics made in Section 2.1 and the closure assumption about X made in Equation 5 . Simplifying the lower bound in Equation 7 , ( 7 ) = T−1∑ t=0 Q ( st , at ) − V ( st ) + ∑ ( s , a ) ∈X − 1 2σ2 ( ( Q ( s , a ) − ( E [ R ( s , a ) ] + γV ( s′ ) ) ) 2 +Var [ R ( s , a ) ] ) − log √ 2πσ2 . ( 10 ) Applying the assumption from Section 2.1 that the rewards have zero mean and unit variance , ( 10 ) = T−1∑ t=0 Q ( st , at ) − V ( st ) + ∑ ( s , a ) ∈X − 1 2σ2 ( ( Q ( s , a ) − γV ( s′ ) ) 2 + 1 ) − log √ 2πσ2 . ( 11 ) Second bound . Next , the bound is intentionally loosened , in order to arrive at an objective that is maximized by the SQIL algorithm.2 By the MaxCausalEnt assumptions in Section 2 , V ( s0 ) = Eτ∼p ( τ |π ) [ T−1∑ t=0 γt ( Q ( st , at ) − γV ( st+1 ) +H ( π ( ·|st ) ) ) ] , ( 12 ) 2The experiments in Section 5 compare the performance of SQIL , which optimizes the bound in Equation 13 , to the performance of SQIL-T , which optimizes the tighter bound in Equation 11 . See Section A.1 in the appendix for a more detailed description of SQIL-T. where H denotes entropy . Maximizing Equation 11 with respect to Q yields a solution in which Q ( s , a ) − γV ( s′ ) ≤ f ( σ2 ) , where f is an increasing function of σ2 . Thus , V ( s0 ) ≤ ( f ( σ2 ) + log |A| ) / ( 1− γ ) , and ( 11 ) ≥ ( 11 ) + V ( s0 ) − ( f ( σ2 ) + log |A| ) / ( 1− γ ) . ( 13 ) The soft Q function that maximizes the lower bound in Equation 13 is Q∗ , argmax Q V ( s0 ) + T−1∑ t=0 Q ( st , at ) − V ( st ) − 1 2σ2 ∑ ( s , a ) ∈X ( Q ( s , a ) − γV ( s′ ) ) 2 . ( 14 ) Practical implementation . For environments with a small , discrete state space S and known dynamics , X can be set to S × A . For environments with a continuous state space S and unknown dynamics , we setX to be the set of state-action pairs in the demonstration rolloutsDdemo and rollouts Dsamp of the imitation policy , which are periodically sampled during training.3 We fit a parameterized soft Q function Qθ to minimize the loss , ` ( θ ) , α ∑ τ∈Ddemo ( −Vθ ( s0 ) + T−1∑ t=0 − ( Qθ ( st , at ) − Vθ ( st ) ) ) + λdemoBθ ( Ddemo , 0 ) + λsampBθ ( Dsamp , 0 ) , ( 15 ) where α , λdemo , λsamp ∈ R≥0 are constant hyperparameters , Vθ denotes the soft value function given by Qθ and Equation 2,4 and Bθ denotes the sum of squared soft Bellman errors , Bθ ( D , r ) , ∑ τ∈D T−1∑ t=0 ( Qθ ( st , at ) − ( r + γVθ ( st+1 ) ) ) 2 , ( 16 ) where r ∈ R is a constant that does not depend on the state or action . The experiments in Section 5 use a convolutional neural network or multi-layer perceptron to model Qθ , where θ are the weights of the neural network ( Schmidhuber , 2015 ) . The loss in Equation 15 bounds the log-posterior of the soft Q values given the demonstrations . Searching for a soft Q function that optimizes this objective is equivalent to performing approximate inference under the MaxCausalEnt model of expert behavior . The objective combines a behavioral cloning loss , which does not reason about the consequences of actions , with a regularization term that incorporates information about state transitions . The regularization term performs soft Bellman backups on a reward of zero – a consequence of the prior distribution over rewards , which has an expected value of zero – and can be seen as imposing a prior distribution on the soft Q values that is independent of the expert demonstrations .

Summary: This paper proposes an imitation learning algorithm framed as an off-policy RL scenario. They introduce all demonstrations into a replay buffer, with positive reward r=1. Subsequent states derived from the agent will be also introduced in the replay buffer but with reward r=0. There is no additional reward involved. The hope is that agents will learn to match the expert in states appearing in the replay buffer and not try to do much else.