Learning the Dynamical Response of Nonlinear Non-Autonomous Dynamical Systems with Deep Operator Learning Neural Networks

High-fidelity numerical schemes are the primary computational tools for simulating and predicting complex dynamical systems, such as climate modeling, robotics, or the modern power grid. However, these schemes may be too expensive for control, optimization, and uncertainty quantification tasks, which often require a large number of forward simulations and consume significant computational resources. Therefore, there is a growing interest in developing tools that can accelerate the numerical simulation of complex dynamical systems without compromising accuracy.

By providing faster alternatives to traditional numerical schemes, machine learning-based computational tools hold the promise of accelerating the rate of innovation for complex dynamical systems. Hence, a recent wave of machine and deep learning-based tools has demonstrated the potential of using observational data to construct fast surrogates of complex systems Efendiev et al. (2021); Qin et al. (2021) and providing efficient solutions to complex engineering tasks such as fault diagnosis Choudhary et al. (2023); Mishra et al. (2022b, a) and power engineering Moya and Lin (2023); Moya et al. (2023). These tools aim to (i) learn the governing equations of a dynamical system or (ii) learn to predict the system’s response from data.

On the one hand, several works Brunton et al. (2016a, b); Schaeffer (2017); Sun et al. (2020) use observational data to discover the unknown governing equations of the underlying system. For example, Brunton et al.  Brunton et al. (2016a) proposed identifying nonlinear systems using sparse schemes and a large set of dictionaries. They then extended their work in Brunton et al. (2016b) to identify input/output mappings that describe control systems. In Sun et al. (2020), the authors used sparse approximation schemes to recover the governing partial or ordinary differential equations of unknown systems.

On the other hand, there is growing interest in predicting the future response of a dynamical system using time-series data Qin et al. (2019); Raissi et al. (2018); Qin et al. (2021); Proctor et al. (2018). For instance, the authors of Efendiev et al. (2021); Chung et al. (2021) trained a transformer with data from the early stages of time-dependent partial differential equations to recursively predict the solution of future stages. Qin et al.  Qin et al. (2019) used a recurrent residual neural network (ResNet) to approximate the mapping from the current state to the next state of an unknown autonomous system. Similarly, in Raissi et al. (2018), the authors used feed-forward neural networks (FNN) as a building block of a multi-step scheme that predicts the response of an autonomous system. In Qin et al. (2021), the authors extended their previous work Qin et al. (2019) to non-autonomous systems with time-dependent inputs. To this end, they parametrized the input locally within small time intervals.

Many of the works mentioned above require large amounts of training data to avoid overfitting. However, acquiring this data may be prohibitively expensive for complex dynamical systems. Additionally, if one fails to avoid overfitting when using traditional neural networks for next-state methods, the predicted response may drift and accumulate errors over time. It is therefore crucial to develop deep learning-based frameworks that can efficiently handle the infinite-dimensional nature of predicting the response of non-autonomous systems with time-dependent inputs for long time horizons, as studied in this paper.

Operator Learning Chen and Chen (1995); Lu et al. (2019); Li et al. (2020); Zhang et al. (2022) was first addressed by the seminal work of Chen and Chen (1995). Unlike classical function learning, operator learning involves approximating an operator between two function spaces. Recently, the seminal paper Lu et al. (2021) extended the work of Chen and Chen (1995) and proposed the Deep Operator Neural Network (DeepONet) framework. DeepONet exhibits small generalization errors and learns with limited training data Zhang et al. (2022). This has been demonstrated in many application areas, such as power engineering Moya et al. (2023), multi-physics problems Cai et al. (2021), and turbulent combustion problems Ranade et al. (2021). Extensions to the original DeepONet Lu et al. (2021) have enabled the incorporation of physics-informed models Wang et al. (2021b), handling noisy data Lin et al. (2023), and designing novel optimization methods Lin et al. (2023, 2021).

DeepONet is a prediction-free and domain-free tool Zhang et al. (2022). Specifically, DeepONet can predict the target function value at any point in its domain, and the input and output functions are not necessarily required to share the same domain. This property allows for handling incomplete datasets during the training phase. In Zhang et al. (2022), it was further relaxed the assumptions on input function discretization, improving flexibility in data preparation and prediction accuracy for operator learning.

In this paper, we focus on extending the This paper focuses on extending the original DeepONet framework to learn the solution operator of non-autonomous systems with time-dependent inputs for long-time horizons. By using such a data-driven operator framework, control policies can be designed for continuous nonlinear systems, among other applications.

In particular, one of our motivations behind learning to approximate the solution operator of a non-autonomous system is its application in control, and Model-Based Reinforcement Learning (MBRL) Wang et al. (2019); Sutton (1990). In MBRL, one learns to approximate the system to control (from data) and then uses the learned model to seek an optimal policy without extensive interaction with the actual system. A common approach is to learn a discrete-time forward model that predicts the next state using the current state and selected control action. Such an MBRL framework has delivered successful and efficient results in discrete-time problems such as games Kaiser et al. (2019). However, most control systems in science and engineering (e.g., robotics Nagabandi et al. (2020), unmanned vehicles Hafner et al. (2019), or laminar flows Fan et al. (2020)) are continuous. Of course, one can always discretize the continuous dynamics and apply discrete-time MBRL using traditional neural networks (e.g., see Brockman et al. (2016)). However, if one fails to handle the inherent epistemic uncertainty Deisenroth et al. (2013), such a strategy may lead to error accumulation (due to model bias) and poor asymptotic performance. To alleviate these drawbacks, we build on the original DeepONet Lu et al. (2021) to design an effective and efficient framework that can learn the solution operator of a nonlinear non-autonomous system with time-dependent inputs, which one can then apply within the framework of control or continuous MBRL Du et al. (2020).

Formally, the objectives of this paper are twofold.

1. 1.

   Approximation of the local solution operator: Our goal is to create a neural operator-based framework that can learn to map (i) the current state of the dynamical system with control and (ii) a local approximation of the input to the next state of the non-autonomous dynamical system.

2. 2.

   Long/Medium-term simulation: We aim to design an efficient scheme that uses the operator learning framework to simulate the dynamical system’s response to a given input over a long or medium-term horizon.

The contributions that will achieve the above objectives are summarized below.

1. 1.

   We model the non-autonomous/control dynamical system in the deep operator learning framework and propose a DeepONet-based numerical scheme that effectively and recursively simulates the solution of the system.

2. 2.

   We also propose a novel data-driven Runge-Kutta (RK) DeepONet-based method that reduces the cumulative error and improves recursive prediction accuracy in long-term simulations.

3. 3.

   We provide an estimation of the cumulative error for the proposed numerical scheme based on DeepONet. Our estimated error bounds are tighter than those presented in Qin et al. (2021). Additionally, we demonstrate that the novel data-driven RK DeepONet-based scheme achieves better stability bounds than the DeepONet-based scheme.

4. 4.

   We test the proposed frameworks on state-of-the-art models, such as the predator-prey system, the pendulum, the cart-pole system, and a power engineering task. The effectiveness of the methods is observed in all experiments.

We organize the rest of this paper as follows. In Section 2, we review operator learning and introduce the proposed DeepONet-based algorithms. Next, we present the novel data-driven Runge-Kutta DeepONet-based scheme and estimate its corresponding error bound in Section 3. We present the numerical experiments in Section 4. We provide a discussion of our results and future work in Section 5, and conclude the paper in Section 6.

We consider the problem of learning from data the solution operator of the following continuous-time non-autonomous system with time-dependent inputs

where $x(t)\in\mathcal{X}\subseteq\mathbb{R}^{n}$ is the state vector, $u(t)\in\mathcal{U}\subseteq\mathbb{R}^{p}$ the input vector, and $f:\mathcal{X}\times\mathcal{U}\to\mathcal{X}$ an unknown function. Additional assumptions on the input function u and the function $f$ will be discussed later. Also, throughout this paper, we assume $u(t)$ is a scalar input, i.e., $p=1$. Extending the proposed framework to the vector-valued case is straightforward.

Solution Operator. Let $\mathcal{F}$ denote the solution operator of (1). $\mathcal{F}$ takes as input the initial condition $x(a)=x_{0}\in\mathcal{X}$ and the sequence of input function values $u_{[a,t)}:=\{u(\tau)\in\mathcal{U}:\tau\in[a,t)\}$, and outputs the state $x(t)\in\mathcal{X}$ at time $t\in[a,b]$. We compute the solution operator via

Approximate System. In practice, to learn the operator $\mathcal{F}$, we only have access to an approximate/discretized representation of the input function $u(t)$. Let $u_{m}$ denote this approximate representation, which yields the following approximate system

whose solution operator is

In the above, with a slight abuse of notation, we denoted $u_{m}$ as the input discretized using $m\geq 1$ sensors or interpolation points within the interval $[a,b]$, and the approximate state function as $\tilde{x}(t)$.

In this section, we propose a data-driven Runge-Kutta explicit DeepONet-based scheme Iserles (2009) that predicts the new state vector $\hat{x}(t_{n}+h_{n})$ using the current state value $\hat{x}(t_{n})$, i.e.,

Here $k_{1}$ and $k_{2}$ are, respectively, the estimates of $f$ at $t_{n}$ and $t_{n+1}$. We compute these estimates (see equation (13)) using (i) the forward pass of trained DeepONet $\mathcal{F}_{\theta^{*}}$ and (ii) automatic differentiation. Note that in (13b), we use the notation $\bar{x}(t_{n+1})$ for the estimate of the state at $t_{n+1}$ obtained using the DeepONet’s $\mathcal{F}_{\theta^{*}}$ forward pass.

We detail the proposed data-driven RK explicit DeepONet-based scheme in Algorithm 2 and Figure 3. Two remarks about our algorithm are provided next. (i) For simplicity, we only present our scheme for the improved Euler method or RK-2 Iserles (2009). However, we remark that we can extend our idea to any explicit Runge-Kutta scheme. (ii) If $u_{m}^{n+1}$ is available at $t_{n}$, then we can compute $k_{2}$ as follows:

Then, equations (13a) and (12) will work as a predictor-corrector scheme with the updated input information. Other strategies can also be adopted within the proposed RK scheme. However, we let the design of such strategies for our future work.

To evaluate our framework, we tested the proposed DeepONets on five tasks: the autonomous Lorentz 63 system (in Section 4.1), the predator-prey dynamics with control (in Section 4.2), the pendulum swing-up (in Section 4.3), the cart-pole system (in Section 4.4), and a power engineering application (in Section 4.5). For the three control tasks, we used only $n_{s}=1$ sensor. The reasons for selecting only one sensor are twofold. First, we want to show that DeepONet is effective even when the input signal is encoded with minimal information. For reference, in Qin et al. (2021), the authors encoded the input signals (used in their experiments) with at least $n_{s}=3$ interpolation points (sensors). Second, the $n_{s}=1$ sensor scenario resembles the scenario of sampled-data control systems Wittenmark et al. (2002) or reinforcement learning tasks Sutton and Barto (2018) with continuous action space. Finally, for the power engineering application task, we used $n_{s}=2$ sensors.

Training dataset. For each of the three continuous control tasks (predator-prey, pendulum, and cart-pole systems), we generate the training dataset $\mathcal{D}_{\text{train}}$ as follows. We use Runge-Kutta (RK-4) Iserles (2009) to simulate $N_{\text{train}}$ trajectories of size two. For each trajectory, the input to RK-4 is the initial condition $x(t_{n})$ uniformly sampled from $\mathcal{X}$ and the input $u(t_{n})$ uniformly sampled from the set $\mathcal{U}$. The output from the RK-4 algorithm is the state $x(t_{n}+h)$, where $h$ is uniformly sampled from the interval $[0,0.25]$. Such a procedure gives the dataset:

Training protocol and neural networks. We implemented our framework using JAX Bradbury et al. (2018)⁵⁵5We will publish the code on GitHub after publication. The neural networks for the Branch and Trunk Nets are the modified fully-connected networks proposed in Wang et al. (2021a) and used in our previous paper Moya et al. (2023). We trained the parameters of the networks using Adam Kingma and Ba (2014). Moreover, we selected (i) the default hyper-parameters for the Adam algorithm and (ii) an initial learning rate of $\eta=0.001$ that exponentially decays every 2000 epochs.