Data-driven discovery of non-Newtonian astronomy
via learning non-Euclidean Hamiltonian

Oswin So∗
Massachusetts Institute of Technology, USA
oswinso@mit.edu

Gongjie Li
Georgia Institute of Technology, USA
gongjie.li@physics.gatech.edu

Evangelos A. Theodorou
Georgia Institute of Technology, USAevangelos.theodorou@gatech.edu

Molei Tao
Georgia Institute of Technology, USA
mtao@gatech.edu

Abstract
Incorporating the Hamiltonian structure of physical dynamics into deep learningmodels provides a powerful way to improve the interpretability and prediction
accuracy. While previous works are mostly limited to the Euclidean spaces, theirextension to the Lie group manifold is needed when rotations form a key component of the dynamics, such as the higher-order physics beyond simple point-mass
dynamics for N -body celestial interactions. Moreover, the multiscale nature ofthese processes presents a challenge to existing methods as a long time horizon isrequired. By leveraging a symplectic Lie-group manifold preserving integrator, we
present a method for data-driven discovery of non-Newtonian astronomy. Preliminary results show the importance of both these properties in training stability andprediction accuracy.

1

Introduction

Deep Neural Networks (DNN) have been demonstrated to beeffective tools for learning dynamical systems from data. One
important class of systems to be learned have dynamics describedby physical laws, whose structure can be exploited by learning the
Point
Hamiltonian of the system instead of the vector field [1, 2]. An
Rigidappropriately learned Hamiltonian can endow the learned system
Lie T2
with properties such as superior long prediction accuracy [3] andapplicability to chaotic systems [3–5].
Figure 1: One planet’s orbit
To learn continuous dynamics from discrete data, one important around a star: rigid body correcstep is to bridge the continuous and discrete times. Seminal tion results in a precession, i.e.
work initially approximated the time derivative via finite differ- slow rotation of the orbital axis.
ences and then matched it with a learned (Hamiltonian) vector Our method, ‘Lie T2’, learns V
field[1, 2]. Recent efforts avoid the inaccuracy of finite difference from data and predicts a trajecby numerically integrating the learned vector field. Especially tory that matches the ground truth
relevant here is SRNN [6], which uses a symplectic integrator with the rigid body potential.
to ensure the learned dynamics is symplectic (a necessity for
Hamiltonian systems). Although SRNN only demonstrated learning separable Hamiltonians, breakthrough in symplectic integration of arbitrary Hamiltonians [7] was used to extend SRNN [8]. Further
efforts on improving the time integration error have also been made [9–11]. Meanwhile, alternative
approaches based on learning a symplectic map instead of the Hamiltonian also demonstrated efficacy[3, 12], although these approaches have not been extended to non-Euclidean problems.
∗

Work was done while Oswin was at Georgia Tech.

Machine Learning and the Physical Sciences workshop, NeurIPS 2022.

In fact, one relatively under-explored area is learning Hamiltonian dynamics on manifolds like the Liegroup manifold family2 . One important member of this family is SO(n), which describes isometries
in Rn and is important for, e.g., dynamical astronomy. The evolution of celestial bodies correspondto a mechanical system, and the 2- and 3-body problems have been a staple problem in works onlearning Hamiltonian (e.g., [1, 3, 6, 15]); however, the Newtonian (point-mass) gravity considered isalready well understood. Practical problems in planetary dynamics are complicated by higher-order
physics such as planet spin-orbit interaction, tidal dissipation, and general relativistic correction.
While it is unclear what would be a perfect scientific model for these effects, planetary rotation is anecessary component to account for spin-orbit interaction and tidal forcing, creating an SO(3)⊗Ncomponent of the configuration space. To learn these physics from data, we need to learn on the Liegroup.
Rigid body dynamics also play important roles in other applications such as robotics. In a seminal
work [16], Hamiltonian dynamics on SO(3) are used to learn rigid body dynamics for a quadrotor.
In that work, Runge-Kutta 4 integrator is used. Consequently, the method is applicable to short
time-horizon (see Sec.3 and last paragraph of Sec.2).
For our problem of learning non-Newtonian astronomy, the time-horizon has to be long. Hence,
we use a different approach by leveraging a Lie-group preserving symplectic integrator. Structurepreserving integration of dynamical systems on manifolds has been extensively studied in literature,for example for Lie groups [17–21] and more broadly, geometric integration [22–25].
In summary, we propose a deep learning methodology for performing data-driven discovery of
non-Newtonian astronomy. By leveraging the use of a symplectic Lie-group manifold preserving
integrator, we show how a non-Euclidean Hamiltonian can be learned for accurate prediction of
non-Newtonian effects. Moreover, we provide insights that show the importance of both symplecticity
and exact preservation of the Lie-group manifold in training stability.

2

MethodGiven observations of a dynamically evolving system, our goal is to learn the physics that governs its
evolution from the data. Denote by (qk,l , Rk,l , pk,l , Πk,l )K
k=1 , l ∈ [L] a dataset of snapshots of L
continuous-time trajectories of a system with N interacting rigid bodies. That is,(qk,l , Rk,l , pk,l , Πk,l ) = ql (k∆t), Rl (k∆t), pl (k∆t), Πl (k∆t) ,where ql (t), Rl (t), pl (t), Πl (t) is a solution of some latent Hamiltonian ODE to be learned corre⊗N
sponding to mechanical dynamics on T ∗ SE(3) . ∆t is a (possibly large) observation timestep,⊗N
R ∈ SO(3)is the rotational configuration of the N rigid bodies, and Π ∈ so(3)⊗N denotes
each’s angular momentum in their respective body frames.
⊗N

Importantly, since the configuration space Q = SE(3)is not flat, the mechanical dynamics are not∂Hgiven by q̇ = ∂H
∂ p , ṗ = − ∂ q for some Hamiltonian H that depends on the generalized coordinatesq ∈ Q and generalized momentum p ∈ T q∗ Q. Instead, the equations of motion can be derived via
either Lagrange multipliers [20, 26] or a Lie group variational principle [20, 27], which will beq̇i = pi /mi ,(1a)∂Vṗi = −
+ Fpi (1b)∂pi

−1Ṙi = Ri J\
i Πi

(1c)




⊺

∨

∂V
⊺ ∂V
Π̇i = Πi × J−1
−
Ri
+ FΠi
(1d)
i Πi − Ri
∂Ri
∂Ri
PN
P
assuming a physical Hamiltonian H(q, R, p, Π) = i=1 12 pTi pi /mi + i 21 ΠTi J−1
i Πi + V (q, R)
that sums total (translation and rotational) kinetic energy and interaction potential V , where
mi , Ji denote the mass and inertial tensor of the ith body, and Fp , FΠ are forcing terms to model
nonconervative forces. Πi ∈ R3 is a vector, ∧ is the map from R3 to Skew3 and ∨ is its inverse ([20]for more details). By learning the potential V , external forcing Fp and torque FΠ , we can learn thephysics of the system.
2

We note extensions to include holonomic constraints in [13] and to handle contact in [14]).

2

2.1

Machine Learning Challenges Posed by Dynamical Astronomy

We study this setup because it helps answer scientific questions like: what physics governs the
motions of celestial bodies, such as planets in a planetary system? The leading order physics isof course already well known, namely these bodies can be approximated by point masses that areinteracting through a 1/r gravitational potential. However, planets are not point masses, and theirrotations matter because they shape planetary climates [28, 29] and even feedback to their orbits[30]. This already starts to alter V even if one only considers classical gravity. P
For example, thegravitational potential V for interacting bodies of finite sizes should be V (q, R) = i<j Vi,j , where!
Z Z
Gρ(xi )ρ(xj )Gmi mj
1
Vi,j (q, R) =−dxi dxj = −+O.
2∥qi + Ri xi − qj − Rj xj ∥∥qi − qj ∥∥qi − qj ∥
Bi Bj
|
{z
} |
{z
}
Vi,j,point

Vi,j,resid(2)
Working with the full potential is complicated since Bi is not known and the integral is not analyticallyknown. Can we directly learn Vresid from time-series data?
Classical gravity (i.e. Newtonian physics) is not the only driver of planetary motion — tidal forcesand general relativity (GR) matter too. The former provides a dissipation mechanism and plays
critical roles in altering planetary orbits [31, 32]; the latter doesn’t need much explanation andhas been demonstrated by, e.g., Mercury’s precessing orbit [33]. Tidal forces depend on celestialbodies’ rotations [34] and thus is a function of both q, R. GR’s effects cannot be fully characterizedwith classical coordinates q, R, p, Π, but post-Newtonian approximations based purely on thesecoordinates are popular [35]. Can we learn both purely from data if we did not have theories for
either?
In addition to the scientific questions, there are also significant machine learning challenges:
Multiscale dynamics. Rigid-body correction (Vresid ), tidal force, and GR correction are all muchsmaller forces compared to point-mass gravity. Consequently, their effects do not manifest until longtime. Thus, one challenge for learning them is that the dynamical system exhibits different behaviors
over multiple timescales. It is reasonable to require long time series data for the small effects tobe learned; meanwhile, when observations are expensive to make, the observation time step ∆t canbe much longer than the smallest timescales. Can we still learn the physics in this case? We willleverage symplectic integrator and its mild growth of error over long time [7, 26] to provide a positive
answer.
Respecting the Lie group manifold. However, even having a symplectic integrator is not enough⊗Nbecause the position variable of the latent dynamics (i.e. truth) stays on SE(3) . If the integrated⊺
solution falls off this manifold such that R R = I no longer holds, it is not only incorrect but likely
misleading for the learning of V (q, R). Popular integrators such as forward Euler, Runge-Kutta 4(RK4) and Leapfrog [1, 6, 36] unfortunately do no maintain the manifold structure.
2.2

Learning with Lie Symplectic RNNs

Our method can be viewed as a Lie-group generalization of the seminal work of SRNN [6], where agood integrator that is both symplectic and Lie-group preserving is employed as a recurrent block.
Lie T2 : A Symplectic Lie-Group Preserving Integrator. To construct an integrator that achieves
both properties, we borrow from [20] the idea of Lie-group and symplecticity preserving splitting, andsplit our Hamiltonian as H = HKE +HPE +Hasym , which contains the axial-symmetric kinetic energy,
potential energy and asymmetric kinetic energy correction terms. This enables computing the exact
[KE]
[PE]
[asym]
integrators ϕt , ϕt and ϕt(see App B for details). We then construct a 2nd-order symplectic
integrator Lie T2 by applying the Strang composition scheme. To account for non-conservative
forces, the corresponding non-conservative momentum update ϕ[force] : (p, Π) ← F (q, R, p, Π) is[Lie T ]inserted in the middle of the composition [20]. This gives ϕh 2 for stepsize h as
[KE]

[PE]

[asym]

T2:= ϕh/2 ◦ ϕh/2 ◦ ϕh/2
ϕLieh[force]◦ ϕh[asym]◦ ϕh/2[PE][KE]◦ ϕh/2 ◦ ϕh/2

(3)

While the implementation of our integrator in the case of dynamical astronomy is specific to the⊗NSE(3)manifold, we note that the use of a splitting method to construct Lie-group preserving3..
.
(q1 , R1 , p1 , Π1 )(q0 , R0 , p0 , Π0 )T2
ϕθ,LiehT2ϕθ,Lie
h

...
...

T2
ϕθ,Lieh(q̂2 , R̂2 , p̂2 , Π̂2 )T2ϕθ,Lieh(q̂1 , R̂1 , p̂1 , Π̂1 )L(θ)

H layer SRNN

Figure 2: Inputs are fed through a recurrent layer with Lie T2 . Prediction error is used as a loss on θ.
integrators can be applied to construct symplectic integrators for any Hamiltonian system defined ona Lie group [26].
A Recurrent Architecture for Nonlinear Regression. Given the simplicity of Vpoint , we assume
θthis is known and learn Vresid and F θ with multi-layer perceptron (MLP) Vresid
and F θ without
assuming any pairwise structure (see App C for discussion). We then use ϕθ,Lie T2 to integrate
dynamics forward, where θ denotes the dependence on the networks. However, when the temporal
T2
spacing between observations ∆t is large, using a single ϕθ,Liewill result in large errors for the∆tθ,Lie T2
T2
fast timescale dynamics. Instead, we compose ϕh
H times as ( q̂k+1,l , p̂k+1,l ) = ϕθ,Lie◦hθ,Lie T2
· · · ◦ ϕh
( qk,l , pk,l ), where H = h/∆t ∈ Z determines the integration stepsize h, and we use
q = (q, R) and p = (p, Π) to denote the position and momentum related terms respectively. We
perform training by minimizing the following empirical loss over random minibatches of size Nb
Nb X
K n
o
1 X
2
2
θ
θ
p
−
p̂
L(θ) :=
qk,l − q̂k,l
+
(4)
k,l
k,l 2
2
Nb K
l=1 k=1

Note that we do not assume access to the true derivatives q̇k,j and ṗk,j used in the loss function ofsome works [1, 37, 38]. The forward pass of the training process is summarized in Fig. 2 (see App C
for details).

Results

We aim to answer two questions. Q1 Can we learn
multiscale physics? Q2 How important are symplecticity (S) and Lie-group preservation (L) for learning? The closest baseline for our problem is work of[16], which learns short timescale rigid-body Hamiltonian dynamics for robotics. Placed in our framework, their work corresponds to using RK4 for the recurrent block, which is neither S nor L. Therefore, toinvestigate Q2, we vary the choice of integrator in ourframework as follows: Normal: Explicit Euler, RK4.
S: Verlet. L: Lie RK2(CF2) and Lie RK4(CF4) [21].
We leave the precise details to Apps C and D.

10−2
10−6

10−10
0

10
H error3

kRT R − Ik2

Benefit. Learning an accurate V θ , F θ requires accurate numerical simulation which also leads toa trainable model. Without preservation of the manifold structure, training can lead to ‘shortcuts’
outside the manifold that seemingly match the data but completely mislead the learning. Symplecticity
also plays a vital role in controlling the long time integration error — under reasonable conditions, apth-order symplectic integrator has linear O(∆thp ) error bound, whereas a pth-order nonsymplectic
one has an exponential O(eC∆t hp ) error bound [7, 26]. While these bounds do not matter for small∆t, they are significant for multiscale problems where ∆t is macroscopic but h is microscopic.
Consequently, improving error estimates for a nonsymplectic integrator by reducing h makes the
RNN exponentially deep — this often renders training difficult [39] and is not desirable.

100

200
300
Integrator steps

400

500

−1

10−4
Euler
RK4

10−7
0

20

VerletLie RK24060
Integrator steps

Lie RK4
Lie T280

100

Figure 3: Results on TRAPPIST-1 with short
Toy Two-Body Problem. We consider an illustrative
data separation. Top: SO(3) manifold ertwo-body problem to demonstrate the effects of Vrigid .
ror. Bottom: Hamiltonian error over the inteIn Fig. 1 ‘Point’ & ‘Rigid’ denote exact solutions for
grated trajectory. Only Lie T2 achieves low
a point-mass and rigid-body potential, and ‘Lie T2’
errors in both metrics.
the prediction of our method based on a V learned
from data. Compared to ‘Point’, ‘Rigid’ induces anapsidal precession (rotation of the orbital axis) due to spin-orbit couplings. Our method successfully
predicts this interaction and matches the trajectory of ‘Rigid’.
4

Table 1: TRAPPIST-1 with long data separation
evaluated after 500 integrator steps. All methods
except for Lie T2 diverge during training. ∥∆R∥
is computed as the geodesic between the predicted Table 2: TRAPPIST-1 with short data separation evaluatedafter 500 integrator steps.
and true rotations (see App E for details).
∥∆q∥2 ∥∆R∥ ∥∆ ∂V∥∂q∥∆q∥2 ∥∆R∥ ∥∆ṗ∥2 ∥∆Π̇∥ 2
No Corrections 8.6e−6 1.6
3.3
Euler, RK4,
Verlet, Lie RK2, ∞∞∞Lie RK4
Lie T2 (Ours) 6.1e−8 4.0e−2 1.2e−1

4.7e−2∞3.1e−3Euler∞
RK4
1.6e−5
S Verlet
8.6e−8
LLie RK2
5.3e−5
LLie RK4
1.6e−5SLLie T2 (Ours) 8.0e−8∞
3.9e−12.0e+7
3.3e−13.5e−12.4e−1

2

∂V
∥∆ ∂R
∥2

∞
1.5e+3
1.6e+12.6e+2
1.5e+3
1.6e+1∞7.9e−14.1e−17.6e−18.4e−14.0e−1

We next test our method by learning the dynamics of the TRAPPIST-1 system [40] which consists ofseven earth-sized planets and is notable for potential habitability for terrestrial forms of lives.
TRAPPIST-1, Large ∆t. To answer Q1, we choose a large data timestep ∆t = 2.4e−3 yr. The
closest planet has an orbital period of ∼ 2∆t (4.1e−3 yr), while the rigid body correction, tidal
force and GR correction act on much longer scales. Only Lie T2 successfully trains. All othermethods diverge during training (denoted by ∞) despite attempts at stabilization with techniques suchas normalization (LayerNorm [41], GroupNorm [42]). Reducing h improves integration accuracy,
but increases the RNN depth and makes training more unstable. We compare with the solution
for point-mass potential only (No Correction). Our method reduces the error up to two orders ofmagnitude in measures of trajectory error and potential gradients (Table 1). We evaluate each methodby computing the norm of the error of the predicted positions and forces. The geodesic on SO(3) isused to compute errors between rotations ∥∆R∥. See App E for more details on how each metric isdefined and computed.
TRAPPIST-1, Small ∆t. To gain more insight on Q2, we shrink ∆t until almost all methods canconverge and only consider conservative forces (i.e. no tidal force or GR). The mean errors in thepredicted trajectory and derivatives of the learned potential V after 500 integrator steps are shown in
Table 2. Both S methods achieve small errors in position related terms. Verlet has a large rotational
error since it does not integrate on the rotation manifold. L methods achieve lower rotational errors
but are worse elsewhere. Lie T2 being both SL achieves the lowest error on both fronts.

4

Broader Impact

Our work presents an approach for learning multiscale, higher order physics on the Lie-groupmanifold in the context of non-Newtonian astronomy. This research, though directly applicable toastronomy, can also be applied to perform data-driven discovery multiscale phenomena on Lie groups
in other fields.

References[1] Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advancesin neural information processing systems, 32, 2019.
[2] Tom Bertalan, Felix Dietrich, Igor Mezić, and Ioannis G Kevrekidis. On learning hamiltoniansystems from data. Chaos: An Interdisciplinary Journal of Nonlinear Science, 29(12):121107,
2019.
[3] Renyi Chen and Molei Tao. Data-driven prediction of general hamiltonian dynamics via
learning exactly-symplectic maps. In Marina Meila and Tong Zhang, editors, Proceedings of the
38th International Conference on Machine Learning, volume 139 of Proceedings of Machine
Learning Research, pages 1717–1727. PMLR. URL https://proceedings.mlr.press/v139/chen21r.html.
[4] Anshul Choudhary, John F Lindner, Elliott G Holliday, Scott T Miller, Sudeshna Sinha, andWilliam L Ditto. Physics-enhanced neural networks learn order and chaos. Physical Review E,101(6):062207, 2020.
5[5] Chen-Di Han, Bryan Glaz, Mulugeta Haile, and Ying-Cheng Lai. Adaptable hamiltonian neuralnetworks. Physical Review Research, 3(2):023156, 2021.
[6] Zhengdao Chen, Jianyu Zhang, Martin Arjovsky, and Léon Bottou. Symplectic recurrent neuralnetworks. arXiv preprint arXiv:1909.13334, 2019.
[7] Molei Tao. Explicit symplectic approximation of nonseparable hamiltonians: Algorithm andlong time performance. Physical Review E, 94(4):043303, 2016.
[8] Shiying Xiong, Yunjin Tong, Xingzhe He, Shuqi Yang, Cheng Yang, and Bo Zhu. Nonseparable
symplectic neural networks. arXiv preprint arXiv:2010.12636, 2020.
[9] Daniel DiPietro, Shiying Xiong, and Bo Zhu. Sparse symplectically integrated neural networks.
Advances in Neural Information Processing Systems, 33:6074–6085, 2020.
[10] Marco David and Florian Méhats. Symplectic learning for hamiltonian neural networks. arXiv
preprint arXiv:2106.11753, 2021.
[11] Frederik Baymler Mathiesen, Bin Yang, and Jilin Hu. Hyperverlet: A symplectic hypersolver
for hamiltonian systems. In Proceedings of the AAAI Conference on Artificial Intelligence,
volume 36, pages 4575–4582, 2022.
[12] Pengzhan Jin, Zhen Zhang, Aiqing Zhu, Yifa Tang, and George Em Karniadakis. Sympnets:
Intrinsic structure-preserving symplectic networks for identifying hamiltonian systems. Neural
Networks, 132:166–179, 2020.
[13] Marc Finzi, Ke Alexander Wang, and Andrew G Wilson. Simplifying hamiltonian and lagrangianneural networks via explicit constraints. Advances in neural information processing systems,
33:13880–13889, 2020.
[14] Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Extending lagrangian andhamiltonian neural networks with differentiable contact models. Advances in Neural Information
Processing Systems, 34:21910–21922, 2021.
[15] Yupu Lu, Shijie Lin, Guanqi Chen, and Jia Pan. Modlanets: Learning generalisable dynamics
via modularity and physical inductive bias. In International Conference on Machine Learning,
pages 14384–14397. PMLR, 2022.
[16] Thai Duong and Nikolay Atanasov. Hamiltonian-based neural ode networks on the se (3)manifold for dynamics learning and control. arXiv preprint arXiv:2106.12782, 2021.
[17] Arieh Iserles, Hans Z Munthe-Kaas, Syvert P Nørsett, and Antonella Zanna. Lie-group methods.
Acta numerica, 9:215–365, 2000.
[18] Nawaf Bou-Rabee and Jerrold E Marsden. Hamilton–pontryagin integrators on lie groups part i:
Introduction and structure-preserving properties. Foundations of computational mathematics, 9(2):197–219, 2009.
[19] Elena Celledoni, Håkon Marthinsen, and Brynjulf Owren. An introduction to lie group
integrators–basics, new developments and applications. Journal of Computational Physics, 257:
1040–1061, 2014.
[20] Renyi Chen, Gongjie Li, and Molei Tao. Grit: A package for structure-preserving simulationsof gravitationally interacting rigid bodies. The Astrophysical Journal, 919(1):50, 2021.
[21] Elena Celledoni, Ergys Çokaj, Andrea Leone, Davide Murari, and Brynjulf Owren. Lie group
integrators for mechanical systems. International Journal of Computer Mathematics, 99(1):58–88, 2022.
[22] Ernst Haier, Christian Lubich, and Gerhard Wanner. Geometric Numerical integration: structurepreserving algorithms for ordinary differential equations. Springer, 2006.
[23] Benedict Leimkuhler and Sebastian Reich. Simulating hamiltonian dynamics. Number 14.
Cambridge university press, 2004.
6

[24] Sergio Blanes and Fernando Casas. A concise introduction to geometric numerical integration.
CRC press, 2017.
[25] Jesus-Maria Sanz-Serna and Mari-Paz Calvo. Numerical hamiltonian problems. Courier Dover
Publications, 2018.
[26] E. Hairer, Christian Lubich, and Gerhard Wanner. Geometric Numerical Integration: StructurePreserving Algorithms for Ordinary Differential Equations. Number 31 in Springer Series in
Computational Mathematics. Springer, 2nd ed edition.
[27] Taeyoung Lee, N Harris McClamroch, and Melvin Leok. A lie group variational integrator forthe attitude dynamics of a rigid body with applications to the 3D pendulum. In Proceedings of2005 IEEE Conference on Control Applications, 2005. CCA 2005., pages 962–967. IEEE. doi:
10.1109/CCA.2005.1507254. URL http://ieeexplore.ieee.org/document/1507254/.
[28] Billy Quarles, Gongjie Li, and Jack J Lissauer. Milankovitch cycles for a circumstellar earthanalogue within α centauri-like binaries. Monthly Notices of the Royal Astronomical Society,509(2):2736–2757, 2022.
[29] Renyi Chen, Gongjie Li, and Molei Tao. Low spin-axis variations of circumbinary planets.
Monthly Notices of the Royal Astronomical Society, 515(4):5175–5184, 2022.
[30] Sarah Millholland and Gregory Laughlin. Obliquity-driven sculpting of exoplanetary systems.
Nature Astronomy, 3:424–433, March 2019. doi: 10.1038/s41550-019-0701-7.
[31] Rosemary A. Mardling and D. N. C. Lin. Calculating the Tidal, Spin, and Dynamical Evolutionof Extrasolar Planetary Systems. 573(2):829–844. ISSN 0004-637X, 1538-4357. doi: 10.1086/
340752. URL https://iopscience.iop.org/article/10.1086/340752.
[32] Smadar Naoz, Will M Farr, Yoram Lithwick, Frederic A Rasio, and Jean Teyssandier. Hot
jupiters from secular planet–planet interactions. Nature, 473(7346):187–189, 2011.
[33] Gerald Maurice Clemence. The relativity effect in planetary motions. Reviews of Modern
Physics, 19(4):361, 1947.
[34] P Hut. Tidal evolution in close binary systems. Astronomy and Astrophysics, 99:126–140, 1981.
[35] Luc Blanchet. Gravitational radiation from post-newtonian sources and inspiralling compact
binaries. Living reviews in relativity, 17(1):1–187, 2014.
[36] Nate Gruver, Marc Anton Finzi, Samuel Don Stanton, and Andrew Gordon Wilson. Deconstructing the inductive biases of hamiltonian neural networks. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=EDeVYpT42oS.
[37] Sam Greydanus and Andrew Sosanya. Dissipative hamiltonian neural networks: Learning
dissipative and conservative dynamics separately. arXiv preprint arXiv:2201.10085, 2022.
[38] Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley
Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.04630, 2020.
[39] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent
neural networks. In International conference on machine learning, pages 1310–1318. PMLR,
2013.
[40] Michaël Gillon, Emmanuël Jehin, Susan M Lederer, Laetitia Delrez, Julien de Wit, Artem
Burdanov, Valérie Van Grootel, Adam J Burgasser, Amaury HMJ Triaud, Cyrielle Opitom,
et al. Temperate earth-sized planets transiting a nearby ultracool dwarf star. Nature, 533(7602):
221–224, 2016.
[41] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint
arXiv:1607.06450, 2016.
[42] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conferenceon computer vision (ECCV), pages 3–19, 2018.
7

[43] Jerrold E Marsden and Tudor S Ratiu. Introduction to mechanics and symmetry: a basic
exposition of classical mechanical systems, volume 17. Springer Science & Business Media,
2013.
[44] Darryl D Holm, Tanya Schmah, and Cristina Stoica. Geometric mechanics and symmetry: fromfinite to infinite dimensions, volume 12. Oxford University Press, 2009.
[45] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal
Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, andQiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://github.com/google/jax.
[46] Tom Hennigan, Trevor Cai, Tamara Norman, and Igor Babuschkin. Haiku: Sonnet for JAX,
2020. URL http://github.com/deepmind/dm-haiku.
[47] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network
function approximation in reinforcement learning. Neural Networks, 107:3–11, 2018.
[48] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh.
Set transformer: A framework for attention-based permutation-invariant neural networks. In
International conference on machine learning, pages 3744–3753. PMLR, 2019.
[49] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
arXiv:1711.05101, 2017.
[50] Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ODE-Net: Learning Hamiltonian Dynamics with Control. URL http://arxiv.org/abs/1909.12077.
[51] Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Benchmarking energyconserving neural networks for learning dynamics from data. In Learning for Dynamics andControl, pages 1218–1229. PMLR, 2021.
[52] Elena Celledoni, Arne Marthinsen, and Brynjulf Owren. Commutator-free Lie group methods.
19(3):341–352. ISSN 0167739X. doi: 10.1016/S0167-739X(02)00161-9. URL https://linkinghub.elsevier.com/retrieve/pii/S0167739X02001619.

8

Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on
how to answer these questions. For each question, change the default [TODO] to [Yes] , [No] , or
[N/A] . You are strongly encouraged to include a justification to your answer, either by referencing
the appropriate section of your paper or providing a brief inline description. For example:• Did you include the license to the code and datasets? [Yes] See Section ??.
• Did you include the license to the code and datasets? [No] The code and the data areproprietary.
• Did you include the license to the code and datasets? [N/A]Please do not modify the questions and only use the provided macros for your answers. Note that the
Checklist section does not count towards the page limit. In your paper, please delete this instructions
block and only keep the Checklist section heading above along with the questions/answers below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes] See App F. Also see App C
for a discussion on how the assumption of no prior knowledge on the structure of thepotential limits the scalability of our method in comparison to exploiting the pairwise
nature of the gravitational potential.
(c) Did you discuss any potential negative societal impacts of your work? [N/A](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section 2
for the assumptions of the underlying theoretical framework.
(b) Did you include complete proofs of all theoretical results? [Yes] See Section 2 andApps A and B for the derivation of the equations.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [No](b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] See App C.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes] See App C.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We, we cite all
existing code used in App C.
(b) Did you mention the license of the assets? [Yes] See footnotes in App C.
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A](d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A](e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A]
9

(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]10

A

Rigid Body Equations of Motion

We present two derivations from Chen et al. [20]. The first derivation derives a constrained Hamiltonian system via Lagrange multipliers, while the second uses a variational principle for mechanics on
Lie groups. We also note that the same equations can be derived via the use of the Port-Hamiltonian
framework as in Duong and Atanasov [16], though they choose to express the linear momentum
coordinates in the moving body frame rather than the inertial frame as is done in this work.
A.1

Constrained Hamiltonian System

Since the manifold SE(3)⊗N is a product space of the individual manifolds SE(3), we consider
the latter for brevity and drop indices i for each body. Furthermore, SE(3) can be seen as the
product space of R3 × SO(3) for which the former is unconstrained, so we focus our attention
on the latter. We can view R to be in the embedded Euclidean space R3×3 ←−- SO(3) and use
R ∈ SO(3) := h(R) = {R⊺ R − I3×3 = 03×3 } as a holonomic constraint. Using Lagrange
multipliers Λ [26] for the constraint h gives us the following Lagrangian
h
i
1
1
(A1)
L(R, Ṙ) = Tr ṘJd Ṙ − V (R) − Tr [Λ⊺ (R⊺ R − I3×3 )]2
2
where Jd denotes the nonstandard moment of inertia [20, 27]Z
Jd :=ρ(x)xx⊺ dx ,Jd = Tr[J]I3×3 − J,
J = Tr[Jd ]I3×3 − Jd ,
(A2)
and

B

#
λ1 λ4 λ6
(A3)
Λ = λ4 λ2 λ5 ∈ R3×3
λ6 λ5 λ3is a 6-dimensional symmetric matrix of Lagrange multipliers. Performing the Legendre transform for(A1) gives us the conjugate rotational momentum P as"

P=

∂L(R, Ṙ)
= ṘJd∂ Ṙ

(A4)

and corresponding Hamiltonian1
1
Tr PJd −1 P + V (R) + Tr [Λ⊺ (R⊺ R − I3×3 )](A5)
2
2
The constraint for P can be obtained by taking the time derivative of h(R) = 0 according to Haier⊺⊺1
et al. [22], i.e., J−1
d P R + R PJd = 03×3 . Hence, our equations of motion so far looks like

∂H

Ṙ =
= PJ−1d∂P(A6)Ṗ = − ∂H = − ∂V − RΛ∂R∂R
on the manifold
H(R, P) =

−1
⊺
⊺
J−1
d P R + R PJd = 03×3

M := {(R, P)|R⊺ R = I3×3 ,(A7)Let Ω̂ = R⊺ Ṙ denote the body’s angular velocity in the inertial frame, where ∧ denotes the
cross-product operation ûv = u × v. Then, it can written using P asΩ̂ = R⊺ PJ−1
d

(A8)



∂V (R)˙⊺Ω̂ = Jd−1 P⊺ PJ−1
+
R
−
−
RΛ
J−1
d
d
∂R

(A9)

Taking the time derivative of Ω̂ gives

Since our goal is to represent the dynamics of Π, we convert the above to the body frame momentum
using Π = JΩ. Using properties of the hat map [20, Appendix A.1] gives us⊺
ˆ = Tr[Jd ]Ω̂ − JdΠ̂ = JΩd Ω = Ω̂Jd − Jd Ω̂ .

11

(A10)

Taking the time derivative of the above then gives˙˙˙Π̂ = Ω̂Jd − Jd Ω̂,

(A11)


⊺ 


∂V (R)−1 ⊺−1⊺⊺ ∂V (R)R − (Λ − Λ⊺ ) ,+= Jd P P − P PJd − R
∂R
∂R

⊺ 


∂V (R)−1⊺⊺⊺ ∂V (R)
+
R ,
= J−1
P
P
−
P
PJ
−
R
d
d
∂R
∂R

(A12)
(A13)

where the symmetric Λ vanishes on the last line.
Since P = RΩ̂Jd , applying properties of the hat map [20, Appendix A.1] further simplifies theabove as  ∂V (R)  ∂V (R) ⊺ ˆ⊺⊺Π̇ = Ω̂ Ω̂Jd − Jd Ω̂ Ω̂ − R⊺
R ,
(A14)
+
∂R
∂R


⊺ 
∂V (R)∂V (R)= Ω\× Jd Ω − R⊺
+
R .
(A15)
∂R
∂R
Thus, applying the vee map and converting back to J and Π gives


⊺ ∨
∂V (R)⊺ ∂V (R)Π̇ = Ω × Jd Ω − R
+
R
,
∂R
∂R


⊺ ∨
∂V (R)∂V (R)= Ω × () Tr[J] − J)Ω − R⊺
+
R
,
∂R
∂R


⊺ ∨
∂V (R)⊺ ∂V (R),= −Ω × JΩ − R
+
R
∂R
∂R


⊺ ∨
∂V (R)∂V (R)= Π × J−1 Π − R⊺
+
R
.
∂R
∂R

(A16)
(A17)
(A18)
(A19)

Hence, the rotational equations of motion on SO(3) are
d1
Ṙ = RJ Π,



⊺ ∨
∂V (R)∂V (R)
Π̇ = Π × J−1 Π − R⊺
.
−
R
∂R
∂R
A.2

(A20)

Variational Principle for Mechanics on the Lie Group

We can also apply the Euler-Lagrange equations for Hamilton’s variational principle on a Lie group,
a topic that has been well studied (e.g., Marsden and Ratiu [43], Holm et al. [44]) [20], and wesummarize the results for the special case of rigid bodies from the expository part of Lee et al. [27].
Denote the infinitesimally varied rotation by Rϵ = R exp(ϵη̂), with ϵ ∈ R and η ∈ R3 , where exp(·)
is the exponential map from so(3) to SO(3). The varied angular velocity Ω̂ϵ isΩ̂ϵ = R⊺ϵ Ṙϵ = exp(−ϵη̂) Ṙ exp(ϵη̂) + R exp(ϵη̂)ϵη̂ ,(A21)ˆ= exp(−ϵη̂)Ω̂ exp(−ϵη̂) + exp(−ϵη̂)η̇,= Ω̂ + ϵ η̇ˆ + Ω̂η̂ − η̂ Ω̂ + O(ϵ2 ).

(A22)
(A23)

Consider the action
Z t1
S(Ω, R) =Z t1
L(Ω, R) dt =

t0

t0

12

i
1 h
Tr Ω̂Jd Ω̂ − V (R) dt .
2

(A24)

Taking the variation Sϵ of the action S, we haveSϵ (Ω, R) = S(Ωϵ , Rϵ ),Z t1 h
1Tr − η̇ˆ Jd Ω̂ + Ω̂Jd
2
t0





i
∂V (R)dt+ η̂ Ω̂ Jd Ω̂ + Ω̂Jd − η̂ Jd Ω̂ + Ω̂Jd Ω̂ + Tr η̂R⊺
∂R

= S(Ω, R) + ϵ(A25)+ O(ϵ2 ).
d
S = 0, i.e.,
Using Hamilton’s Principle, we have dϵ
ϵ=0 ϵ
 

Z
∂V (R)1 t1Ω̇ + Ω\
× JΩ + 2R⊺
Tr η̂ Jc
dt = 0,2 t0∂R(R)must be skew-symmetric, giving us
for any η ∈ R3 . Hence, JcΩ̇ + Ω\
× JΩ + 2R⊺ ∂V∂R⊺
∂V (R)c⊺ ∂V (R)\JΩ̇ = −Ω × JΩ +
R−R
.
∂R
∂R

Thus, by definition of Π and applying the vee map, we recover the same update∨⊺∂V (R)∂V (R)
Π̇ = Π × J−1 Π +R − R⊺
.
∂R
∂R

B

(A26)

(A27)

(A28)

Details of the Lie T2 splitting integrator

The full Hamiltonian H takes the form
H(q, R, p, Π) =

N
X
1

2
i=1

1
p⊺i pi /mi + Π⊺i J−1
i Πi + V (q, R)
2(B1)

We borrow from Chen et al. [20] the idea of Lie-group and symplecticity preserving splitting andsplit our Hamiltonian as H = HKE + HPE + Hasym , where
HKE :=

N
X
1

2
i=1

1
p⊺i pi /mi + Π⊺i J−1
i,sym Πi2

HPE := V (q, R)
Hasym :=

(B3)

N
X

1 ⊺ −1Π J
Πi
2 i i,asym
i=1

where we assume that Ji is axis-symmetric, i.e.,

(1)
J
 i
Ji :=  0
0

(B2)

0
(2)

Ji
0

(B4)


0

0 

(B5)

(3)
Ji

−1
−1
and J−1such that J−1i,sym and Ji,asym denote the axial-symmetric and residual terms of Ji
i,sym +
−1
−1Ji,asym = Ji , i.e.,




(1)
1/Ji
0
0
0
0
0


(1)
(2)
(1)

J−1
J−1
(B6)
0
1/Ji
0 ,
− 1/Ji0 .
i,sym := i,asym := 0 1/Ji
(3)
0
0
0
0
0
1/Ji

Then, each of HKE + HPE + Hasym can be integrated exactly.
13

Exact integration of HKE
V = 0)For HKE , we have the equations of motion (using (1a)-(1d) but withq̇i = pi /mi ṗi = 0

(B7a)
(B7b)

\

Ṙi = Ri J−1i,sym ΠiΠ̇i = Πi × J−1 Πi
i,sym

(B7c)
(B7d)

The equation for Π̇i (B7d) is the Euler equation for a free rigid body [20]. It is exactly solvable witha simple expression for axial-symmetric bodies, since in this case Π̇ simplifies as
 


(3)
(1)
1/Ji − 1/Ji
Πi,y Πi,z
 




(3)
(1)
Π̇i = Πi × J−1(B8)i,sym Πi = − 1/J
− 1/J
Π Π 
i

i

i,y

i,z

0

Consequently, Πi,z (t) = Πi,z (0), meaning that we can express the above as the linear differential
equation




(3)
(1)
0
1/Ji − 1/Ji
Πi,z (0) 0
 



(3)
(1)
Π̇i (t) = 
(B9)
Π (0)
0
0 Πi (0)
− 1/J − 1/J
i

i,z

i

0

0

0

"d#
0
= −θt 0 Πi (0)
1


(3)
(1)
where θ := 1/Ji − 1/Ji
Πi,z (0) and has the solution

(B10)



"d#
0
Πi (t) = exp −θt 0  Πi (0) = Rz⊺ (θt)Πi (0)
1

(B11)

where Rz denotes the rotation matrix around the z axis. Taking the above back to (B7c) then gives us[KE]
the solution for Ri as well, giving the flow ϕh of HKE asqi (h) = qi (0) + pi /mi h,pi (h) = pi (0),


(B12)
(1)

R
(h)
=
R
(0)R
∥Π
(0)∥h/JRz (θh),iii
Πi (0)
i
Πi (h) = Rz⊺ (θh)Πi (0).
Exact integration of HPE For HPE , the equations of motion are [20]∂Vq̇i = 0, ṗi = − ∂qi ,


⊺ ∨

∂V
∂V

Ṙi = 0, Π̇i = − R⊺i
−
Ri
.
∂Ri
∂Ri

(B13)

[PE]

Since qi and pi stay constant, pi and Πi change at constant rates. Hence, the flow ϕhgiven byqi (h) = qi (0)




∂V



pi (h) = pi (0) − ∂qi h,
Ri (h) = Ri (0),





⊺ ∨



∂V
∂V

Πi (h) = Πi (0) − R⊺i
−
Ri
h.
∂Ri
∂Ri
14

of HPE is(B14)

Exact integration of Hasym Finally, the equations of motion of Hasym are given by (again adapting(1a)-(1d))(
q̇i = 0, ṗi = 0,
(B15)
−1
\
Ṙi = Ri J−1
i,asym Πi , Π̇i = Πi × Ji,asym Πi .
[asym]

which can be solved to obtain the flow ϕhof Hasym asqi (h) = qi (0)pi (h) = pi (0),(2)

Ri (h) = Ry (δΠi h)Ri (0),




(2)
Πi (h) = Ry (−δΠi h)Πi (0).
(2)

where δ := 1/Ji

(B16)

(1)

− 1/Ji .

Having obtained analytical solutions for each of the flows ϕ[KE] , ϕ[PE] , ϕ[asym] , we then combine themwith the non-conservative momentum update from the non-conservative forcing terms Fpi and FΠi[force]with flow ϕhqi (h) = qi (0)pi (h) = pi (0) + Fpi h, Ri (h) = Ri (0), Πi (h) = Πi (0) + FΠi h.

(B17)

Consequently, the full Lie T2 integrator is obtained by applying the Strang composition scheme toobtain[KE][PE][asym][force][asym][PE][KE]
T2:= ϕh/2
ϕLie◦ ϕh/2 ◦ ϕh/2 ◦ ϕh
◦ ϕh/2 ◦ ϕh/2 ◦ ϕh/2
(B18)
h

C

Training Details

We implement our method using Jax3 [45] and use the Haiku framework [46]4 for constructing thedeep neural networks. In all our experiments, we use a multilayer perceptron (MLP) with 3 hidden
θlayers each of size 256 with the SiLU activation [47] for all networks (Vresid, Fpθ , FΠθ ). Each method
is run until convergence. Other common training hyperparameters used are summarized in Table 3.
θStructure of Vresid
Note that we do not assume prior knowledge on the pairwise structure of thepotential function Vresid or of the forcing terms Fp , FΠ . More specifically, since the true Gravitational
potential (2) only acts pairwise between rigid bodies, the rigid body correction potential also acts
pairwise and has the structure
X
Vresid (q, R) =Vi,j,resid (qi , qj , Ri , Rj ).
(C1)
i<j

For the forcing terms, the tidal forcing term also acts pairwise with coupled effects on pi and Πi due
the relationship between forces and torques, while the post-Newton general-relativity correction term
acts on each planet individually and affects only pi . However, in this work, we assume that none ofthis prior knowledge is available and aim to learn everything purely from data. Hence, we choose tolearn the high dimensional forms of Vresid , Fp , FΠ . The fact that we were able to obtain improvements
despite the high dimensionality of the input space (8 planets each with (qi , Ri ) ∈ R3 × R3×3 ≈ R12 ,i.e., R96 for V and R144 for Fp , FΠ ) is an indication of the generality of our approach. We
believe that the proposed approach can be made more scalable and have better generalization if
we do assume some knowledge about the structure of the forces at play and instead, e.g., learnθθ
Vresid(qi , qj , Ri , Rj , ϕi , ϕj ) instead of the more general Vresid(q, R) and similarly for the forcing
terms, where ϕi , ϕj contains known or potentially learned information about physical properties
about each rigid body which are needed. Here, since the ordering of i and j is not important, some
type of permutation-invariant encoding such as [48] should be used to further improve generalization.
34

https://github.com/google/jax. The repository is licensed under Apache-2.0.
https://github.com/deepmind/dm-haiku. The repository is licensed under Apache-2.0.

15

Table 3: Common hyperparameters used for training in all experiments.
Name

Value

Batch size
Optimizer
Learning Rate

256
AdamW [49]4e−4Training All experiments are done on a single RTX 3090 locally. The length of each run isdependent on the complexity of the integrator. For example, explicit Euler is the simplest and has thefastest per-iteration time, while Lie RK4 has the slowest per-iteration time. The Lie T2 integrator
has a per-iteration time between the two extremes with each experiment taking approximately threehours.
Data generation We use the GRIT5 simulator [20] for generating the dataset in all cases. The GRIT
simulator uses a high order Lie-group symplectic integrator derived using a tailored splitting scheme
and can handle tidal forces (between all pairs of bodies) and general relativistic effects (between eachplanet and the host star, assuming that the host star is much more massive than the other planets)
using a post-Newtonian correction term. Each dataset consists of 32 different trajectories, which isthen split in a 80-20 training-validation split. Each trajectory is then further subsampled to provide
128 tuples (q, R, p, Π) of datapoints. Each trajectory is generated by adding small multiplicative
Gaussian noise to the coordinates of each body.
C.1

Toy Two-Body Problem

The parameters for this system were hand-picked to provide intuition on how the additional corrections
differ from the pure point-mass potential that is widely used in other Hamiltonian learning literature.
C.2

TRAPPIST-1

The initial conditions for this system were taken directly from the TRAPPIST-1 example from GRIT.

D

Details on Integrators used for Comparison in Section 3

To answer the question Q2 of how important are symplecticity (S) and Lie-group preservation (L)for learning, we vary the choice of integrator in our experiments. The integrators used can be broadly
split into four categories:
Neither Symplectic S nor Lie-group preserving L: This category contains the popular explicit
Euler and Runge-Kutta 4 integrators, neither of which are symplectic or Lie-group preserving. The
finite-difference scheme from Greydanus et al. [1], Greydanus and Sosanya [37] can be interpreted as
an application of the explicit Euler integrator David and Méhats [10]. Explicit Euler is also used in
the works of ]. Runge-Kutta 4 is a popular fourth-order integrator due to its implementation simplicity
and is used in Finzi et al. [13], Zhong et al. [14], Duong and Atanasov [16], Zhong et al. [50, 51].
Symplectic S but not Lie-group preserving L: This category contains the verlet integrator, which
we define loosely in the current work as the integrator obtained by using classical splitting and Strang
composition, but without performing exact integration on the manifold. In the Euclidean case where
H(q, p) = 12 p⊺ p/m + V (q), this takes the form
p(0)
h,
m
∂V
p(h) = p(0) −
h,
∂q
p(1)
q(h) = q(h/2) +h,mq(h/2) = q(0) +5https://github.com/GRIT-RBSim/GRIT. The repository is licensed under Apache-2.0.

16

(D1)
(D2)
(D3)

and corresponds to the “leapfrog” proposed in SRNN Chen et al. [6]. In our setting where the space
is not Euclidean, we interpret a “naive” implementation to look like the following qi (h/2) = qi (0) + pi (0) h,
(D4a)
m


\
R (h/2) = R (0) + R (0) J−1
Π (0)h,
(D4b)
i

i

i

i

i


∂V


 pi (h) = pi (0) −
h,


∂qi



Π (h) = Πi (0) +

 i

(D5a)
∨ !



⊺
∂V
∂V
⊺
Ri (0)
Πi (0) × J−1
Π
(0)
−
R
(0)
−
i
i
i
∂Ri
∂Ri


 qi (h) = qi (h/2) + pi (h) h,m\R (h) = R (h/2) + R (h)J−1
Π (h)h,
i

i

i

i

i

h,

(D5b)

(D6a)
(D6b)

Not Symplectic S but is Lie-group preserving L: This category contains the Lie RK2 (CF2) and
Lie RK4 (CF4) commutator-free Lie-group preserving integrators from Celledoni et al. [52]. These
integrators are Lie-group preserving but are not symplectic. Specifically, the flow of Lie RK2 isdescribed by

"
#
"
# 
q̇i (0)qi (h/2)
qi (0)p (h/2)= pi (0) +  ṗi (0)  h, i
Πi (h/2)
Πi (0)
Π̇i (0)
(D7)


−1
\

F
= Ji Πi (0),
 1
Ri (h/2)= exp(F1 h/2)Ri (0),

"
#
"
# 
q̇i (h/2)
q
(h)
q
(0)
i
i




p (h)
= pi (0) +  ṗi (h/2)  h, i
Πi (h)
Πi (0)
Π̇i (h/2)
(D8)


−1\

F
=
J
Π
(h/2),
ii 2Ri (h)= exp(F2 h)Ri (0).
For Lie RK4, the coordinates in Euclidean space (i.e., q, p, Π) follow the normal RK4 integration ina similar fashion as above, while the flow of the R coordinate is described by(\
F1 = J−1i Πi (0),(D9)
R1i = exp(F1 h/2)Ri (0),(−1 1
F2 = J\i Πi ,(D10)
2
Ri = exp(F2 h/2)Ri (0),(−1 2
F3 = J\i Πi ,(D11)
3
Ri = exp((F3 − F1 /2)h)R2i ,

−1 3
\


F4 = Ji Πi ,






h
(3F1 + 2F2 + 2F3 − F4 ) Ri (0),Ri (h/2) = exp
(D12)
12






h


(−F1 + 2F2 + 2F3 + 3F4 ) Ri (h/2),
Ri (h) = exp
12
where the superscripts denote the intermediate outputs of each stage.
Both Symplectic S and Lie-group preserving L: This category contains our proposed Lie T2
integrator which is both symplectic and Lie-group preserving using a splitting technique that allows
for exact integration, borrowed from Chen et al. [20]. See App B for a detailed derivation.
17

E

Definition of Evaluation Metrics

The errors ∥∆q∥2 , ∥∆R∥ in Table 1 and Table 2 are computed by predicting a trajectory with 500
integrator steps, and then computing∆qk,l := qk,l − q̂k,l ,∆Rk,l := Rk,l − R̂k,l ,and the norm for R ∈ SO(3)⊗N is the geodesic computed asNX
tr(R⊺i Ri ) − 1∥R∥ :=∥Ri ∥,∥Ri ∥ := cos−1
2
i=1

(E1)
(E2)

(E3)

The numbers shown in the table are taken to be the mean across all LK samples.
The errors ∥∆ṗ∥ and ∆Π̇ in Table 1 are intended to measure how well the forces (both conservativeand non-conservative) are learned. Consequently, these are computed along the dataset and not thepredicted trajectory and are defined asp̂((k + 1)h) − p(kh) p((k + 1)h) − p(kh)−,hhΠ̂((k + 1)h) − Π(kh) Π((k + 1)h) − Π(kh)∆Π̇k =−,
hh∆ṗk =(E4)(E5)

where p̂ and Π̂ in the equation above denote the one-step predictions using the learned integrator ϕh ,i.e.
q̂(h), p̂(h), R̂(h), Π̂(h), := ϕh q(0), p(0), R(0), Π(0) .
(E6)Similarly, the errors ∆ ∂Vand ∆ ∂V
in Table 2 measure how well the potential (and hence
∂q
∂q
2
2
the conservative forces) have been learned. Hence, these are again computed along the dataset andnot the predicted trajectory and are defined as22θ∂Vresid∂Vresid
−
,
(E7)
∂q
∂q 2
2


⊺ ∨ ⊺ ∨ 2
N
2
θ
θ
X
∂V
∂Vresid∂Vresid⊺ ∂Vresid⊺ ∂Vresid
:=
∆
Ri
−
Ri
− Ri
−
Ri
,
∂R 2
∂Ri
∂Ri
∂Ri
∂Ri
i=1
2
(E8)

∆

∂V
∂q

:=

∂V
∂V
where the additional manipulations for ∆ ∂R
denote the projection of R⊺i ∂R
on the skew-symmetrici⊺ ∂Vmatrices. This is done because only the skew-symmetric part of Ri ∂Ri is used in the equations of
motion (1a)–(1d) and hence the learned symmetric part can be arbitrarily defined.

F

Limitations and Future Directions

One big assumption we make in our work is that the masses and inertial tensors mi , Ji are known.
This may be a restricting assumption when applying this method to learning physics for rigid bodies
where this information is not available and must be learned jointly with the physics, an approachtaken in many works on learning with Hamiltonian structure (e.g., [16, 51]). Moreover, as discussed
in App C, our framework does not assume that prior knowledge on the structure (e.g., pairwise,
independent) of the potential or forcing terms is known. If we make this assumption, then additional
physical properties for each rigid body may need to be provided to fully specify the structured physics(e.g., time lag and tide constants for tidal forcing). Extending this framework to handle the learningof per-body physical properties at the same time is left for future work.
Another direction that we have not been able to explore in this work is the robustness of our method
to noise in the dataset. In this work, we use a dataset generated using the GRIT [20] simulator (seeApp C). Consequently, the dataset used is clean and not corrupted by any noise. As noted in SRNN
18

Chen et al. [6], integrating for multiple consecutive timesteps may allow the network to better discern
the true noiseless trajectory when computing the loss function. Moreover, the benefit from performing
this multi-step training is integrator dependent, with improvements not observed when using the naive
explicit Euler integration. Given that, 1. our configuration space now lies on a Lie-group manifold,and 2. the physics of our system are mutiscale, exploring whether those insights are applicable to the
problem considered in this work is a future interesting direction that will help inform practitioners
wishing to apply this methodology to real world data.

19