Robust Imitation via Decision-Time Planning

The goal of imitation learning is to mimic expert behavior from demonstrations , without access to an explicit reward signal . A popular class of approach infers the ( unknown ) reward function via inverse reinforcement learning ( IRL ) followed by maximizing this reward function via reinforcement learning ( RL ) . The policies learned via these approaches are however very brittle in practice and deteriorate quickly even with small test-time perturbations due to compounding errors . We propose Imitation with Planning at Test-time ( IMPLANT ) , a new algorithm for imitation learning that utilizes decision-time planning to correct for compounding errors of any base imitation policy . In contrast to existing approaches , we retain both the imitation policy and the rewards model at decision-time , thereby benefiting from the learning signal of the two components . Empirically , we demonstrate that IMPLANT significantly outperforms benchmark imitation learning approaches on standard control environments and excels at zero-shot generalization when subject to challenging perturbations in test-time dynamics . 1 INTRODUCTION . The objective of imitation learning is to optimize agent policies directly from demonstrations of expert behavior . Such a learning paradigm sidesteps reward engineering , which is a key bottleneck for applying reinforcement learning ( RL ) in many real-world domains , e.g. , autonomous driving , robotics . In the presence of a finite dataset of expert demonstrations however , a key challenge with current approaches is that the learned policies can quickly deviate from intended expert behavior and lead to compounding errors at test-time ( Osa et al. , 2018 ) . Moreover , it has been observed that imitation policies can be brittle and drastically deteriorate in performance with even small perturbations to the dynamics during execution ( Christiano et al. , 2016 ; de Haan et al. , 2019 ) . A predominant class of approaches to imitation learning is based on inverse reinforcement learning ( IRL ) and involve successive application of two steps : ( a ) an IRL step where the agent infers the ( unknown ) reward function for the expert , followed by ( b ) an RL step where the agent maximizes the inferred reward function via a policy optimization algorithm . For example , many popular IRL approaches consider an adversarial learning framework ( Goodfellow et al. , 2014 ) , where the reward function is inferred by a discriminator that distinguishes expert demonstrations from roll-outs of an imitation policy [ IRL step ] and the imitation agent maximizes the inferred reward function to best match the expert policy [ RL step ] ( Ho & Ermon , 2016 ; Fu et al. , 2017 ) . In this sense , reward inference is only an intermediary step towards learning the expert policy and is discarded posttraining of the imitation agent . We introduce Imitation with Planning at Test-time ( IMPLANT ) , a new algorithm for imitation learning that incorporates decision-time planning within an IRL algorithm . During training , we can use any standard IRL approach to estimate a reward function and a stochastic imitation policy , along with an additional value function . The value function can be learned explicitly or is often a byproduct of standard RL algorithms that involve policy evaluation , such as actor-critic methods ( Konda & Tsitsiklis , 2000 ; Peters & Schaal , 2008 ) . At decision-time , we use the learned imitation policy in conjunction with a closed-loop planner . For any given state , the imitation policy proposes a set of candidate actions and the planner estimates the returns for each of actions by performing fixedhorizon rollouts . The rollout returns are estimated using the learned reward and value functions . Finally , the agent picks the action with the highest estimated return and the process is repeated at each of the subsequent timesteps . Conceptually , IMPLANT aims to counteract the imperfections due to policy optimization in the RL step by using the reward function ( along with a value function ) estimated in the IRL step for decision-time planning . We demonstrate strong empirical improvements using this approach over benchmark imitation learning algorithms in a variety settings derived from the MuJoCo-based benchmarks in OpenAI Gym ( Todorov et al. , 2012 ; Brockman et al. , 2016 ) . In default evaluation setup where train and test environments match , we observe that IMPLANT improves by 16.5 % on average over the closest baseline . We also consider transfer setups where the imitation agent is deployed in test dynamics that differ from train dynamics and the test dynamics are inaccessible to the agent during both training and decision-time planning . In particular , we consider the following three setups : ( a ) “ causal confusion ” where the agent observes nuisance variables in the state representation during training ( de Haan et al. , 2019 ) , ( b ) motor noise which adds noise in the executed actions during testing ( Christiano et al. , 2016 ) , and ( c ) transition noise which adds noise to the next state distribution during testing . In all these setups , we observe that IMPLANT consistently and robustly transfers to test environments with improvements of 35.2 % on average over the closest baseline . 2 PRELIMINARIES . Problem Setup . We consider the framework of Markov Decision Processes ( MDP ) ( Puterman , 1990 ) . An MDP is denoted by a tupleM = ( S , A , T , p0 , r , γ ) , where S is the state space , A is the action space , T : S × A × S → R≥0 are the stochastic transition dynamics , p0 : S → R≥0 is the initial state distribution , r : S ×A → R is the reward function , and γ ∈ [ 0 , 1 ) is the discount factor . We assume an infinite horizon setting . At any given state s ∈ S , an agent makes decisions via a stochastic policy π : S × A → R≥0 . We denote a trajectory to be a sequence of state-action pairs τ = ( s0 , a0 , s1 , a1 , · · · ) . Any policy π , along with MDP parameters , induces a distribution over trajectories , which can be expressed as pπ ( τ ) = p ( s0 ) ∏∞ t=0 π ( at|st ) T ( st+1|st , at ) . The return of a trajectory is the discounted sum of rewards R ( τ ) = ∑∞ t=0 γ tr ( st , at ) . In reinforcement learning ( RL ) , the goal is to learn a parameterized policy πθ that maximizes the expected returns w.r.t . the trajectory distribution . Maximizing such an objective requires interaction with the underlying MDP for simulating trajectories and querying rewards . However , in many highstakes scenarios , the reward function is not directly accessible and hard to manually design . In imitation learning , we sidestep the availability of the reward function . Instead , we have access to a finite set of D trajectories τE ( a.k.a . demonstrations ) that are sampled from an expert policy πE . Every trajectory τ ∈ τE consists of a finite length sequence of state and action pairs τ = ( s0 , a0 , s1 , a1 , · · · ) , where s0 ∼ p0 ( s ) , at ∼ πE ( ·|st ) , and st+1 ∼ T ( ·|st , at ) . Our goal is to learn a parameterized policy πθ which best approximates the expert policy given access to τE . Next , we discuss the two major families of techniques for imitation learning . 2.1 BEHAVIORAL CLONING . Behavioral cloning ( BC ) casts imitation learning as a supervised learning problem over state-action pairs provided in the expert demonstrations ( Pomerleau , 1991 ) . In particular , we learn the policy parameters by solving a regression problem with states st and actions at as the features and target labels respectively . Formally , we minimize the following objective : ` BC ( θ ) : = ∑ ( st , at ) ∈τE ‖at − πθ ( st ) ‖22 . ( 1 ) In practice , BC agents suffer from distribution shift in high dimensions , where small deviations in the learned policy quickly accumulate during deployment and lead to a significantly different trajectory distribution relative to the expert ( Ross & Bagnell , 2010 ; Ross et al. , 2011 ) . 2.2 INVERSE REINFORCEMENT LEARNING . An alternative indirect approach to imitation learning is based on inverse reinforcement learning ( IRL ) . Here , the goal is to infer a reward function for the expert and subsequently maximize the inferred reward to obtain a policy . For brevity , we focus on adversarial imitation learning approaches to IRL ( Goodfellow et al. , 2014 ) . These approaches represent the state-of-the-art in imitation learning and are also relevant baselines for our empirical evaluations . Generative Adversarial Imitation Learning ( GAIL ) is an IRL algorithm that formulates imitation learning as an “ occupancy measure matching ” objective w.r.t . a suitable probabilistic divergence ( Ho & Ermon , 2016 ) . GAIL consists of two parameterized networks : ( a ) a policy network πθ ( generator ) which is used to rollout agent trajectories ( assuming access to transition dynamics ) , and ( b ) a discriminator Dφ which distinguishes between “ real ” expert demonstrations and “ fake ” agent trajectories . Given expert trajectories τE and agent trajectories τθ , the discriminator minimizes the cross-entropy loss : ` IRL ( φ ) : = −EτE [ logDφ ( τE ) ] − Eτθ [ log ( 1−Dφ ( τθ ) ) ] . ( 2 ) We then feed the discriminator output− log ( 1−Dφ ( s , a ) ) as the inferred reward function to the generator policy . The policy parameters θ can be updated via any regular policy optimization algorithm for the RL objective , e.g. , Ho & Ermon ( 2016 ) use the TRPO algorithm ( Schulman et al. , 2015 ) . By simulating agent rollouts , GAIL seeks to match the full trajectory state-action distribution of the imitation agent with the expert as opposed to BC which greedily matches the conditional distribution of individual actions given the states . In practice , GAIL and its variants ( Li et al. , 2017 ; Fu et al. , 2017 ) outperform BC but might need excessive interactions with the training environment for sampling rollouts during training . Crucially , both BC and IRL approaches tend to fail catastrophically in the presence of small perturbations and nuisances at test-time ( de Haan et al. , 2019 ) . 3 THE IMPLANT FRAMEWORK . In the previous section , we showed that current IRL algorithms consider reward inference as an auxiliary task for imitation learning . Once the agents have been trained , the reward function is discarded and the learned policy is deployed.1 Indeed , if the RL step post reward inference ( e.g. , generator updates in GAIL ) were optimal , then the reward function provides no additional information about the expert relative to the imitation policy . However , this is far from reality , as current RL algorithms can fail to return optimal solutions due to either representational or optimization issues . For example , there might be a mismatch in the architecture of the policy network and the expert policy , and/or difficulties in optimizing non-convex objective functions . In fact , the latter challenge gets exacerbated in adversarial learning scenarios due to a non-stationary reward . Building off these observations , we propose Imitation with Planning at Test-time ( IMPLANT ) , an imitation learning algorithm that employs the learned reward function for decision-time planning . The pseudocode for IMPLANT is shown in Algorithm 1 . We can dissect IMPLANT into two sequential phases : a training phase and a planning phase . Training phase : We can invoke any IRL algorithm , e.g. , GAIL to optimize for a stochastic imitation policy πθ by optimizing for some inferred reward function rφ . Additionally , we also train a parameterized value function Vψ at this stage . Value function estimation is often a subroutine for many RL algorithms including those which are used to update the policy within the IRL setup , such as actor-critic methods ( Konda & Tsitsiklis , 2000 ) . For such algorithms , learning a value function does not incur any additional computation . Planning phase : At decision-time , we use the imitation policy along with the learned value and reward functions for closed-loop planning . We build our planner based on model-predictive control ( MPC ) ( Camacho & Alba , 2013 ) . At any given state st and time t ≥ 0 , we are interested in choosing action sequences for trajectories which maximizes the following objective : at , at+1 , · · · , = argmax at , at+1 , ··· R ( τ ) = ∞∑ t′=t γt ′−tr ( st′ , at′ ) ( 3 ) where s0 ∼ p0 and st+1 ∼ T ( ·|st , at ) for all t ≥ 0 . 1In some cases , the reward function is transferred to a new environment and a new policy is learned using the reward function and additional interactions with the new environment . See Section 5 for further discussion . Algorithm 1 : Imitation with Planning at Test-time ( IMPLANT ) 1 Input : available dynamics T̂ , expert demonstrations τE , rollout budget B , rollout policy π , horizon H , test start state s0 2 Note : For brevity , we omit relevant MDP parameters in the list of arguments 3 Function Train ( τE ) : 4 Learn a policy πθ and a reward function rφ with any existing IRL algorithm given access to demonstrations τE , e.g. , GAIL 5 Estimate a value function Vψ for πθ 6 return πθ , rφ , Vψ ; 7 Function Plan ( s , πθ , Vψ , rφ , π , H , B , T̂ ) : 8 Set s = s0 9 while agent is alive do 10 // Agent planning 11 Sample B trajectories { τ ( 1 ) , τ ( 2 ) , ... , τ ( B ) } of max length H starting from s using dynamics T̂ ; sample the first action a ( i ) 0 ∼ πθ , and sample subsequent actions from π as a ( i ) > 0 ∼ π , for i ∈ { 1 , 2 , ... , B } 12 Estimate trajectory returns R̂φ , ψ ( τ ( i ) ) using Vψ and rφ ( see Eq . 4 ) 13 Pick best action index i∗ = argmaxi R̂φ , ψ ( τ ( i ) ) and execute the best action a ( i ∗ ) 0 14 // Environment feedback 15 Observe true reward r ( s , a ( i ∗ ) 0 ) and true next state s ∼ T ( ·|s , a ( i∗ ) 0 ) 16 end This objective has also been applied for model-based RL with a learned dynamics model and blackbox access to the rewards function ( Nagabandi et al. , 2018 ; Chua et al. , 2018 ) . Unlike the RL setting however , we do not know the reward function for imitation learning . The true dynamics model may be available for planning ( i.e. , T̂ = T ) as in Ho & Ermon ( 2016 ) or can be estimated from expert demonstrations or online interactions ( Baram et al. , 2016 ) . Hence , we can do rollouts as before in regular model-based RL but need to rely on learned estimates for the reward function . In particular , we use the learned reward function rφ up to a fixed horizon H and a terminal value function Vψ thereafter to estimate the trajectory return as : R ( τ ) ≈ t+H−1∑ t′=t γt ′−trφ ( st′ , at′ ) + γ HVψ ( sH ) : = R̂φ , ψ ( τ ) . ( 4 ) Substituting Eq . 4 in Eq . 3 , we obtain a surrogate objective for optimization . To optimize this surrogate , we propose a variant of the random shooting optimizer ( Richards , 2005 ) that works as follows . At the current state st , we first sample a set of B candidate actions independently from the imitation policy . For each candidate action , we estimate a score based on their expected returns by performing rollout ( s ) of fixed-length H . The rollout policy π from which we sample all subsequent actions could be random ( potentially high variance ) or the imitation policy πθ ( potentially high bias ) or a mixture . In our experiments , we obtained consistently better performance with using πθ as the rollout policy π . For each trajectory , we estimate its return via Eq . 4 and finally , pick the action with the largest return . Consistent with the closed-loop nature of MPC , we repeat the above procedure at the next state st . Doing so helps correct for errors in estimation and optimization in the previous time step , albeit at the expense of additional computation . The algorithm has two critical parameters that induce similar computational trade-offs . First , we need to specify a budget B for the total number of rollouts . The higher the budget , larger is our search space for the best action . Second , we need to specify a planning horizon H . For larger lengths , we need extra computation that also involves interactions with the dynamics of the environment and rely more on the learned reward function than the value function for estimating returns in Eq . 4 . However , since the rollouts are independent , we can mitigate additional computational costs by parallelizing the rollouts . While this parallelization is indeed bottlenecked by the rollout with the largest horizon , in all of our experiments , we perform rollouts of fixed length and the horizon that corresponds to the optimal performance is relatively small ( 10 ∼ 50 ) . Thus , the gains due to parallelization are significant . In the next section , we present our empirical validation that also investigates the effect of planning horizon on the performance of the algorithm in greater detail .

Summary: The paper considers imitation learning problem in the presence of small perturbations and nuisances at test-time. In particular, it proposes Imitation with Planning at Test-time (IMPLANT), a new algorithm for imitation learning that incorporates decision-time planning within an inverse reinforcement learning algorithm. To counteract the imperfection due to policy optimization in RL step, IMPLANT uses the reward function estimated in the IRL step for decision-time planning. The effectiveness of the proposed method has been empirically evaluated on two kinds of setups, i.e., the default ‘no-transfer’ setting and the ‘transfer’ setting where the test dynamics is a perturbated version of the training dynamics.