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

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.

• The paper introduces SG-JEPA, a model that improves multi-step prediction accuracy by 31-48% in long-horizon planar dynamics and(projectile prediction by 30-50% across various gravity levels through improved representation learning.
• SG-JEPA's success in zero-shot physics generalization is attributed to its encoder, which captures reusable motion structures from planar dynamics, even in unseen composite shapes.
• Allowing gravity as a conditioning variable in SG-JEPA enables effective transfer to robotic control tasks, enhancing performance by 2-1.5x in tasks involving dynamic projectile motion or robotic interception and impact.

Research question and contribution

The paper studies whether a latent world model trained from visual observations can learn a parameterized physical law rather than merely fit the transition statistics present during training. The controlled parameter is gravity, which is supplied explicitly to the model but sampled from a narrow distribution during training and from a substantially wider grid at evaluation. This construction separates two forms of generalization: extrapolation across a known physical parameter and stability under long autoregressive rollouts.

The proposed model, Semigroup-JEPA (SG-JEPA), extends LeWorldModel by jointly training a visual encoder and a gravity-conditioned temporal predictor with a recursive latent rollout objective. The model is evaluated on passive rigid-body motion, projectile dynamics, and robotic manipulation. The central empirical claim is that multi-step training improves both long-horizon prediction and downstream control at held-out gravity values, with much of the improvement attributable to the representation learned by the encoder rather than to the predictor architecture alone (2609.10464).

The paper’s framing is technically important because one-step latent prediction does not directly constrain the behavior of a model when its own predictions become future inputs. A representation can yield low teacher-forced error while discarding variables required for stable composition. SG-JEPA instead optimizes the consequences of repeatedly applying a shared latent transition operator, thereby making predictive closure and error propagation explicit training concerns.

SG-JEPA architecture and objective

SG-JEPA encodes each RGB observation with a ViT-Tiny encoder and maps the final CLS token into a 256-dimensional latent. Gravity is appended to the action vector and passed through an action encoder. This design makes gravity available to both passive and action-conditioned tasks without requiring the model to infer it from trajectories. The temporal predictor receives a history of $H=20$ latent-action pairs and predicts the next latent. The primary experiments compare GRU and SSM predictors, with a Transformer used for LeWM-style baselines and architectural ablations.

The key methodological change is the replacement of teacher-forced one-step prediction with a $K$-step autoregressive latent rollout. The model begins from a true context, predicts the first future latent, inserts that prediction into the context window, and recursively generates subsequent latents. The rollout objective weights prediction steps geometrically, using $K=5$ and $\gamma=0.95$ in the principal configuration. Unlike the standard DINO-WM and Original LeWM baselines, SG-JEPA therefore receives direct gradient pressure from errors caused by feeding predictions back into the model.

Figure 1: SG-JEPA jointly trains the visual encoder and gravity-conditioned predictor through recursive latent rollouts, then uses the frozen representation for gravity-conditioned diffusion-policy control.

The objective combines the autoregressive rollout loss with SIGReg, which regularizes random one-dimensional projections of the encoded latents toward a standard Gaussian. Targets are generated by the same trainable encoder rather than a stop-gradient target network. The use of SIGReg is intended to prevent representational degeneracy while leaving the rollout loss responsible for selecting dynamics-relevant information.

The experimental configuration also includes a hybrid Muon/AdamW optimizer. On the planar square task, Muon produces lower best validation objectives than pure AdamW for the GRU, SSM, and Transformer predictors: $4.990$ versus $5.864$, $5.403$ versus $6.922$, and $5.278$ versus $6.384$, respectively. These correspond to relative reductions of $K$0, $K$1, and $K$2. This result is useful for reproducibility, although it does not isolate whether the optimizer improves the learned representation, optimization speed, or both.

Experimental design

The authors construct eight MuJoCo datasets. Four contain planar rigid bodies undergoing free fall and collisions; Approach Ball isolates 3D projectile dynamics; and three manipulation tasks require a robot arm to catch, repeatedly bounce, or redirect a ball. Episodes contain 64 frames sampled at 16 Hz. Training gravity is drawn from a narrow Gaussian centered at either $K$3 or $K$4, depending on the task family. Evaluation uses a 25-value grid that extends well beyond the training distribution, including values corresponding approximately to planetary gravities.

Figure 2: The evaluation suite spans planar collisions, projectile motion, interception, repeated paddle contact, and paddle-to-basket control.

Prediction is evaluated using frozen MLP probes trained to recover physical state from latent windows. The reported metric subtracts the probe error on encoded ground-truth latents from the probe error on autoregressive rollouts. This excess-error construction is important: it attempts to measure degradation introduced by latent dynamics rather than conflating it with the probe’s inability to decode the representation.

Control uses an independently trained gravity-conditioned Diffusion Policy for each frozen encoder. The policy receives a history of 20 latent features, predicts 16 actions, executes only the first 4 or 8, and replans from new observations. Thus, the control experiments evaluate frozen representations rather than online use of the SG-JEPA predictor. This distinction strengthens the interpretation that representation quality, not merely model-based planning, drives the control gains.

Long-horizon planar prediction

The strongest and most systematic results occur on the square dataset. At a 44-step rollout horizon, both SG-JEPA variants substantially outperform DINO-WM in position, velocity, and cumulative rotation. The GRU variant reduces the three error types relative to DINO-WM by approximately $K$5–$K$6. Original LeWM remains close to DINO-WM despite using a jointly trained encoder.

Figure 3: At 44 rollout steps, SG-JEPA produces lower normalized position, velocity, and cumulative-rotation errors than DINO-WM across the planar evaluation suite.

The advantage increases with rollout horizon. On the square, an SG-JEPA variant leads every reported metric from five through 44 forecast steps. On the triangle, DINO-WM retains an advantage for some short-horizon position and velocity measurements, but SG-JEPA obtains the lowest position and velocity errors at the 44-step horizon. DINO-WM remains best for cumulative rotation on the triangle, which the paper attributes to error amplification at asymmetric contacts.

The per-gravity results support the claim that the gain is not restricted to the center of the training distribution. SG-JEPA is strongest across most of the held-out square gravity grid, with especially clear advantages in velocity and rotation at large gravity values. However, the performance is not uniformly superior at every gravity or for every physical quantity. The triangle’s rotation result and several high-gravity velocity reversals are material exceptions.

The paper also tests transfer to an unseen composite “house” shape constructed from a square and right triangle. Translational dynamics transfer more successfully than rotation. On the Small encoder, the model trained on component shapes has $K$7 the position error and $K$8 the velocity error of a house-trained model near the training gravity range, while the rotation error is substantially worse after autoregressive rollout. This result indicates that the learned representation captures some reusable motion structure, but does not establish shape-independent physical modeling.

Projectile prediction and robotic control

On Approach Ball, SG-JEPA reduces mean position error by approximately $K$9–$K=5$0 relative to DINO-WM and approximately $K=5$1 relative to Original LeWM. The detailed aggregate results are:

Method	Mean position error	Horizon-44 position error	Mean velocity error	Horizon-44 velocity error
Original LeWM	0.0986 m	0.1229 m	1.1361 m/s	1.1400 m/s
SG-JEPA (GRU)	0.0491 m	0.0764 m	0.5391 m/s	0.7829 m/s
SG-JEPA (SSM)	0.0495 m	0.0785 m	0.5664 m/s	0.8861 m/s
DINO-WM	0.0705 m	0.1152 m	0.5672 m/s	0.7411 m/s

SG-JEPA is best at 22 of 25 held-out gravity values for Approach Ball, with the exceptions concentrated at very low gravity, where Original LeWM is sometimes stronger. The velocity results are more qualified than the position results: SG-JEPA is competitive and often better at intermediate horizons, but DINO-WM has lower velocity error at the final horizon. Consequently, the paper’s strongest 3D prediction claim concerns position, not uniformly improved physical-state prediction.

Figure 4: SG-JEPA improves long-horizon projectile prediction and transfers to closed-loop control across held-out gravity values, although individual tasks and physical quantities exhibit reversals.

The representation-level gains transfer to control when the world-model predictor is discarded and a separate diffusion policy is trained. On Arm Catcher Ball, SG-JEPA raises capture success from $K=5$2 to $K=5$3, more than doubling the baseline rate. At the training gravity, the reported rates are $K=5$4 for SG-JEPA and $K=5$5 for DINO-WM. The advantage remains significant across much of the range from $K=5$6 to $K=5$7, although performance declines at extreme gravity values.

On Arm Paddle Ball, success increases from $K=5$8 to $K=5$90.119$\gamma=0.95$00.114$\gamma=0.95$10.635$\gamma=0.95$20.466$$m/s. The larger velocity improvement is relevant to repeated contact, because errors in flight velocity alter the phase and geometry of subsequent impacts.</p> <p>The Franka Paddle-to-Basket task presents a more nuanced result. SG-JEPA increases strict basket-entry success from approximately$$\gamma=0.95$3 to $\gamma=0.95$430.36\%$\gamma=0.95$530.00\%$\gamma=0.95$628.80\%$\gamma=0.95$727.60\%$ for DINO-WM. This localizes the advantage to post-impact trajectory control rather than interception.

There are also pronounced task-specific reversals. At $\gamma=0.95$8, DINO-WM achieves $\gamma=0.95$9 strict entry while SG-JEPA obtains $4.990$0g=16.0$4.990$158.6\%$4.990$20.8\%$$for DINO-WM. The authors appropriately treat these as reproducible but narrow reversals rather than evidence for a general mechanism. They demonstrate that aggregate success rates conceal substantial dependence on gravity and contact regime.</p> <h2 class='paper-heading' id='what-the-model-learns-representation-versus-predictor'>What the model learns: representation versus predictor</h2> <p>The paper’s most informative analysis separates local transition error from its recursive amplification. A linear feature model assumes feature dynamics of the form</p> <p>$$4.990$3

with learned representation $4.990$4. The local defect decomposes into predictor error on retained features and closure error caused by information discarded by the representation. In this formulation, gravity conditioning alone is insufficient: transfer depends on the predictor’s fit, the representation’s predictive closure, and coverage of the relevant law directions in the training distribution.

For affine gravity dependence, the proposed coverage factor grows as

$4.990$5

This makes explicit why narrow gravity sampling creates a difficult extrapolation problem. The result is conditional on a shared finite law basis for the true dynamics, projected dynamics, and learned predictor. The neural predictor is not constrained to obey this basis exactly, so the bound should be interpreted as a sufficient result for an idealized model rather than as a direct guarantee for SG-JEPA.

The empirical crossover study provides stronger evidence about the source of the observed performance. The authors freeze encoders trained with GRU and Transformer predictors, discard the original predictors, and train fresh GRU and Transformer predictors on each representation.

Frozen encoder source	Fresh predictor	Mean rollout error
GRU-trained	GRU	1.376
GRU-trained	Transformer	1.269
Transformer-trained	GRU	1.555
Transformer-trained	Transformer	1.453

The GRU-trained representation remains better under both fresh predictors, with approximately a $4.990$6 reduction relative to the Transformer-trained representation. This is a significant result because it contradicts a predictor-centric interpretation of the main comparison: the benefit follows the representation produced during joint training, not simply the recurrent predictor used at inference.

Teacher-forced evaluation shows that the advantage is already present before autoregressive feedback. At far-OOD gravity values, the GRU-trained representation reduces local error by approximately $4.990$7 relative to DINO-WM. Once predictions are recursively fed back, the gap becomes much larger, peaking near $4.990$8 around horizon 20 while the local gap remains below $4.990$$9. The amplification is nonmonotonic and decreases at longer horizons, consistent with cancellation and regime-dependent error propagation rather than an assumption of uniformly divergent dynamics.</p> <p><img src="https://images.emergentmind.com/paper-images/2609-10464/theory_claim_compact_3panel.png" alt="Figure 5" title="" class="markdown-image" loading="lazy"></p> <p><p class="figure-caption">Figure 5: Encoder–predictor crossover and teacher-forced diagnostics indicate that SG-JEPA’s advantage originates in the learned representation and is amplified by recursive feedback.</p></p> <p>The history ablation gives a related result. Shortening the context from 20 to 4 frames imposes only a$$5.864$0 additional penalty on the Transformer-trained source during the first five forecasts but a $5.864$1 penalty over the full 44-step rollout relative to the GRU-trained source. The result suggests that the GRU-trained representation is more nearly sufficient for recursive prediction with shorter context. It does not prove that GRUs intrinsically produce superior physical representations; the representation is shaped jointly by architecture, optimization, conditioning, and rollout training.

Ablations and adaptation

The hyperparameter studies support the selected configuration but also expose objective mismatch. On the square task, the AdamW learning rate producing the lowest validation objective does not produce the best probe representation. The authors therefore select $5.864$2 for AdamW parameters based on downstream readout quality rather than training loss alone. This supports the broader methodological point that latent objective values are imperfect proxies for physical prediction.

The SIGReg sweep identifies $5.864$3 as the best planar tradeoff. It produces effective rank $5.864$4, compared with $5.864$5 for coefficient $5.864$6, and correlates strongly with later probe quality at early checkpoints. Nevertheless, the paper explicitly notes that SIGReg’s raw loss is not a direct measure of dimensional collapse and that the analysis does not establish SIGReg as the cause of the dynamics advantage.

The history and rollout sweeps favor $5.864$7 and $5.864$8. For Approach Ball, increasing the rollout horizon to $5.864$9 produces a marginally lower 32-step rollout NMSE of $5.403$0 versus $5.403$1, but increases peak host memory from $5.403$2 to $5.403$3 GiB. Longer histories are both less accurate and more expensive in the reported ablations.

Counterfactual gravity interventions verify that the predictor uses the supplied parameter. Holding the physical trajectory fixed while changing only the gravity input increases rollout error, and correct gravity minimizes the error for both GRU and Transformer predictors. This rules out the trivial explanation that gravity is merely an unused conditioning channel.

Sparse post-training experiments further test adaptation. Post-training on support gravities $5.403$4 and evaluating on 13 unseen interpolation values improves SG-JEPA GRU error by $5.403$5 on triangles and $5.403$6 on squares. Across all eight model-shape combinations, mixed post-training with source replay yields a mean reduction of $5.403$7, compared with $5.403$8 for target-only post-training. The improvement is therefore not restricted to the observed support values, although only four of eight mixed configurations satisfy the source-range retention criterion.

Limitations and open questions

The experimental generalization axis is one-dimensional. Gravity is explicitly supplied, held constant within each episode, and varied while the scene interface remains fixed. The results therefore do not establish generalization to unknown physical parameters, vector-valued parameter changes, friction, mass, restitution, or combinations of interacting laws.

The benchmark is also simulation-based and comparatively structured. Contacts create hybrid transitions, and the theory’s clean linear law basis does not fully describe branch changes, event-time errors, frictional effects, or nonlinear history dependence. The authors acknowledge that the theory is a surrogate analysis rather than an identification of the neural encoder or GRU.

Shape transfer is uneven. Translational dynamics transfer from triangles and squares to the unseen house shape, but rotation degrades substantially after recursive rollout. This indicates that the method does not automatically produce object-compositional physical representations.

Finally, several headline comparisons use a single world-model training seed, while control results average multiple policy rollouts rather than multiple independently trained world models. The reported confidence intervals therefore characterize episode or rollout variability more directly than optimization variability. The paper leaves open whether the representation advantage persists across substantially different visual complexity, data scales, predictor capacities, and random initializations.

Conclusion

SG-JEPA combines gravity conditioning, joint encoder–predictor optimization, SIGReg, and recursive latent rollout training. Across planar dynamics, projectile prediction, and robotic control, it generally improves long-horizon performance and extrapolation beyond the training gravity distribution. The strongest numerical gains include a reduction of Approach Ball mean position error from $5.403$9 to $6.922$0 m relative to DINO-WM and an increase in Arm Catcher Ball capture success from $6.922$1 to $6.922$2.

The paper’s principal explanatory result is that the advantage is primarily representation-level: a GRU-trained encoder remains superior after replacement with fresh predictors, and its advantage is visible under teacher forcing before recursive feedback. Autoregressive composition then amplifies this local difference. The evidence supports recursive rollout training as a mechanism for shaping latents toward predictive closure, while the gravity-coverage analysis clarifies that conditioning on a physical parameter is not itself sufficient for extrapolation.