Hamiltonian Graph Networks with ODE Integrators

TL;DR

The problem addresses learning physical system dynamics with strong structural biases. The authors propose a framework that combines graph networks with differentiable ODE integrators and a Hamiltonian internal representation, introducing DeltaGN, OGN, and HOGN variants. Key findings show improved predictive accuracy, energy conservation, and zero-shot generalization across unseen time-steps and integrator orders, with HOGN often best aligning with true Hamiltonian dynamics under RK4. This approach advances learned simulation and holds promise for broad applicability beyond traditional physics domains.

Abstract

We introduce an approach for imposing physically informed inductive biases in learned simulation models. We combine graph networks with a differentiable ordinary differential equation integrator as a mechanism for predicting future states, and a Hamiltonian as an internal representation. We find that our approach outperforms baselines without these biases in terms of predictive accuracy, energy accuracy, and zero-shot generalization to time-step sizes and integrator orders not experienced during training. This advances the state-of-the-art of learned simulation, and in principle is applicable beyond physical domains.

Figures (12)

• Figure 1: (a) The data reflects a system's temporal dynamics, which is governed by physics. (b) The baseline DeltaGN model takes as input a state represented by a graph and a time-step, and uses a $\mathrm{GN}_V$ to predict state changes. (c) Our OGN and HOGN models take as input an input state, time-step, and a function to evaluate the ODE's time derivatives, $f_{\dot{\mathbf{q}}, \dot{\mathbf{p}}}$, and uses an integrator, which queries $f_{\dot{\mathbf{q}}, \dot{\mathbf{p}}}$ at different points, to predict the state after the time-step. (d) The $f_{\dot{\mathbf{q}}, \dot{\mathbf{p}}}$ takes as input any state and outputs its time derivatives. (e) The OGN model uses a $\mathrm{GN}_V$ as $f^{\textrm{OGN}}_{\dot{\mathbf{q}}, \dot{\mathbf{p}}}$. (f) The HOGN model's $f^{\textrm{HOGN}}_{\dot{\mathbf{q}}, \dot{\mathbf{p}}}$ uses a $\mathrm{GN}_\mathbf{u}$ to predict the Hamiltonian, which is then differentiated w.r.t. the input state.
• Figure 2: (a) Predictive accuracy (on 20-step trajectory) across models using RK4 for the ODE based models. The HOGN is most accurate. (b) Energy accuracy across the same models. (c-g) Last 300 steps of a 500-step trajectory of a 6-particle system, where dots indicate the final position of the particles (colors fade into the past). (c) Ground truth trajectory for all particles. (d-g) Trajectory for one of the particles (blue) superimposed by the trajectory obtained (red) when integrating the True Hamiltonian, or (e-g) using the best seed of the different learned models, at a time step of 0.1. While errors of all models are very small at the beginning of the trajectory, the HOGN is the only learned model still indistinguishable from ground truth at the end of the long 500-step trajectory (https://tinyurl.com/hogn-example-trajectories).
• Figure 3: (a-b) Time-step size generalization error in 20-step-long trajectories when trained with (a) a fixed time-step of 0.1 or (b) variable time-steps (0.02-0.2). (c-d) Predictive accuracy and energy conservation across models and integrators. (e-f) Results when varying the integrator used at test time, where points that share the same train and test integrator are highlighted with black circles.
• Figure D.1: (a-b) Predictive accuracy and energy conservation across models and integrators including symplectic integrators (analogous to Fig. \ref{['fig:generalization']}c-d).
• Figure D.2: (a-d) Predictive accuracy when varying the integrator used at test time, where points that share the same train and test integrator are highlighted with black circles (analogous to Fig. \ref{['fig:generalization']}e-f).