Learning to Simulate Complex Physics with Graph Networks

Alvaro Sanchez-Gonzalez; Jonathan Godwin; Tobias Pfaff; Rex Ying; Jure Leskovec; Peter W. Battaglia

Source-linked AI summary

Learning to Simulate Complex Physics with Graph Networks

Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, Peter W. Battaglia

arXiv:2002.09405v2Feb 21, 2020cs.LGphysics.comp-phstat.ML

TL;DR

Complex physical simulation is expensive and difficult for traditional simulators and standard end-to-end learning methods. The paper introduces GNS, which represents physical systems as interacting particles in graphs and learns their dynamics through message-passing. A single architecture simulates fluids, rigid solids, and deformable materials while generalizing to larger systems and longer rollouts than seen during training.

Problem
Traditional simulators are costly and often narrow, while standard learning approaches struggle with large state spaces and complex dynamics.
Method
GNS represents physical states as particle graphs and approximates their dynamics with learned message-passing in a single architecture.
Results
A single GNS architecture simulates interacting fluids, rigid solids, and deformable materials using tens of thousands of particles over thousands of timesteps.
Takeaways & Limitations
GNS provides a simpler, more accurate, and better-generalizing approach to learned simulation across diverse material systems.
Takeaways & Limitations
The approach focuses on particle-based simulation, while future work must explore efficient parameterizations and implementations to realize advantages over traditional simulators.

Abstract

from arXiv · show

Here we present a machine learning framework and model implementation that can learn to simulate a wide variety of challenging physical domains, involving fluids, rigid solids, and deformable materials interacting with one another. Our framework---which we term "Graph Network-based Simulators" (GNS)---represents the state of a physical system with particles, expressed as nodes in a graph, and computes dynamics via learned message-passing. Our results show that our model can generalize from single-timestep predictions with thousands of particles during training, to different initial conditions, thousands of timesteps, and at least an order of magnitude more particles at test time. Our model was robust to hyperparameter choices across various evaluation metrics: the main determinants of long-term performance were the number of message-passing steps, and mitigating the accumulation of error by corrupting the training data with noise. Our GNS framework advances the state-of-the-art in learned physical simulation, and holds promise for solving a wide range of complex forward and inverse problems.

1. Introduction

Traditional simulators are costly, computationally demanding, and often limited in generality or accuracy, while standard learning methods struggle with large state spaces and complex dynamics. GNS addresses this challenge with particle graphs and learned message-passing, enabling broad physical simulation and generalization beyond training conditions.

Motivation: Traditional simulators require years of engineering effort and trade off generality against accuracy in narrow settings.They also require substantial computational resources, making scaling up prohibitive.
Motivation: Standard end-to-end learning approaches struggle to overcome the large state spaces and complex dynamics of physical systems.
GNS Framework: GNS represents physical states as graphs of interacting particles and approximates dynamics through learned message-passing among nodes.
Results: A single deep learning architecture accurately simulates interacting fluids, rigid solids, and deformable materials.
Results: GNS generalizes to longer time scales than training and performs well across dozens of experiments with robustness to hyperparameter choices.Long-term performance depends on long-range interactions, spatial-invariance biases, and procedures mitigating error accumulation.

2. Related Work

Prior work includes particle-based, differentiable, learned, and graph-network approaches to physical simulation. GNS builds on these lines while targeting a simpler, more general, and more accurate approach across fluid, rigid, and deformable materials.

Particle-Based Simulation: Particle-based simulation represents physical states as particles and computes dynamics from local particle interactions, with SPH as a prominent fluid example.
Differentiable Simulation: Differentiable particle simulators such as DiffTaichi, PhiFlow, and Jax-MD backpropagate gradients through their architectures.
Learned Simulation: Learned simulators can predict complex phenomena more efficiently than engineered simulators, including learned components for faster fluid prediction.
Graph Networks: Graph Networks learn forward dynamics through message-passing, propagating latent information between interacting entities via edges.
GNS Positioning: GNS generalizes related approaches as a simpler, more accurate, and more broadly applicable method across fluid, rigid, and deformable systems.

3. GNS Model Framework

GNS learns a simulator that maps particle states to dynamics information and iteratively updates predicted states. Its graph-network dynamics model encodes particles, performs learned message-passing, and decodes per-particle dynamics information.

General Learnable Simulation: A simulator maps preceding states to future states by repeatedly applying a learned or numerical dynamics update.Rollouts are generated iteratively from the simulator’s predicted states.
General Learnable Simulation: The learnable simulator uses dθ to compute dynamics information, which an update mechanism combines with the current state to predict the next state.The framework uses an Euler integrator as its update mechanism.
Particle Representation: GNS represents the physical system as particles whose interactions approximate exchanges such as energy and momentum among neighbors.The interaction model determines the effects and materials the simulator can represent.
Graph Representation: Particle-based simulation corresponds to message-passing on a graph, with particles as nodes and pairwise relations as edges.
Encode-Process-Decode: The GNS dynamics model uses an ENCODER, PROCESSOR, and DECODER to transform particle states into dynamics information.
Encode-Process-Decode: The ENCODER constructs a latent graph from particle and pairwise features, while the PROCESSOR performs M learned message-passing steps before the DECODER extracts node dynamics.The number of message-passing steps is expected to scale with interaction complexity.

4. Experimental Methods

The experiments evaluate GNS across diverse materials and simulation engines using particle trajectories, graph construction, neural components, and noise-corrupted training. Models are assessed through long rollouts and held-out test MSE.

Datasets: The study evaluates water, sand, and goop, whose distinct fluid, granular, and plastically deformable behaviors usually require separate material models or algorithms.
Datasets: The experiments use PBD, SPH, and MPM engines to generate scenarios spanning fluids, obstacles, and multiple material behaviors.
Datasets: Datasets typically contain 1000 training, 100 validation, and 100 test trajectories lasting 300–2000 timesteps.
Representations: Each particle input includes position, C = 5 previous velocities, and static material features, while targets are per-particle average accelerations computed from positions.
GNS Implementation: The ENCODER connects particles within a fixed-resolution connectivity radius and recomputes edges with nearest neighbors during rollouts.
GNS Implementation: The PROCESSOR stacks M graph networks, and the DECODER’s MLP outputs accelerations used by an Euler integrator to update positions and velocities.
Training: Training corrupts input velocities with random-walk noise to make the training distribution closer to noisy rollout inputs and reduce error accumulation.The noise uses N(0, σv = 0.0003).
Evaluation: Models are selected using full-length validation rollouts and evaluated on held-out test-set mean squared error.

5. Results

GNS accurately simulates diverse interacting materials and generalizes to longer, larger, and out-of-distribution settings. Performance depends especially on message-passing capacity, neighborhood connectivity, parameter sharing, encoder choice, and training noise.

Scaling: Up to 19k particles and thousands of rollout timesteps remain within GNS’s demonstrated operating range.The models were trained for one-step predictions but produced plausible long-term trajectories at high resolution.
Material simulation: GNS simulates fluids, rigid obstacles, deformable materials, and multiple materials with one architecture.The model handles water, sand, goop, obstacle interactions, and jointly trained multimaterial systems.
Generalization: GNS generalizes to 28k particles after 2500 steps, 85k particles in a domain 32x larger in area, and unseen material shapes and interactions.These tests use inflows, larger spatial domains, unseen objects, and multimaterial interactions beyond training distributions.
Key Architectural Choices: More message-passing steps improve one-step and rollout accuracy, although computation time scales linearly with the step count.The authors attribute the improvement to longer-range and more complex particle interactions.
Key Architectural Choices: Unshared processor parameters, larger connectivity radii, relative encoding, and intermediate training noise generally improve rollout performance.Larger radii increase computation and memory, while excessive noise reduces one-step accuracy.

6. Conclusion

The paper presents GNS as a particle-based, graph-message-passing framework for simulating interacting physical materials. Its experiments support accurate, scalable, and generalizable simulation, while future work must extend beyond mesh-free particle representations and improve computational efficiency.

Conclusion: A single GNS architecture simulates fluids, rigid solids, and deformable materials using tens of thousands of particles over thousands of timesteps.The framework combines particle-based physical representations with learned graph message-passing.
Conclusion: GNS is reported as simpler, more accurate, and better generalizing than previous approaches.This comparison is stated as a conclusion across the paper’s evaluated settings.
Limitations and Future Work: The current work focuses on mesh-free particle methods, though the framework may also apply to mesh-based data such as finite-element methods.Future work also includes more efficient GNS computations and use of parallel hardware.
Conclusion: The framework is presented as a step toward more sophisticated generative models with greater capacity for physical reasoning.This consequence is stated within the paper’s broader conclusion.

Supplementary Material: Learning to Simulate Complex Physics with Graph

The GNS simulator predicts accelerations and uses semi-implicit Euler integration to update particle states. Training uses one-step loss, imposing a Markovian inductive bias, while whole-trajectory optimization can exploit prediction memory in ways that reduce generalization.

Update mechanism: Semi-implicit Euler integration updates the next state using accelerations predicted by GNS.With Δt = 1, predicted acceleration directly influences the next position update rather than only the velocity update.
Optimizing parameters of learnable simulator: The general simulator objective optimizes parameters over a distribution of state trajectories generated from initial conditions.The objective considers the rollout produced by the simulator across K timesteps.
Optimizing parameters of learnable simulator: GNS is trained with a one-step loss, rather than directly optimizing whole simulated trajectories.This choice focuses training on one-step prediction accuracy.
Optimizing parameters of learnable simulator: The one-step objective imposes the inductive bias that dynamics are Markovian and operate consistently at every trajectory time.The model is therefore trained to use the current state rather than trajectory-specific history.
Optimizing parameters of learnable simulator: Whole-trajectory optimization can favor models that encode initial conditions in predictions as a memory buffer, improving long-rollout accuracy but weakening generalization.The pendulum example illustrates how accumulated intermediate errors can motivate such hidden encoding.
Future work: Memory mechanisms may help simulators recognize and respect symmetries such as conservation of energy and momentum.The paper identifies memory-based simulation as a direction for future work.

B.2. Implementation Details

The implementation represents particle dynamics with finite-difference velocities and accelerations, material and wall features, and radius-based neighborhood graphs. Several architectural choices are intentionally kept simple because performance is relatively insensitive to them.

Input and output representations: Velocity is defined from the position difference between consecutive timesteps, while acceleration is the difference between consecutive velocities.The resulting acceleration uses three consecutive positions: p̈_t^k = p_t^{k+1} − 2p_t^k + p_t^{k−1}.
Input and output representations: Material type is encoded as a learned 16-dimensional particle embedding, while flat walls can be represented by clipped distances to each wall.Clipping wall distances to connectivity radius R preserves spatial translation invariance.
Building the graph: The graph connects particles within a connectivity radius selected to yield roughly 10–20 neighbors, using a standard k-d tree search.The radius was not fine-tuned across same-resolution 2D scenes.
Neural network parametrizations: Replacing deep encoder and decoder MLPs with linear layers produced similar performance.This result indicates that the tested nonlinear parametrization was not essential for performance.

B.3. Training

Training injects temporally accumulated velocity noise to reduce mismatch between training and rollout inputs, while alternatives use short self-generated rollouts. The pipeline also normalizes data, samples sequential trajectories into shuffled pairs, and scales across TPU and GPU hardware.

Noise injection: Noise injection perturbs input velocities with independent Gaussian samples and accumulates them over time as a random walk.Position features are adjusted afterward to preserve the finite-difference velocity definition.
Noise injection: Short self-generated rollouts provide another way to expose the model to its own predictions, but are more expensive and generalized less well to long sequences in the authors’ experience.Noise injection was preferred over the two-step training alternative described from prior work.
Normalization: Normalization statistics are computed from exact accumulated sums, squared sums, and particle counts after input noise is applied.This avoids moving-average shifts during training and supports scale-free learning rates across datasets.
Loss function and optimization procedures: A 1000-step trajectory yields 995 input-target pairs when conditioning on five past states, and pairs are sampled from a 10k shuffle buffer.Targets are computed from clean positions before input noise adjustments.
Loss function and optimization procedures: Optimization uses Adam with a nominal mini-batch size of two examples, up to 20M gradient updates, and an exponentially decaying learning rate.Predictions and targets are normalized before computing the loss.
Hardware: Training runs interchangeably on second-generation TPUs and V100 GPUs, with TPUs accelerating batch parallelism.The largest systems additionally use model parallelism across 16 TPU cores per example.

B.4. Distributional Evaluation Metrics

Evaluation supplements particle-wise MSE with permutation-invariant distributional metrics, because matching particle distributions can matter even when particle identities differ. Ablations show robustness across many hyperparameters, with notable gains from two input velocities and training-stability benefits from LayerNorm on larger datasets.

Distributional metrics: Particle-wise MSE can penalize predictions that match the true particle distribution but swap particle identities, motivating permutation-invariant optimal-transport metrics.The evaluation therefore measures distributional agreement in addition to exact particle correspondence.
Robustness: The GOOP ablation scans show robustness to typical hyperparameter changes.This supports the broader claim that many tested settings have limited effect on performance.
Input history: Two recent velocities significantly outperform conditioning on one, while larger input histories show similar performance; the models use C = 5.Because velocity is a position difference, C + 1 recent positions are required.
Architecture: Latent and MLP hidden sizes of 64 or larger produce little performance change.The ablation varies both latent size and MLP hidden size.
Architecture: Except with zero hidden layers, MLP depth has little effect on performance.The zero-layer case corresponds to a linear layer.
Architecture: Replacing encoder and decoder MLPs with linear layers produces no significant performance change.This agrees with the separate parametrization experiment reported in the implementation details.
Normalization: LayerNorm is slightly worse on small datasets when enabled but improves training stability on larger datasets.The performance effect depends on dataset size rather than being uniformly beneficial.
Graph mechanisms: Including self-edges or updating edge latent states produces similar performance.Neither choice materially changes the reported ablation results.

C.2. Noise-Related Training Parameters

The GOOP ablations examine how noise is introduced, scaled, and handled during training, including whether graph connectivity is recomputed and which noise component targets correct. These variations are compared with the default model using the same qualitative vertical scale.

Noise type: Four noise modes vary whether perturbations affect only the latest velocity, all states coherently, independently, or through a random walk.
Noise standard deviation: Noise standard deviation is varied as a separate training parameter in the GOOP experiments.
Graph reconnection: Recomputing graph connectivity after position noise does not change performance.
Position-noise correction: The fraction of position noise corrected by target accelerations ranges from 0% to 100%, with intermediate 10% and 30% settings.Because Euler updates cannot exactly correct position and velocity noise simultaneously, the model must choose one target or a compromise.
Ablation figures: Additional ablations and noise-related variations are shown in grey against the default model in red, with quartile error bars and shared vertical limits.

C.3. Distributional Evaluation Metrics

The distributional evaluation compares rollout behavior using MSE, Optimal Transport, and MMD. These metrics generally agree with MSE-based analyses, while distributional metrics can better distinguish visually good and bad rollouts.

MSE and distributional metrics generally lead to similar conclusions across the analyses.
Distributional metric differences between qualitatively good and bad rollouts can be more prominent and align more closely with visual impressions.
Rollout error is plotted over time for water, sand, and goop using MSE, Optimal Transport, and MMD across three evaluation trajectories.
MSE tends to increase over time in chaotic systems, whereas Optimal Transport and MMD can decrease for water and sand as particles equilibrate.For goop, separate chaotic clusters can keep Optimal Transport error high.
Figure C.4 reports Optimal Transport and Maximum Mean Discrepancy rollout errors for architectural ablations, averaged across the test dataset.Bars show median-seed performance and error bars show lower and higher quartiles.

C.5. Quantitative Generalization Results on CONTINUOUS Domain

CConv uses displacement-dependent continuous convolutions that can be expressed as graph-network edge updates, but its strengths and failures vary across materials and interactions. GNS more faithfully handles complex multimaterial behavior in the supplied comparisons.

Inference Times: Inference-time comparisons were informative post-hoc, while precomputed-neighborhood measurements showed neighborhood computation dominated graph-network inference.The reported timing results were not the main goal of the work, and optimized neighborhood calculations could improve performance.
CConv as a Graph Network: CConv can be represented as a graph network with a displacement-specific edge update function.Its particle displacements become edge attributes, while neighborhood messages are pooled to update each particle.
Continuous Convolution: CConv interpolates displacement-dependent weight matrices from nearby discrete kernel-grid coordinates for continuous particle positions.In one dimension this uses linear interpolation; in three dimensions it uses trilinear interpolation.
Qualitative Comparisons: CConv simulated sand reasonably well but struggled with complex solids, preserving goop rest states, and preventing local explosions in MULTIMATERIAL.In BOXBATH it simulated the fluid well but struggled to keep the box shape intact.
Qualitative Comparisons: CConv was strongest for water-like fluids, yet did not match GNS and was weaker on other materials and cross-material interactions.The comparison attributes part of this gap to GNS using more message-passing steps, pairwise interaction functions, and higher-capacity multilayer perceptrons.
Comparison with DPI: GNS faithfully reproduced BOXBATH trajectories for both water and solid particles without specialized mechanisms, and also reproduced the fluid layering effect.The comparison used material labels rather than explicit solidity constraints or forced preservation of box-particle displacements.