Evaluating Model-Based Planning and Planner Amortization for Continuous Control

1 INTRODUCTION . In recent years , model-free RL algorithms have improved and scaled to the point where it is feasible to learn adaptive behavior for high-dimensional systems in diverse circumstances ( Schulman et al. , 2017 ; Heess et al. , 2017 ) . Despite these improvements , there is a widespread intuition that model-based methods can further improve data efficiency . This has led to recent advances for continuous control problems that demonstrate improved learning efficiency by leveraging model learning during policy training ( Lowrey et al. , 2019 ; Nagabandi et al. , 2020 ; Hubert et al. , 2021 ) . However , there is also work which urges moderation in interpreting these results , by showing that well-tuned model-free baselines can compare favorably against some model-based approaches ( Springenberg et al. , 2020 ; van Hasselt et al. , 2019 ) . In this paper , we focus on model-based RL with learned models and proposals.1 We use modelpredictive control ( MPC ) to improve the quality of behavior generated from a policy that serves as a proposal for the planning procedure . When acting , the MPC-based search procedure improves the quality of the data collected which serves as a sort of “ active exploration ” . This data can then be used to improve the proposal , in effect consolidating the gains . This hybrid approach interpolates between model-free RL , on one side , and trajectory optimization with a model on the other . Fundamentally , the spectrum on which this hybrid approach is situated reflects a trade-off in terms of reusability / generality versus compute cost at deployment time . Policies obtained by model-free RL often generalize poorly outside of the situations they have been trained for , but are efficient to execute ( i.e. , they are fully amortized ) . Models offer potentially greater generalization , insofar as the model is accurate over a broad domain of states , but it can be computationally costly to derive actions from models . Ultimately , the pure planning approach involving deriving actions from models is usually prohibitive and some form of additional knowledge is required to render the search space tractable for real-time settings . Coming from the model-free RL perspective , the natural approach is to learn a policy that serves as a relatively assumption-free proposal for the planning process . While in most model-based RL publications the ambition is to speed up policy learning for a single task , we believe this is not necessarily the most promising setting , as it is difficult to dissociate the learning of a dynamics model ( which is potentially task independent ) from that of a value function ( which is task specific ) . A different intuition is that model-based RL might not accelerate 1We use the term proposal to mean a `` proposal distribution '' , which is a distribution over the space of actions . learning on a particular problem , but rather enable efficient behavior learning on new tasks with the same embodiment . In such transfer settings , we might hope that a dynamics model will offer complementary generalization to the policy , and we can transfer the policy , model , or both . In this work , we validate these intuitions about the relative value of the different components of this system . For high-dimensional locomotion problems , we find that even with a good learned model successful model predictive control requires good proposal distributions to succeed . This effect is particularly pronounced in domains with multiple goals ( or equivalently multiple different tasks ) . Building on this insight we show that it is sometimes possible to leverage learned models with a limited search budget to boost exploration and learn policies more efficiently . Policies that are trained by off-policy updates from data acquired through planning do not reliably perform well when used without planning in our experiments . To overcome this problem , and enable planner amortization into a policy that does not require planning at test time , we propose training the policy according to a combination of a behavioral cloning objective ( on MPC data ) and an off-policy update with MPO ( Abdolmaleki et al. , 2018b ) . When transferring models and proposals to other tasks we find only marginal improvements in data efficiency relative to a well-tuned model-free baseline . Overall , our results suggest that while MPC with learned models can lead to more data efficiency , and planners can be amortized effectively into compact policies , it is not a silver bullet and model-free methods are strong baselines . 2 BACKGROUND & RELATED WORK . Model-predictive control refers to the use of model-based search or planning over a short horizon for selecting an action . In order to make it computationally tractable , it is common to “ seed ” the search at a given timestep with either an action from a proposal policy or an action obtained by warmstarting from the result of the previous timestep . In simulation , it has been demonstrated that a fairly generic MPC implementation can be effective for control of a relatively high DoF humanoid ( Tassa et al. , 2012 ) and that MPC with learned models ( Chua et al. , 2018 ; Wang & Ba , 2020 ; Nagabandi et al. , 2020 ) can achieve data-efficient solutions to simple control problems . Real RC helicopter control has also been achieved using an MPC approach that made use of a learned model ( Abbeel et al. , 2006 ) . MPC approaches are gaining wider use in robotics ( multicopters ( Neunert et al. , 2016 ; Torrente et al. , 2021 ) , quadruped ( Grandia et al. , 2019 ; Sleiman et al. , 2021 ) , humanoids ( Tassa et al. , 2014 ; Kuindersma et al. , 2016 ) ) , and dexterous manipulation ( Nagabandi et al. , 2020 ) ; but the computational speed of the planner is a bottleneck for hardware deployment . There are different ways around this , with the core practical solution being to plan in lower dimensional reduced coordinate models . Alternatively , POPLIN ( Wang & Ba , 2020 ) explores learning proposals for MPC and planner distillation but is tested on simple tasks and does not leverage model-free RL for proposal learning . Amortized policies map states to actions quickly , but implicitly reflect considerable knowledge by being trained on diverse data . Model-free RL produces policies which amortize the knowledge reflected in rollouts required to produce them . Similarly , it is possible to produce diverse trajectories from a planner and distil them into a single policy ( Levine & Koltun , 2013 ; Mordatch & Todorov , 2014 ; Mordatch et al. , 2015 ; Pan et al. , 2017 ; Xie et al. , 2021 ) . Reinforcement learning approaches with MPC have become more popular recently , and our work fits within this space . As noted in the introduction , previous work often emphasizes the role of amortization through the learning of value functions and models . TreePI ( Springenberg et al. , 2020 ) , MuZero ( Schrittwieser et al. , 2020 ; Hubert et al. , 2021 ) , SAVE ( Hamrick et al. , 2020 ) , DPI ( Sun et al. , 2018 ) and MPC-Net ( Carius et al. , 2020 ) all perform versions of hybrid learning with model-based MCTS or planning being used to gather data which is then used to train the model and value function to accelerate learning . Other recently proposed algorithmic innovations blend MPC with learned value estimates to trade off model and value errors ( Bhardwaj et al. , 2021 ) . Here , we primarily consider learning dynamics models to enable transfer to new settings with different reward functions . Other uses of models unlike hybrid MPC-RL schemes have also been explored in the literature ; however , they are not the focus of this work . Nevertheless we highlight two of these approaches : value gradients can be backpropagated through dynamics models to improve credit assignment ( Nguyen & Widrow , 1990 ; Heess et al. , 2015 ; Amos et al. , 2020 ) and it is possible to train policies on model rollouts to improve data efficiency ( Sutton , 1991 ; Janner et al. , 2019 ) . Recently , there has also been considerable effort towards exploring model-based offline RL methods in order to gain full value from offline datasets ( Yu et al. , 2020 ; Argenson & Dulac-Arnold , 2021 ; Kidambi et al. , 2020 ) . 3 TASKS . In this paper we consider a number of locomotion tasks of varying complexity , simulated with the MuJoCo ( Todorov et al. , 2012 ) physics simulator . We consider two embodiments : an 8-DoF ant from dm_control ( Tunyasuvunakool et al. , 2020 ) and a model of the Robotis OP3 robot with 20 degrees of freedom . For each embodiment , we consider three tasks : walking forward , walking backward and “ go-to-target-pose ” ( GTTP ) , a challenging task that is the focus of our evaluation . In all tasks , the agent receives egocentric proprioceptive observations ( joint angles , velocities and end-effector positions ) and additional task observations . In the walking forward and backward tasks the agent is rewarded for maximizing forward ( or backward ) velocity in the direction of a narrow corridor . For the OP3 robot we also include a small pose regularization term . The task is specified through a relative target direction observation . The GTTP task , which builds on existing motion tracking infrastructure ( Hasenclever et al. , 2020 ) , consists of either body on a plane , with a target pose in relative coordinates as a task-specific observation and proximity to the target pose rewarded . When the agent is within a threshold distance of the target pose ( i.e . it achieves the current target ) , it gets a sparse reward and a new target is sampled . For the ant we use target poses from policies trained on a standard go-to-target task ( Tunyasuvunakool et al. , 2020 ) . For the OP3 , we use poses from the CMU mocap database ( cmu ) ( retargeted to the robot ) . We use thousands of different target poses ; the agent has to learn to transition between them . Thus the GTTP task can be seen as either a single highly diverse task or as a multi-task setting with strongly related tasks and consistent dynamics . We believe the GTTP task should be particularly amenable to model-based methods : it combines a high-dimensional control problem with a diverse goal distribution . This makes it hard to solve quickly with model-free methods . However , since the dynamics are shared between all goal poses , a dynamics models should be beneficial in leveraging the common structure . See Sec . A for an in-depth motivation and Sec . F.5 for results on tasks from the DeepMind Control Suite ( Tassa et al. , 2018 ) . 4 METHOD . 4.1 MODEL PREDICTIVE CONTROL ( MPC ) FOR BEHAVIOR GENERATION . We consider control obtained by model-based planning that refines initial action-samples from a proposal distribution . MPC executes the first action from the planned action sequence and iteratively re-plans after each timestep ( c.f . actor loop in Alg . 1 ) . In this scheme , actions are sampled either from the proposal πθ ( with parameters θ ) or the planner ( PLANNER ) , forming a mixture distribution that is specified by the probability of choosing planner actions ( pplan ) . Setting pplan = 0 results in using only samples from the proposal and vice-versa for pplan = 1 . Setting pplan = 0.5 leads to interleaved execution of proposal and planner actions , like the stochastic mixing of student and expert actions in DAGGER ( Ross et al. , 2011 ) . The PLANNER subroutine takes in the current state st , the proposal πθ ( with parameters θ ) , a model mφ ( with parameters φ ) that predicts next state st+1 given current state st and action at , the known reward function r ( st , at , st+1 ) . Optionally , a learned state-value function Vψ ( s ) ( with parameters ψ ) that predicts the expected return from state s can be provided . While we consider sample based planners , primarily a Sequential Monte Carlo based non-iterative planner ( SMC , Alg . 2 ) ( Piché et al. , 2019 ; Gordon et al. , 1993 ) and the Cross-Entropy Method ( CEM , Alg . 3 ) ( Botev et al. , 2013 ) , other planners ( Williams et al. , 2017 ) could be used . See Sec . B for more details . The hybrid agent in Alg . 1 has several learnable components : the model mφ , the proposal πθ as well as the ( optional ) value function Vψ ; these are learned based on data generated by the MPC loop on the actor and we measure learning progress w.r.t environment steps collected . Note that we assume the reward function is known since it is under the control of the practitioner in most robotics applications . Algorithm 1 Agent combining MPC with model-free RL Given : Randomly initialized proposal πθ , ( pre-trained or random ) model mφ , random critic Qψ , optionally ( pre-trained or random ) task-agnostic proposal ρω . // Modules to be learned Given : Known reward r , planning probability pplan , replay buffer B , MPO loss weight α , BC loss weight β , learning rates & optimizers ( ADAM ) for the different modules . // Known modules and parameters // MPC loop – Asynchronously on the actors while True do Initialize ENV and observe state s0 . while episode is not terminated do // Choose between planner and proposal action depending on pplan // Use mixture of ( pre-trained ) task-agnostic proposal ( ρω ) and learned proposal ( πθ ) as proposal for proposal transfer exps ( Sec . 5.4 ) . Use pre-trained model for model transfer exps ( Sec . 5.3 ) . at ∼ { PLANNER ( st , πθ , mφ , r , Vψ ) if x ≤ pplan πθ ( st ) otherwise . where x ∼ U [ 0 , 1 ] Step ENV ( st , at ) → ( st+1 , rt ) and write transition to replay buffer B end while end while // Asynchronously on the learner while True do Sample batch B of trajectories , each of sequence length T from the replay buffer B Update action-value function Qψ based on B using Retrace ( Munos et al. , 2016 ) . Update model mφ based on B using multi-step loss in Equation 20 ( Sec . 4.2 ) Update proposal πθ based on B using Equation 4 ( Sec . 4.3 ) Optionally , for from-scratch experiments , update task-agnostic proposal ρω using behavioural cloning ( Equation 1 ) on transitions in B. end while

Summary: The paper provides a thorough, detailed evaluation of model-based RL, as realized via MPC-based planning. In particular, it examines: (a) how it can be improved via learned policies that act as MPC proposal distributions; and (b) to understand and disentangle the benefits of (i) model transfer and (ii) proposal transfer, when using model-based RL. Comparisons are performed using both CEM and SMC samplers.  For the locomotion-related benchmarks considered, MPC with a learned model & proposal can provide modest improvements wrt a well-tuned model-free baseline (MPO). However, for a more challenging go-to-target-pose (GTTP) task, they offer significantly improved learning speed and final performance. The benefits of model-transfer are rather marginal, even when transferred to the same task,  counter-to-common-belief-and-intuition.