Exploring Model-based Planning with Policy Networks

1 INTRODUCTION . A model-based reinforcement learning ( MBRL ) agent learns its internal model of the world , i.e . the dynamics , from repeated interactions with the environment . With the learnt dynamics , a MBRL agent can for example perform online planning , interact with imaginary data , or optimize the controller through dynamics , which provides significantly better sample efficiency ( Deisenroth & Rasmussen , 2011 ; Sutton , 1990 ; Levine & Abbeel , 2014 ; Levine & Koltun , 2013 ) . However , MBRL algorithms generally do not scale well with the increasing complexity of the reinforcement learning ( RL ) tasks in practice . And modelling errors in dynamics that accumulate with time-steps greatly limit the applications of MBRL algorithms . As a result , many latest progresses in RL has been made with model-free reinforcement learning ( MFRL ) algorithms that are capable of solving complex tasks at the cost of large number of samples ( Schulman et al. , 2017 ; Heess et al. , 2017 ; Schulman et al. , 2015 ; Mnih et al. , 2013 ; Lillicrap et al. , 2015 ; Haarnoja et al. , 2018 ) . With the success of deep learning , a few recent works have proposed to learn neural network-based dynamics models for MBRL . Among them , random shooting algorithms ( RS ) , which uses modelpredictive control ( MPC ) , is shown to have good robustness and scalability ( Richards , 2005 ) . In shooting algorithms , the agent randomly generates action sequences , use the dynamics to predict the future states , and choose the first action from the sequence with the best expected reward . However , RS usually has worse asymptotic performance than model-free controllers ( Nagabandi et al. , 2017 ) , and the authors of the the PETS algorithm ( Chua et al. , 2018 ) suggest that the performance of RS is directly affected by the quality of the learnt dynamics . They propose a probabilistic ensemble to capture model uncertainty , which enables PETS algorithm to achieve both better sample efficiency and better asymptotic performance than state-of-the-art model-free controllers in environments such as Cheetah . However , PETS is not as effective on environments with higher dimensionality . 1https : //github.com/WilsonWangTHU/POPLIN . -200.000 -200.000 -200.000 -190.000 -190.000 -190.000 -190.000 -190.000 -190.0 00 -190.000 -190.000 -19 0.0 00 -180 .000 -18 0.0 00 -18 0.0 00 -180.000 -180.000 -18 0.0 00 -17 0.00 0 -170.000 -170.000 -170.000 -160.000 ( a1 ) Reward Surface ( a2 ) PETS Iter 1 ( a3 ) PETS Iter 2 ( a4 ) PETS Iter 3 -525.000 -525.00 0 -500.000 -500.000 -500.000 -500.000 -475.000 -475.000 -47 5.0 00 -475.000 -475.000 -450.000 -450.000 -425.000 -425.000 -425.000 -400.000 ( b1 ) Reward Surface ( b2 ) POPLIN Iter 1 ( b3 ) POPLIN Iter 2 ( b4 ) POPLIN Iter 3 Figure 1 : We transform each planned candidate action trajectory with PCA into a 2D blue scatter . The top and bottom figures are respectively the visualization of PETS ( Chua et al. , 2018 ) and our algorithm . The red area has higher reward . From left to right , we show how candidate trajectories are updated , across different planning iterations within one time-step . As we can see , while both reward surface is not smooth with respect to action trajectory . POPLIN , using policy networks , has much better search efficiency , while PETS is stuck around its initialization . The details are in section 5.3 . In this paper , we explore MBRL algorithms from a different perspective , where we treat the planning at each time-step as an optimization problem . Random search in action space , as what is being done in state-of-the-art MBRL algorithms such as PETS , is insufficient for more complex environments . On the one hand , we are inspired by the success of AlphaGo ( Silver et al. , 2016 ; 2017 ) , where a policy network is used to generate proposals for the Monte-Carlo tree search . On the other hand , we are inspired by the recent research into understanding deep neural networks ( Nguyen & Hein , 2017 ; Li et al. , 2018 ; Soudry & Hoffer , 2017 ) . Deep neural networks , frequently observed in practices , is much less likely to get stuck in sub-optimal points . In Figure 1 , we apply principal component analysis ( PCA ) on the action sequences generated in each planning iteration within one time-step . The reward surface of the action space is not smooth and prone to local-minimas . We argue that optimization in the policy network ’ s parameter space will be more efficient . Furthermore , we note that the state-of-the-art MBRL algorithm with MPC can not be applied real-time . We therefore experiment with different policy network distillation schemes for fast control without MPC . To sum up , the contribution of this paper is three-fold : • We apply policy networks to generate proposals for MPC in high dimensional locomotion control problems with unknown dynamics . • We formulate planning as optimization with neural networks , and propose policy planning in parameter space , which obtain state-of-the-art performance on current bench-marking environments , being about 3x more sample efficient than the previous state-of-the-art algorithm , such as PETS ( Chua et al. , 2018 ) , TD3 ( Fujimoto et al. , 2018 ) and SAC ( Haarnoja et al. , 2018 ) . • We also explore policy network distillation from the planned trajectories . We found the distilled policy network alone achieves high performance on environments like Cheetah without the expansive online planning . 2 RELATED WORK . Model-based reinforcement learning ( MBRL ) has been long studied . Dyna ( Sutton , 1990 ; 1991 ) algorithm alternately performs sampling in the real environments and optimize the controllers on the learned model of the environments . Other pioneering work includes PILCO ( Deisenroth & Rasmussen , 2011 ) , where the authors model the dynamics using Gaussian Process and directly optimize the surrogate expected reward . Effective as it is to solve simple environments , PILCO heavily suffers the curse of dimensionality . In ( Levine & Abbeel , 2014 ; Levine & Koltun , 2013 ; Levine et al. , 2016 ; Chebotar et al. , 2017 ; Zhang et al. , 2018 ) , the authors propose guided policy search ( GPS ) . GPS uses iLQG ( Li & Todorov , 2004 ; Todorov & Li , 2005 ; Tassa et al. , 2012 ) as the local controller , and distill the knowledge into a policy neural network . In SVG ( Heess et al. , 2015 ) , the authors uses stochastic value gradient so that the stochastic policy network can be optimized by back-propagation with off-policy data . Recently with the progress of model-free algorithms such as TRPO and PPO ( Schulman et al. , 2015 ; 2017 ) , Kurutach et al . ( 2018 ) ; Luo et al . ( 2019 ) propose modern variants of Dyna , where TRPO ( Schulman et al. , 2015 ) is used to optimize the policy network using data generated by the learnt dynamics . Concurrent to this work , Janner et al . ( 2019 ) further use SAC ( Haarnoja et al. , 2018 ) to train the policy network , and gets state-of-the-art performance on many tasks . At the same time , random shooting methods proposed by Nagabandi et al . ( 2017 ) ; Chua et al . ( 2018 ) have shown its robustness and effectiveness on benchmarking environments . PETS algorithm ( Chua et al. , 2018 ) is considered by many to be the state-of-the-art shooting algorithm , which we discuss in detail in section 3 . Dynamics is also used to obtain better value estimation to speed up training ( Gu et al. , 2016 ; Feinberg et al. , 2018 ; Buckman et al. , 2018 ) . Latent dynamics models using VAE ( Kingma & Welling , 2013 ) are commonly used to solve problems with image input ( Ha & Schmidhuber , 2018a ; b ; Hafner et al. , 2018 ; Kaiser et al. , 2019 ) . 3 BACKGROUND . 3.1 REINFORCEMENT LEARNING . In reinforcement learning , the problem of solving the given task is formulated as a infinite-horizon discounted Markov decision process . For the agent , we denote the action space and state space respectively as A and S. We also denote the reward function and transition function as r ( st , at ) and f ( st+1|st , at ) , where st ∈ S and at ∈ A are the state and action at time-step t. The reward function is assumed known to the agent in this work . The agent maximizes its expected total reward J ( π ) = Eπ [ ∑∞ t=0 r ( st , at ) ] with respect to the agent ’ s controller π . 3.2 RANDOM SHOOTING ALGORITHM AND PETS . Our proposed algorithm is based on the random shooting algorithm ( Richards , 2005 ) . In random shooting algorithms ( Nagabandi et al. , 2017 ; Chua et al. , 2018 ) , a data-set of D = { ( st , at , st+1 ) } is collected from previously generated real trajectories . The agent learns an ensemble of neural networks denoted as fφ ( st+1|st , at ) , with the parameters of the neural networks denoted as φ . In planning , the agent randomly generates a population of K candidate action sequences . Each action sequence , denoted as a = { a0 , ... , aτ } , contains the control signals at every time-steps within the planning horizon τ . The action sequence with the best expected reward given the current dynamics network fφ ( st+1|st , at ) is chosen . RS , as a model-predictive control algorithm , only executes the first action signal and re-plan at time-step . In PETS ( Chua et al. , 2018 ) , the authors further use cross entropy method ( CEM ) ( De Boer et al. , 2005 ; Botev et al. , 2013 ) to re-samples sequences near the best sequences from the last CEM iteration . 4 MODEL-BASED POLICY PLANNING Algorithm 1 General POPLIN Framework 1 : while Training iterations not Finished do 2 : for ith time-step of the agent do 3 : CEM planning as in section 4.1 , 4.2 4 : Execute the first action from CEM . 5 : end for 6 : Dynamics update and policy distillation . 7 : end while In this section , we describe two variants of POPLIN : model-based policy planning in action space ( POPLIN-A ) and model-based policy planning in parameter space ( POPLIN-P ) . Following the notations in section 3.2 , we define the expected planning reward function at time-step i as follows : R ( si , ai ) = E [ i+τ∑ t=i r ( st , at ) ] , where st+1 ∼ fφ ( st+1|st , at ) . ( 1 ) The action sequence ai = { ai , ai+1 , ... , ai+τ } is generated by the policy search module , as later described in Section 4.1 and 4.2 . The expectation of predicted trajectories { si , si+1 , ... , si+τ } is estimated by creating P particles from the current state . The dynamics model fk , tφ ( st+1|st , at ) used by kth particle at time-step t is sampled from deterministic or probabilistic ensemble models . To better illustrate , throughout the paper we denote this dynamics as a fixed deterministic model , i.e . fk , tφ ≡ fφ . In practice the dynamics uses probabilistic ensemble models , which requires some trivial modifications to the math and we refer readers to PETS Chua et al . ( 2018 ) for details .

Summary: This paper presents POPLIN, a novel model-based reinforcement learning algorithm, which trains a policy network to improve model-prediction control. The paper studies extensively how to utilize the policy, by planning in action space or planning in parameter space and how to train the policy, by behavioral cloning, by GAN or by averaging the results of CEM. The experiments show that the proposed algorithm can perform very well in MuJoCo tasks.