Neural Dynamical Systems: Balancing Structure and Flexibility in Physical Prediction

1 INTRODUCTION . The use of function approximators for dynamical system modeling has become increasingly common . This has proven quite effective when a substantial amount of real data is available relative to the complexity of the model being learned ( Chua et al. , 2018 ; Janner et al. , 2019 ; Chen et al. , 1990 ) . These learned models are used for downstream applications such as model-based reinforcement learning ( Nagabandi et al. , 2017 ; Ross & Bagnell , 2012 ) or model-predictive control ( MPC ) ( Wang & Ba , 2019 ) . Model-based control techniques are exciting as we may be able to solve new classes of problems with improved controllers . Problems like dextrous robotic manipulation ( Nagabandi et al. , 2019 ) , game-playing ( Schrittwieser et al. , 2019 ) , and nuclear fusion are increasingly being approached using model-based reinforcement learning techniques . However , learning a dynamics model using , for example , a deep neural network can require large amounts of data . This is especially problematic when trying to optimize real physical systems , where data collection can be expensive . As an alternative to data-hungry machine learning methods , there is also a long history of fitting models to a system using techniques from system identification , some of which include prior knowledge about the system drawn from human understanding ( Nelles , 2013 ; Ljung et al. , 2009 ; Sohlberg & Jacobsen , 2008 ) . These models , especially in the gray-box setting , are typically data-efficient and often contain interpretable model parameters . However , they are not well suited for the situation where the given prior knowledge is approximate or incomplete in nature . They also do not generally adapt to the situation where trajectories are drawn from a variety of parameter settings at test time . This is an especially crucial point as many systems of interest exhibit path-dependent dynamics , which we aim to recover on the fly . In total , system identification methods are sample efficient but inflexible given changing parameter settings and incomplete or approximate knowledge . Conversely , deep learning methods are more flexible at the cost of many more samples . In this paper , we aim to solve both of these problems by biasing the model class towards our physical model of dynamics . Physical models of dynamics are often given in the form of systems of ordinary differential equations ( ODEs ) , which are ubiquitious and may have free parameters that specialize them to a given physical system . We develop a model that uses neural networks to predict the free parameters of an ODE system from the previous timesteps as well as residual terms added to each component of the system . To train this model , we integrate over the ODE and backpropagate gradients from the prediction error . This particular combination of prior knowledge and deep learning components is effective in quickly learning the dynamics and allows us to adjust system behavior in response to a wide variety of dynamic parameter settings . Even when the dynamical system is partially understood and only a subset of the ODEs are known , we find that our method still enjoys these benefits . We apply our algorithm to learning models in three synthetic settings : a generic model of ballistics , the Lorenz system ( Lorenz , 1963 ) , and a generalized cartpole problem , which we use for control as well . We also learn a high-level model of plasma dynamics for a fusion tokamak from real data . The contributions of this paper are • We introduce Neural Dynamical Systems ( NDS ) , a new class of model for learning dynamics that can incorporate prior knowledge about the system . • We show that these models naturally handle the issue of partial or approximate prior knowledge , irregularly spaced data , and system dynamics that change across instantiations , which generalizes the typical system identification setting . We also show that these advantages extend to control settings . • We demonstrate this model ’ s effectiveness on a real dynamics problem relevant to nuclear fusion and on synthetic problems where we can compare against a ground truth model . 2 RELATED WORK . System Identification and Deep Learning with Structure There is a long tradition of forecasting physical dynamics with either machine learning or techniques based on domain knowledge of the dynamics , especially in the field of system identification , where Ljung ( 2010 ) , Schoukens & Ljung ( 2019 ) and Cressie & Wikle ( 2015 ) are good summaries . Often , this space is discussed as a spectrum from a purely prior-knowledge-based system ( white-box ) to a purely data-driven system ( black-box ) with several shades of gray in between . White-box models use prior knowledge to precisely give the relationship between quantities of interest over time and there is extensive literature on solving them ( Brenan et al. , 1995 ) . ‘ Shades of gray ’ may distinguish between levels of prior knowledge or how equations cover subsets of the state space ( Ljung , 2010 ) . Other prior work focuses on online parameter estimation ( Vahidi et al. , 2005 ) , but this relies on an ongoing trajectory through the system and is difficult to use in our setting . In nonlinear black-box settings , there are a variety of techniques used to solve system identification models . Volterra series , a generalization of Taylor series which respects dependency on the past , have been used for system identification ( Rugh , 1981 ) . Block models such as the Hammerstein ( 1930 ) and Weiner ( Billings , 1980 ) models and their combination can also identify systems . Feedforward and recurrent neural networks have been widely used to model dynamical systems ( Chua et al. , 2018 ; Nagabandi et al. , 2017 ; Hochreiter & Schmidhuber , 1997 ) , with additional constraints on stability ( Manek & Kolter , 2020 ) or the Hamiltonian ( Chen et al. , 2019 ) and many others added . Nonlinear autoregressive moving average models with exogenous variables ( NARMAX ) have also been used widely to model dynamical systems and this class is a superset of nearly everything else discussed ( Brunton et al. , 2015 ; Rahrooh & Shepard , 2009 ) . Broadly , none of these algorithms are well-suited to a setting where the dynamic parameters of the system change across rollouts . There have also been several approaches for including physical structure in deep models . Raissi et al . ( 2019 ) use automatic partial derivative computation to force a neural network to fit a given ODE or PDE solution . de Avila Belbute-Peres et al . ( 2018 ) uses a linear complementarity problem to differentiate through 2d physics simulations however their method is not general to more dimensions or other types of problems besides mechanics . Cranmer et al . ( 2019 ) uses graph networks to discover physical laws . Chen et al . ( 2019 ) , Sanchez-Gonzalez et al . ( 2019 ) and Cranmer et al . ( 2020 ) force the network to respect Hamiltonian and Lagrangian constraints but without specific problem data on the system . Psichogios & Ungar ( 1992 ) predicts physical parameters for a given ODE model and Rico-Martinez et al . ( 1994 ) predict residuals . Thompson & Kramer ( 1994 ) similarly builds a hybrid parameter prediction function into a dynamical model . These last three works are especially similar to ours , though they use tiny networks , are problem-specific in their setup , and don ’ t take advantage of backpropagation through a numerical ODE solver . Neural Ordinary Differential Equations As most numerical ODE solvers are algorithms involving differentiable operations , it has always been possible in principle to backpropagate through the steps of these solvers dating back to at least Runge ( 1895 ) . However , since each step of the solver involves calling the derivative function , naı̈ve backpropagation incurs an O ( n ) memory cost , where n is the number of derivative calls made by the solver . Historically Pontryagin ( 2018 ) and recently Chen et al . ( 2018 ) demonstrated that by computing gradients through the adjoint sensitivity method , the memory complexity of backpropagating through a family of ODE solvers can be reduced to O ( 1 ) for a fixed network , as opposed to the naive O ( n ) . However , this work only used generic neural networks as the derivative function and did not consider dynamics modeling . They also provide a PyTorch package which we have built off of in our work . There has been some work using neural ordinary differential equations to solve physical problems . Portwood et al . ( 2019 ) used a fully-connected neural ODE with an RNN encoder and decoder to model Navier-Stokes problems . Rudy et al . ( 2019 ) used a neural network integrated with a RungeKutta method for noise reduction and irregularly sampled data . There has also been work learning the structure of dynamical systems , first with a convolutional warping scheme inspired by advectiondiffusion PDEs ( de Bezenac et al. , 2018 ) , then with a Neural ODE which was forced to respect boundary conditions and a partial observation mechanism ( Ayed et al. , 2019 ) . Machine Learning for Nuclear Fusion As far back as 1995 , van Milligen et al . ( 1995 ) showed that by approximating the differential operator with a ( single-layer , in their case ) neural network , one could fit simple cases of the Grad-Shafranov equation for magnetohydrodynamic equilibria . Recent work has shown that plasma dynamics are amenable to neural network prediction . In particular , Kates-Harbeck et al . ( 2019 ) used a convolutional and LSTM-based architecture to predict possible plasma disruptions ( when a plasma instability grows large and causes a loss of plasma containment and pressure ) . There has also been work in the field of plasma control : a neural network model of the neutral beam injection for the DIII-D tokamak has been deployed in order to diagnose the effect of controls on shots conducted at the reactor ( Boyer et al. , 2019b ) . Additionally , ( Boyer et al. , 2019a ) used classic control techniques and a simpler model of the dynamics to develop a controller that allows characteristics of the tokamak plasma to be held at desired levels . Others have used contextual Bayesian optimization to choose single-state controls which direct the plasma to desirable states ( Char et al. , 2019 ; Chung et al. , 2020 ) . 3 PROBLEM SETTING . Typically , a dynamical system ẋ = fφ ( x , u , t ) with some parameters φ is the conventional model for system identification problems . Here , state is x ∈ X , control is u ∈ U , and time is t ∈ R. The objective is to predict future states given past states , past and future controls , and prior knowledge of the form of f . We denote x ( φ , t , u , x0 ) = x0 + ∫ t 0 fφ ( x , u , t ) dt as the state obtained by integrating our dynamical system around f to time t. We consider in this work a more general setting and address the problem of prediction and control over a class of dynamical systems , which we define as the set { ẋ = fφ ( x , u , t ) | φ ∈ Φ } , where Φ is the space of parameters for the dynamical system ( e.g . spring constant or terminal velocity ) . We can generate a trajectory from a class by sampling a φ ∼ P ( Φ ) for some distribution P and choosing initial conditions and controls . In real data , we can view nature as choosing , but not disclosing , φ . For a particular example j , we sample φ ∼ P ( Φ ) and x0 ∼ P ( X0 ) and are given controls u indexed as u ( t ) and input data { x ( φ , ti , u , x0 ) } Ti=0 during training . At test time , we give a shorter , prefix time series { x ( φ , ti , u , x0 ) } T ′ i=0 but assume access to future controls . Then the prediction objective for a class of systems for N examples for timesteps { ti } TT ′+1 is x̂ = arg min x̂ Eφ∼P ( Φ ) , x0∼P ( X0 ) [ T∑ i=T ′+1 ||x ( φ , ti , u , x0 ) − x̂ti || 2 2 ] . ( 1 ) This objective differs from the traditional one in that implicitly , identifying φ for each trajectory needs to be done from the problem data in order to be able to predict the data generated by fφ . Similarly , the control problem is u = min u Eφ∼P ( Φ ) , x0∼P ( X0 ) [ ∫ t 0 c ( u ( t ) , x ( t ) ) dt ] , s.t . x ( t ) = x0 + ∫ t 0 fφ ( x , u , t ) dt ( 2 ) for some cost functional c. We will primarily explore the prediction problem in this setting , but as secondary considerations , we explore robustness to noise , the ability to handle irregularly spaced input data , and the ability to recover the parameters φ which generated the original trajectories . We will also consider the control problem in a simple setting .

Summary: The paper proposes a neural-network architecture for modeling dynamical systems that incorporates prior domain knowledge of the system's dynamics. More specifically, the main contributions are the mechanisms for incorporating such knowledge, in terms of fully or partially known structure (differential equations) of the system, which in turn positively affects the modeling performance. The results from the experimental evaluation (on 3 synthetic and one real-world experiments), in general, show that the proposed Neural Dynamical Systems (NDS), and in particular the ones trained with partial prior knowledge, have better performance than several standard benchmarks (such as NeuralODEs, LSTMs, Sparse Regression etc.).