TANGO: Time-Reversal Latent GraphODE for Multi-Agent Dynamical Systems
Zijie Huang Wanjia Zhao Jingdong Gao Ziniu Hu Xiao Luo Yadi Cao Yuanzhou Chen Yizhou Sun Wei Wang
Learning complex multi-agent system dynamics from data is crucial across many
domains, such as in physical simulations and material modeling. Extended from
purely data-driven approaches, existing physics-informed approaches such as
Hamiltonian Neural Network strictly follow energy conservation law to introduce
inductive bias, making their learning more sample efficiently. However, many
real-world systems do not strictly conserve energy, such as spring systems with
frictions. Recognizing this, we turn our attention to a broader physical
principle: Time-Reversal Symmetry, which depicts that the dynamics of a system
shall remain invariant when traversed back over time. It still helps to
preserve energies for conservative systems and in the meanwhile, serves as a
strong inductive bias for non-conservative, reversible systems. To inject such
inductive bias, in this paper, we propose a simple-yet-effective
self-supervised regularization term as a soft constraint that aligns the
forward and backward trajectories predicted by a continuous graph neural
network-based ordinary differential equation (GraphODE). It effectively imposes
time-reversal symmetry to enable more accurate model predictions across a wider
range of dynamical systems under classical mechanics. In addition, we further
provide theoretical analysis to show that our regularization essentially
minimizes higher-order Taylor expansion terms during the ODE integration steps,
which enables our model to be more noise-tolerant and even applicable to
irreversible systems. Experimental results on a variety of physical systems
demonstrate the effectiveness of our proposed method. Particularly, it achieves
an MSE improvement of 11.5 % on a challenging chaotic triple-pendulum systems.