Amortized Equation Discovery in Hybrid Dynamical Systems

Prior works focus on the simplest hybrid dynamical models, i.e. piece-wise affine systems with linear transition rules (Vidal et al., 2003; Ferrari-Trecate et al., 2003; Roll et al., 2004; Juloski et al., 2005; Paoletti et al., 2007; Ozay et al., 2008; Bako, 2011; Ohlsson & Ljung, 2013). Recently,  Mangan et al. (2019); Novelli et al. (2022) relieve these constraints and propose methods for non-linear hybrid systems. Among them, Hybrid-SINDy (Mangan et al., 2019) proposes a two-stage method, i.e. it first uses k-NN to group time series into small cluster fragments and then discovers governing equations separately in each fragment by models proposed in non-hybrid systems, e.g. SINDy (Brunton et al., 2016). Based on Hybrid-SINDy, Novelli et al. (2022) also follows a two-stage paradigm while introducing the number of discontinuities in hybrid systems as a known prior for better performance. Although effective, these two-stage methods learn the mode of each segment individually and do not leverage the commonalities across segments. In this paper, we reformulate the problem and propose an amortized end-to-end learning framework to jointly categorize modes, discover equations, and learn mode-switching behaviors.

Many methods have been proposed for equation discovery in non-hybrid dynamical systems. Bongard & Lipson (2007) and Schmidt & Lipson (2009) leverage genetic programming (Koza et al., 1994) to discover nonlinear differential equations from data. SINDy (Brunton et al., 2016) uses symbolic sparse regression on a library of candidate model terms to select the fewest terms required to describe the observed dynamics. Several methods extend this approach to new settings, e.g. identifying partial differential equations (Rudy et al., 2017), considering additional physical constraints (Loiseau & Brunton, 2018), including control signals (Kaiser et al., 2018), and introducing integral terms for denoising (Schaeffer & McCalla, 2017). These methods cannot be directly applied to hybrid systems because they cannot model an unknown number of modes and unknown mode-switching behaviors.

Switching dynamical systems refer to the same systems as hybrid dynamical systems, but highlight different aspects in the literature. Hybrid systems denote systems with a mixture of continuous and discrete states, while switching dynamical systems highlight the switching behaviors of system states. Many methods focus on switching linear dynamical systems where they set matrix calculations to model linear state transitions (Ackerson & Fu, 1970; Ghahramani & Hinton, 2000; Oh et al., 2005). Recently, switching non-linear dynamical systems model state transitions as neural networks, e.g. SNLDS (Dong et al., 2020), REDSDS (Ansari et al., 2021), and GRASS (Liu et al., 2023). While effective in modeling state-switching behaviors, they cannot discover closed-form equations from data. To categorize dynamics, our method is inspired by previous switching dynamical systems (Dong et al., 2020; Ansari et al., 2021; Liu et al., 2023) to infer latent mode variables. Differently, our method can jointly discover parsimonious closed-form equations to characterize dynamics and infer the values of the mode variables.

To model modes and mode-switching behaviors in hybrid systems, inspired by REDSDS (Ansari et al., 2021), we introduce latent categorical variables, i.e. mode variables, to learn categorical distributions of modes and index each set of equations representing each type of dynamics. Specifically, mode variables are discrete variables $\mathbf{z}\coloneqq\mathbf{z}_{1:T}$, where each $z_{t}\in\{1,\dots,K\}$ categorizes the mode at time step $t\in\{1,\dots,T\}$.

Besides mode variables, we also model latent count variables to learn the duration distributions of each mode. Count variables can help us avoid frequent mode switching, thanks to the fact that mode durations typically follow a geometric distribution (Ansari et al., 2021). They are modeled as discrete categorical variables $\mathbf{c}\coloneqq\mathbf{c}_{1:T}$, where each state $c_{t}\in\{1,\dots,d_{\rm max}\}$ explicitly models the run-length of the currently active mode at time $t$ and $\rm d_{max}$ is the maximal number of steps before a mode switch. Count variables $c_{t}$ are incremented by 1 when the mode $z_{t}=k$ stays the same at the next time step $z_{t+1}=k$, or they reset to 1 when mode $z_{t}=k$ switches to another one $z_{t}\neq k$.

Each mode $k\in[K]$ is assigned its own set of candidate basis functions $\Theta_{k}(\mathbf{y}_{t})$ and learnable coefficient weights $\mathbf{w}_{k}$, which we will use to discover its equation. For instance, at time $t$, the mode variable $z_{t}=k$ indexes the candidate basis function $\Theta_{k}(\mathbf{y}_{t})$ and the learnable weights $\mathbf{w}_{k}$, which together define the equation representing the dynamics of mode $k$. In practice, different modes share the same candidate basis functions, i.e. $\Theta_{j}(\mathbf{y}_{t})=\Theta_{k}(\mathbf{y}_{t}),\forall j,k\in\{1,\cdots,K\}$, unless we have some prior knowledge of the hybrid system. However, the learnable coefficient weights of different modes are individual and never shared, i.e. $\mathbf{w}_{j}\neq\mathbf{w}_{k},\forall j,k\in\{1,\cdots,K\}$. We collect all the candidate basis functions in a single vector $\Theta(\mathbf{y}_{t})=(\Theta_{1}(\mathbf{y}_{t}),\cdots,\Theta_{K}(\mathbf{y}_{t}))$ and similarly we collect all learnable coefficient weights $\mathbf{w}=(\mathbf{w}_{1},\cdots,\mathbf{w}_{K})$.

Assuming Markovian dynamics, the joint generative probability of hybrid systems in our model is described as

In the initial states, count variables are ignored as they are always 1 when starting. $p(z_{1})$ is the initial distribution over all possible modes. $p(\mathbf{y}_{1}|z_{1})$ models the initial observation states conditioned on the initial modes. For later time steps $t\geq 2$, the count transition probability $p(c_{t}|c_{t-1},z_{t-1})$ models how the count variables at time $t$ change over time depending on their previous values and mode variables at time $t-1$. The mode transition probability $p(z_{t}|z_{t-1},c_{t},\mathbf{y}_{t-1})$ models mode-switching behaviors on how modes switch at time $t$ conditioned on the previous mode and observation states at time $t-1$ as well as the updated count state $c_{t}$ at time $t$. The observation transition probability $p(\mathbf{y}_{t}|\mathbf{y}_{t-1},z_{t})$ models how the observations at time $t$ are influenced by their previous values at time $t-1$ conditioned on the updated mode variables at time $t$. Equations are amortized and learned at $p(\mathbf{y}_{t}|\mathbf{y}_{t-1},z_{t})$ by all segments of each mode to characterize mode dynamics. More specifically, conditioned on motion mode $z_{t}=k$, $p(\mathbf{y}_{t}|\mathbf{y}_{t-1},z_{t})$ first indexes a set of candidate basis functions $\Theta_{k}$ and coefficient weights $\mathbf{w}_{k}$, which are used together to obtain derivatives $\dot{\mathbf{y}}_{t-1}=\Theta_{k}\cdot\mathbf{w}_{k}$ of $\mathbf{y}_{t-1}$ at time $t-1$. With known time intervals $\Delta_{t}$, we finally achieve $\mathbf{y}_{t}=\dot{\mathbf{y}}_{t-1}\cdot\Delta_{t}+\mathbf{y}_{t-1}$ assuming the dynamics do not change much in short time intervals. For inference of the latent mode and count variables, we conduct exact inference similar to the forward-backward algorithm in HMM (Eddy, 1996). The graphical model for the exact inference is the same as the generative model, which is illustrated in Figure 2. The neural network implementations and details of the inference model are in Appendix A.1 and A.2.

Learnable parameters of AMORE are optimized by maximizing the observation likelihood with sparse regularization on coefficient weights $\mathbf{w}$ of candidate basis functions

The derivatives of the training objective and further expansions over time are detailed in Appendix A.3.

Assume that $N$ objects and $K$ motion modes exist in multi-object hybrid systems. Inspired by Kipf et al. (2018); Liu et al. (2023), we model interactions between objects by latent edge variables $\mathbf{e}=\mathbf{e}_{1:T}^{1:N^{2}}=\{e_{t}^{1},\cdots,e_{t}^{N^{2}}\}_{t=1}^{T}$ including self-loop, thus totally $N^{2}$ for $N$ objects. For each pair of objects, interactions $e_{t}^{m\rightarrow n}$ model whether object $m$ interacts with object $n$ at time $t$. The edge variables are modeled in a latent temporal graph $\mathcal{G}_{t}=(\mathcal{V}_{t},\mathcal{E}_{t})$, where edges $e_{t}^{m\rightarrow n}\in\mathcal{E}_{t}$ and nodes $\mathcal{V}_{t}$ summarize states of objects. For instance, $\mathbf{v}^{m}_{t}=\{z^{m}_{t},c^{m}_{t},\mathbf{y}^{m}_{t}\}$ for $\mathbf{v}^{m}_{t}\in\mathcal{V}_{t}$ defines one graph node summarizing states of object $m$ at time $t$ which includes observation $\{\mathbf{y}^{m}_{t}\}$ and latent states $\{z^{m}_{t},c^{m}_{t}\}$. Edge $e_{t}^{m\rightarrow n}$ signals interaction relationships between node $\mathbf{v}_{t}^{m}$ and node $\mathbf{v}_{t}^{n}$ in graph $\mathcal{G}_{t}$.

We set the number of all possible motion modes as $K$ in multi-object hybrid dynamical systems. The $K$ motion modes are shared across $N$ objects, which are implemented by $K$ sets of candidate basis functions $\Theta(\mathbf{y}_{t})=(\Theta_{1}(\mathbf{y}_{t}),\cdots,\Theta_{K}(\mathbf{y}_{t}))$ and learnable coefficient weights $\mathbf{w}=(\mathbf{w}_{1},\cdots,\mathbf{w}_{K})$. Each mode $k\in[K]$ has its own $\Theta_{k}(\mathbf{y}_{t})$ and $\mathbf{w}_{k}$ as in Sec. 4. Thus there are K sets of learnable weights for learning dynamics of K modes across N objects. Both the time and space complexity of AMORE-MIO regarding learnable weights of basis functions is $\mathcal{O}(K)$. For instance, the mode variable $z_{t}^{n}=k$ of the object $n$ at time $t$ indexes $\Theta_{k}(\mathbf{y}_{t})$ and $\mathbf{w}_{k}$ which together form equations to represent the dynamics of mode $k$. Different from single-object scenarios, the mode-switching behaviors of each object are not only influenced by their own evolving nature but also by external influences of potentially interacting objects. We model the influences of interactions on the mode-switching behaviors between objects, which are detailed in the following generative model of AMORE-MIO.

Assuming Markovian dynamics, the joint generative probability of multi-object hybrid systems in AMORE-MIO is calculated as

where the initial states and count transition probability are defined as in Eq. 4 of single-object scenarios but with $n$ objects. For later time steps $t\geq 2$, the edge transition probability $p(\mathbf{e}_{t}^{1:N^{2}}|\mathbf{e}_{t-1}^{1:N^{2}},\mathbf{c}_{t-1}^{1:N},\mathbf{z}_{t-1}^{1:N},\mathbf{y}_{t-1}^{1:N})$ models how the edge variables evolve depending on all the states at the previous time step. We model the influences of interactions on the mode transition probability $p(\mathbf{z}_{t}^{1:N}|\mathbf{z}_{t-1}^{1:N},\mathbf{y}_{t-1}^{1:N},\mathbf{c}_{t}^{1:N},\mathbf{e}_{t}^{1:N^{2}})$, which characterizes the mode-switching behaviors of multi-object hybrid dynamical systems. Based on the updated modes of each object, the observation transition probability $p(\mathbf{y}_{t}^{1:N}|\mathbf{y}_{t-1}^{1:N},\mathbf{z}_{t}^{1:N})$ can be factorized over objects $\prod_{n=1}^{N}p(\mathbf{y}_{t}^{n}|\mathbf{y}_{t-1}^{n},z_{t}^{n})$ where equations of each mode are amortized and learned by all segments from all objects belonging to the same mode. The further expansion over objects of the joint generative probability is in Appendix B.1. For inference of latent mode, count, and edge variables, we conduct posterior approximate inference for edge variables $q_{\phi_{e}}(\mathbf{e}|\mathbf{y})$ conditioned on observations $\mathbf{y}$, and then conduct exact inference of mode and count variables $p_{\theta}(\mathbf{z},\mathbf{c}|\mathbf{y},\tilde{\mathbf{e}})$ conditioned on observations $\mathbf{y}$ and the approximate edge variables $\tilde{\mathbf{e}}\sim q_{\phi_{e}}(\mathbf{e}|\mathbf{y})$. The generative and inference models of AMORE-MIO are illustrated in Fig. 3. Learnable parameters of AMORE-MIO are optimized by maximizing the evidence lower bound with sparse regularization on the learnable coefficient weights

Neural network implementations, the derivations and the detailed inference model are in Appendix B.2,  B.4. and B.3.

We train all datasets with a fixed batch size of 40 for 20,000 training steps. We use the Adam optimizer with $10^{-5}$ weight-decay and clip gradients norm to 10. The learning rate is warmed up linearly from $5\times 10^{-5}$ to $2\times 10^{-4}$ for the first 2,000 steps, and then decays following a cosine manner with a rate of 0.99. Each experiment is running on one Nvidia GeForce RTX 3090 GPU. $d_{min}$ and $d_{max}$ of the count variables are simply set as 20 and 50, respectively for all datasets. The number of edge types $L$ is set as 2, containing one no-interaction type and one with-interaction type. More details are in Appendix C.2.

For evaluation of discovered equations, following Course & Nair (2023), we use the reconstruction error between the discovered coefficients of equations and ground truth, i.e. ${\rm RER}=\frac{1}{T}\sum_{t=1}^{T}(||\mathbf{w}_{t}-\xi_{t}||_{2}\;/\;||\xi_{t}||_{2})$ where $\mathbf{w}_{t}$ and $\xi_{t}$ are the learned and ground-truth coefficients at time $t$. For evaluation of segmentation, following Ansari et al. (2021), we use frame-wise segmentation accuracy, i.e. Accuracy and $F_{1}$ after matching the labels using the Hungarian algorithm (Kuhn, 1955), Normalized Mutual Information (NMI) and Adjusted Rand Index (ARI) to measure similarities. For evaluation of forecasting, we use Normalized Mean Absolute Error (NMAE) and Normalized Root Mean Squared Error (NRMSE). We conducted each experiment with 5 random seeds. We report the average score of each experiment in the main paper and put the statistics (error bars) in the appendix due to limited space.

Hybrid-SINDy (Mangan et al., 2019) uses a two-stage paradigm and cannot perform forecasting, thus we compare with it on discovered equations and segmentation. To compare with Hybrid-SINDy on forecasting, we continue the value of the last observable time point as forecasting results of Hybrid-SINDy. For forecasting, we compare with other three recent representative methods, i.e. SVI (Course & Nair, 2023) which is designed for equation discovery in non-hybrid systems and can perform forecasting, LLMTime (Gruver et al., 2023) which utilize pre-trained large language models (LLM) to do forecasting, and TimeGPT (Garza & Mergenthaler-Canseco, 2023) which is the first foundation model for time series. GRASS (Liu et al., 2023) does not discover equations, but models multi-object switching dynamical systems, so it is used for comparison in multi-object systems.