Semigroup-JEPA improves physics learning and control in robot simulations
Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization
Machine LearningArtificial IntelligenceComputer Vision and Pattern Recognition
Summary
Understanding how physical forces like gravity affect movement is tricky for AI models that predict how things will behave. The authors introduce Semigroup-JEPA, a method that teaches AI to better capture these physics by working with different gravity settings and learning from multiple steps of prediction. This leads to more accurate predictions and better robot control in simulated environments. They found that the improvement mainly comes from teaching the AI to recognize useful features about the physics rather than just improving the prediction steps.
Joint-Embedding Predictive Architecturelatent representationworld modelphysics simulationgravityautoregressive rolloutencoderpredictorzero-shot generalizationrobot control
Authors
Andy Zeyi Liu, Haoran Sun, Lucas Baker, Randall Balestriero, John Sous
Abstract
Joint-Embedding Predictive Architecture (JEPA) world models learn a compact latent representation of the world that supports prediction and planning, but their capability to learn physics and generate physically realistic dynamics remains hitherto untested. In this work, we introduce SemiGroup-JEPA (SG-JEPA), which extends the LeWorldModel framework by supplying the parameter governing the physics to the temporal model via action-conditioning and jointly training an encoder and predictor through an autoregressive latent rollout. To evaluate the model's ability to generalize out of distribution, we design dynamical tasks under different gravitational fields that, despite obeying the same physical law, exhibit qualitatively different dynamics, ranging from floating motion in weak gravitational fields to rapid bouncing in strong ones. In contrast to DINO-WM, SG-JEPA reduces open-loop prediction error by up to 2 times on two-dimensional datasets, and increases control success rate up to 2.5 times for three-dimensional robotic datasets, for which we train independent diffusion policies. To explain this advantage, we develop a linear feature model that separates local law-conditioned error from its recursive amplification under rollout. Guided by this model, we find that back-propagating the multi-step rollout loss into the representation trains the encoder to keep the features that the predictor can carry forward, and that those are the features the dynamics depend on, so most of the gain comes from the encoder learning better features rather than from the predictor learning better dynamics. See project page at https://sg-jepa.github.io.