Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization
Introduced Semigroup-JEPA for zero-shot physics generalization, achieving 34% error reduction and 23.3% control success improvement.
Key Findings
Methodology
Semigroup-JEPA extends LeWorldModel by integrating gravity as a conditioning parameter and training the encoder and predictor jointly with a discounted K-step rollout loss and SIGReg regularization, improving latent dynamics consistency.
Key Results
- On 2D freefall datasets, SG-JEPA reduced position prediction error by 34% compared to DINO-WM, maintaining accuracy over 44 rollout steps.
- In the 3D Arm Catcher Ball task, control success rate increased from 9.5% to 23.3%, outperforming baselines significantly.
- SG-JEPA achieved the best performance on 22/25 gravity values in out-of-distribution tests, demonstrating strong zero-shot generalization.
Significance
This work addresses limitations in existing latent models for long-horizon and out-of-distribution dynamics prediction. By introducing gravity conditioning and recursive training, SG-JEPA enables zero-shot generalization across diverse gravitational environments, advancing physics modeling and robotics.
Technical Contribution
Proposed a gravity-conditioned latent prediction framework with SIGReg regularization and discounted multi-step loss, significantly improving dynamic fidelity and cross-distribution generalization. SG-JEPA suppresses error accumulation in long-horizon rollouts compared to SOTA.
Novelty
First to integrate gravity as a conditioning variable in latent models, achieving zero-shot physics generalization through recursive rollouts and multi-step loss optimization. Outperforms DINO-WM and LeWM in dynamic prediction and control.
Limitations
- Limited performance on extreme gravity values outside training range, especially under high gravity.
- Struggles with modeling complex contact events like collisions.
- High computational cost, particularly for long-horizon tasks.
Future Work
Future work could explore more efficient latent regularization techniques, improve modeling of complex contact events, and expand to broader physical dynamics scenarios.
AI Executive Summary
Modeling physical dynamics is a critical challenge in AI, but current latent models struggle with long-horizon predictions and out-of-distribution generalization.
Semigroup-JEPA introduces gravity conditioning and recursive latent rollouts to significantly enhance dynamic learning. On 2D freefall and 3D robotic control tasks, it reduced prediction errors by 34% and improved control success rates to 23.3%.
By combining SIGReg regularization and discounted multi-step loss, SG-JEPA ensures latent consistency over long horizons. While challenges remain in extreme gravity scenarios, SG-JEPA offers a promising solution for zero-shot physics generalization with broad applications in robotics and simulation.
Deep Analysis
Background
Physics modeling aims to predict object motion in dynamic environments. Existing methods like DINO-WM and LeWorldModel perform well under fixed dynamics but struggle with out-of-distribution scenarios.
Core Problem
How to achieve zero-shot generalization across gravitational environments with limited training data, while suppressing error accumulation in long-horizon predictions.
Innovation
- �� Introduced gravity as a conditioning variable to enhance dynamic modeling.
- �� Used discounted multi-step rollout loss to optimize long-horizon predictions.
- �� Applied SIGReg regularization to improve latent representation quality.
Methodology
- �� Encoder: Utilized ViT-Tiny to extract latent variables, incorporating gravity as input.
- �� Predictor: Employed GRU or SSM architectures for recursive latent prediction.
- �� Loss function: Combined discounted multi-step rollout loss with SIGReg regularization to optimize dynamic consistency.
Experiments
Evaluated on MuJoCo-generated 2D freefall and 3D robotic control datasets, with gravity values extending beyond the training range. Baselines included DINO-WM and LeWM.
Results
SG-JEPA reduced 2D prediction errors by 34%, improved 3D control success rates to 23.3%, and outperformed baselines on most out-of-distribution gravity values.
Applications
Applicable in robotics, physics simulation, and game engines, particularly in scenarios requiring dynamic adaptation to diverse environments.
Limitations & Outlook
Limited generalization to extreme gravity values, high computational cost, and challenges in modeling complex contact events.
Plain Language Accessible to non-experts
Imagine a robot catching a ball in different gravity environments. SG-JEPA is like a smart coach that not only teaches the robot how to catch but also helps it adapt to the gravity on Earth, the Moon, or even Jupiter. This adaptability is achieved through repeated practice and strategy adjustments, ensuring the robot performs well in unknown settings.
ELI14 Explained like you're 14
Think of playing a video game where your character jumps on different planets. SG-JEPA is like a cheat code that lets your character automatically adjust to each planet's gravity, so they can jump perfectly whether they're on the Moon or Jupiter. Cool, right?
Glossary
Latent Variable
A low-dimensional representation capturing key features of data.
Used as the core representation for dynamic prediction.
Latent Rollout
Generating future latent sequences through multi-step model predictions.
Key for long-horizon dynamic prediction.
SIGReg Regularization
A technique ensuring latent variables follow a standard Gaussian distribution.
Improves latent representation quality.
Gravity Conditioning
Incorporating gravity as an input parameter.
Enhances generalization across different gravitational environments.
Discounted Multi-Step Loss
A loss function that optimizes long-horizon predictions by weighting errors over multiple steps.
Improves dynamic consistency.
Open Questions Unanswered questions from this research
- 1 How can generalization to extreme gravity values be improved?
- 2 Can this approach be extended to model complex dynamics like fluids or soft bodies?
Applications
Immediate Applications
Robotic Control
Optimize robot operations in diverse gravity environments, such as space missions.
Physics Simulation
Enhance dynamic modeling in game engines or virtual reality applications.
Long-term Vision
Interplanetary Missions
Enable robots to perform complex tasks across multiple planetary environments.
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.