World models are supposed to learn dynamics, not just predict the next frame. But does a model that predicts the next latent state accurately over a single step actually understand physics when that prediction is composed over dozens of steps? And can it generalize to physical regimes it has never seen during training? A team from Yale, Brown, and Jump Trading tests these questions with Semigroup-JEPA, a world model that conditions its latent dynamics on gravity and trains through recursive autoregressive rollouts rather than single-step teacher forcing. The results: up to 2x lower prediction error and 2.5x higher control success compared to prior methods, with a surprising insight that the encoder, not the predictor, is where most of the gain lives.
Why composing single-step predictions isn't enough
Prior JEPA world models like DINO-WM and LeWorldModel (LeWM) train a predictor to minimize the distance between its one-step prediction and the true next latent. This is computationally cheap and stable, but it ignores a fundamental issue: at test time the model never sees true latents. It sees its own predictions, fed back as input. Errors compound. A one-step loss that looks small on paper can accumulate into total drift over a 44-step rollout.
The gap between teacher-forced and free-rollout accuracy is especially painful when the model needs to generalize to unseen physical parameters. A model trained only on Earth-strength gravity (g ≈ 9.8) may produce reasonable single-step predictions at g = 1.63 (Moon gravity), but the trajectory differs qualitatively, floating in slow arcs rather than parabolic drops. If the encoder discards any feature the predictor needs to handle that regime, the error multiplies at every step of rollout.
SG-JEPA addresses both problems simultaneously: it trains through K-step recursive rollouts so the encoder learns features the predictor can carry forward, and it conditions the predictor on the physics parameter (gravity) so the same latent representation supports dynamics across a wide range of regimes.
How gravity enters the latent dynamics
SG-JEPA builds on the LeWM architecture. A Vision Transformer (ViT-Tiny: 12 layers, 3 heads, hidden width 192) encodes each observation frame into a CLS token h_t in R^192, which a one-hidden-layer projector maps to a 256-dimensional latent z_t. The key modification is action conditioning: gravity g is concatenated with any external control action to form an augmented action vector. An action encoder (temporal convolution followed by SiLU MLP) projects this to a 256-dimensional conditioning vector c_t.
The predictor is a GRU or state-space model (Mamba/S6-style) rather than a Transformer. The default GRU variant has 3 residual layers of width 512, with the hidden state concatenated with the action embedding at every layer. The Transformer variant (matching LeWM's original design) uses 6 layers with 16 attention heads and Adaptive Layer Normalization for action injection. Both variants use the same encoder and training procedure; the only difference is the temporal backbone.
For a trajectory at fixed gravity with no external forces, iterating the predictor produces a discrete semigroup on the latent history: S_g(k + ℓ) = S_g(ℓ) ∘ S_g(k). The name "Semigroup-JEPA" reflects this algebraic structure.
Training through recursive rollouts
The training loss has two components. The rollout loss feeds the model's own predictions back as input for K = 5 steps, starting from a context window of H = 20 true encoded frames. Each predicted latent enters the next history window, so the model must remain accurate without teacher forcing. A discount factor γ = 0.95 downweights later predictions so accumulated errors don't dominate the gradient signal:
ℒ_roll = Σ_{k=1}^{K} w_k ‖ẑ_{t+k} − z_{t+k}‖²
where w_k = γ^{k-1} / Σ_{j=1}^{K} γ^{j-1}SIGReg regularizes the latent space by encouraging random one-dimensional projections of encoded latents to match a standard Gaussian. Unlike most contrastive methods, targets come from the same trainable encoder without stop-gradient, so the encoder and predictor are jointly optimized through the full rollout.
The optimizer is a hybrid: Muon applies orthogonalized updates to all 2D tensor parameters at learning rate 1e-4, while AdamW handles everything else at 5e-5. This combination consistently beats pure AdamW by 15-22% on validation loss across predictor types.
Tasks that test real physics understanding
The evaluation spans eight MuJoCo environments in two categories. For prediction: four 2D planar shapes (right triangle, square, pentagon, house) that receive a single impulse and then free-fall inside a box, plus a 3D Approach Ball with projectile flight. For control: three robotic tasks (Arm Catcher Ball, Arm Paddle Ball, Franka Paddle-to-Basket) requiring the arm to interact with a moving ball under gravity.
The critical design choice is gravity sampling. For the 2D tasks and Arm Catcher Ball, training gravity is drawn from N(4, 0.5²). The test set uses a 25-point grid from −2 to 10, spanning everything from reversed gravity (objects fly upward) to strong gravity (objects slam down rapidly). For the 3D tasks, training uses N(9.8, 2.0²), with test values including Pluto (0.62), Moon (1.63), Mars (3.72), and Venus (8.87) alongside Earth and beyond up to 20.
Every episode is 64 frames at 16 Hz (4 seconds). The model takes H = 20 context frames and rolls out for 44 steps at test time, probing position, velocity, and rotation through a frozen MLP probe trained only on the training episodes.
Results: halving error, more than doubling success
On 2D Square at a 44-step rollout horizon, SG-JEPA (GRU) reduces position error by 31%, velocity error by 48%, and rotation error by 37% relative to DINO-WM. The GRU and SSM variants perform comparably and both dominate across the full gravity grid. On triangles, DINO-WM leads slightly at short horizons for position and velocity, but SG-JEPA overtakes it at longer rollouts, consistent with the rollout training reducing error accumulation.
The 3D results are stronger. On Approach Ball, SG-JEPA reduces mean position error by about 34% versus DINO-WM and 50% versus Original LeWM. From Figure 4, DINO-WM performs comparably at short horizons but drifts more quickly during rollout, while SG-JEPA variants stay accurate longer. The advantage holds across 22 of 25 test gravity values, with the exceptions at very small gravity where Original LeWM edges ahead.
For control, the model's frozen encoder feeds a separate Diffusion Policy. On Arm Catcher Ball, SG-JEPA (GRU) raises capture success from 9.5% to 23.3%, a 2.45x improvement. The gain peaks near the training gravity and extends across most of the held-out range. On Franka Paddle-to-Basket, SG-JEPA converts more paddle contacts into basket entries (27.4% to 30.5%). On Arm Paddle Ball, success rises from 17.7% to 23.8%.
Where the advantage actually lives
The most interesting finding is not that SG-JEPA works better, but where the improvement comes from. The authors develop a fresh-predictor crossover experiment: starting from a trained checkpoint, they freeze the encoder and replace the predictor with a freshly initialized one, then retrain the predictor alone. They do the reverse as well, replacing the encoder while keeping the trained predictor.
The results are striking. The GRU-trained encoder retains about 12% lower rollout error than the Transformer-trained encoder, regardless of which fresh predictor is attached. Replacing the GRU-trained encoder with a fresh one degrades performance; replacing the Transformer-trained encoder with a fresh one improves it. The GRU-trained encoder's advantage survives even when paired with a fresh Transformer predictor. Meanwhile, the Transformer-trained encoder actually has lower latent MSE on the raw predictions, but this doesn't translate into better physical-state predictions when probed downstream.
The explanation ties back to the linear feature model. The one-step prediction error at an unseen gravity g* decomposes into a closure term (features the encoder discards but the dynamics need) and an operational error (the predictor's own approximation). A coverage bound shows that the closure term scales with L_law(g*) = 1 + (g* − μ_tr)²/σ_tr², which grows quadratically with distance from the training distribution. The encoder's job is to minimize this closure term by retaining the features the dynamics actually depend on. Training through recursive rollouts forces exactly this: if the encoder drops a dynamical feature, the error compounds at every step, and the backpropagated gradient penalizes the encoder for it.
The fresh-predictor crossover confirms this empirically. The GRU-trained encoder already has the better feature selection before any predictor is attached. And when predictions are fed back, the small local advantage (about 32% lower one-step error on far-OOD gravity values) amplifies into a much larger long-horizon gap, peaking near 0.38 around horizon 20 while the local gap stays below 0.06.
Ablations and implementation details
The SIGReg coefficient λ_SIG = 0.72 gives the best balance of effective rank (138), probe accuracy, and rollout quality. Higher values increase rank but hurt downstream prediction. The rollout discount γ = 0.95 outperforms γ = 1.0 (no discount) and γ = 0.90, suggesting that slightly downweighting later predictions prevents the early-error domination that destabilizes training.
History length H = 20 with predictor depth L = 3 is the best performance-cost tradeoff. Going to L = 4 improves metrics by only 2-4% while adding 4.7M parameters. On Approach Ball, doubling K from 5 to 10 gives 1.3% improvement on 32-step rollout accuracy but uses 73.5 GiB versus 61.3 GiB of RAM.
Sparse-gravity post-training, where the pretrained model is fine-tuned on just four gravity values {0, 2, 6, 8}, reduces error by 13.6% on the square and 20.8% on the triangle across 13 held-out interpolation gravities. Mixed post-training (keeping some original training data alongside the new values) is substantially more consistent than target-only adaptation, averaging 14.1% improvement versus 6.9%.
On conventional OGBench and DMC control tasks, the GRU predictor matches the Transformer on success rate (70% vs 68% on Cube, 98% vs 100% on Reacher) with lower planning latency (0.111s vs 0.195s), suggesting the GRU choice is not a compromise.
Limitations and open questions
The experiments vary only gravity as a physical parameter. Vector-valued physical quantities (friction coefficients, material stiffness, mass ratios) or having the model infer dynamical parameters directly from observation rather than receiving them as inputs would be a natural extension. Shape generalization is uneven: training on triangles and squares transfers translational dynamics of the house-shaped composite, but rotation does not transfer well, and the pentagon remains difficult.
The theoretical analysis uses a linear feature model, while the actual predictor is nonlinear and history-dependent. Contacts can switch transition branches, violating the smooth dynamics assumption. A theory incorporating nonlinearities and action-dependent perturbations is needed to close the gap between the linear analysis and the empirical behavior.
The control tasks use separate diffusion policies trained on frozen encoder features, so the world model and the policy are decoupled. Whether joint training of the world model and policy would further improve OOD control remains untested.
For a working developer, the practical takeaway is straightforward: if you're building a world model for planning or control, train it through multi-step recursive rollouts rather than single-step teacher forcing. The cost is modest (K = 5 steps with discount), and the benefit is an encoder that retains the features the dynamics actually need. The frozen encoder then serves as a general-purpose feature extractor for downstream policies, and the GRU backbone is fast enough for real-time replanning at 0.11s per step.