Hamiltonian Neural Networks

Sam Greydanus, Misko Dzamba, Jason Yosinski | 2019

• Paper: arxiv.org/abs/1906.01563
• Blog: greydanus.github.io/2019/05/15/hamiltonian-nns/

Basic usage

To train a Hamiltonian Neural Network (HNN):

• Task 1: Ideal mass-spring system: python3 experiment-spring/train.py --verbose
• Task 2: Ideal pendulum: python3 experiment-pend/train.py --verbose
• Task 3: Real pendulum (from this Science paper): python3 experiment-real/train.py --verbose
• Task 4: Two-body problem: python3 experiment-2body/train.py --verbose
• Task 4b: Three-body problem: python3 experiment-3body/train.py --verbose
• Task 5: Pixel pendulum (from OpenAI Gym): python3 experiment-pixels/train.py --verbose

To analyze results

• Task 1: Ideal mass-spring system: analyze-spring.ipnyb
• Task 2: Ideal pendulum: analyze-pend.ipnyb
• Task 3: Real pendulum: analyze-real.ipnyb
• Task 4: Two-body problem: analyze-2body.ipnyb
• Task 4b: Three-body problem: analyze-3body.ipnyb
• Task 5: Pixel pendulum: analyze-pixels.ipnyb

Summary

Even though neural networks enjoy widespread use, they still struggle to learn the basic laws of physics. How might we endow them with better inductive biases? In this paper, we draw inspiration from Hamiltonian mechanics to train models that learn and respect exact conservation laws in an unsupervised manner. We evaluate our models on problems where conservation of energy is important, including the two-body problem and pixel observations of a pendulum. Our model trains faster and generalizes better than a regular neural network. An interesting side effect is that our model is perfectly reversible in time.

The HNN recipe

1. Make a dataset of pixel-space observations of a physical system where energy is conserved. Here we're working with a pendulum.

2. Train an autoencoder on the dataset. The autoencoder is a bit unusual: its latent representation gets fed to the HNN, which tries to model the system's dynamics in latent space.

3. Since the HNN uses the latent representation to model dynamics, we can think of the latent factors as being analogous to canonical coordinates (e.g. position and velocity).

4. Phase space plots are a common way to visualize Hamiltonians. We can make a phase space plot in the autoencoder's latent space. We can also integrate along the energy contours of phase space to predict the dynamics of the system (in the figure below, we intentionally "add energy" halfway through).

5. After integrating in latent space, we can project back into pixel space to simulate the dynamics of the system.

Here's what it looks like when we add energy halfway through the simulation:

Modeling larger systems

We can also model larger systems -- systems with more than one pair of canonical coordinates. The two-body problem, for example, has four coordinate pairs.

Numbers

Train loss

• Choose data of the form x=[x0, x1,...] and dx=[dx0, dx1,...] where dx is the time derivative of x
• Let dx' = model.time_derivative(x)
• Compute L2 distance between dx and dx'

Test loss

Do the same thing with test data

Energy MSE

• Choose a trajectory [x0, x1,...] from test data
• Use RK4 integration to estimate [x0', x1',...] using the model
• Compute the L2 distance between [energy(x0), energy(x1),...] and [energy(x0'), energy(x1'),...]

Dependencies

• OpenAI Gym
• PyTorch
• NumPy
• ImageIO
• Scipy

This project is written in Python 3.