Learning Modular Simulations for Homogeneous Systems
Abstract
Complex systems are often decomposed into modular subsystems for engineering tractability.
Although various equation based white-box modeling techniques make use of such structure, learning based methods have yet to incorporate these ideas broadly.
We present a modular simulation framework for modeling homogeneous multibody dynamical systems, which combines ideas from graph neural networks and neural differential equations.
We learn to model the individual dynamical subsystem as a neural ODE module.
Full simulation of the composite system is orchestrated via spatio-temporal message passing between these modules.
An arbitrary number of modules can be combined to simulate systems of a wide variety of coupling topologies.
We evaluate our framework on a variety of systems and show that message passing allows coordination between multiple modules over time for accurate predictions and in certain cases, enables zero-shot generalization to new system conﬁgurations.
Furthermore, we show that our models can be transferred to new system conﬁgurations with lower data requirement and training effort, compared to those trained from scratch.
Introduction
Modeling and simulation are key enablers in engineering as viable alternatives to real world systems.
Simulation allows for the testing and evaluation of large/complex systems in a risk-free environment and has been of use in many disciplines.
Simulation also plays a key role in domains such as deep reinforcement learning, which requires a large volume of interactions with the environment to learn effective policies.
However, traditional simulators often require a large amount of engineering effort to create and can be computationally intensive when high ﬁdelity predictions are needed.
This has created interest in data-driven simulations where machine learning algorithms are used to learn system dynamics directly from observations [1, 2, 3, 4, 5].
While machine learning algorithms excel at learning correlations from data, such end-to-end learning approaches result in models that cannot be reused for different system conﬁgurations.
Additionally, end-to-end learning also requires large amounts of data and training effort when trying to model real-world systems.
In this paper we propose a machine learning based framework to learn and compose modular simulators.
We build upon the observations that complex systems comprises of multiple modules (e.g., multibody systems), where each module evolves according to its own set of dynamical equations with corresponding inputs and outputs [6].
Examples include multiple pendulums connected via a single thread and multiple quadrotors carrying a large structure or a beam.
Instead of modeling the entire system jointly, it might be more beneﬁcial to model individual subsystems (such as a single pendulum or one quadrotor) and then consider interaction effects.
The ability to harness such modularity can help reduce the sample complexity for machine learning as well create transferable subsystems that enable rapid prototyping.
The technical methodology we describe here builds upon a class of models referred to as mechanistic models .
Mechanistic models aim to introduce inductive biases related to the underlying mechanisms of the systems being modeled, which are in turn informed by domain knowledge.
Some examples are neural networks motivated by differential equations [7, 8, 9, 10, 11, 12, 13], physics-informed neural networks [14, 15, 16, 17, 18, 19, 20] etc.
As mechanistic models encode structural knowledge of the system at hand, they can be used to learn modular simulations, as opposed to treating the entire system as a singular block.
In this work, we propose an approach to simulate homogeneous systems from data.
Under the assumption that the systems of interest can be modeled as ordinary differential equations, we present a class of neural ODEs which we call message-passing neural ODEs (MP-NODE).
Instead of modeling the entire graph of a system through ODEs, the message passing neural ODE only learns the model of the subsystem that forms represents a node in the graph.
In order to enable interaction between the nodes, the MP-NODE contains additional message variables along with the state variables.
The entire system is thus composed of several MP-NODE modules with the weights shared between them, and the coordination as well as ability to account for model approximation errors over time is learnt through the messages.
We show that MP-NODEs allow for expressive and efﬁcient modular simulations which can be ﬁne-tuned to new system conﬁgurations signiﬁcantly faster than training a model from scratch.
Our key contributions are listed below.
We present the message-passing neural ODE (MP-NODE), an augmented neural ODE formulation aimed at enabling interaction between multiple components through messages.
We test the MP-NODE on various connected systems and demonstrate its ability to learn complex dynamics and generalize to unseen system conﬁgurations.
We also show that this modular framework allows ﬁnetuning for new conﬁgurations with signiﬁcantly lower training effort than training models from scratch.
Related Work
Let us assume X t ∈ X is the state of the world at time t .
Then, dynamically stepping over T timesteps with optional control input U t ∈ U gives a trajectory of states X t 0: T =
( X t 0 , . . .
, X t T ) .
A simulator s : X × U → X , models the dynamics by mapping previous states and control inputs to future states.
Traditional simulators often deﬁne the dynamics as a modular set of differential or differentiable algebraic equations and use appropriate numerical solvers as the update mechanism.
For data-driven simulations, we parameterize the dynamics of the simulator as f θ : X × U → Y .
Here, the parameters θ can be optimized for some training objective.
The semantics of the dynamics information Y ∈ Y is governed by an update function, usually an integrator.
Learning such simulations from data has been quite important in ﬁelds like physics [21] and graphics [22].
Learned simulations can often be faster and more efﬁcient for complex phenomenon prediction than just equation based simulators [23, 24].
Combining graph neural networks (GNN) [25, 26, 27] and differential equations [7, 8] has been explored in Poli et al.
[28], Huang et al.
[29] and Zang and Wang [30].
The dynamics of their test systems are relatively simple.
Importantly, there is no study of transferability of the learned dynamics to changing graph structure.
More recently, graph neural networks have been applied in the context of particle simulations by Sanchez-Gonzalez et al.
[31] and mesh-based simulations by Pfaff et al.
Moreover, unlike other graph neural net based dynamics learning works [28, 30] they also use edge based features.
However, these largely ignore the existence of control inputs, limiting their usefulness for reinforcement learning and control applications.
Some work in the past has examined the idea of message passing for continuous time Bayesian networks [33], or as a combination of generative graphical models and GNNs [34].
The focus of these works was to model full system dynamics for a particular graph structure, where messages play a role in computing the interactions, but do not evaluate transfer performance to different graph structures.
We distinguish our work from such prior methods primarily through the fact that our method is generalizable between different graph structures.
An approach that is close to ours from the neural ODE body of work would be the augmented neural ODEs[8] which also contain extra dimensions in the state, although these additional variables are zero-valued.
In our method, we repurpose those extra dimensions to explicitly pass messages across multiple neural ODEs as per the graph structure.
Methodology
The proposed framework tackles the problem of learning modular simulations where (1) smaller subsystems are individually modeled and (2) can be composed to simulate larger systems and processes.
Speciﬁcally, given observations from a complex system, the methodology seeks to learn both the individual simulation module as well the message passing mechanism that would best explain the data traces.
At the core of this method are Neural Differential Equations that use messages to coordinate amongst themselves.
Such message passing enables the framework to model various interaction effects between the modules as well as account for errors due to model approximations and time discretization.
The next subsection reviews Neural ODEs.
We then present our MP-NODE which extends neural ODEs to incorporate message passing.
Note that we assume that the graphs are homogeneous and the underlying graph structure is given.
However, we do show that the learnt Neural ODEs can be reused across different given graph topologies.
Handling heterogeneous subsystems along with learning the underlying graph structure is a more challenging problem and will be the focus of future work.
3.1 Background: Neural Differential Equations
Many continuous-time dynamical systems which can be found in physics (oscillators, mechanical systems), biology (population dynamics) etc. are best modeled by ordinary differential equations.
Typically, such a system can be described using an ODE as ˙ x = f ( x, u, t ) .
A data-driven way of learning a simulation of this system would thus be to estimate θ such that ˙ x = f θ ( x, u, t ) ﬁts a set of given observations.
Neural ordinary differential equations [7] are a class of neural networks that link the concepts of residual networks and dynamical systems.
Given a set of observations of the system state over time, under the assumption that the system dynamics can be described through an ODE, neural ODEs can be used to approximate the time derivative directly.
By integrating the model output using a black-box ODE solver for an arbitrarily long timestep and backpropagating appropriately, the model parameters θ can be updated to minimize the distance between the predicted states (Equation 2) and the observed ones. [Equation]
[Equation]
3.2 Message Passing with Neural ODEs (MP-NODE)
Now, let us consider a connected system that evolves as dXdt =
f ( G, X, U, t ) , where X ( t ) ∈ R n × d represents the entire system state, constituting of n connected nodes whose state is of length d each.
G = {V , E} captures the graph structure and how the nodes are connected to each other, which we assume is known.
N ( k ) refers to the neighbors of a node k in the graph.
While neural ODEs have been successfully used to model dynamical systems, modeling large systems with complex connectivity and dynamics is still challenging and computationally intensive.
Furthermore, even the slightest change to the connectivity of the system necessitates retraining of the entire system dynamics model.
On the other hand, connected systems can often be decomposed into repeatable subsystems that communicate with each other, hence it can be more efﬁcient to learn these individual blocks rather than attempting to model the connected system as a whole.
The ﬁrst key aspect of our framework is that the focus is placed on the subsystems rather than the entire system.
For a given homogeneous network, we essentially only learn the model corresponding to a single subsystem, or node.
When orchestrating full simulation of the system, the weights are shared across all individual nodes.
Secondly, due to this modular nature, the individual nodes need to learn to coordinate with each other and understand how their long term evolution changes based on interaction according to the graph structure.
We facilitate such communication between the nodes by augmenting each subsystem state with additional variables which we call messages with the augmented state for
node k being
X kt ,U kt
M kt f θ
X kt +
M (cid:48) kt +
MP
M kt + cat
U kt + f θ
X kt +
M (cid:48) kt + (cid:82) dt (cid:82) dt (cid:82) dt (cid:82)
dt m 1 t +1 m 2 t + m 3 t + m 4 t + m 5 t + m 6 t +
MP
X t ,U t
Encoder
G+
GN ···
G M
Decoder
X t +
Processor (cid:82) dt (cid:126)h 1 t (cid:126)h 2 t (cid:126)h 3 t (cid:126)h 4 t (cid:126)h 5 t (cid:126)h 6 t
GN
[Equation]
m k (cid:21) .
With the neural ODE operating on the augmented state, each neural subsystem outputs a message along with its predicted state at each timestep.
Similarly, each node also receives a message at the input, which results from an aggregation operation between all output messages of the node neighbors which are identiﬁed according to G .
We call this model the message passing neural ODE (MP-NODE) and describe the structure in Figure 1(Left).
The full formulation of the MP-NODE can thus be written as below, where m 0 k ( t ) is the input message at node k and m k is the output message.
[Equation]
While the outputs from the ODE solver contain both the predicted state and a message, the supervision and loss computation only happen on the predicted states, allowing the messages to evolve freely.
Over time, the model is able to use these additional variables to potentially encode the effects of different aggregations and interactions between neighbors, ﬁnally reaching an optimum conﬁguration that is applicable for all the nodes.
Learning such a generalizable subsystem allows the MP-NODE framework to be directly applicable for another system conﬁguration with higher or lower number of nodes, as opposed to attempting to learn dynamics for the entire system state such as is done with a conventional neural ODE, which needs retraining for new conﬁgurations.
The length of the message can be varied, and we treat it as a hyperparameter that changes the expressivity of the model.
In our implementation, we use the mean as the aggregator for incoming messages from the neighbors of a node.
Experiments
We evaluate the performance of MP-NODE on ﬁve different systems: two of them selected to focus on robotic regimes and three of them as representative of standard graph neural network benchmarks, for all of which we generate datasets of trajectories.
We implement our method in Julia [35] and make use of the SciML ecosystem [36, 37].
We list training-speciﬁc details such as hyperparameters, learning rates, optimizers in the appendix.
Source code for our experiments is available at https://github.com/microsoft/MPNODE.jl . We list the systems used below, and provide additional details in the appendix.
• Coupled Pendulum: This system is a combination of two identical simple pendula connected by a string, behaving in an undamped manner.
Hence, we train an MP-NODE that models a single pendulum, which is instantiated twice and the communication is enabled through the messages.
• Lorenz Attractor: A Lorenz system is a simpliﬁed model of atmospheric convection, described through a 3-dimensional system that exhibits chaotic behavior in speciﬁc cases [38].
In our case, we consider multiple Lorenz systems coupled in a fully connected manner. We denote Lorentz3 and Lorentz10 as fully connected systems with 3 and 10 nodes respectively.
20 − [Figure]
−
[Figure] [Figure]
{+θ 1 ( r a d ) datapred     − 
+}
[Figure]
s) θ
[Figure]
( r a d )    {+
+}
[Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
time(s) θ 1 ( r a d ) datapred    20 − 1 .
5 − 1 . 0 − 0 . 01 .time(s) θ 2 ( r a d )    20 0 . 0 0 .0 .
[Figure]
(s) mmsgnorm    20 0 . 0 0 .0 .
2 0 .0 .time(s)
m
[Figure] [Figure]
[Figure]
[Figure]
[-denote Lorentz3 and Lorentz
 as fully connected systems with 3 and 10 nodes respectively.-][Figure]
• Gene Dynamics: This system simulates a continuous time system from biology: the gene regulatory dynamics governed by the Michaelis-Menten equation[39] over grid-like conﬁgurations with varying connectivity.
Similar to one of the test cases in [30], we generate data for three types of graph topologies: a) Erdós-Rényi (ER) [40]; b) Barabási-Albert (BA) [41] and c) Wattz-Strogatz (WS) [42].
We use two conﬁgurations in our experiments, where the dynamics are simulated over a 4 × 4 grid and an 8 × 8 grid.
• Kuramoto System: The Kuramoto model [43] describes the behavior of large sets of coupled oscillators.
We simulate the Kuramoto dynamics over the same three types of networks mentioned above, for 10 nodes, which we refer to as Kuramoto10.
We use Lindner et al.
[44] to build our reference simulators to generate our datasets.
• Quadrotor swarm: This system involves a more challenging dynamics learning task of multiple quadrotors carrying a heavy load (that a single quadrotor cannot lift) while passing through a narrow doorway [45].
We consider two versions of this: one with 3 quadrotors, another with 6.
We compare our proposed approach to the one proposed by Sanchez-Gonzalez et al.
[31] that generalized different graph based dynamics learning methods, which we refer to as the L2S baseline.
We do not use edge features in our implementation of L2S (Figure 1(Right)).
4.1 Effect of messages
When the MP-NODE method of modeling is applied to systems with connected nodes, we ﬁnd that over time, the messages evolve to enable coordination between the nodes.
Figure 2 compares predictions (red) with observations (blue) of state evolution for the coupled pendulum system.
Each column corresponds to a pendulum and we plot two different trajectories (rowsand 3) in Figure 2.
We observe that the predictions closely match the ground truth observations.
We also plot the norm of outgoing messages (green) (rows 2 and 4) in Figure 2 and notice that they dynamically change in order to appropriately coordinate across the two nodes. [Figure] [Figure] [Figure]
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure]
We also investigate the importance of message passing by comparing performance of MP-NODE with and without messages.
Speciﬁcally, we turn the messages off by clamping the message outputs to zero.
Note that such forcing the messages to zero is similar to the augmented ODE applied to individual nodes.
Figure 3a shows that we achieve lower training loss with messages on (red) vs when off (blue).
We again test the idea of forcing messages to zero during training (similar to augmented neural ODE) on the
Kuramoto system.
As can be seen from Figure 4, using the additional dimensions for message-passing between nodes as in MP-NODE leads to substantial improvements in learning performance.
This is consistent with the ﬁndings for the coupled pendulum system.
Similarly, Figure 3b highlights that if no messages are passed between the nodes at test time, each node starts behaving like a simple pendulum that is unconnected to any neighbors.
This indicates that the message parameters are important and help the model by encoding the interaction between the nodes.
4.2 Finetuning
The modular nature of the MP-NODE also allows for easy transfer to different conﬁgurations of systems.
We examine the ability and performance of ﬁnetuning in several cases.
Increasing number of nodes: A common use case would be to take a trained MP-NODE module, representing a subsystem, and ﬁnetune it for a larger graph than what it was originally trained on.
We use the MP-NODE that was trained on Lorenz3, and ﬁnetune this model for the Lorenz conﬁguration.
We compare this training effort with that of MP-NODE models trained from scratch just for Lorenz10.
We show in Figure 5a that even with a small dataset of 50 trajectories, ﬁnetuning a pretrained MP-NODE reaches a low error signiﬁcantly faster than the models trained from scratch with larger training sets.
Similarly, we put the modular nature of the MP-NODE to test in the gene dynamics system by taking
MP-NODE models trained with smaller grid sizes and ﬁnetuning them on data from larger grid sizes.
In Figure 5b, we show the ﬁnetuning performance when a model trained on a 4 × 4 grid is transferred
NodeX
NodeY
V a l u e
GroundTruth 10301 2 3 4 5 6
NodeX
NodeY
V a l u e
Prediction 1030t =1 2 3 4 5 6
NodeX
NodeY
V a l u e
GroundTruth 20301 2 3 4 5 6
NodeX
NodeY
V a l u e
Prediction 2030t = 1 .
25 s 1 2 3 4 5 6
NodeX
NodeY
V a l u e
GroundTruth 27 .
NodeX
NodeY
V a l u e
Prediction 27 3036t = 5 .
s
NodeX
NodeY
V a l u e
GroundTruth 10301 2 3 41 2
NodeX
NodeY
V a l u e
Prediction 1030t = 0 s 1 2 3 41 2
NodeX
NodeY
V a l u e
GroundTruth 7 . 512 .
NodeX
NodeY
V a l u e
Prediction 7 . 512 .
517 .t
= 1 .
25 s 1 2 3 41 2
NodeX
NodeY
V a l u e
GroundTruth 6 79 101 2 3 41 2
NodeX
NodeY
V a l u e
Prediction 6 79 10t = {+5 . 
+}
s
[Figure]
[Table]
to an 8 × 8 grid, compared to one trained from scratch on the 8 × 8 grid. Similar to the trend with the Lorenz systems, we see that ﬁnetuning a trained MP-NODE allows for accurate predictions faster than training a model from scratch.
Furthermore, in Figure 6a, we show a qualitative comparison of ground truth and predicted state over the ﬁnetuned Gene 8 × 8 grid at different time steps.
This shows that due to the modular characteristic of our approach, ﬁnetuning a pretrained MP-NODE module allows for rapid adaptation to arbitrary number of nodes in a graph.
Higher system complexity: We also evaluate the possibility of ﬁnetuning MP-NODEs for different system parameters.
As an example, we ﬁnetune a MP-NODE model trained on Lorenz3 to Lorenz but with different coupling intensity and show the results in Figure 5d.
We observe that ﬁnetuning requires a lot less data to achieve better performance than training from scratch.
Similarly, when we ﬁnetune the MP-NODE model trained on Lorenz3 to Lorenz10 for longer time horizon of 10 s in Figure 5e, we again ﬁnd that a lot less data is required to achieve better performance than training from scratch.
Changing graph structure: We also examine the ability of MP-NODE to be ﬁnetuned for different graph structures.
To this end, we train an MP-NODE on a Kuramoto10 system connected according to the Barabasi-Albert (BA) network and attempt to ﬁnetune it for systems of other network types.
Similar to above, we see that ﬁnetuning is more efﬁcient at adapting to new network types than training models from scratch.
A comparison of training effort is shown in Figure 5c when transferring the Kuramoto10 model trained on the BA network to the ER network.
4.3 Zero-shot generalization
Considering that the modular nature of MP-NODE allows us to connect arbitrary number of subsystems together for a given graph structure, we can investigate the performance for new conﬁgurations without explicitly ﬁnetuning the model.
When tested on the Lorenz system, we ﬁnd that an MPNODE trained only on Lorenz3 exhibits reasonable zero-shot generalization performance on a higher number of Lorenz attractors, such as Lorenz7 and Lorenz10 (see appendix) without requiring any additional training.
We notice that this ability holds for not only changing numbers of nodes, but also different graph topologies.
We train the MP-NODE on the Gene 4 × 4 system connected according to the BarabasiAlbert (BA) topology, with the dataset comprising of data corresponding to ﬁve different adjacency matrices in the BA topology.
By training on data coming from different types of connections, the
[Figure]
MP-NODE learns how node dynamics are are affected by different graph structures.
After training, we observe that the MP-NODE model that was trained only on BA topology generalizes to the other unseen topologies as well (ER and WS).
Table 1 shows the MP-NODE’s performance on both seen and unseen network topologies.
We observe that MP-NODE exhibits low test error on unseen topologies as well, while achieving better validation performance than L2S.
In Figure 6b, we take a qualitative look at the predictions for gene dynamics on the ER network using the zero-shot transferred model.
4.4 Effect of message size
[Figure] In order to further improve the performance when modeling generally complex dynamics of systems such as Kuramoto or the quadrotor swarm, we investigate whether having a larger message size helps.
We train MP-NODE models on the Kuramoto10 system with message dimensionalities of 1 , 3 , 7 , and 13 .
By default, each Kuramoto node has a state length of 3.
We show this comparison in Figure 8.
We notice that larger messages result in higher accuracy, but after a certain threshold (messsage size =in this case), the improvement becomes minimal.
In the case of the quadrotor swarm, we choose message sizes of 7 , 13 , 20 for analysis.
Through Figure 7, we observe that all tested message sizes for MP-NODE lead to similar ﬁnal test errors.
Although L2S seems to do slightly better on the 3 Quadrotor System, it performs worse on the [-
-]6 Quadrotor System.
Similar to the trend in other datasets, we ﬁnd in Figure 9 that ﬁnetuning MP-NODE on even limited data allows learning more accurate simulations much faster than learning from scratch, and this holds true for different message sizes as well.
Although, we note that higher message dimensions can lead to higher data requirements.
4.5 Summary
To better demonstrate the advantages of MP-NODE when it comes to ﬁnetuning for different system conﬁgurations, we consider a metric: for a given test conﬁguration, we compute the ‘number of epochs taken to reach minimum test error multiplied by the value of the minimum test error observed’.
Intuitively, we can see that a lower number is better.
We summarize our ﬁndings in Table 2, where we compare MP-NODE ﬁnetuning vs. training from scratch, and for some conﬁgurations, we also compare with L2S. In the case of Lorenz and Kuramoto systems, we evaluate the performance under varying number of trajectories in the ﬁnetuning dataset. Through Figure 9, we observe that MP-NODE ﬁnetuning consistently achieves a lower combination of test error and number of epochs required to reach it, even in the low-data regime.
[Figure]
[Table]
Conclusions
We present a framework for modular modeling of complex homogeneous systems called the message passing neural ODE (MP-NODE).
MP-NODE builds upon neural ODEs, and learns a model for the individual node within a connected homogeneous system, which can then be shared across all instances in the system.
Augmenting the system state with additional message variables allows the model to encode the features required for coordination and accurate prediction over long horizons.
This modular way of thinking about complex systems opens up several possibilities such as zero-shot generalization to unseen conﬁgurations and rapid ﬁnetuning to more complex conﬁgurations.
We apply our framework to a wide variety of systems such as coupled pendulum, Lorenz attractors, biological networks, Kuramoto systems and a swarm of quadrotors, and showcase its ability to learn complex dynamics in a modular fashion while outperforming existing methods.
We also show that pretrained MP-NODE models can be ﬁnetuned for new conﬁgurations signiﬁcantly faster than training from scratch which has the potential to greatly amortize data and compute requirements for data-driven simulations.
While we focused on neural ODEs, general principles of parameter sharing and message passing as used in this work would be applicable to modern state-space models [46] too.
Limitations: Current implementation of MP-NODE is limited to homogeneous systems only.
For better ability to adapt to more complex systems, an extension that can handle heterogeneous systems would be needed.
Although encoder-decoder architectures could help in extending to nodes with different node state dimensions, different time-stepping requirements for the nodes might require a different messaging setup.
Moreover, the approach currently requires full knowledge of the graph structure which may not always be available for large complex systems.
In our current problem setup, we also do not consider potential issues such as delays in communication between nodes, which we leave for future work.
Broader Impact: Pretrained simulation nodes that are reusable can minimize data, compute, and energy requirements.
Data-driven simulations can often speed up conventional slow simulations while also being capable of capturing real world phenomena ignored in traditional simulations.
We show through our analysis of ﬁnetuning and generalization that our approach of modeling subsystems that are inherently reusable as opposed to speciﬁc conﬁgurations of systems has the potential to alleviate data and compute requirements.
However, limited data and lack of physical constraints can lead to arbitrarily bad predictions for unseen simulation regimes.
Nevertheless, building upon this approach, there is a potential to enable large variety of existing simulations to work with heterogeneous compute hardware and data via surrogate modeling [47].