Learning to Simulate Complex Physics

Chapter 1: Particles as Graphs

The first design decision: how to turn a physical system into a graph. The answer is beautifully simple. Each particle is a node. Two particles become connected by an edge if their Euclidean distance is below a threshold r — the interaction radius. This creates a connectivity radius graph, similar to a k-nearest-neighbors graph.

The interaction radius r is a hyperparameter. Too small: neighbors miss important interactions. Too large: every particle connects to every other, removing sparsity and defeating the purpose of a GNN. In practice, r is set based on the physics: for water, the interaction length is roughly 2× the average particle spacing.

Boundary Handling

Physical systems have walls, floors, and containers. These aren't particles, but they exert forces. GNS handles this with a trick: add the distance to each boundary as an extra dimension of the node feature vector. If the box has 4 walls, add 4 distances: [d_left, d_right, d_top, d_bottom]. The GNN learns from data how much walls affect particle motion — no hard-coded reflection equations.

Chapter 2: The Interaction Network — Live Sim

GNS uses a three-part GNN architecture: an encoder that embeds raw node and edge features into latent vectors, a processor (M message-passing layers) that propagates interactions through the graph, and a decoder that reads out each node's acceleration from its final latent vector.

python
class GNS(nn.Module):
    def __init__(self, input_dim, hidden=128, M=10):
        # Encoder: raw features → latent
        self.node_enc = MLP(input_dim, hidden)     # [d*C + types] → 128
        self.edge_enc = MLP(3, hidden)              # [rel_disp(2) + dist(1)] → 128
        # Processor: M rounds of message passing
        self.processors = nn.ModuleList([
            InteractionLayer(hidden) for _ in range(M)
        ])
        # Decoder: latent → acceleration
        self.decoder = MLP(hidden, 2)  # → (ax, ay)

    def forward(self, pos, vel_hist, particle_type, edges):
        # Build features
        node_feat = torch.cat([vel_hist, particle_type], dim=-1)
        i, j = edges
        rel_disp = pos[j] - pos[i]
        dist = rel_disp.norm(dim=-1, keepdim=True)
        edge_feat = torch.cat([rel_disp, dist], dim=-1)
        # Encode
        h = self.node_enc(node_feat)   # [N, 128]
        e = self.edge_enc(edge_feat)   # [E, 128]
        # Process (M rounds, with residuals)
        for layer in self.processors:
            h, e = layer(h, e, edges)
        # Decode acceleration
        acc = self.decoder(h)  # [N, 2] — predicted (ax, ay)
        return acc

Chapter 3: Rollout

Training GNS to predict one-step acceleration is straightforward: minimize the L2 loss between the predicted and ground-truth acceleration at each time step. But the real challenge is rollout — autoregressive generation of long trajectories where each step's output becomes the next step's input.

GNS addresses exposure bias with noise injection during training. Before computing each training example, add small random Gaussian noise to the particle positions. This mimics the corrupted inputs the model will see during rollout. The model learns to be robust to small positional errors because it saw them during training.

The noise level σ is a crucial hyperparameter. Too small: model is still brittle at rollout. Too large: training is destabilized. The paper uses σ chosen to match the typical one-step positional error of a preliminary trained model — a principled data-driven calibration.

Training Target: Acceleration, not Position

Why train to predict acceleration rather than position directly? Because acceleration is bounded and smooth. Positions can drift arbitrarily. If you train to predict the next position and the current position has a small error, the model must output a large correction — amplifying instability. Acceleration prediction gives the physics integrator more natural control over trajectory correction.

Chapter 4: Generalization

The most striking claim of the GNS paper: a model trained on 1000-particle simulations generalizes to 5000-particle simulations. This is remarkable because the graph structure changes completely — different number of nodes, different connectivity, different dynamics. How?

This is why the GNN architecture is not just a design choice but a philosophical commitment. A fully connected network takes the entire state as input and cannot generalize to different system sizes. An LSTM processes a fixed-length sequence and has the same problem. The GNN processes each node's local neighborhood independently — and that operation scales to any number of nodes.

Material Generalization

GNS also generalizes across materials using the particle type feature. Each particle's feature vector includes a one-hot encoding of its material type (water, sand, goop, rigid). When a new simulation mixes materials, the GNN learns interactions between different type codes. The MultiMaterial scenario — water + sand + rigid blocks — is trained jointly and the GNN learns all pairwise interaction types from data.

Container Generalization

Training containers are rectangular boxes. At test time, the model encounters L-shaped containers, ramps, and bowls (in the ramp scenarios). Since boundary information is encoded as distances to boundary planes (not as explicit geometry), the model generalizes to different container shapes as long as the local boundary distance is still meaningful.

Chapter 5: Results

GNS was evaluated on 5 physics domains with rollout lengths of 300-1000 steps. Metrics: mean squared error on particle positions (MSE) and qualitative visual fidelity assessed by human raters.

Ablation studies confirmed three key design choices:

• Noise injection: Without noise, MSE degrades 3-5× on long rollouts. With noise, rollout is stable.
• Velocity history (C steps): C=1 (current velocity only) gives 2× worse MSE than C=5. Memory matters.
• Number of processor layers M: M=10 is significantly better than M=3. Deeper message passing captures longer-range interactions within the interaction radius.

Chapter 6: vs Traditional Simulation

Traditional physics simulators are built on first principles. They solve the Navier-Stokes equations for fluids, the Mohr-Coulomb yield criterion for sand, constitutive relations for elastic and plastic solids. Every material requires new equations, new numerical methods, new implementations. Each simulator is a specialized tool.

What GNS Cannot Do

Several failure modes of GNS are worth knowing:

• Conservation violations: Traditional simulators can enforce energy and momentum conservation by construction. GNS may slightly violate conservation laws in long rollouts — the model never explicitly was taught to conserve energy.
• Extrapolation: If the test scenario has particle velocities much higher than anything in the training data, GNS can fail catastrophically. Traditional solvers handle novel regimes by solving the right equations.
• Interpretability: You cannot inspect a trained GNS and recover the effective viscosity, yield stress, or equation of state it has learned. It is a black box.
• Data requirements: GNS needs thousands of reference trajectories. For exotic materials or multi-physics phenomena, generating training data from a reference simulator may itself be expensive.

Chapter 7: Connections

GNS is part of a broader movement: neural simulation, or replacing expensive physics solvers with learned surrogates. Understanding its lineage and successors maps the field.