Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization

2026-09-09

AI Summary

研究通过引入Semigroup-JEPA模型，利用动作条件和自回归潜变量展开方法改进了JEPA框架，以提高物理动态学习能力和泛化性能。

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.

Problem

• JEPA
• Physics Generalization
• Latent Dynamics
• Gravitational Fields
• Zero-Shot

Innovation

• Semigroup-JEPA
• action-conditioning
• autoregressive latent rollout
• zero-shot physics generalization
• multi-step rollout loss