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 paper investigates different choices of designing the planning model which predict the next state on the observation space. The paper compares (1) prediction with deterministic or stochastic models, (2) using 1-step forward prediction loss or multi-step forward prediction loss to train the model, (3) prediction with single network or network ensemble, and (4) using perfect observation or add gaussian noise to the observation.  Through the above experiments, the paper suggests that deterministic models should be trained with multi-step loss, while stochastic models work better when the 1-step log-likelihood loss is used and noise is added to the observation space. In both cases, ensemble networks tend to show better performance.