PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Andy Zeyi Liu, Haoran Sun, Lucas Baker, Randall Balestriero, John Sous

arXiv:2609.10464v1cs.LGcs.AIcs.CV

TL;DR

JEPA world models had not been established as learners of physically realistic, OOD dynamics. SG-JEPA adds gravity conditioning and recursive encoder–predictor training, then evaluates prediction and control across gravity ranges; it outperforms prior baselines, with the analysis attributing most gains to better encoder features. The study is limited by its single scalar physical variable, uneven shape transfer, and a linear theory for nonlinear, history-dependent dynamics.

  • Problem

    JEPA world models had not been tested for learning physical dynamics that remain accurate under long-horizon rollout and OOD gravity.

  • Method

    SG-JEPA gravity-conditions LeWM and jointly trains its encoder and predictor through recursive latent rollouts, evaluated on prediction and robot control tasks.

  • Results

    SG-JEPA outperforms Original LeWM and DINO-WM on long-horizon and OOD prediction, while frozen SG-JEPA representations improve success across all three manipulation tasks.

  • Takeaways & Limitations

    The analysis indicates that multi-step rollout training mainly improves the encoder’s retention of dynamics-relevant features rather than the predictor’s dynamics modeling.

  • Takeaways & Limitations

    The study varies only scalar gravity, finds uneven transfer across object shapes, and analyzes nonlinear history-dependent models with a linear theory.

Abstract

from arXiv · show

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.

1 Introduction

SG-JEPA tests whether latent world models learn physical dynamics that remain accurate under long rollouts and unseen gravity values. It combines gravity conditioning with recursive latent training and finds that representation quality is central to OOD generalization.

  • Research questions: The study asks whether learned physical dynamics remain accurate when model updates are composed over long horizons and transferred to OOD gravity values.Gravity is sampled from a narrow training range and evaluated across a wider grid.
  • Significance: A world model that genuinely learns dynamics could help agents adapt to new environments, such as transferring Earth-trained dynamics to Martian conditions.This consequence motivates testing physical generalization beyond the training regime.
  • Motivation and approach: SG-JEPA extends LeWM with gravity conditioning and jointly trains its encoder and predictor through discounted autoregressive latent rollouts.SIGReg regularizes the latent space, and evaluation spans 2D rigid-body, 3D projectile, and robot-arm tasks.
  • Analysis: The analysis separates one-step, gravity-conditioned transition error from its recursive amplification during rollout.A linear feature model is used to identify where SG-JEPA’s long-horizon advantage arises.
  • Contribution: Recursive latent rollout loss addresses both OOD gravity generalization and long-horizon error accumulation.Unlike one-step prediction, the loss back-propagates through repeated latent predictions.

2 Problem Setting

The problem setting combines gravity-conditioned latent dynamics with evaluation on passive motion, projectile motion, and robot control. Prediction is measured through physical-state probes over open-loop rollouts, while frozen representations support separately trained diffusion policies.

  • Model architecture: SG-JEPA uses an encoder and gravity-conditioned predictor, with gravity supplied through the action-conditioning pathway even for action-free trajectories.The encoder maps observations to low-dimensional latents, while predictor variants include Transformer, GRU, and SSM architectures.
  • Semigroup dynamics: Repeated application of the shared history update forms a discrete semigroup for fixed gravity in action-free trajectories.With controls, successive update blocks still compose, but each block is modified by its action.
  • Training objective: The training objective recursively predicts K future latents from H encoded frames, feeding each prediction into the next history window instead of using only teacher-forced one-step prediction.The rollout loss is combined with SIGReg, which regularizes encoded latents.
  • Environments: Experiments cover eight MuJoCo datasets spanning freefall, projectile motion, and robot-arm control, with narrow gravity training distributions and wider-grid testing that includes OOD values.Five datasets evaluate prediction and three evaluate control.
  • Prediction evaluation: Prediction uses frozen world models and MLP probes to estimate physical states from latent windows at open-loop horizons up to 44 frames.The reported excess error subtracts probe error measured on encoded ground-truth latents.
  • Control evaluation: Control freezes each visual encoder and trains a separate gravity-conditioned diffusion policy on its features, executing only the first E predicted actions before replanning.The policy uses a recurrent feature summary and gravity conditioning.

3 Empirical results

SG-JEPA improves long-horizon physical prediction and closed-loop control across 2D and 3D gravity-conditioned tasks, especially at held-out gravity values. Gains vary by task because gravity changes the dominant control bottleneck.

  • 2D prediction: 31–48%: SG-JEPA reduces the three long-horizon error types on the 2D square dataset relative to DINO-WM.An SG-JEPA variant leads every square metric from 5 through 44 rollout steps.
  • 3D prediction: 34%: SG-JEPA reduces mean Approach Ball position error relative to DINO-WM, with the advantage persisting across longer rollouts.SG-JEPA is best at 22 of 25 tested gravity values, except at very small gravities where Original LeWM performs best.
  • 3D control: 9.5% to 23.3%: SG-JEPA raises Arm Catcher Ball capture success across a range of in-distribution and OOD gravity values.The gain is largest and most consistent for this task, extending from g = 0 to g = 9.
  • 3D control: 27.4% to 30.5%: SG-JEPA increases Franka Paddle-to-Basket average success by converting more paddle contacts into basket entries.DINO-WM has a slightly higher paddle-hit success rate, so the main difference occurs after impact.
  • 3D control: 17.7% to 23.8%: SG-JEPA raises Arm Paddle Ball success, while gravity exposes distinct bottlenecks across interception, post-impact targeting, and repeated contact regulation.The 3D control evaluations use task-specific Diffusion Policies trained on frozen world-model features and are averaged over five seeds.

4 Where Does the Long-Horizon Advantage Come From?

The analysis separates local law-conditioned representation error from its recursive amplification during rollout. Evidence indicates that SG-JEPA’s advantage chiefly comes from the encoder learning a representation that preserves dynamics, while feedback magnifies a smaller one-step advantage.

  • Theoretical model: A linear feature model separates one-step local transition error from recursive rollout amplification, allowing the two sources of OOD advantage to be measured separately.The model treats the closure term as information about the next latent that depends on features discarded by the representation.
  • Theoretical model: Predictive closure requires the current latent to contain all information needed for the conditional mean of the next latent at a given gravity.When C_W(g) = 0, the representation is predictively closed at gravity g.
  • Theoretical model: Unseen-gravity accuracy depends on law coverage, predictor fit, and representation closure, not gravity conditioning alone.For affine free-flight dynamics, the bound grows with the distance of the test gravity from the training-gravity mean.
  • Rollout mechanism: One-step errors of equal size can have different rollout consequences because learned updates transform and propagate them; multi-step training penalizes this distinction directly.The free-rollout error is governed by repeated composition of the true and learned discrete evolution operators.
  • Empirical tests: 12%: the GRU-trained encoder lowers mean rollout error with either a fresh GRU or Transformer predictor, showing the gain follows the representation rather than the original predictor.Fresh-GRU errors are 1.376 versus 1.555, and fresh-Transformer errors are 1.269 versus 1.453.
  • Empirical tests: 0.38: the free-rollout gap peaks near horizon 20, while the teacher-forced local gap remains below 0.06.The amplification begins from horizon 2 and shrinks again at long horizons, so it is not monotone.
  • Empirical tests: 32%: SG-JEPA lowers far-OOD one-step error before feedback, supporting the conclusion that its OOD advantage lies in the encoder representation learned during GRU training.The advantage survives replacing the predictor and is already present under teacher forcing.

5 Conclusion

SG-JEPA combines gravity conditioning with recursive latent-rollout training and improves long-horizon, out-of-distribution prediction and downstream control across the evaluated 2D and 3D tasks. Its gains are attributed primarily to the learned encoder representation, while the study identifies limits in physical-variable scope, shape transfer, and theoretical modeling.

  • SG-JEPA combines gravity conditioning with joint encoder–predictor training through recursive latent rollouts.
  • Across 2D and 3D prediction and robot control, SG-JEPA outperforms Original LeWM and DINO-WM at longer horizons and out-of-distribution gravity values.
  • The same frozen representation transfers to closed-loop control, where diffusion policies improve success rates on all three manipulation tasks.
  • Controlled retraining studies indicate that SG-JEPA’s advantage lies primarily in the encoder representation rather than the original predictor.
  • The experiments vary only scalar gravity, while shape transfer is uneven and the linear theory does not model neural nonlinearities, history dependence, or contact-perturbed branches.

B Datasets

The benchmark contains eight MuJoCo datasets spanning planar rigid-body prediction, projectile prediction, and robotic control. Episodes record observations, actions, physical states, and task-specific outcomes under sampled and held-out gravity conditions.

  • The benchmark evaluates eight MuJoCo datasets: four planar rigid-body prediction tasks, one 3D projectile task, and three 3D robotic control tasks.
  • The 2D tasks use rigid bodies with randomized initial poses and a single planar impulse inside a box, while recording an eight-dimensional physical state.
  • Approach Ball predicts free projectile motion, whereas Arm Catcher Ball, Arm Paddle Ball, and Franka Paddle-to-Basket require arm-mediated interception, bouncing, or basket placement.
  • Gravity is constant within an episode, sampled from narrow training distributions and evaluated on wider held-out grids, and is supplied as a normalized action coordinate.
  • Each episode contains 64 frames at 16 Hz, with 2D videos at 128×128 and 3D videos at 256×256 resolution.

C Model and training details

SG-JEPA uses jointly trained visual encoders and temporal predictors with recursive latent rollouts, SIGReg, and action conditioning. The reported implementation fixes rollout and history settings while selecting optimization and regularization configurations through validation studies.

  • The encoder is a ViT-Tiny model, and the temporal model receives action information, including gravity, throughout prediction.
  • During training, predicted latents are recursively fed back into a length-H history window for K steps, with H = 20, K = 5, and discount γ = 0.95 in final models.
  • The baseline LeWM uses teacher-forced one-step latent MSE, whereas DINO-WM freezes a pretrained DINOv2 encoder and trains a causal Transformer predictor.
  • The hybrid Muon/AdamW optimizer lowers best validation objectives relative to AdamW for GRU, SSM, and Transformer predictors.
  • A SIGReg coefficient of 0.72 provides the best balance of effective rank, frozen-probe accuracy, and rollout quality in the reported 2D experiments.
  • History length H = 20 gives the best long-horizon accuracy, while H = 30 worsens real-latent control errors despite lower training objectives.

D.2.2 Approach Ball: history and training rollout horizon

Approach Ball ablations favor a moderate history and rollout configuration that balances prediction accuracy with resource use. Gravity interventions further show that predictors use the supplied physical parameter, while representation-level analyses attribute recurrent advantages to the learned encoder.

  • Approach Ball ablation: (H, K, λSIG) = (20, 5, 0.18) gives the lowest probe-validation NMSE, direct-probe test NMSE, and 20-step rollout NMSE for Approach Ball.
  • Rollout discount: At horizon 44, γ = 0.95 gives the lowest joint NMSE and velocity error on Arm Catcher Ball, with position error 0.122 versus the best 0.118 at γ = 1.00.
  • Gravity conditioning: For both GRU and Transformer predictors, rollout error is minimized when supplied gravity matches true gravity, and larger mismatches generally increase error.
  • Representation analysis: The recurrent advantage follows the source representation: crossover studies show that the GRU-trained representation yields lower physical-state error with either freshly fitted predictor.
  • Downstream control: SG-JEPA GRU and SSM predictors preserve downstream control success while achieving lower planning latency and higher effective rank on both evaluated datasets.

E.2 Shape generalization

SG-JEPA transfers translational dynamics from triangle and square components to an unseen house shape, with position generalizing more reliably than velocity or rotation. Recursive prediction, especially for rotation, remains the main source of error under shape and gravity shifts.

  • Transfer results: SG-JEPA transfers substantial house translational dynamics from triangle and square training shapes, while rotation remains the hardest quantity to predict.Position transfers most consistently, whereas velocity is less stable; direct house training remains better near the training distribution.
  • OOD gravity: Across much of both gravity tails, T+S closely follows the house-trained reference on position despite never seeing a house during training.With the Small encoder, T+S matches or outperforms the house-trained model at several extreme gravity values.
  • Error localization: 1.40× position error, 1.70× velocity error, and 3.36× cumulative-rotation error occur after rollout for Small T+S relative to the house-trained reference.The corresponding true-latent probing ratios are 0.90, 1.23, and 1.12, showing that recursive prediction amplifies the remaining gap, especially for rotation.
  • Evaluation setup: Figure 19 reports house rollout error at h44 for Tiny and Small encoders across position, velocity, and cumulative-rotation error.

F.1 Approach-Ball results

On Approach-Ball, SG-JEPA provides its clearest advantage for position prediction across long rollouts and held-out gravities, while velocity performance is competitive but not uniformly better than DINO-WM. The evaluation reports errors over 44 forecast steps and 25 held-out gravity values.

  • Task scope: Approach-Ball contains only position and velocity as physical quantities because it has no ball rotation or friction.
  • Position prediction: SG-JEPA GRU and SSM achieve the two lowest position errors and reduce mean position error by roughly 30% relative to DINO-WM.They are best both on average and at the final forecast step.
  • Velocity prediction: Velocity performance is not uniformly better: GRU is slightly lower than DINO-WM on rollout-average error, while DINO-WM is more accurate at h44.GRU or SSM is lower over much of the intermediate rollout, but DINO-WM closes the gap near the end.
  • Gravity generalization: GRU is strongest across most central and moderate-gravity values, DINO-WM at the high-gravity tail, and Original LeWM at the few lowest values.These gravity-specific results support a clear SG-JEPA advantage for position but only competitiveness for velocity.
  • Evaluation setup: Figure 20 evaluates mean ball-velocity L2 error at each recursive forecast step and mean h1–h44 error across 25 held-out gravity values.The training distribution is g ∼N(9.8, 2.02).

F.2 Arm-Catcher-Ball control results

SG-JEPA substantially improves Arm-Catcher-Ball control over DINO-WM across gravity values, with gains that extend beyond the training distribution. The task’s phase-sensitive dynamics make velocity and contact timing important for successful capture.

  • Evaluation protocol: The evaluation tests 23 gravity values with five paired simulator-rollout seeds and requires the ball to enter and be latched in the catcher.Each checkpoint receives a separate diffusion policy conditioned on a 20-frame latent history.
  • Control success: SG-JEPA more than doubles overall capture rate and improves on DINO-WM by 13.8%.At training gravity, SG-JEPA reaches 42.3% versus 16.0% for DINO-WM; gains remain positive across every gravity.
  • Evaluation setup: Figure 21 reports overall capture rate and method comparisons by gravity, averaged over five rollout executions.
  • Method comparison: Both SG-JEPA predictor variants remain ahead of Original LeWM and DINO-WM overall.The comparison includes GRU and SSM SG-JEPA variants alongside the two baselines.
  • Task mechanism: Successful capture requires preserving ball velocity and contact phase, not merely the current position.As gravity increases, shorter bounce intervals amplify timing errors and can shift the trajectory to the wrong contact phase.

F.3 Franka paddle-to-basket results

For Franka paddle-to-basket control, SG-JEPA achieves stronger strict-entry outcomes than the baselines, especially because success depends on the post-impact ball state. Open-loop probes also show lower long-horizon position and velocity errors, while per-gravity results reveal sharp task-specific reversals.

  • Metrics: Table 11 distinguishes relaxed success, strict entry, and strict entry conditioned on valid paddle-blade contact.All values are percentages.
  • Control outcomes: Both SG-JEPA pipelines achieve the highest strict-entry rates, although DINO-WM records slightly more hits and the highest relaxed-success rate.The SG-JEPA advantage appears after contact, where the policy must produce the appropriate outgoing position and velocity.
  • Caveat: The two large per-gravity reversals are reproducible across five rollout seeds, but two gravity values do not support a broader mechanistic conclusion.
  • Task mechanism: High paddle-hit rates indicate that interception is not the main bottleneck; post-impact timing, orientation, and incoming velocity determine whether the ball reaches the basket.Weak gravity can prolong flight and cause overshoot, whereas strong gravity demands a larger, more precisely directed impulse.
  • Open-loop evaluation: Figure 23 reports mean ball-position and ball-velocity Euclidean error over 5,000 held-out episodes, with SG-JEPA lower in position from h9 onward and in velocity at every step.
  • Open-loop prediction: SG-JEPA has lower position error at 36 of 44 forecast steps and lower velocity error at all 44.At h44, position error falls from 0.119 to 0.114 m and velocity error from 0.635 to 0.466 m/s.

G Sparse-Gravity Post-Training Ablation

The ablation tests whether sparse-gravity post-training improves long-horizon prediction at unseen gravities and whether source replay preserves broader gains. It also develops a linear feature analysis of local error, rollout amplification, and representation choice.

  • Sparse-gravity post-training: Mixed post-training reduces aggregate error in all eight shape–model combinations and outperforms target-only post-training in every combination.Mean reduction rises from 6.85% to 14.05%, a gain of 7.21 percentage points.
  • Sparse-gravity post-training: 20.84% improvement occurs for SG-JEPA (GRU) on triangles, while square SG-JEPA (GRU) improves by 13.60%.The largest mixed improvement for squares is 20.81% for Original LeWM.
  • Sparse-gravity post-training: Post-training reduces error at many unobserved gravities between support values, with the largest separations generally near low- and high-gravity ends.The evaluation uses 13 held-out interpolation gravity values and tests source-range retention against each pretrained baseline.
  • Linear feature analysis: The linear feature model separates one-step error into predictor error, representation-discarded dynamical information, and a conditionally zero-mean residual.Its rollout analysis tracks how local error and transition residuals are transformed by later predicted dynamics.
  • Linear feature analysis: A multi-step loss applies a different quadratic geometry to one-step transition error, allowing representations to exchange rank as the training horizon changes.The analysis is exact for its linear model but does not identify the neural encoder or GRU with linear operators.

H.3 How gravity coverage controls test error

The theory relates unseen-gravity error to coverage of the gravity-dependent transition law under shared-basis assumptions. For affine free-flight dynamics, it yields an explicit coverage factor, while emphasizing that the result is a sufficient bound rather than a general guarantee.

  • Gravity coverage: The semigroup law composes evolution over time at fixed gravity, whereas generalization across gravity depends on how the one-step operator varies with gravity and which law directions training covers.Low error at training gravities alone gives no bound at an unseen gravity.
  • Gravity coverage: The law-coverage factor is large when the test gravity lies in a weakly covered direction of the training design.The analysis assumes the training distribution covers every direction in the chosen gravity-law basis.
  • Local test error: Theorem H.4 bounds unseen-gravity local error using training-law errors under shared-basis assumptions and bounded latent and physical feature states.The proof applies transfer bounds separately to representation non-closure and predictor transition error.
  • Affine gravity: For collision-free ballistic motion, gravity enters the transition affinely, yielding a two-dimensional law basis and an explicit mean–variance coverage factor.The planar experiments use a clipped Gaussian gravity distribution, so the exact factor uses moments of that clipped distribution.
  • Scope of the result: The affine free-flight result applies only before contact and to displayed physical-coordinate states; arbitrary learned features and contact-induced branch changes need not preserve it.Tangential components have no gravity forcing in the idealized free-flight model.

H.3.1 When the coverage result does not apply

The coverage result fails when training does not observe relevant law directions, when analytic conditioning is mistaken for extrapolation, or when rollout composition changes which representation is preferred. The resulting bounds must therefore account for unseen-law error, recursive propagation, residuals, and horizon-dependent defect weighting.

  • Missing law directions: A singular law-feature covariance permits zero average training error but nonzero error at an unseen gravity when the test law has an uncovered direction.This failure is avoided when the test feature lies in the covariance’s covered range.
  • Finite-design coverage: Finite observed gravity values can identify a residual family only under a full-rank feature design, with transfer controlled by the corresponding leverage factor.The finite-design result uses the observed feature matrix and its inverse Gram matrix to bound prediction at the target gravity.
  • Analytic conditioning: Analytic conditioning alone provides no extrapolation guarantee because two analytic conditional models can agree at every observed gravity yet differ arbitrarily at an unseen one.Supplying gravity to the model is therefore insufficient without additional structural assumptions.
  • Recursive rollout: The exact rollout recursion shows that each local defect is transformed by later learned iterates, so propagated errors may reinforce or cancel rather than grow monotonically.Teacher forcing measures local error before predictions are reused, whereas free rollout exposes recursive amplification or cancellation.
  • Physical error bounds: The physical-error bound separates readout residual, unseen-law error, recursive propagation, and transition residuals, but remains an upper bound rather than a stand-alone law-coverage guarantee.Its statement retains trajectory bounds, rollout gain, and realized transition residuals.
  • Horizon-dependent preference: Multi-step objectives can change representation preference with rollout horizon: one-step profiling prefers WP, while horizon-two and full profiling can prefer WF.The construction proves horizon-dependent ranking changes, not that longer rollouts always help or that the toy mechanism is learned by SG-JEPA.
  • Scope limits: The theoretical constructions have limited scope: the history bridge does not establish global nonlinear GRU behavior, and the hinge result excludes discontinuous, frictional, restitution-dependent, and multi-contact dynamics.The SIGReg calculation likewise does not prove noncollapse or attribute the dynamics difference to SIGReg.
Loading 2609.10464v1…