NEURAL DYNAMICAL SYSTEMS: BALANCING STRUC-TURE AND FLEXIBILITY IN PHYSICAL PREDICTION

Abstract

We introduce Neural Dynamical Systems (NDS), a method of learning dynamical models in various gray-box settings which incorporates prior knowledge in the form of systems of ordinary differential equations. NDS uses neural networks to estimate free parameters of the system, predicts residual terms, and numerically integrates over time to predict future states. A key insight is that many real dynamical systems of interest are hard to model because the dynamics may vary across rollouts. We mitigate this problem by taking a trajectory of prior states as the input to NDS and train it to dynamically estimate system parameters using the preceding trajectory. We find that NDS learns dynamics with higher accuracy and fewer samples than a variety of deep learning methods that do not incorporate the prior knowledge and methods from the system identification literature which do. We demonstrate these advantages first on synthetic dynamical systems and then on real data captured from deuterium shots from a nuclear fusion reactor. Finally, we demonstrate that these benefits can be utilized for control in small-scale experiments.

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.

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. (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.



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

