Extracting Strong Policies for Robotics Tasks from Zero-Order Trajectory Optimizers

1 INTRODUCTION . The general purpose of model-based and model-free reinforcement learning ( RL ) is to optimize a trajectory or find a policy that is fast and accurate enough to be deployed on real robotic systems . Policies optimized by model-free RL algorithms achieve outstanding results for many challenging domains ( Heess et al. , 2017 ; Andrychowicz et al. , 2020 ) , however , in order to converge to the final performance , they require a large number of interactions with the environment and can hardly be used on real robots , which have a limited lifespan . Moreover , real robotic systems are high-dimensional and have a highly non-convex optimization landscape , which makes policy gradient methods prone to converge to locally optimal solutions . In addition , model-free RL methods only gather task-specific information , which inherently limits their generalization performance to new situations . On the other hand , recent advances in model-based RL show that it is possible to match modelfree performance by learning uncertainty-aware system dynamics ( Chua et al. , 2018 ; Deisenroth & Rasmussen , 2011 ; Du et al. , 2019 ) . The learned model can then be used within a model-predictive control framework for trajectory optimization . Zero-order optimizers are gaining a lot of traction in ∗equal contribution . We acknowledge the support from the German Federal Ministry of Education and Research ( BMBF ) through the Tübingen AI Center ( FKZ : 01IS18039B ) and from the Max Planck ETH Center for Learning Systems . the model-based RL community ( Chua et al. , 2018 ; Wang & Ba , 2020 ; Williams et al. , 2015 ) since they can be used with any choice of model and cost function , and can be surprisingly effective in finding high-performance solutions ( Pinneri et al. , 2020 ) ( close to a global optimum ) in contrast to their gradient-based counterparts , which are often highly dependent on hyperparameter tuning ( Henderson et al. , 2017 ) . One of the most popular optimizers is the Cross-Entropy Method ( CEM ) , originally introduced in the 90s by Rubinstein & Davidson ( 1999 ) . Despite their achievements , using zero-order methods for generating action sequences is time consuming in complex high-dimensional environments , due to the extensive sampling , making it hard to deploy them for real-time applications . Extracting a policy from powerful zero-order optimizers like CEM would bridge the gap between model-based RL in simulation and real-time robotics . As of today , this is still an open challenge ( Wang & Ba , 2020 ) . We analyze this issue and showcase several approaches for policy extraction from CEM . In particular , we will use the sample-efficient modification of CEM ( iCEM ) presented in Pinneri et al . ( 2020 ) . Throughout the paper , we will call these optimizers “ experts ” as they provide demonstration trajectories . To isolate the problem of bringing policy performance close to the expert ’ s one , we consider the true simulation dynamics as our forward model . Our contributions can be summarized as follows : • pinpointing the issues that arise when trying to distill a policy from a multimodal , stochastic teacher ; • introducing APEX , an Adaptive Policy EXtraction procedure that integrates iCEM with DAgger and a novel adaptive variant of Guided Policy Search ; • our specific integration of methods produces an improving adaptive teacher , with higher performance than the original iCEM optimizer ; • obtaining strong policies for hard robotic tasks in simulation ( HUMANOID STANDUP , FETCH PICK & PLACE , DOOR ) , where model-free policies would usually just converge to local optima . Videos showing the performance of the extracted policies and other information can be found at https : //martius-lab.github.io/APEX . 2 RELATED WORK . Our objective is to extract high-performing policies from CEM experts that can operate with a few planning samples to make iterative learning fast . Other kinds of zero-order optimizers have been used to generate control sequences ( Williams et al. , 2015 ; Lowrey et al. , 2019 ) but they still have to evaluate thousands of trajectories for each time step . Even simple random shooting has been used as a trajectory optimizer to bootstrap a model-free policy ( Nagabandi et al. , 2018 ) . To train policies from optimal control solutions , it was shown that the expert optimizers need to be guided towards the learning policy – known as guided policy search ( GPS ) ( Levine & Koltun , 2013 ; Levine & Abbeel , 2014 ) . In our work , the expert does not come from optimal control but is the stochastic iCEM optimizer , which we will also refer to as teacher . We apply GPS in a model model-predictive control setting , as done in Levine & Koltun ( 2013 ) ; Mordatch & Todorov ( 2014 ) ; Mordatch et al . ( 2015 ) ; Zhang et al . ( 2016 ) ; Kahn et al . ( 2017 ) ; Sun et al . ( 2018 ) using local trajectory optimization to generate a dataset for training a global policy through imitation learning . These approaches alternate between training a policy and creating new data with a model-based supervisor guided towards the learner , which was formalized in Sun et al . ( 2018 ) . Stochastic experts require particular guidance strategies , such as an adaptive cost formulation that we propose here , together with expert warm-starting via distribution initialization and additional samples from the policy . A simple form of warm-starting was already done in Wang & Ba ( 2020 ) . Recently , approaches like simple point-to-point supervised training such as Behavioral Cloning ( BC ) , or Generative Adversarial Network training ( GAN ) have been explored ( Wang & Ba , 2020 ) for policy distillation from CEM , but only largely sub-optimal policies could be extracted . When the policy is used alone at test time and not in combination with the MPC-CEM optimizer , its performance drops significantly , becoming almost random for some environments . We argue that the reason behind the difficulty in distilling a policy from CEM expert data is the multimodality of the CEM solution space and its inherent stochasticity due to the sampling , which we address below . Another possibility to train higher-performance policies from experts is DAgger ( Ross & Bagnell , 2010 ) , which is an on-policy method , asking the expert to relabel/correct policy actions . Nevertheless , as we will show in the following sections , DAgger alone is not sufficient to extract high-performing policies in our setting . To solve this problem , we use a guiding cost in combination with DAgger . A combination of GPS and DAgger-style relabeling was proposed in PLATO ( Kahn et al. , 2017 ) , however to create unbiased training data from iLQG experts . Since unguided DAgger is not appropriate with CEM , PLATO is not successful in our setting either . The components of our algorithm will be explained in the following sections . 3 METHODS . Trajectory optimization aims to find a suitable action sequence ~at = ( at , at+1 , . . . , at+h ) of horizon h that minimizes a cost function f ( ~at , st ) , where st is the current state of the system . i.e . ~a ? t ← argmin ~at f ( ~at , st ) . ( 1 ) The cost function f encodes the task . Optimal control is obtained if f is the trajectory cost until step h plus the cost-to-go under the optimal policy for an infinite-horizon problem , evaluated at the last state in the finite horizon trajectory . Typically , the cost-to-go is replaced by a proxy cost such as the distance to a target . In our case , the cost f is given by the sum of negative environment rewards up to the planning horizon . 3.1 IMPROVED CROSS-ENTROPY METHOD FOR MPC . To optimize Eq . 1 , we use a variant of the Cross-Entropy Method . Although originally introduced as an adaptive importance sampling procedure to estimate rare-event probabilities , it was recently employed in model-based RL ( Chua et al. , 2018 ; Wang & Ba , 2020 ) as a trajectory optimizer . A practical implementation of CEM uses a Gaussian proposal distribution over the optimization variables ~a , N ( µ , σ ) , evaluates the cost function f ( ~a ) , and refits this distribution iteratively by using the top-k samples . It finds low-cost regions of f and high-performing samples of the variable ~a . The cost function can be either expressed through a learned value function or through the dynamics model of the system , which in turn can be learned from data or in an analytical form . In the model-predictive control ( MPC ) framework , CEM can be used at every time-step to generate the optimal action plan , of which only the first action is executed , and the whole procedure is repeated for the next steps until task completion . In this paper , we will make use of iCEM , a sample-efficient improvement of CEM by Pinneri et al . ( 2020 ) that makes use of colored-noise and memory . In particular , it generates correlated action sequences with a non-flat frequency spectrum , differently from the flat Gaussian noise of CEM . This , together with the memory addition , results in an order of magnitude sample reduction . 3.1.1 USING A POLICY TO INFORM THE OPTIMIZATION : ICEMπ Unguided search for action sequences can be hard , in particular , if the cost function has large flat regions ( sparse rewards ) . Since we aim at extracting a policy π ( s ) , we can expect that after some examples the policy can be used to guide the search into the right region . Thus , we warm-start the mean µ of the iCEM Gaussian distribution with the policy actions . In particular , at time t = 0 where no prior information exists , the mean is initialized from rolling out the trajectory with the policy until the planning horizon : µ ← ( π ( s0 ) , π ( s1 ) , . . . π ( sh ) ) . Whenever the action for a new step is computed , iCEM uses shift initialization of the mean , and only the mean-action at the end of the horizon is initialized from the policy . Since iCEM is sample-based we also provide it with samples from the policy directly . More concretely , the actions performed by the policy are added in the last iteration if iCEM . The policy-informed iCEM algorithm is called iCEMπ and is shown in Alg . 1 . Algorithm 1 : Improved Cross-Entropy Method with warm-starting and samples from policy π , denoted as iCEMπ . Blue marks policy warm-starting and red is adding policy samples . Input : f ( ~a , s ) : cost function ; π : policy ; N : # of samples ; h : planning horizon ; k : elite-set size ; β : colored-noise exponent ; σinit : noise strength ; iter : # of iterations ; γ : reduction factor . 1 τ ← ∅ 2 for t = 0 to T−1 do 3 s← get the current state from the environment 4 if t == 0 then 5 µ0← h-steps rollout of model with π starting from s 6 else 7 µt← shifted µt−1 with last time-step action given by π 8 σt← constant vector in Rd×h with values σinit 9 for i = 0 to iter−1 do 10 Ni←max ( N · γ−i , 2 · k ) 11 samples← Ni samples from clip ( µt + Cβ ( d , h ) σ2t ) ; // with Cβ ( d , h ) colored-noise Normal distribution with noise-exponent β and dimension ( d , h ) 12 add a fraction of the shifted/reused elite-set to samples 13 if i == last-iter then 14 add µi to samples 15 add policy actions to samples ; // h-step rollout with policy from current state s 16 costs← cost function f ( ~a , s ) for ~a in samples 17 elite-sett← best k samples according to costs 18 µt , σt← fit Gaussian distribution to elite-sett with momentum 19 a← first action of best elite sequence 20 add ( s , a ) to τ 21 execute a 22 return τ ; // return the trajectory

Summary: This paper presents an approach to distill a model-based planning expert into a policy to enable real-time execution on robotic systems. An improved version of the Cross-Entropy Method, iCEM is used to generate trajectories via model-based optimisation where the forward model is the ground-truth simulator dynamics. The expert is further improved by warm-starting using the learned policy. The policy is learned via imitation learning of trajectories from the expert. Several approaches for this step are presented, including the vanilla approach of behaviour cloning (BC), dataset aggregation (DAGGER) to reduce covariate shift and Guided Policy Search (GPS) where the planner is encouraged to keep close to the policy distribution via a trust region KL loss. The paper discusses the merits of each approach and proposes Adaptive Policy Extraction (APEX), an approach that learns the policy from the expert through a combination of DAGGER and GPS where the tradeoff between the cost terms on planner exploration and policy trust region KL is set adaptively. This approach is tested on four continuous control tasks and achieves good performance compared to baselines — the learned policy gets significant improvements compared to the baselines. Additionally, a study that ablates different parts of the APEX algorithm is also presented.