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 work provides a novel model-based reinforcement learning algorithm for continuous domains (Mujoco) dubbed POPLIN. The presented algorithm is similar in vein to the state-of-the-art PETS algorithm, a planning algorithm that uses state-unconditioned action proposal distributions to identify good action sequences with CEM in the planning routine. The important difference compared to PETS is the incorporation of a parametric state-conditioned policy (trained on real data) in the planning routine to obtain better action-sequences (CEM is used to learn the "offset" from the parametric policy). The paper presents two different algorithmic ablations where CEM either operates in action space or parameter space (POPLIN-A and POPLIN-P respectively), in combination with different objectives to learn the parametric policy. The method is evaluated on 12 continuous benchmarks and compared against state-of-the-art model-based and model-free algorithms, indicating dominance of the newly proposed method.