Semigroup-JEPA

Latent Dynamics Consistency for Zero-Shot Physics Generalization

Andy Zeyi Liu^1,*, Haoran Sun^1,*, Lucas Baker², Randall Balestriero³, John Sous^1,†

¹Yale University ²Jump Trading ³Brown University

^*Equal contribution. ^†Correspondence: john.sous@yale.edu

TL;DR: Compared with LeWorldModel and DINO-WM, SG-JEPA's gravity-conditioned rollout training improves long-horizon prediction and closed-loop control under unseen gravity.

Abstract

JEPA world models learn compact latent representations for prediction and planning, but whether they learn physics and generate realistic dynamics remains untested. We introduce Semigroup-JEPA (SG-JEPA), a gravity-conditioned extension of LeWorldModel that jointly trains an encoder and temporal predictor through autoregressive latent rollouts. We evaluate out-of-distribution generalization on systems that obey the same physical law under different gravitational fields, ranging from floating motion in weak gravity to rapid bouncing in strong gravity.

Compared with DINO-WM, SG-JEPA reduces open-loop prediction error by up to 2x on 2D tasks and improves control success by up to 2.5x on 3D robotic tasks with independently trained diffusion policies. A linear feature model separates local law-conditioned error from its recursive amplification during rollout. The analysis attributes most of the gain to the encoder learning dynamical features that the predictor can carry forward.

Method

Semigroup-JEPA trains a gravity-conditioned recursive latent predictor with rollout and SIGReg losses, then freezes the encoder for receding-horizon diffusion-policy control — SG-JEPA jointly trains the encoder and gravity-conditioned predictor over \(K\) recursive steps. For control, a separate diffusion policy uses the frozen encoder and replans after \(E\) actions.

SG-JEPA extends LeWorldModel by adding normalized gravity to the action and training the encoder and predictor over a discounted \(K\)-step latent rollout. Each predicted latent enters the next history window. Targets come from the same trainable encoder, while SIGReg prevents collapse without stop-gradient.

For action-free motion at fixed gravity, repeated updates satisfy \(S_g(k+\ell)=S_g(\ell)\circ S_g(k)\), the discrete semigroup behind the name. During control, a task-specific Diffusion Policy uses the frozen encoder to predict \(A\) actions, execute the first \(E\), and replan from new observations.

Results

At horizon 44 on 2D Square, SG-JEPA (GRU) reduces position, velocity, and rotation error by 31% to 48% relative to DINO-WM. On Approach Ball, the GRU and SSM variants reduce mean position error by about 34% relative to DINO-WM and 50% relative to Original LeWM. With separate policies trained on frozen features, SG-JEPA (GRU) also improves success on all three control tasks.

Square dynamics rollout comparing observed RGB frames, decoded target latents, and decoded recursively predicted latents

Franka robot completing the basket task with closed-loop receding-horizon diffusion-policy control

Prediction error and closed-loop success across rollout horizon and held-out gravity — Long-horizon prediction and Diffusion Policy control across held-out gravity values.

Where the gain comes from

Fresh-predictor crossover separates the encoder from the temporal model. The GRU-trained encoder retains about 12% lower rollout error than the Transformer-trained encoder with either fresh predictor. Against DINO-WM, SG-JEPA has lower teacher-forced local error at all 25 test gravities and about 32% lower error far outside the training range. Recursive feedback widens this gap, although it shrinks again at long horizons.

Frozen encoder-predictor crossover, local error by gravity, and local versus free-rollout gap — Frozen-encoder crossover, teacher-forced physical error, and the difference between local and free-rollout gaps on 2D Square.

Ablations

These ablations examine the SIGReg weight and adaptation from a few new gravity values.

In the first three panels, raw SIGReg loss is separated from effective rank and linear-probe accuracy. The remaining panels show rollout error at horizons 5, 20, and 44 across gravity; the yellow band marks the training range, and lower is better except for rank. Overall, \(\lambda_{\mathrm{SIG}}=0.72\) gives the best balance of representation rank, probe accuracy, and rollout quality.

SIGReg coefficient sweep showing training loss, effective rank, linear-probe error, and rollout error at three horizons — SIGReg coefficient sweep on 2D Square: training loss, effective rank, linear-probe error, and rollout error at horizons 5, 20, and 44.

Each row below is one world model, while the columns report position, velocity, and rotation error after 44 rollout steps. Target-only post-training uses four gravity values, \(g\in\{0,2,6,8\}\); mixed post-training also retains data from the original training range. The mixed setting is more consistent at unseen gravities between the four support values and reduces SG-JEPA (GRU)'s aggregate square error by 13.6%.

Square per-gravity errors for four predictors under pretrained, target-only, and mixed post-training — Sparse-gravity post-training on 2D Square at horizon 44. Curves compare the pretrained model with target-only and mixed adaptation across held-out gravity values.

Limitations

The experiments vary one supplied scalar, gravity. Transfer across shapes is uneven: house translation transfers better than rotation, and the pentagon remains difficult. The theory is linear, while the learned predictor is nonlinear and history dependent, and contacts can change transition branches.

BibTeX

@misc{liu2026semigroupjepa,
  title         = {Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization},
  author        = {Liu, Andy Zeyi and Sun, Haoran and Baker, Lucas and Balestriero, Randall and Sous, John},
  year          = {2026},
  eprint        = {XXXX.XXXXX},
  archivePrefix = {arXiv}
}