SyMetric: Measuring the Quality of Learnt Hamiltonian Dynamics Inferred from Vision

1 Introduction . If you want to understand how the world will change , a good place to start is with simple physical principles , such as the laws of motion . One way to achieve good performance in predicting dynamics forward and backward in time is by building a model using the Hamiltonian formalism from classical mechanics . This formalism aims to model physical systems which conserve energy . Apart from real physical problems , other important dynamical systems can also be modelled using this formulation , including GAN optimisation [ 35 , 38 , 34 , 36 ] , multi-agent learning in non-transitive zero-sum games [ 5 , 6 , 53 ] , or the transformations induced by flows in generative models [ 41 ] . While many naturally observed dynamics do not preserve energy – for example , energy is often added or dissipated – they can 1The code for reproducing all results is available on https : //github.com/deepmind/ deepmind-research/tree/master/physics_inspired_models . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . still be modelled using the Hamiltonian formalism by augmenting it with additional terms capturing the deviations from the energy-conserved state [ 37 , 23 , 43 , 55 , 56 ] . Recently , a body of work has emerged that brings these well-established principles of modelling dynamics from physics – such as the conservation of energy , and numerical formulations from the theory of differential equations – to neural network architectures [ 49 , 20 , 9 , 3 , 58 , 8 , 4 , 31 , 50 , 43 , 56 , 12 , 17 , 24 , 14 , 32 , 11 , 45 , 57 , 54 , 48 , 21 ] . Most of these papers , however , learn dynamics directly from an abstract state , such as the positions and momenta of the physical objects . This may be a suitable choice for some applications , for example in molecular simulations , but in situations where inference at test time has to be made from high-dimensional observations , like images – a setup common to many robotics , reinforcement learning or autonomous driving challenges – it is important to be able to have models that can both correctly infer the abstract state space and faithfully capture the underlying Hamiltonian dynamics without access to the ground truth state space information at inference time . Currently only a handful of methods exist for learning dynamics with physical priors from pixels [ 49 , 9 , 45 , 3 , 58 ] . These models learn in a completely unsupervised manner using pixel observations of the dynamics , however , currently there is no good way to know whether these models have indeed managed to recover the underlying dynamics faithfully through learning . In this paper we empirically demonstrate that reconstruction based measures typically used in the field are poor indicators of how well the model has learnt the underlying dynamics and propose a better method for evaluating such models . In summary , our contributions are as follows : 1 . We introduce performance measures for identifying whether a model trained with the Hamiltonian prior has managed to capture the true dynamics of an energy-conserving system when learning from high-dimensional observations . 2 . We use these measures to identify a set of hyperparameters and architectural modifications that significantly improves the performance of Hamiltonian Generative Networks ( HGN ) [ 49 ] , an existing state of the art model for recovering Hamiltonian dynamics from pixel observations , both in terms of long time-scale predictions , and interpretability of the learnt latent space . 2 Preliminaries . Hamiltonian dynamics The Hamiltonian formalism is a mathematical formulation of Newton ’ s equations of motion for describing energy-conservative dynamics ( see [ 25 , 29 ] ) . It describes the continuous time evolution of a system in an abstract phase space with state s= ( q , p ) ∈R2n , where q∈Rn is a vector of positions , and p∈Rn is the corresponding vector of momenta . The time evolution of the system in phase space is then given by the Hamiltonian equations of motion : q̇= ∂H ∂p , ṗ=−∂H ∂q ( 1 ) where the HamiltonianH : R2n→R maps the state s= ( q , p ) to a scalar representing the energy of the system . Based on the above equations of motion , the energy conservation property is easy to see , since dH dt = ∂H ∂q q̇+ ∂H ∂p ṗ=0 . Hamiltonian dynamics as a symplectic map The time evolution of Hamiltonian systems is described using a symplectic map . To define a symplectic map , we first need to define a symplectic matrix . A matrixM is called symplectic ifMTAM =A , whereA= [ 0 I −I 0 ] , and I∈Rn is the identity matrix . Then , a differentiable mapF ( s ) : R2n→R2n is called symplectic if its Jacobian matrix J= ∂F∂s is symplectic everywhere , i.e . JTAJ =A . Indeed , that Hamiltonian dynamics are governed by a symplectic map dsdt =A ∂H ∂s ( see Section 16.3 in [ 18 ] for the derivation ) . Canonical transformations Physical systems can be equivalently described in any number of arbitrary coordinate systems . For example , the dynamics of a simple mass-spring system can be described equally well using either canonical Cartesian phase space coordinates or the action-angle coordinates ( see Fig . 1 ) . A canonical transformation is a mapping that moves between equivalent state spaces in a way that preserves the form of the Hamiltonian dynamics . For example , there is a canonical transformation between the Cartesian phase space coordinates or the action-angle coordinates shown in Fig . 1 , since both preserve the correct underlying Hamiltonian dynamics . More formally , let us consider a Hamiltonian system with a phase-space state ( q , p ) , a Hamiltonian H ( q , p ) and equations of motion given by Eq . 1 . We can transform the state of this system according to new variables2 Q = Q ( q , p ) and P = P ( q , p ) . In general , such a transformation can not be expected to preserve the equations of motion in the new variables ( Q , P ) . However , if there exists a new Hamiltonian Ĥ ( Q , P ) that describes the same dynamics as the ones described by the original HamiltonianH ( q , p ) , where Q̇= ∂Ĥ∂P and Ṗ =− ∂Ĥ ∂Q , then such a transformation is called canonical . Canonical transformations are also symplectic maps [ 2 ] . We provide an illustration of a canonical transformation in Fig . 1 and refer to Stewart [ 46 ] for more details . 3 Measuring the quality of learnt Hamiltonian dynamics . To measure whether a model has learnt the underlying Hamiltonian dynamics faithfully , we propose measuring whether there exists a canonical transformation between the ground truth phase space and a subspace or a submanifold within the latent space discovered by the model . Note that a direct comparison between the inferred state space and the ground truth phase space ( e.g . L2 distance ) is impossible because physical systems can be equivalently described in any number of arbitrary coordinate systems and there is no reason why the models should converge to the phase space arbitrarily chosen as the “ ground truth ” when learning from pixels . As discussed in Sec . 2 , symplectic maps ( or canonical transformations ) are typically defined over two spaces of the same dimension . In our case , however , the latent space of the model is most often overparametrised to have more dimensions than the ground truth state . Saying this , if the model has learnt to capture the underlying Hamiltonian dynamics , the dynamics in its latent space should “ mimic ” the dynamics in the ground truth state space , in the following sense : given a latent space trajectory S ( t ) ∈ R2m and a ground state trajectory s ( t ) ∈ R2n where m ≥ n there should exist a map F : R2m → R2n such that the energy landscape defined by the learnt Hamiltonian Hm ( S ) is a scaled version of the energy landscape defined by the original Hamiltonian Hn ( F ( S ) ) , i.e . Em =Hm ( S ) = cHn ( F ( S ) ) = cEn ( where c is a constant , and E is the energy of the system , see Sec . A.5 in the Appendix for more motivation for the constant ) . One can think of the map F as the pseudo-inverse of a symplectic map between the space R2n of the ground truth state and another space R2n which is embedded in the higher dimensional space R2m of the model latent space . This pseudo-inverse preserves the Hamiltonian dynamics of the ground truth state space in the model latent space . If we expand the time derivative of the projection of the learnt state space back to the ground truth state space ŝ ( t ) =F ( S ( t ) ) , we see that ∂ŝ ( t ) ∂t = ∂F ∂S ∂S ( t ) ∂t = ∂F ∂S Am ∂Hm ∂S =c ∂F ∂S Am ∂F ∂S T ∂Hn ∂s =c ( JAmJ T ) ∂Hn ∂s . ( 2 ) Hence , if we can find a mapping F where cJAmJT =An holds , then it suggests the dynamics in the latent space do in fact “ mimic ” the dynamics in the ground truth state space , and hence the model has recovered the underlying Hamiltonian dynamics of the system . 2These can also depend on time but we are not considering time-dependent transformations here . Hence , we will only consider restricted canonical transformations in this paper . To measure how well the model has learnt to mimic the original dynamics , we first learn the mapping F : S 7→ s that explains as much variance in the original state space s as possible . We then check whether this mapping is symplectic by measuring how much JAmJT deviates fromAn up to a constant scaling factor , where J= ∂F∂S is the Jacobian of the learnt mapping , andAk= [ 0 I −I 0 ] is a 2k×2k anti-symmetric matrix . We next discuss each of the two steps in more detail . Learning to explain the ground truth dynamics We first learn the mapping F between the model state space S and the ground truth phase space s. A naive implementation would use an MLP to approximate F , however one has to be very careful not to overfit to the training data . In the low data regimes we found that even small MLPs trained with L1 regularization often overfit by learning an unnecessarily complicated map that ends up not being symplectic even if a symplectic map exists ( see Sec . A.3 in the Appendix ) . Hence , we make sure to use at least 1000x more datapoints for training the MLP than the number of its parameters . Given that learning the map F relies on the availability of ground truth phase space data , which may be hard to obtain in certain domains , we also propose a more data efficient way of learning F that only requires labeling 100 training trajectories . We do so by going back to the basics of machine learning and avoiding overfitting by controlling the expressivity of the mapping F by training a Lasso regularised linear regression over a polynomial expansion over the learnt state space features S , similarly to DiPietro et al [ 14 ] . We start with the expansion order of 1 and progressively increase it up to an upper bound determined by the hyperparameter κ , until enough variance in s is explained . In our experiments we use κ=5 ( see Alg . 1 in Appendix for more details ) . To limit the exponential computational cost of computing higher order polynomial expansions of larger latent spaces , where possible we first filter out “ uninformative ” state dimensions defined as those dimensions that have average KL from the prior less than the thresholdKL [ qφ ( Si ) |p ( Si ) ] < 0.01 as in Duan et al [ 16 ] , however this step is not necessary , and we have successfully applied our approach without this step in practice to a 1024-dimensional latent space . The MLP and polynomial regression ( PR ) instantiations of learning the mapping F have different trade-offs . As discussed above , MLP requires orders of magnitude more ground truth state data ( in our experiments at least 640x ) than PR , however because it is exposed to more data , it can be more accurate than PR , especially in terms of estimating how much variance in the ground truth dynamics can be explained well by the latent dynamics . The MLP can be arbitrarily expressive , while PR calculations get exponentially more expensive as the expressivity of F is increased . PR , however , has the benefit of interpretability . Having learnt F using either the MLP or PR method , we calculate its goodness of fit , R2 , which measures how much information about the ground truth phase space dynamics can be recovered from the learnt state space according to R2 = 1− ∑ ( F ( S ) −s ) 2∑ ( s−s ) 2 , where s is the mean of the ground truth state variables , F ( S ) are the ground truth state variables predicted from the model state variables S , and the summation is over the available data points [ 13 ] . R2 serves as a useful estimate of how much information about the ground truth dynamics is preserved in the latent trajectories . We use it instead of estimation error based alternatives ( e.g . MSE , which worked equally well in practice ) , because it is not sensitive to the magnitude of the data , and hence allows for a more direct comparison between scores obtained from different models trained on different datasets . Measuring symplecticity of the mapping We check the symplecticity of the learnt mapping F by calculating Ân = JAmJT . As the next step we could measure directly whether this product is equal to the anti-symmetric matrix An , however depending on the model it is possible that the product will instead be equal toATn ( we are going to drop the subscript for easier readability where the dimensionality can be inferred from context ) , which may happen if the Hamiltonian learnt by the model is related to the original Hamiltonian through a negative constant c in Eq . 2 ( also see Sec . A.3 in the Appendix ) . Hence , we instead calculate ÂÂT =JAJTJATJT and measure how far it deviates from the identity matrix : Sym=MSE ( cÂÂT , I ) ( 3 ) where c is a normalisation factor as per Sec . 3 ( see Sec . A.5 in the Appendix for more motivation ) , and I ∈R2n . Note that c will be different between different models , since it depends on the learnt Hamiltonian . We estimate it by finding a constant that minimises Sym for each model separately . Since for symplectic maps the right-hand side in Eq . 3 has to be zero everywhere , we need to evaluate it across all avaliable data . In practice we sample t points across k trajectories to evaluate the Jacobian and average Sym score over . The normalisation factor is calculated as c= 1/mean ( max ( |ÂÂT | ) ) , where the absolute maximum value of ÂÂT is taken at every evaluation point , and their average is taken across all the points from a single trajectory . We calculate a different constant for each trajectory to ensure that on datasets where different trajectories are sampled from different Hamiltonians ( e.g . mass-spring +c or pendulum +c ) , Sym is still meaningful . This does not affect Sym for datasets generated from a single Hamiltonian , where in practice the constant is the same across all trajectories .

Summary: The paper suggests novel measures for evaluating latent dynamics model by measuring how close the latent dynamics is to being governed by a symplectic map, a property of Hamiltonian dynamics. The paper analyses variants of the original Hamiltonian Generative Network [48] using the model and identifies better variants. The measure is also compared to related measures such as the MSE image reconstruction error, or the valid prediction time (VPT).