Learning Constrained Lagrangian Dynamics in Keypoint Coordinates from Images
Rembert Daems, Jeroen Taets, Francis wyffels, Guillaume Crevecoeur
Published in Neurocomputing 573 (2024): 127175, and presented (oral, top 1.6%) at the Machine Learning and the Physical Sciences Workshop at NeurIPS 2023.
We present KeyCLD, a framework to learn Lagrangian dynamics from images. Learned keypoints represent semantic landmarks in images and can directly represent state dynamics. Interpreting this state as Cartesian coordinates coupled with explicit holonomic constraints allows expressing the dynamics with a constrained Lagrangian. KeyCLD is trained unsupervised end-to-end on sequences of images. Our method explicitly models the mass matrix, potential energy and the input matrix, thus allowing energy-based control. We demonstrate learning of Lagrangian dynamics from images on the dm_control pendulum, cartpole and acrobot environments, whether they are unactuated, underactuated or fully actuated. Trained models are able to produce long-term video predictions, showing that the dynamics are accurately learned.
KeyCLD learns Lagrangian dynamics from images. (a) An observation is processed by a keypoint estimator. (b) Positions are represented as spatial probability heatmaps. (c) Cartesian coordinates are extracted via spatial softmax and used as state to learn Lagrangian dynamics. (d) A learned renderer reconstructs the original observation — including background, reflections and shadows — from the keypoint bottleneck. All models are jointly learned unsupervised on image sequences.
KeyCLD predicts future frames for the pendulum, cartpole and acrobot environments. Each predicted sequence is based on the first three frames of the ground-truth sequence to estimate velocities. KeyCLD makes accurate long-term predictions, including reflections and shadows. See the paper for comparisons with ablated models and related work.
Learning explicit energy models allows simple and robust energy-shaping control. Below, we reach a target state by leveraging the learned potential-energy models.