Learning Integrable Dynamics with Action-Angle
Networks
Ameya Daigavane∗
Massachusetts Institute of Technology
Cambridge, MA 02139, USA

Arthur Kosmala
Ludwig-Maximilians-Universität München
80333 Munich, Germany

ameyad@mit.edu

a.kosmala@physik.uni-muenchen.de

Miles Cranmer
Princeton University
Princeton, NJ 08544, USA

Tess Smidt
Massachusetts Institute of Technology
Cambridge, MA 02139, USAmcranmer@princeton.edu

tsmidt@mit.edu

Shirley HoCentre for Computation Astrophysics, Flatiron Institute
New York, NY 10010, USAshirleyho@flatironinstitute.org

Abstract
Machine learning has become increasingly popular for efficiently modelling the
dynamics of complex physical systems, demonstrating a capability to learn effective models for dynamics which ignore redundant degrees of freedom. Learned
simulators typically predict the evolution of the system in a step-by-step mannerwith numerical integration techniques. However, such models often suffer frominstability over long roll-outs due to the accumulation of both estimation and integration error at each prediction step. Here, we propose an alternative construction
for learned physical simulators that are inspired by the concept of action-angle
coordinates from classical mechanics for describing integrable systems. We propose Action-Angle Networks, which learn a nonlinear transformation from input
coordinates to the action-angle space, where evolution of the system is linear.
Unlike traditional learned simulators, Action-Angle Networks do not employ any
higher-order numerical integration methods, making them extremely efficient at
modelling the dynamics of integrable physical systems.

1

Introduction

1.1

Modelling Hamiltonian Systems

Hamiltonian systems are an important class of physical systems whose dynamics are governedby a scalar function H, called the Hamiltonian. The state of a Hamiltonian system is a 2-tupleu(t) = (q(t), p(t)), where q ∈ Rn are the positions and p ∈ Rn are the canonical momenta. We
are interested in predicting the time evolution of a Hamiltonian system as a function of time t. In∗

Work performed when affiliated with Google Research and visiting Flatiron Institute.

Machine Learning and the Physical Sciences workshop, NeurIPS 2022.

particular, we seek to learn a model Mϕ (parameterized by ϕ) that can accurately predict the futurestate u(t + ∆t) given the current state u(t):Mϕ (u(t), ∆t) ≈ u(t + ∆t).
(1)

Previous efforts towards modelling Hamiltonian systems [Greydanus et al., 2019, Cranmer et al.,2020, Jin et al., 2020, Chen et al., 2018, Kidger, 2022] have seen success with neural networks trained
via backpropagation on the trajectory prediction objective:Xϕ∗ = arg min∥Mϕ (u(t), ∆t) − u(t + ∆t))∥2 .
ϕtSuch models are often constrained in some way to match the underlying physical evolution, which
improves their accuracy. However, they suffer from several drawbacks that have restricted theireffectiveness: (1) Their predictions tend to be unstable over long roll-outs (when ∆t is large). (2)They tend to require many parameters and need long training times. Further, the inference times ofthese models often scale with ∆t.
Based on the principle of action-angle coordinates [Vogtmann et al., 1997] from classical mechanics,
we introduce a new paradigm for learning physical simulators which incorporate an inductive bias
for learning integrable dynamics. Our Action-Angle Network learns an invertible transformationto action-angle coordinates to linearize the dynamics in this space. In this sense, the Action-Angle
Network can be seen as a physics-informed adaptation of DeepKoopman [Lusch et al., 2018]. The
Action-Angle Network is efficient both in parameter count and inference cost; it requires much
fewer parameters than previous methods to reach similar performance and enjoys an inference time
independent of ∆t.
1.2

Action-Angle Coordinates

The time evolution of any Hamiltonian system is given by Hamilton’s equations:
dq
∂H
dp
∂H
=
,
=−
(2)
dt
∂p
dt
∂q
where H is the Hamiltonian of the system. The complexity for learning these dynamics arises
solely because H is not known and must be inferred; only samples from a trajectory over which
H is conserved is available. The core issue for all physical simulators is the non-linearity of the
dynamics when expressed in the canonical coordinates (q, p). However, for integrable systems [Tong,
2004] which possess significant symmetry, the dynamics are actually linear in a set of coordinates
termed the action-angle coordinates [Vogtmann et al., 1997]. The actions I and angles θ are related
to the canonical coordinates (q, p) via a symplectic transformation T [Meyer and Hall, 2013] as(I, θ) = T (q, p). In the action-angle coordinates, the Hamiltonian only depends on the actions I, notthe angles θ. Thus, Hamilton’s equations in this basis tell us that:
dI
∂H
dθ
∂H
=−
= 0,
=
.
(3)
dt
∂θ
dt
∂I
Thus, the actions I are always constant across a trajectory, while the angles evolve linearly withnconstant rate θ̇ = ∂H
∂I . The action-angle space can be thought of as a torus T , where the actions I
are a function of the radii, and the angles describe the individual phases living in [0, 2π).
Learning the mapping T from canonical coordinates (q, p) to action-angle coordinates (I, θ) for
several physical systems was first explored in [Bondesan and Lamacraft, 2019], which did not
focus on the complete dynamics of integrable systems. We leverage their framework to additionallyincorporate a dynamics model F to learn a complete physical simulator.

Ishikawa et al. [2021] attempt to model the time evolution of actions for integrable systems bylearning the Hamiltonian in action-angle coordinates. However, their overall objective and training
setup are different from ours: they have samples of the true action-angle coordinates (I, θ) as input,
while we only observe the canonical coordinates (q, p).

2

Action-Angle Networks

Suppose we observe the system at time t, and we query our model Mϕ to predict the state at a future
time t + ∆t. The Action-Angle Network performs the following operations, illustrated in Figure 1:
2

Figure 1: An overview of the Action-Angle Network.
• Encode: Convert the current state in canonical coordinates q(t), p(t) to action-angle coordinates
I(t), θ(t) via a learned map T .

• Evolve: Compute the angular velocities θ̇(t) at this instant with the dynamics model F:θ̇(t) ← F(I(t)). Evolve the angles to time t + ∆t with these angular velocities:

θ(t + ∆t) ← θ(t) + θ̇(t) · ∆t (mod 2π)(4)−1• Decode: Convert the new state back to canonical coordinates q(t + ∆t), p(t + ∆t) via T .
Symplectic Normalizing Flows: The transformation T from canonical coordinates to action-angle
coordinates is guaranteed to be symplectic [Vogtmann et al., 1997]. To constrain our encoder to learn
only symplectic transformations, we compose layers of symplectic normalizing flows [Li et al., 2020,
Bondesan and Lamacraft, 2019, Jin et al., 2020]. In particular, we found that the G-SympNet layers[Jin et al., 2020] with our modification described below performed well empirically. G-SympNet was
shown to be universal; they can represent any symplectic transformation given sufficient width anddepth. Similar to affine coupling layers [Dinh et al., 2014], each G-SympNet layer ψi operates on
only one of the coordinates keeping the other fixed, depending on the parity of i:ψ2k (q, p) = (q, p + f (q)), ψ2k+1 (q, p) = (q + f (p), p)(5)
where f is of the following form f (x) = Cx + W T diag(A)σ(W x + B) with learnable parameters
A ∈ Rdo , B ∈ Rdo , C ∈ R and W ∈ Rdo ×n , where do is a hyperparameter. Our minor modification
above allows each layer to model the identity transformation, enabling the training of deeper models.
Action-Angles via Polar Coordinates: We found that learning a mapping from canonical coordinates
which live in R2n directly to action-angle coordinates which live in the torus Tn was challenging forthe network. We hypothesize that the differing topology of these spaces is a major obstacle becauseof the resulting singularities. To bypass this, we borrow a trick from Bondesan and Lamacraft [2019];
we have the G-SympNet instead output the components of the actions (I (x) , I (y) ) in the Cartesian
coordinate basis. These are then converted to action-angles (I, θ) in the polar coordinate basis, by(x) (y)the standard transformation Tpolar applied to each pair of (Ii , Ii ) coordinates:
q
(x)
(y)
(y)
(x)
Ii = (Ii )2 + (Ii )2 , θi = arctan(Ii /Ii )(6)
Our encoder can thus be described as T := TPolar ◦ G-SympNet.

Evolve: From Equation 3, we know that the angular velocities θ̇ are only a function of the actions I
at the instant t0 , which we model as a simple multi-layer perceptron (MLP) F. Since the angles θ
are guaranteed to evolve linearly, we do not need to use higher-order numerical integration schemes.
Instead, a single call to the forward Euler method is exact:θ̇(t0 ) ← F(I(t0 )) θ(t0 + ∆t) ← θ(t0 ) + θ̇(t0 ) · ∆t (mod 2π)(7)−1Decode: As T is symplectic, T is invertible. Thus, our decoder is T −1 := G-SympNet
3−1◦ Tpolar.

2.1

Training

We generate a trajectory Tr = {(q(t), p(t))}Tt=1 of T = 1000 time steps, by either evaluating theclosed-form dynamics equations or via numerical integration. In practise, this trajectory can also be
obtained from raw observations, and does not need to be regularly sampled in time. We train models
upto the first 500 time steps and evaluate performance on the last 500 time steps, for each trajectory.
State Prediction Error: We wish to learn parameters ϕ := (ϕT , ϕF ) that minimize the predictionP
loss Lpredict :1
2(8)
Lpredict = 1+∆tt0 ∈Tr ∥Mϕ (u(t0 ), ∆t) − u(t0 + ∆t)∥ .
Regularization of Predicted Actions: Additionally, we wish to enforce the fact that the predictedactions I are constant across the trajectory. We do this by adding a regularizer Laction on the varianceof the predicted actions:2Pbb = 1Pd
I(t0 ) − E[I]where E[I]
I(t0 ).
(9)
Laction = Var(I)= 1T

t0 ∈TrT

t0 ∈Tr

We find that this global regularizer works better than the loss proposed in Bondesan and Lamacraft[2019] that minimizes only the local pairwise differences between the predicted actions across the trajectory. As the Action-Angle Network is completely differentiable, the parameters ϕ can be obtained
via gradient-descent-based minimization of the total loss: ϕ∗ = argminϕ (Lpredict + λLaction ) , where
λ is a hyperparameter that controls the strength of the regularization.
Schedule for ∆t: Increasing ∆t over the course of training helped, as it corresponded to gradually
increasing the complexity of the prediction task. We set ∆tmax = 10 and sampled ∆t according to:∆t ∼ Uniform (0, (Current training step)/(Maximum training steps)) · ∆tmax(10)

3

Experiments

We compare the Action-Angle Network to three strong baseline models: the Euler Update Network(EUN), the Neural Ordinary Differential Equations (Neural ODE) [Chen et al., 2018], and the physicsinspired Hamiltonian Neural Networks (HNN) [Greydanus et al., 2019]. A comparison of thesemodels can be found in Table 1. The baseline models are further described in Section A.1.
Table 1: Comparing different models.

Parameter countInference time
Learns conservation laws
Learns linear dynamics

Action-Angle Network

EUN

Neural ODE

HNN≈ 8.5K
O(1)
✓
✓

≈ 9K
O(1)

≈ 100KO(∆t)≈ 200KO(∆t)✓✓

As detailed in Section A.2, we simulate a system of coupled harmonic oscillators.2 Figure 2 depicts
the prediction error for each of the models as a function of training samples, showing that the
Action-Angle Network is much more data-efficient than the other baselines. Figure 3a shows that
the Action-Angle Network can be queried much faster than the Neural ODE and the HNN, with aninference time that is independent of ∆t. Figure 3b depicts the prediction error as a function of ∆t,showing that the Action-Angle Network also scales much better with the jump size ∆t, even for jumpsizes larger than those seen during training ∆tmax = 10. Finally, Figure 4 shows that the predictedangular frequencies from the Action-Angle Network closely match the true angular frequencies.

4

Conclusion

Our preliminary experiments indicate that Action-Angle Networks can be a promising alternative
for learning efficient physical simulators. Action-Angle Networks overcome many of the obstacles– instability and inefficiency – faced by state-of-the-art learned simulators today. That being said,
our experiments here only model very simple integrable systems. Next, we plan to analyse theperformance of Action-Angle Networks on modelling large-scale non-integrable systems (such ascosmological simulations and molecular dynamics trajectories) to better understand their tradeoffs.
2

We have created animations of the system trajectory and model predictions at this webpage. Our codebase
to run all experiments and analyses is available at https://github.com/ameya98/ActionAngleNetworks.

4

100

10−2

Mean Prediction Error

Mean Prediction Error

Jump Size
1
2
5
10
20
50

102

10−4
0

100

200

300

400

102

Jump Size
1
2
5
10
20
50

100

10−2

10−4

500

0

100

Training Samples

Jump Size
1
2
5
10
20
50

100

10−2

10−4
0

100

200

300

300

400

500

(b) Euler Update Network
Mean Prediction Error

Mean Prediction Error(a) Action-Angle Network
102

200

Training Samples

400

Jump Size
1
2
5
10
20
50

102

100

10−2

10−4

500

0

100

Training Samples

200

300

400

500

Training Samples(c) Neural ODE(d) Hamiltonian Neural Network

Figure 2: Prediction errors on test data as a function of training samples.

100

Model
Action-Angle Network
Euler Update Network
Neural ODE
HNN

Prediction Error

Inference Time(seconds)

103

Model
Action-Angle Network
Euler Update Network
HNN
Neural ODE

10−1

10−2

101

10−1

10−3

0

10

20

30

40

100

50

101

Jump Size

102

Jump Size(a) Inference times as a function of jump ∆t.

(b) Prediction error as a function of jump ∆t.

Figure 3: Evaluating each model on (a) inference time and (b) prediction error.
300

6000

250

4000

DensityDensity

5000

3000

200
150

2000

100

1000

50

0
0.0498 0.0500 0.0502 0.0504 0.0506 0.0508 0.0510

0Angular Frequency

0.772

0.774

0.776

0.778

0.780

0.782

Angular Frequency

Figure 4: Kernel density plot of angular frequencies predicted by the Action-Angle Network. True
angular frequencies are indicated by dashed vertical lines.

5

Broader Impact Statement
In this paper, we have proposed a novel method for modelling the dynamics of Hamiltonian systems,with improvements in computational efficiency. There has been much interest in accelerating
simulations of various kinds across scientific domains that could have impact on human lives (forexample, simulating biological cell cycles and immune system responses). However, given thetechniques here apply to highly-constrained physical systems, we do not anticipate any negativesocial implications of our work.

Acknowledgments and Disclosure of Funding
The authors acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Centerfor providing (HPC, database, consultation) resources that have contributed to the research results
reported within this paper.

References
Roberto Bondesan and Austen Lamacraft. Learning Symmetries of Classical Integrable Systems,
2019. URL https://arxiv.org/abs/1906.04645.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural Ordinary
Differential Equations, 2018. URL https://arxiv.org/abs/1806.07366.
Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho.
Lagrangian Neural Networks, 2020. URL https://arxiv.org/abs/2003.04630.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear Independent Components
Estimation, 2014. URL https://arxiv.org/abs/1410.8516.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski.
Hamiltonian Neural Networks.
In Advances in Neural Information Processing Systems, volume 32. Curran
Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper/2019/file/
26cd8ecadce0d4efd6cc8a8725cbd1f8-Paper.pdf.
Fumihiro Ishikawa, Hidemaro Suwa, and Synge Todo. Neural Network Approach to Construction ofClassical Integrable Systems. Journal of the Physical Society of Japan, 90(9):093001, sep 2021.
doi: 10.7566/jpsj.90.093001. URL https://doi.org/10.7566%2Fjpsj.90.093001.
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. ISSN 0893-6080. doi: https://doi.org/10.1016/j.neunet.2020.08.017. URL
https://www.sciencedirect.com/science/article/pii/S0893608020303063.
Patrick Kidger. On Neural Differential Equations, 2022. URL https://arxiv.org/abs/2202.
02435.
Shuo-Hui Li, Chen-Xiao Dong, Linfeng Zhang, and Lei Wang. Neural Canonical Transformation
with Symplectic Flows. Phys. Rev. X, 10:021020, Apr 2020. doi: 10.1103/PhysRevX.10.021020.
URL https://link.aps.org/doi/10.1103/PhysRevX.10.021020.
Bethany Lusch, J. Nathan Kutz, and Steven L. Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature Communications, 9(1):4950, 2018. doi:
10.1038/s41467-018-07210-0. URL https://doi.org/10.1038/s41467-018-07210-0.
K. Meyer and G. Hall. Introduction to Hamiltonian Dynamical Systems and the N-Body Problem.
Applied Mathematical Sciences. Springer New York, 2013. ISBN 9781475740738. URL https://books.google.st/books?id=KITdBwAAQBAJ.
David Morin.
Normal Modes, 2022.
URL https://scholar.harvard.edu/files/
david-morin/files/waves_normalmodes.pdf. Accessed: 6th September, 2022.
6

David Tong. Lectures on Classical Dynamics, 2004. URL https://www.damtp.cam.ac.uk/
user/tong/dynamics.html.
K. Vogtmann, A. Weinstein, and V.I. Arnol’d. Mathematical Methods of Classical Mechanics.
Graduate Texts in Mathematics. Springer New York, 1997. ISBN 9780387968902. URL https://books.google.com/books?id=Pd8-s6rOt_cC.

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes] We have designed and evaluated Action-Angle
Networks, showing how they can efficiently learn the dynamics of integrable systems.
(b) Did you describe the limitations of your work? [Yes] Action-Angle Networks can beapplied to only integrable systems. Their performance on non-integrable Hamiltoniansystems still needs investigation.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We havementioned an overview in the Broader Impact statement. We do not expect any negative
societal impacts of our work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms
to them? [Yes] Our paper does not utilise or expose any human-derived data. Alldata for the physical systems discussed here are synthetically generated via numericalintegration.
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
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)? [Yes] Yes, we haveadded a link to our code in Section 3.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] Yes, the entire training pipeline and configuration is available atthe link for the code in Section 3.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No] We plan to address this in a future version.
(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] Experiments were run on the CPU
nodes of MIT SuperCloud, and can be easily reproduced with minimal resources.
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? [N/A](b) Did you mention the license of the assets? [Yes] Our code linked in Section 3 islicensed under the Apache License 2.0.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]We have released code and simulation tools at the link in Section 3.
(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](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]7

A

Appendix

A.1

Baseline Models

Euler Update Network: This model updates the latent state via the forward Euler method:û(t) ← Encode(u(t))
û(t + ∆t) ← û(t) + ∆t · F(û(t))
u(t + ∆t) ← Decode(û(t + ∆t))
Neural Ordinary Differential Equations: A generalization of Euler Update Networks, Neural
ODEs [Chen et al., 2018] have demonstrated state-of-the-art performance on several time-series
forecasting problems. They usually use a higher-order numerical integration scheme to update thelatent state:û(t) ← Encode(u(t))
Z t+∆tû(t + ∆t) ← û(t) +F(û(s))dstu(t + ∆t) ← Decode(û(t + ∆t))
Hamiltonian Neural Network: Hamiltonian Neural Networks (HNNs) [Greydanus et al., 2019]models the Hamiltonian in a latent space explicitly and updates the latent coordinates via Hamilton’s
equations.
q̂(t), p̂(t) ← Encode(u(t))
Z t+∆t∂H(q̂(s), p̂(s))dsq̂(t + ∆t) ← q̂(t) +∂ p̂t
Z t+∆t∂p̂(t + ∆t) ← p̂(t) −H(q̂(s), p̂(s))ds∂q̂
tu(t + ∆t) ← Decode(q̂(t + ∆t), p̂(t + ∆t))
Wherever applicable, we used the Dormand-Prince 5(4) solver, a 5th order Runge-Kutta method fornumerical integration.
All of these baselines can technically simulate the Action-Angle Network by encoding the canonicalcoordinates into the action-angle coordinates to linearize the dynamics.
A.2

Harmonic Oscillators

We model a system of n point masses, which are connected to a wall via springs of constant kw andto each other via springs of constant kp . This can be described by the set of n differential equations:
m

X
d2 qi
= −kw qi +kp (qj − qi ).
2
dt

(11)

j̸=i

where qi (t) is the position of the ith point mass. This forms a system of coupled harmonic oscillators,
where kp controls the strength of the coupling. Any solution to Equation 11 is a linear combinationof the ‘normal modes’ of the system [Morin, 2022]:
Pn
q(t) = r=1 Ar cr cos (ωr t + ϕr ).
The n angular frequencies ω and coefficients c are found by solving the following eigenvalueeigenvector equation (M + ω 2 In )c = 0, where M is the matrix defined as:
k + (n − 1)kp− w
if i = jmi(12)
Mij = k p
if i ̸= jmi8