Learning Dynamics Models for Model Predictive Agents

1 INTRODUCTION . Recently reinforcement learning ( RL ) ( Sutton & Barto , 2018 ) , in particular actor-critic approaches ( Lillicrap et al. , 2015 ) were shown to successfully solve a variety of continuous control problems ( Schulman et al. , 2017 ; Lillicrap et al. , 2015 ; Fujimoto et al. , 2018 ; Haarnoja et al. , 2018 ) . The simplicity of this approach has led to an explosion of research demonstrating the effectiveness of these methods ( Vecerik et al. , 2019 ; Kalashnikov et al. , 2018 ) . However , model-free RL suffers from two key disadvantages ( Dulac-Arnold et al. , 2021 ) . First , model-free RL is sample inefficient , often requiring millions or billions of environment interactions . Second , the learned policies are tied to a specific task , making transfer of learned knowledge in multi-task settings or across related tasks difficult ( Finn et al. , 2017 ; Andrychowicz et al. , 2017 ) . Model-Based RL holds the promise of overcoming these drawbacks while maintaining the benefits of model-free methods . Model-based RL involves learning a dynamics model from data , and then using this model to optimise behaviour , most often with an online planner . Since the dynamics model is ideally independent of the reward and state-action distribution , multi-task settings are simple : the reward function can be changed at will ( Rajeswaran et al. , 2020 ; Argenson & Dulac-Arnold , 2021 ; Belkhale et al. , 2021 ) . For the same reason , offline learning with off-policy data is possible ( Yu et al. , 2020 ; Kidambi et al. , 2020 ; Argenson & Dulac-Arnold , 2021 ) . Moreover , model-based RL can achieve a better sample efficiency ( Janner et al. , 2019 ; Nagabandi et al. , 2018 ; 2020 ; Byravan et al. , 2020 ; Buckman et al. , 2018 ; Hafner et al. , 2019b ; Argenson & Dulac-Arnold , 2021 ) . In this paper we investigate the impact of the different model learning design choices on the planner performance . For this evaluation we focus on learning models using feed-forward neural networks as these are the dominant models in the recent model-based RL literature . A large portion of these papers mainly proposes different methods to train the networks and variations on the planner . Unfortunately , each new paper proposes a series of evolutions , and so individual design choices are rarely ablated fully . This is the goal of our work . We look specifically at the effects of various model learning design choices on the planning performance . We consider 4 different design choices , deterministic vs. stochastic models , multi-step vs. 1-step losses , use of ensembles , and input noise . These design choices have been proposed to obtain better long term prediction ( Abbeel et al. , 2005 ; Chua et al. , 2018 ; Hafner et al. , 2019b ) and planning performance ( Venkatraman et al. , 2015 ; SanchezGonzalez et al. , 2020 ; Pfaff et al. , 2020 ) . To derive best practices for each choice-parameter , we perform ablation studies evaluating the impact of the different choices . In addition we investigate the qualitative performance of the learned models , by evaluating their ability to describe complex trajectories with reasonable fidelity , and to be consistent . The contributions of this paper are • We perform a systematic comparison evaluating the planning performance of different design choices for model learning , i.e. , stochastic models , multi-step loss , network ensembles and input noise , on five DeepMind Control Suite environments . • We observe that mean squared error ( 1-step or multi-step ) is not a good predictor of planning performance when using a learned model . • We find that multiple model learning design choices can obtain high reward and consistent long-term predictions for the evaluated systems excluding the humanoid . The differences in reward between the different combinations is not significant . • We characterise best practices for learning models using feed-forward networks . For our experiments , deterministic models need to be trained using the multi-step loss and combined with ensembles . Stochastic models require more ensemble components and input noise . The paper is structured as follows . First , we cover related work on model-based RL for continuous control in Sec 2 . Afterwards , we review the naive approach to learning dynamics models ( Sec . 3 ) . Section 4 introduces the design choices of model learning and evaluates the learned models using model-predictive control . The subsequent discussion ( Sec . 5 ) summarizes the insights , results and highlights general challenges of model learning for planning . 2 RELATED WORK . Learning dynamics model has a long tradition in robotics . The seminal work of Atkeson et al . ( 1986 ) proposed to learn the dynamics parameters of a rigid-body kinematic chain from data . To remove the limitations of analytic rigid-body models ( Atkeson et al. , 1986 ; Lutter et al. , 2020 ; 2021 ) , black-box function approximation was leveraged ( Schaal et al. , 2002 ; Choi et al. , 2007 ; Calinon et al. , 2010 ; Nguyen-Tuong et al. , 2009 ; Nguyen-Tuong & Peters , 2010 ; Lutter et al. , 2019 ; Lutter & Peters , 2019 ) to learn system models able to learn the peculiarities of real-world robotic systems . The majority of model-based RL focuses on learning dynamics models , which predict the next step using the current state and action . The model learning for control literature refers to these models as forward models ( Nguyen-Tuong & Peters , 2011 ) . Model-based RL algorithms use the model as a simulator to generate additional data ( Sutton , 1991 ; Janner et al. , 2019 ; Morgan et al. , 2021 ) , to evaluate the reward of an action sequence ( Chua et al. , 2018 ; Nagabandi et al. , 2018 ; Lambert et al. , 2019 ; Nagabandi et al. , 2020 ) , to improve the value function estimate ( Feinberg et al. , 2018 ; Buckman et al. , 2018 ) or backpropagate the policy gradients through time ( Miller et al. , 1995 ; Deisenroth & Rasmussen , 2011 ; Heess et al. , 2015 ; Byravan et al. , 2020 ; Amos et al. , 2021 ) . Current modelbased RL methods use 1-step predictions to compute the optimal trajectory . Most current modelbased RL approaches use different architectures of deep networks to approximate the model , e.g. , deterministic networks ( Nagabandi et al. , 2018 ; Byravan et al. , 2020 ; Feinberg et al. , 2018 ; Kurutach et al. , 2018 ) , stochastic networks ( Heess et al. , 2015 ; Lambert et al. , 2019 ) , recurrent networks ( Hafner et al. , 2019a ; b ; 2020 ; Ha & Schmidhuber , 2018 ) , stochastic ensembles ( Chua et al. , 2018 ; Nagabandi et al. , 2020 ; Janner et al. , 2019 ; Kidambi et al. , 2020 ; Yu et al. , 2020 ; Buckman et al. , 2018 ; Rajeswaran et al. , 2020 ; Lambert et al. , 2020b ) and graph neural networks ( Sanchez-Gonzalez et al. , 2018 ) . In this work we focus on feed-forward neural networks for learning dynamics models with continuous states and actions . The learned models are used with model predictive control ( MPC ) ( Garcia et al. , 1989 ; Allgöwer & Zheng , 2012 ; Pinneri et al. , 2020 ) to solve continuous control tasks . Simple feed-forward neural networks are the standard of most model-based RL and already offer many different variations that affect the planning performance . For the considered environments , MPC enables the direct evaluation of the learned models without requiring a policy or value function approximation . Therefore , the model performance for planning can be directly measured and compared to planning with the ground-truth model . 3 LEARNING DYNAMICS MODELS . We concentrate on 1-step forward models and review this setup in more detail . Dynamics models predict the next observation xi+1 using the current observation xi and action ui . These models can be used as a simulator for planning without interacting with the real system . For example , some sample-based planning algorithms sample action sequences close to the current plan , evaluates the reward of each action sequence by simulating the trajectories using the dynamics model and updates the plan using the action sequences with the highest reward . The simplest approach to learn such a forward model is to use a deterministic model that minimizes the prediction error . This objective can be achieved by minimising the 1-step mean squared error between the predicted and observed next state . This optimization loss is described by θ∗ = arg min θ 1 ND ∑ x , u∈D ( xi+1 − x̂i+1 ) T ( xi+1 − x̂i+1 ) with x̂i+1 = f ( xi , ui ; θ ) , with the model parameters θ , the dataset D with ND samples . Dataset Current model-based RL algorithms re-use the data collected by the agent to learn the model . Therefore , the data is acquired in the vicinity of the optimal policy . This exploration leads to a specialized model for the considered task as only regions of high reward are explored . Hence , the current learned models are not independent of the state-action distribution and the reward function can not be interchanged at will . In contrast to model-based RL , system identification focuses purely on obtaining the optimal data for learning the model cover most of the state space . Integrator Instead of predicting the next observation , one can predict the change of the observation . Then the system dynamics are modelled as an integrator described by x̂i+1 = xi + ∆t f ( xi , ui ; θ ) with the time step ∆t . This technique has proven to result in more stable long term prediction compared to predicting the next state and has been widely adapted ( Chua et al. , 2018 ; Janner et al. , 2019 ; Yu et al. , 2020 ; Amos et al. , 2021 ; Lambert et al. , 2020a ) . Normalization Optimizing the MSE can be problematic if the scales of the different observations are different . In this case , the model overfits to the dimensions with larger scales . For forward models this results in overfitting to the velocities as the position is usually an order of magnitude smaller . To mitigate this overfitting the MSE can be normalized using the diagonal variance of the dataset ΣD . The normalized MSE is described by θ∗ = arg min θ 1 ND ∑ x , u∈D ( xi+1 − x̂i+1 ) T Σ−1D ( xi+1 − x̂i+1 ) . The normalized MSE is comparable to standardization of the model inputs and outputs . In our experiments we did not observe any significant benefits using standardization instead of the normalized MSE . 4 RESULTS . In the following we present our quantitative and qualitative evaluation of the model learning design choices . We evaluate the planning performance for 4 different design choices , i.e. , stochastic models ( Sec . 4.2 ) , multi-step loss ( Sec . 4.3 ) , network ensembles ( Sec . 4.4 ) and input noise ( Sec . 4.5 ) . These different design choices have been previously proposed within the literature to improve the model for planning ( Abbeel et al. , 2005 ; Chua et al. , 2018 ; Hafner et al. , 2019b ) and generate better long-term predictions ( Venkatraman et al. , 2015 ; Sanchez-Gonzalez et al. , 2020 ; Pfaff et al. , 2020 ) . For each of these proposed models we first introduce the approach , summarize the results and conclude the best practices . The best hyperparameter for model learning were identified first using random search and the ablation studies compare to the optimal hyperparameter . The humanoid is only used for the initial ablation study as no model-learning approach learns a model that could be used for planning . A more detailed experimental setup , the hyperparameters and additional ablation studies are contained in the appendix . The links within the results point to the referenced section contained in https : //sites.google.com/view/learning-better-models/ .

Summary: This work ablates some of the design choices that go into learning a dynamics model for control-based environments. They ablate 4 choices: use of deterministic vs. stochastic models, multistep losses, network ensembles, and input noise. The authors study this in a few of the DeepMind control suite environments.