SG-JEPA and Physics-Aware World Models
SG-JEPA improves world model prediction by incorporating physical parameters as explicit inputs to the encoder, enabling more stable latent dynamics.
SemiGroup-JEPA (SG-JEPA) improves the predictive accuracy of world models by treating physical parameters as explicit inputs during training. By extending the LeWorldModel framework with action-conditioning and a multi-step autoregressive rollout, the model achieves a two-fold reduction in open-loop prediction error compared to DINO-WM architectures on 2D datasets. This matters because it moves world modeling away from purely observation-based learning toward architectures that treat governing dynamics as learnable, manipulatable parameters.
The Bottleneck is the Encoder
Most world models learn to predict future states based on past observations. When environments change in ways the model has not seen, performance degrades. This failure is often blamed on the predictor, but the authors argue that the bottleneck is the encoder. The encoder in previous architectures often blends environmental noise with core dynamic features, making it impossible for the predictor to generalize across varying physical conditions. SG-JEPA forces consistency by requiring the encoder to accept physical parameter inputs alongside observations. This explicit conditioning forces the latent space to map observations to a representation that is structurally dependent on the chosen parameter, preventing the model from collapsing different physical regimes into a single, inaccurate latent projection. The authors' analysis using a linear feature model suggests that the performance gain comes from the encoder learning specific, stable features that the predictor can reliably carry forward through time.
How Multi-Step Loss Filters Dynamics
To force this representation, SG-JEPA uses a multi-step autoregressive rollout loss back-propagated into the encoder. In standard models, single-step prediction often relies on correlations that vanish over time. By forcing the model to predict multiple steps into the future based on a fixed set of physical parameters, the system suppresses transient noise. The loss function measures the discrepancy between predicted and actual states across a long horizon, creating a gradient that punishes features that do not evolve according to the conditioning parameters. Only features that maintain a stable mapping under the provided physical variables survive this process. If a feature fails to follow the consistent dynamic path defined by the input parameter over multiple steps, the loss increases. This ensures that the latent representation remains sensitive to the underlying governing dynamics.
| Metric | DINO-WM | SG-JEPA |
|---|---|---|
| 2D Open-Loop Prediction Error | 1.0x (Baseline) | 0.5x |
| 3D Robotic Control Success Rate | 1.0x (Baseline) | 2.5x |
Implementation in Robotic Systems
For those building robotic systems, this result signals that physics should not be treated as a hidden variable for the model to discover. Supplying physics parameters as explicit conditionings changes the latent structure. Imagine a robotic arm moving through water versus air. Instead of the model seeing two different pixel-level trajectories and struggling to reconcile them, it sees the same motion path conditioned on a viscosity variable. This allows the model to learn the concept of drag as a structural rule rather than a visual curiosity. The latent space becomes a set of steerable dynamics, where changing the conditioning variable adjusts the predicted rollout accordingly.
Consider the example of a simulated cart-pole system where you vary the mass of the pole. Without SG-JEPA, a model might struggle to predict the swing frequency as the mass changes, treating the variance as a loss of signal. With SG-JEPA, providing the mass as an input constant allows the latent transition function to account for the inertia shift directly. The system effectively learns the mass-dependence of the motion as a functional rule, rather than trying to memorize every possible mass value as an independent visual sequence.
This approach relies on the assumption that physics can be cleanly separated from environmental noise. In complex real-world scenarios, where contacts and sensor noise are non-linear, the separation of error and recursive amplification is much harder to maintain. It is not yet clear if the encoder can isolate these variables when the underlying dynamics are not perfectly synthetic. We also do not yet know how this architecture performs when the provided physical parameters are noisy or only partially observable, which is the standard state for most robotics hardware today.