Hamiltonian Generative Networks

1 INTRODUCTION Any system capable of a wide range of intelligent behaviours within a dynamic environment requires a good predictive model of the environment ’ s dynamics . This is true for intelligence in both biological ( Friston , 2009 ; 2010 ; Clark , 2013 ) and artificial ( Hafner et al. , 2019 ; Battaglia et al. , 2013 ; Watter et al. , 2015 ; Watters et al. , 2019 ) systems . Predicting environmental dynamics is also of fundamental importance in physics , where Hamiltonian dynamics and the structurepreserving transformations it provides have been used to unify , categorise and discover new physical entities ( Noether , 1915 ; Livio , 2012 ) . Hamilton ’ s fundamental result was a system of two first-order differential equations that , in a stroke , unified the predictions made by prior Newtonian and Lagrangian mechanics ( Hamilton , 1834 ) . After well over a century of development , it has proven to be essential for parsimonious descriptions of nearly all of physics . ∗Equal contribution . 1More results and video evaluations are available at : http : //tiny.cc/hgn Hamilton ’ s equations provide a way to predict a system ’ s future behavior from its current state in phase space ( that is , its position and momentum for classical Newtonian systems , and its generalized position and momentum more broadly ) . Hamiltonian mechanics induce dynamics with several nice properties : they are smooth , they include paths along which certain physical quantities are conserved ( symmetries ) and their time evolution is fully reversible . These properties are also useful for machine learning systems . For example , capturing the time-reversible dynamics of the world state might be useful for agents attempting to account for how their actions led to effects in the world ; recovering an abstract low-dimensional manifold with paths that conserve various properties is tightly connected to outstanding problems in representation learning ( see e.g . Higgins et al . ( 2018 ) for more discussion ) ; and the ability to conserve energy is related to expressive density modelling in generative approaches ( Rezende & Mohamed , 2015 ) . Hence , we propose a reformulation of the well-known image manifold hypothesis by extending it with a Hamiltonian assumption ( illustrated in Fig . 1 ) : natural images lie on a low-dimensional manifold embedded within a high-dimensional pixel space and natural sequences of images trace out paths on this manifold that follow the equations of an abstract Hamiltonian . Given the rich set of established tools provided by Hamiltonian descriptions of system dynamics , can we adapt these to solve outstanding machine learning problems ? When it comes to adapting the Hamiltonian formalism to contemporary machine learning , two questions need to be addressed : 1 ) how should a system ’ s Hamiltonian be learned from data ; and 2 ) how should a system ’ s abstract phase space be inferred from the high-dimensional observations typically available to machine learning systems ? Note that the inferred state may need to include information about properties that play no physical role in classical mechanics but which can still affect their behavior or function , like the colour or shape of an object . The first question was recently addressed by the Hamiltonian Neural Network ( HNN ) ( Greydanus et al. , 2019 ) approach , which was able to learn the Hamiltonian of three simple physical systems from noisy phase space observations . However , to address the second question , HNN makes assumptions that restrict it to Newtonian systems and appear to limit its ability to scale to more challenging video datasets . In this paper we introduce the first model that answers both of these questions without relying on restrictive domain assumptions . Our model , the Hamiltonian Generative Network ( HGN ) , is a generative model that infers the abstract state from pixels and then unrolls the learned Hamiltonian following the Hamiltonian equations ( Goldstein , 1980 ) . We demonstrate that HGN is able to reliably learn the Hamil- tonian dynamics from noisy pixel observations on four simulated physical systems : a pendulum , a mass-spring and two- and three- body systems . Our approach outperforms HNN by a significant margin . After training , we demonstrate that HGN produces meaningful samples with reversible dynamics and that the speed of rollouts can be controlled by changing the time derivative of the integrator at test time . Finally , we show that a small modification of our architecture yields a flexible , normalising flow-based generative model that respects Hamiltonian dynamics . We show that this model , which we call Neural Hamiltonian Flow ( NHF ) , inherits the beneficial properties of the Hamiltonian formalism ( including volume preservation ) and is capable of expressive density modelling , while offering computational benefits over standard flow-based models . 2 RELATED WORK . Most machine learning approaches to modeling dynamics use discrete time steps , which often results in an accumulation of the approximation errors when producing rollouts and , therefore , to a fast drop in accuracy . Our approach , on the other hand , does not discretise continuous dynamics and models them directly using the Hamiltonian differential equations , which leads to slower divergence for longer rollouts . The density model version of HGN ( NHF ) uses the Hamiltonian dynamics as normalising flows along with a numerical integrator , making our approach somewhat related to the recently published neural ODE work ( Chen et al. , 2018 ; Grathwohl et al. , 2018 ) . What makes our approach different is that Hamiltonian dynamics are both invertible and volume-preserving ( as discussed in Sec . 3.3 ) , which makes our approach computationally cheaper than the alternatives and more suitable as a model of physical systems and other processes that have these properties . Also related is recent work attempting to learn a model of physical system dynamics end-to-end from image sequences using an autoencoder ( de Avila Belbute-Peres et al. , 2018 ) . Unlike our work , this model does not exploit Hamiltonian dynamics and is trained in a supervised or semi-supervised regime . 2.1 HAMILTONIAN NEURAL NETWORK . One of the most comparable approaches to ours is the Hamiltonian Neural Network ( HNN ) ( Greydanus et al. , 2019 ) . This work , done concurrently to ours , proposes a way to learn Hamiltonian dynamics from data by training the gradients of a neural network ( obtained by backpropagation ) to match the time derivative of a target system in a supervised fashion . In particular , HNN learns a differentiable function H ( q , p ) that maps a system ’ s state ( its position , q , and momentum , p ) to a scalar quantity interpreted as the system ’ s Hamiltonian . This model is trained so thatH ( p , q ) satisfies the Hamiltonian equation by minimizing LHNN= 1 2 [ ( ∂H ∂p − dq dt ) 2+ ( ∂H ∂q + dp dt ) 2 ] , ( 1 ) where the derivatives ∂H∂q and ∂H ∂p are computed by backpropagation . Hence , this learning procedure is most directly applicable when the true state space ( in canonical coordinates ) and its time derivatives are known . Accordingly , in the majority of the experiments presented by the authors , the Hamiltonian was learned from the ground truth state space directly , rather than from pixel observations . The single experiment with pixel observations required a modification of the model . First , the input to the model became a concatenated , flattened pair of images ot= [ xt , xt+1 ] , which was then mapped to a low-dimensional embedding space zt= [ qt , pt ] using an encoder neural network . Note that the dimensionality of this embedding ( z∈R2 in the case of the pendulum system presented in the paper ) was chosen to perfectly match the ground truth dimensionality of the phase space , which was assumed to be known a priori . This , however , is not always possible . The latent embedding was then treated as an estimate of the position and the momentum of the system depicted in the images , where the momentum was assumed to be equal to the velocity of the system – an assumption enforced by the additional constraint found necessary to encourage learning , which encouraged the time derivative of the position latent to equal the momentum latent using finite differences on the split latents : LCC= ( pt− ( qt+1−qt ) ) 2 . ( 2 ) This assumption is appropriate in the case of the simple pendulum system presented in the paper , however it does not hold more generally . Note that our approach does not make any assumptions on the dimensionality of the learned phase space , or the form of the momenta coordinates , which makes our approach more general and allows it to perform well on a wider range of image domains as presented in Sec . 4 . 3 METHODS . 3.1 THE HAMILTONIAN FORMALISM . The Hamiltonian formalism describes the continuous time evolution of a system in an abstract phase space s = ( q , p ) ∈ R2n , where q ∈ Rn is a vector of position coordinates , and p ∈ Rn is the corresponding vector of momenta . The time evolution of the system in phase space is given by the Hamiltonian equations : ∂q ∂t = ∂H ∂p , ∂p ∂t =−∂H ∂q ( 3 ) where the HamiltonianH : R2n→R maps the state s= ( q , p ) to a scalar representing the energy of the system . The Hamiltonian specifies a vector field over the phase space that describes all possible dynamics of the system . For example , the Hamiltonian for an undamped mass-spring system is H ( q , p ) = 12kq 2+ p 2 2m , wherem is the mass , q∈R 1 is its position , p∈R1 is its momentum and k is the spring stiffness coefficient . The Hamiltonian can often be expressed as the sum of the kinetic T and potential V energiesH=T ( p ) +V ( q ) , as is the case for the mass-spring example . Identifying a system ’ s Hamiltonian is in general a very difficult problem , requiring carefully instrumented experiments and researcher insight produced by years of training . In what follows , we describe a method for modeling a system ’ s Hamiltonian from raw observations ( such as pixels ) by inferring a system ’ s state with a generative model and rolling it out with the Hamiltonian equations . 3.2 LEARNING HAMILTONIANS WITH THE HAMILTONIAN GENERATIVE NETWORK . Our goal is to build a model that can learn a Hamiltonian from observations . We assume that the data X= { ( x10 , ... , x1T ) , ... , ( xK0 , ... , xKT ) } comes in the form of high-dimensional noisy observations , where each xi=G ( si ) =G ( qi ) is a non-deterministic function of the generalised position in the phase space , and the full state is a non-deterministic function of a sequence of images si= ( qi , pi ) =D ( xi0 , ... , x i t ) , since the momentum ( and hence the full state ) can not in general be recovered from a single observation . Our goal is to infer the abstract state and learn the Hamiltonian dynamics in phase space by observing K motion sequences , discretised into T+1 time steps each . In the process , we also want to learn an approximation to the generative process G ( s ) in order to be able to move in both directions between the high dimensional observations and the low-dimensional abstract phase space . Although the Hamiltonian formalism is general and does not depend on the form of the observations , we present our model in terms of visual observations , since many known physical Hamiltonian systems , like a mass-spring system , can be easily observed visually . In this section we introduce the Hamiltonian Generative Network ( HGN ) , a generative model that is trained to behave according to the Hamiltonian dynamics in an abstract phase space learned from raw observations of image sequences . HGN consists of three parts ( see Fig . 2 ) : an inference network , a Hamiltonian network and a decoder network , which are discussed next . The inference network takes in a sequence of images ( xi0 , ... xiT ) , concatenated along the channel dimension , and outputs a posterior over the initial state z ∼ qφ ( ·|x0 , ... xT ) , corresponding to the system ’ s coordinates in phase space at the first frame of the sequence . We parametrise qφ ( z ) as a diagonal Gaussian with a unit Gaussian prior p ( z ) = N ( 0 , I ) and optimise it using the usual reparametrisation trick ( Kingma & Welling , 2014 ) . To increase the expressivity of the abstract phase space s0 , we map samples from the posterior with another function s0=fψ ( z ) to obtain the system ’ s initial state . As mentioned in Sec . 3.1 , the Hamiltonian function expects the state to be of the form s= ( q , p ) , hence we initialise s0 ∈R2n and arbitrarily assign the first half of the units to represent abstract position q and the other half to represent abstract momentum p. The Hamiltonian network is parametrised as a neural network with parameters γ that takes in the inferred abstract state and maps it to a scalarHγ ( st ) ∈R . We can use this function to do rollouts in the abstract state space using the Hamiltonian equations ( Eq . 3 ) , for example by Euler integration : st+1= ( qt+1 , pt+1 ) = ( qt+ ∂H ∂pt dt , pt− ∂H∂qt dt ) . In this work we assume a separable Hamiltonian , so in practice we use a more sophisticated leapfrog integrator to roll out the system , since it has better theoretical properties and results in better performance in practice ( see Sec . A.6 in Supplementary Materials for more details ) . The decoder network is a standard deconvolutional network ( we use the architecture from Karras et al . ( 2018 ) ) that takes in a low-dimensional representation vector and produces a high-dimensional pixel reconstruction . Given that each instantaneous image does not depend on the momentum information , we restrict the decoder to take only the position coordinates of the abstract state as input : pθ ( xt ) =dθ ( qt ) . The objective function . Given a sequence of T+1 images , HGN is trained to optimise the following objective : L ( φ , ψ , γ , θ ; x0 , ... xT ) = 1 T+1 T∑ t=0 [ Eqφ ( z|x1 , ... xT ) [ log pψ , γ , θ ( xt | qt ) ] ] − KL ( qφ ( z ) || p ( z ) ) , ( 4 ) which can be seen as a temporally extended variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) objective , consisting of a reconstruction term for each frame , and an additional term that encourages the inferred posterior to match a prior . The key difference with a standard VAE lies in how we generate rollouts – these are produced using the Hamiltonian equations of motion in learned Hamiltonian phase space .

Summary: The paper introduces a novel way of learning Hamiltonian dynamics with a generative network. The Hamiltonian generative network (HGN) learns the dynamics directly from data by embedding observations in a latent space, which is then transformed into a phase space describing the system's initial (abstract) position and momentum. Using a second network, the Hamiltonian network, the position and momentum are reduced to a scalar, interpreted as the Hamiltonian of the system, which can then be used to do rollouts in the phase space using techniques known from, e.g., Hamiltonian Monte Carlo sampling. Finally, a decoder network can use the system's phase space location at any rollout step to generate images of the system. The HGN can further be modified, leading to a flow-based model, the Neural Hamiltonian Flow (NHF).