Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization
Andy Zeyi Liu, Haoran Sun, Lucas Baker, Randall Balestriero, John Sous
cs.LG, cs.AI, cs.CV
2026-09-10
SG-JEPA conditions JEPA on gravity and trains encoder plus predictor with discounted K-step rollouts. 2D open-loop error drops up to 2x vs DINO-WM; catch success 9.5% to 23.3%.
JEPA world models skip pixel reconstruction and predict the next latent instead. DINO-WM freezes a pretrained DINOv2 encoder and trains only an action-conditioned predictor. LeWorldModel (LeWM) trains the encoder and predictor together, but still with one-step teacher forcing: the predictor always sees true encoded latents, never its own outputs.
Both setups assume a fixed transition law. Change gravity and the same physics produces qualitatively different motion, from floating to rapid bouncing. The question is whether the model learned "the transition at this g" or "how g enters the dynamics."
SG-JEPA changes LeWM in two places.
Gravity is concatenated onto the action and encoded with a temporal convolution plus a SiLU MLP into a 256-d condition. For free-fall there is no external action, so the condition is just g. Training z-scores g with train-split statistics.
The loss is a discounted K-step autoregressive rollout rather than one-step teacher forcing. From H=20 true encodings, the predictor emits K=5 steps and feeds each prediction back into the history window. Weights decay with γ=0.95 so late-step error does not dominate. SIGReg pushes random 1-d projections of the latents toward a standard Gaussian. Targets come from the same trainable encoder, with no stop-gradient.
The encoder is a ViT-Tiny (12 layers, width 192) trained from scratch; the CLS token is projected to 256-d. Predictors tried: Transformer, GRU, and a Mamba-style SSM. The paper treats GRU or SSM as the default SG-JEPA. Matrices are updated with Muon, everything else with AdamW, for 20 epochs.
All data is MuJoCo, 64 frames at 16 Hz. Training gravity is N(4, 0.5²) for planar shapes and Arm Catcher, N(9.8, 2.0²) for the other 3D sets. The test grid is much wider, including negative g and the surface gravities of Pluto, the Moon, Mars, and Venus. Downstream: freeze the world model and probe position, velocity, and rotation with an MLP; freeze the encoder and train a gravity-conditioned Diffusion Policy for catching, paddling, and striking a ball into a basket.
The name is literal. At fixed g with no actions, iterating the shared update is a discrete semigroup. With actions the blocks still compose, each modified by its control.
2D square, 44-step open loop (2.75 s) versus DINO-WM:
| Metric | DINO-WM | SG-JEPA (GRU) |
| Position L2 (m) | 2.006 | 1.034 |
| Velocity L2 (m/s) | 3.082 | 1.915 |
| Cumulative rotation (turns) | 0.341 | 0.236 |
Errors drop 31–48%; position is close to half. The triangle is harder. SG-JEPA still wins position and velocity at step 44, but DINO-WM keeps a slight edge on cumulative rotation (0.875 vs 0.904), which the paper blames on asymmetric contacts amplifying small mistakes.
On Approach Ball, mean position error over h1–h44 is 0.0491 m for GRU, 0.0705 m for DINO-WM, 0.0986 m for LeWM: about 30% below DINO and 50% below LeWM. SG-JEPA wins 22 of 25 test gravities; the losses are at very small g. Velocity is not uniformly better. DINO-WM is more accurate at the final step.
Frozen-encoder Diffusion Policies:
| Task | DINO-WM | SG-JEPA (GRU) |
| Arm Catcher capture | 9.5% | 23.3% |
| Franka paddle-to-basket | 27.4% | 30.5% |
| Arm Paddle qualified bounce | 17.7% | 23.8% |
Catching is about 2.5×, and 42.3% vs 16.0% near the training gravity. Both hit the ball more than 95% of the time on Franka; the gap is the outgoing velocity after contact.
The mechanism sits in the encoder. Freeze each encoder and train a fresh predictor: the GRU-trained representation is about 12% better with either a new GRU (1.376 vs 1.555) or a new Transformer (1.269 vs 1.453). On far-OOD gravity, teacher-forced local error is already about 32% lower. Free rollout then stretches that crack: the gap peaks near 0.38 around horizon 20, while the local gap stays below 0.06. Back-propagating the multi-step loss trains the encoder to keep features the predictor can carry, which are the features the dynamics actually use.
For world-model work this is a recipe you can copy: put the physical parameter in the action channel, replace one-step loss with discounted multi-step rollout, and expect the gain in the representation rather than in the temporal backbone. GRU and SSM trade the lead across horizons, so the story is not "use a GRU."
For control, a frozen encoder plus a Diffusion Policy already transfers catching outside the training gravity band. Gravity is supplied, not inferred from pixels. Deploying on an unknown planet still needs someone to tell you g, or a separate parameter estimator.
This is a controlled simulation study of one mechanism, not a generalist world model.
The authors list three. Only a scalar gravity is varied; there is no vector-valued physics and no attempt to infer parameters from observation. Shape transfer is uneven: triangle-plus-square carries much of a composite "house" rigid body's translation, but not its rotation, and pentagons remain hard. The theory is a linear feature model; the neural predictor is nonlinear and history-dependent, and contacts switch transition branches.
A few more caveats. An intro paragraph on a swap-and-retrain experiment contradicts the abstract and Section 4 on whether the encoder or the predictor holds the gain; the frozen-encoder crossover in Section 4 is the one with numbers. Triangle rotation still loses to DINO-WM. The Franka lift is thin, 27.4% to 30.5%. Everything is MuJoCo, no real robot. The training gravity band is narrow; a wider test grid still shares the same contact model and renderer.