Learning to Simulate

Complex Physics

with Graph Networks

Alvaro Sanchez-Gonzalez ^{* 1} Jonathan Godwin ^{* 1} Tobias Pfaff ^{* 1} Rex Ying ^{* 1 2} Jure Leskovec ² Peter W. Battaglia ¹

Table of contents

01

03

02

04

05

06

Introduction

Related Work

GNS Model Framework

Experimental Methods

Results

Conclusion

2

Introduction

01

3

Idea

machine learning simulators

• Can train simulators directly from observed data.
• The large state spaces and complex dynamics have been difficult for standard end-to-end learning approaches to overcome.

• High-quality simulators require substantial computational resources.
• Even the best are often inaccurate due to insufficient knowledge of, or difficulty in approximating, the underlying physics and parameters.

traditional simulators

4

Graph Network-based Simulators (GNS)

The framework imposes strong inductive biases, where rich physical states are represented by graphs of interacting particles, and complex dynamics are approximated by learned message-passing among nodes.

5

Related Work

02

6

Traditional Ways

Differentiable Particle-based Simulators

smoothed particle hydrodynamics

(SPH)

position-based dynamics

(PBD)

Evaluates pressure and viscosity forces around each particle, and updates particles’

velocities and positions accordingly.

​

Incompressibility and collision dynamics involve resolving pairwise distance constraints between particles, and directly predicting their position changes.

Discretize a block of material into a large number of particles, and compute spatial derivatives and solve momentum equations.

material point method

(MPM)

DiffTaichi, PhiFlow, Jax-MD :

Backpropagate gradients through the architecture.

​

7

Graph Networks

• A type of graph neural network.
• A GN maps an input graph to an output graph with the same structure but potentially different node, edge, and graph-level attributes, and can be trained to learn a form of messagepassing , where latent information is propagated between nodes via the edges.
• GNs and their variants, e.g., "interaction networks", can learn to simulate rigid body, mass-spring, n-body, and robotic control systems , as well as non-physical systems.

GNS framework is a general approach to learning simulation, is simpler to implement, and is more accurate across fluid, rigid, and deformable material systems.

​

8

GNS Model Framework

03

9

Model Structure

10

Graph Networks

11

Experimental Methods

04

12

GNS Implementation Details

Input

ENCODER

x_i^tk = [p_i^tk , p_i˙ ^tk-C+1 , . . . , p_i˙ ^tk,f_i ],

C = 5,

p_i = position

p_i˙ = previous velocities

f_i = features that capture static material properties

constructs the graph structure G⁰

​

MLP encode node features and edge features into the latent vectors, v_i and e_i,j , of size 128.

13

the components of the GNS framework using standard deep learning building blocks,

and used standard nearest neighbor algorithms to construct the graph.

GNS Implementation Details

PROCESSOR

DECODER

Uses a stack of M GNs with identical structure, MLPs as internal edge and node update functions.

Decoder’s learned function, δ^v ,

is an MLP.

After the DECODER, the future position and velocity are updated using an Euler integrator, so the y_i corresponds to accelerations, p¨_i

14

Neural network parameters

MLP

Loss function

All MLPs have

• Two hidden layers
• ReLU activations
• Layer size = 128

followed by a LayerNorm layer

​

​

Randomly sampled particle state pairs (x_i^tk, x_i^tk+1) from training trajectories, calculated target accelerations p¨ ^tk

​

L(x_i^tk , x_i^tk+1 ; θ) = ||dθ(x_i^tk) − p¨ ^tk||²

15

Results

05

16

Simulating Complex Materials

SPH

GNS

MSE

17

• Substantially greater than demonstrated in previous methods—GNS can operate at resolutions high enough for practical prediction tasks and high-quality 3D renderings.
• Although our models were trained to make one-step predictions, the long-term trajectories remain plausible even over thousands of rollout timesteps.

Multiple Interacting Materials

18

Train a single architecture with a single set of parameters to simulate all of our different materials, interacting with each other in a single system.

Generalization

19

The GNS generalizes well even beyond its training distributions, which suggests it learns a more general-purpose understanding of the materials and physical processes experienced during training.

trained:

1x1 domain,

2.5k particles, 600 steps

inference:

8x4 domain,

85k particles, 5000 steps

​

Key Architectural Choices

20

While our GNS model was generally robust to architectural and hyperparameter settings, we also identified several factors which had more substantial impact:

1. the number of message-passing steps,
2. shared vs. unshared PROCESSOR GN parameters,
3. the connectivity radius,
4. the scale of noise added to the inputs during training,
5. relative vs. absolute ENCODER.

​

Comparisons to Previous Models

21

• DPI:

While DPI uses hard-coded constraints to keep the box shape consistent, our model achieves this without any special treatment of the solid particles.

• CConv:

CConv performs well for domains like water, which it was built for, but struggles with some of our more complex materials.

Similarly, in a CConv rollout of the BOXBATH DOMAIN the rigid box loses its shape.

Conclusion

06

22

Conclusion

Simpler

More accurate

Better generalization

23

• GNS approach may also be applicable to data represented using meshes.
• There are also natural ways to incorporate stronger, generic physical knowledge into the framework.