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

Whether JEPA world models learn physical dynamics that remain accurate under recursive rollout and out-of-range gravity has been untested. SG-JEPA conditions latent prediction on gravity and trains the encoder and predictor through recursive rollouts; it outperforms Original LeWM and DINO-WM across 2D and 3D prediction and robot control at longer horizons and out-of-distribution gravity values.

  • Problem

    Whether JEPA world models learn physical dynamics that remain accurate under recursive rollout and out-of-range physical parameters has not been established.

    Source: §1 · pp. 1, 3

  • Method

    SG-JEPA conditions latent prediction on gravity and jointly trains its encoder and predictor through recursive latent rollouts.

    Source: §1 · p. 2, §5 · p. 11

  • Results

    Across 2D and 3D prediction and robot control, SG-JEPA outperforms Original LeWM and DINO-WM at longer horizons and OOD gravity, with largest gains far from training gravity.

    Source: §5 · p. 11

  • Takeaways & Limitations

    Similar performance from the GRU and SSM variants suggests SG-JEPA's improvement does not depend on a single temporal architecture.

    Source: §3 · p. 8

    The experiments vary only scalar gravity, leaving generalization to vector-valued physical variables untested.

    Source: §5 · p. 11

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 a gravity-conditioned JEPA world model can learn dynamics that remain accurate over long rollouts at out-of-distribution gravity values. The introduction presents this as jointly addressing physics generalization and recursive rollout error.

1 Introduction

  • The motivating question is whether learned updates remain physical when composed over long horizons and generalize beyond training dynamics, enabling adaptation to new environments such as Mars.Source: §1 · pp. 1–2
  • Training on a narrow gravity range, SG-JEPA’s learned dynamics remain accurate over long-horizon rollouts at OOD gravity values, demonstrating zero-shot physics generalization.The evaluation spans a wider grid of gravity values than the training range.Source: §1 · p. 2
  • SG-JEPA extends LeWorldModel by conditioning on gravity alongside action and jointly training the encoder and predictor with discounted autoregressive latent rollout loss and SIGReg.The rollout loss is back-propagated through both encoder and predictor.Source: §1 · pp. 2–3
  • Experiments use MuJoCo 2D rigid-body freefall, 3D projectile motion, and robot-arm control, training on a narrow gravity range and testing a wider grid.The study evaluates both transfer to OOD gravity and long-horizon rollout accuracy.Source: §1 · p. 2
  • Replacing the trained predictor with a fresh one degrades performance, while replacing the encoder does not; the reported advantage widens with gravity distance and rollout horizon.The introduction attributes the advantage to what the GRU predictor learns.Source: §1 · p. 2
  • The work addresses OOD physics generalization and rollout error jointly; for action-free trajectories at fixed gravity, repeated updates form a discrete semigroup.Source: §1 · p. 3

2 Problem Setting

SG-JEPA conditions latent prediction on gravity and trains through recursive rollouts, then evaluates physical-state prediction and learned control on MuJoCo tasks spanning varied gravitational fields.

2.1 Model architecture

  • SG-JEPA encodes frames into projected CLS latents and supplies signed episode gravity as an action coordinate to a temporal predictor, using GRU or SSM variants by default.Source: §2.1 · p. 3
  • Its shared history updates form a discrete semigroup, while training recursively feeds K predicted latents into later windows with a K-step rollout loss and SIGReg.Source: §2.1 · p. 4

2.2 Datasets

  • Eight MuJoCo datasets cover passive rigid-body motion, projectile motion, and robot-arm control; training uses narrow gravity distributions, while held-out tests span a wider grid including OOD values.Source: §2.2 · p. 5

2.3 Downstream evaluation

  • For prediction, an MLP probe maps latent windows to physical state, and held-out encoded frames and open-loop rollouts are assessed up to 44 frames using excess state error.Subtracting probe error on encoded ground-truth latents isolates error introduced by rollout.Source: §2.3 · p. 5
  • For control, each model’s frozen visual features train a separate gravity-conditioned diffusion policy; policies infer moving-ball motion from observation history and gravity without a future target frame.Source: §2.3 · pp. 5–6

3 Empirical results

SG-JEPA improves long-horizon prediction and control over LeWM and DINO-WM, especially at out-of-distribution gravity values. Its gains span 2D and 3D tasks, though performance varies by shape, gravity, and task.

2D datasets

  • On squares, SG-JEPA (GRU) reduces all three 2D prediction error types by 31–48% relative to DINO-WM, and an SG-JEPA variant leads every metric from 5 through 44 rollout steps.The advantage is clearest over longer rollouts; SG-JEPA also has the lowest error across most held-out square gravity values, especially for velocity and rotation at large gravity.Source: §3 · pp. 6–7

3D datasets

  • SG-JEPA reduces Approach Ball mean position error by approximately 34% relative to DINO-WM and raises Arm Catcher Ball capture success from 9.5% to 23.3%.The results cover prediction and closed-loop control across gravity values.Source: Figure 4 · p. 7, §3 · pp. 7–9

3D datasets — Control

  • SG-JEPA (GRU) raises success from 27.4% to 30.5% on Franka Paddle-to-Basket and from 17.7% to 23.8% on Arm Paddle Ball.On Franka, it converts more paddle contacts into basket entries; on Arm Paddle Ball, it more often avoids floor contact and completes qualifying bounces. High gravity shortens flight time and lowers apex height, making the 0.28 m threshold harder to reach.Source: Figure 4 · p. 7, §3 · p. 8

2D datasets

  • On triangles, SG-JEPA has the lowest 44-step position and velocity errors despite DINO-WM’s short-horizon edge, but DINO-WM has slightly lower cumulative-rotation error.On an unseen house shape, frozen SG-JEPA representations support reasonable 44-step position and velocity predictions, particularly with a larger ViT encoder.Source: §3 · p. 7

3D datasets — Prediction

  • On Approach Ball, SG-JEPA performs best at 22 of 25 test gravity values; Original LeWM leads only at very small gravity, while DINO-WM drifts faster over longer rollouts.GRU and SSM perform similarly and alternate as best across horizons; SG-JEPA also retains an advantage for velocity generalization.Source: §3 · p. 8

3D datasets — Control

  • Gravity exposes task-specific control bottlenecks: interception timing, post-impact targeting, and repeated-contact regulation.For Arm Catcher Ball, the capture gain extends from g = 0 to g = 9; for Franka Paddle-to-Basket, SG-JEPA converts more contacts into basket entries despite DINO-WM’s slightly higher paddle-hit rate.Source: §3 · p. 8

4 Where Does the Long-Horizon Advantage Come From?

SG-JEPA’s OOD advantage follows the representation learned during GRU training: it appears before recursive feedback and grows under rollout. The linear feature model attributes transfer to predictor fit, law coverage, and representation closure, not gravity conditioning alone.

Theoretical model

  • The linear feature model separates one-step transition error from its recursive propagation and identifies discarded features as a source of teacher-forced error when the representation is not predictively closed.A multi-step objective penalizes how learned updates transform one-step errors, unlike a one-step objective.Source: §4 · p. 9
  • Gravity conditioning alone does not guarantee transfer: unseen-gravity accuracy also depends on law coverage, predictor fit, and representation closure.For free flight, the bound grows with the test gravity’s distance from the training gravity mean.Source: §4 · p. 9

Empirical observations

  • With fresh predictors, the GRU-trained encoder lowers mean rollout error by about 12%: 1.376 versus 1.555 with a fresh GRU, and 1.269 versus 1.453 with a fresh Transformer.The ranking persists across predictor architectures, indicating that the gain follows the representation rather than the original GRU predictor.Source: §4 · p. 10, Figure 5 · p. 10
  • On far-OOD gravity values g ≤2 or g ≥6, SG-JEPA reduces teacher-forced local error by about 32% relative to DINO-WM, before recursive feedback.Its local-error point estimate is lower at all 25 test gravity values.Source: §4 · p. 10, Figure 5 · p. 10
  • From horizon 2 onward, free-rollout gaps are several times larger than teacher-forced gaps, peaking near 0.38 around horizon 20 while the local gap remains below 0.06, then shrinking at longer horizons.The amplification is therefore not monotone.Source: §4 · pp. 10–11, Figure 5 · p. 10

Ablations

  • In the appendices, correct gravity minimizes rollout error under counterfactual gravity inputs, while sparse-support post-training reduces error by 14.1% on average versus 6.9% for target-gravity-only post-training.Post-training reduces SG-JEPA (GRU) error by 13.6% on the square and 20.8% on the triangle.Source: §4 · p. 11

5 Conclusion · A Additional related work

SG-JEPA generalizes physical dynamics across unseen gravity values, with recursive rollout training improving representations and downstream control. The related-work appendix situates this evaluation among latent world models, physical-system generalization, and rollout-training approaches.

5 Conclusion

  • SG-JEPA outperforms Original LeWM and DINO-WM at longer rollout horizons and OOD gravity values, while its frozen representation improves success rates on all three manipulation tasks.The evaluation uses eight MuJoCo datasets, with gravity sampled from a narrow training band and a wider test grid; gains are largest at long horizons and far from training gravity.Source: §5 · p. 11
  • The linear feature model separates law-conditioned local error from rollout amplification, and experiments attribute SG-JEPA’s advantage to encoder representations that recursive rollout preserves.The analysis bounds one-step error at unseen gravity by training-gravity coverage; ablations find that the predictor uses supplied gravity and sparse post-training at new gravities improves accuracy between them.Source: §5 · p. 11
  • Experiments vary only scalar gravity, and transfer across shapes is uneven: triangle and square training transfers house translation but not rotation, while pentagons remain challenging.The authors identify vector-valued physical variables, inferring parameters from observations, and generalist models trained on diverse data as future directions.Source: §5 · pp. 11–12
  • The theory uses a linear feature model, unlike the nonlinear, history-dependent predictor, and contact can perturb transition branches, motivating theory that incorporates nonlinearities and action-perturbing effects.Source: §5 · p. 12

A Additional related work

  • Latent world models predict features rather than pixels, while physical benchmarks range from event and human-judgment evaluation to trajectory and physical-quantity measurement.DINO-WM fits dynamics over frozen DINOv2 features [Zhou et al., 2025, Oquab et al., 2024], whereas LeWM jointly trains encoder and predictor with SIGReg [Balestriero and LeCun, 2025, Maes et al., 2026]; related work also studies self-supervised video, particle simulators, physical-property inference, and held-out configurations [Garrido et al., 2025] [Sanchez-Gonzalez et al., 2020] [Li et al., 2020] [Xue et al., 2026, Tan et al., 2026].Source: A Additional related work · pp. 17–18
  • Prior work adapts to changing dynamics through context inference, universal policies, or domain randomization; SG-JEPA instead supplies gravity and isolates interpolation and extrapolation along that known coordinate.The experiment keeps the scene and control interface fixed, so system identification is removed.Source: A Additional related work · p. 18
  • Unlike prior rollout remedies, SG-JEPA backpropagates rollout loss through encoder and predictor, then compares prediction and policy performance under fixed data and policy training.A frozen-encoder crossover separates predictor contributions; the accompanying lifted-linear and piecewise-affine contact analyses are surrogate analyses, not claims that learned predictors are linear.Source: A Additional related work · p. 18

B Datasets

The benchmark comprises eight MuJoCo datasets spanning planar rigid-body and 3D projectile prediction, plus 3D robotic control, with gravity variation used to evaluate generalization.

Planar rigid-body prediction

  • The benchmark spans eight MuJoCo datasets: four planar rigid-body prediction tasks, one 3D projectile task, and three 3D robotic-control tasks.The planar tasks use a red body shaped as a right triangle, square, pentagon, or house, moving and colliding inside a planar box after a single planar impulse.Source: B Datasets · p. 18

3D prediction and control

  • Approach Ball isolates free-flight projectile prediction, while the robotic tasks use a Z1 arm to catch or repeatedly bounce a ball and a Franka arm to strike one into a basket.Franka success requires blade contact followed by the ball entering the basket; its scripted outcomes include 70% successes and 5% for each of six informative failure modes.Source: B Datasets · p. 19

Gravity sampling and input encoding

  • Training uses 8,000 episodes per dataset, while held-out tests evaluate task-specific gravity grids with 23 or 25 values.Planar tasks and Arm Catcher Ball train with g = max(N(4, 0.52), 0.1); the other 3D datasets train with g = max(N(9.8, 22), 0).Source: B Datasets · p. 19
  • Gravity is provided as an action coordinate and z-scored using training-split statistics that remain fixed for validation, test, and post-training evaluation.Planar actions also include two impulse channels nonzero only at t = 0; Approach Ball receives only gravity, and arm tasks append scripted Cartesian commands, with paddle-orientation commands for Arm Paddle Ball and Franka Paddle-to-Basket.Source: B Datasets · p. 19

Format and storage

  • Episodes contain 64 frames at 16 Hz over 4 seconds, with 128×128 2D videos and 256×256 3D videos; actions align to the transition from frame t to t + 1.Records include encoded RGB, simulator state, aligned actions, rewards, physics data, raw gravity, and episode metadata for controlled tasks.Source: B Datasets · p. 19, Table 1 · p. 19

C Model and training details … D.2.1 2D square: history length and predictor depth

The ablations select H = 20 and L = 3 for 2D square, λSIG = 0.72 for 2D experiments, and a hybrid Muon/AdamW optimizer. The implementation uses autoregressive latent rollouts and frozen-encoder diffusion policies for downstream control.

C.1 World Model

  • SG-JEPA uses a projected 256-dimensional ViT-Tiny CLS token and GRU or SSM predictors with three residual layers, width 512, MLP width 2048, and dropout 0.1.Final models use H = 20, K = 5, γ = 0.95, and 20 epochs; gravity is z-scored and concatenated as an action at every predictor layer. LeWM instead uses a six-layer causal Transformer with teacher-forced one-step latent MSE, while DINO-WM uses frozen DINOv2 patch tokens and a teacher-forced Transformer predictor.Source: §C.1 · pp. 20, 22

C.2 Downstream model

  • A separate diffusion policy is trained for each frozen encoder and control task; the visual encoder stays frozen, and policies do not query a world-model predictor online.State probes are separately fit on encoded ground-truth sequences, frozen, and then applied to predicted rollouts.Source: §C.2 · pp. 22–23

D.1.1 Learning rate sweep

  • The 5 × 10−5 AdamW setting nearly matches the best MLP-probe validation NMSE and gives the best linear-probe NMSE, despite 5 × 10−4 yielding the lowest validation objective.Held-out test results show the same trend, so the study uses 5 × 10−5 for AdamW parameters.Source: §D.1.1 · pp. 23–24, Figure 8 · p. 23

D.1.2 Muon versus AdamW

  • Muon lowered best validation objective versus AdamW for GRU (4.990 vs 5.864), SSM (5.403 vs 6.922), and Transformer (5.278 vs 6.384).The selected checkpoints occurred at epochs 11, 6, and 7 with Muon, versus 12, 8, and 13 with AdamW, respectively.Source: §D.1.2 · p. 24, Figure 9 · p. 24

D.1.3 SIGReg coefficient sweep

  • At λSIG = 0.72, effective rank is 138.0 versus 131.0 at 1.44, and it gives the lowest linear-probe validation NMSE.It also gives the best short- and medium-horizon readout over most of the gravity range and most clearly preserves square geometry and trajectory in decoded rollouts; the study selects it for 2D experiments.Source: §D.1.3 · pp. 25, 27, Figure 10 · p. 25, Figure 11 · p. 26

D.2 History length and rollout horizon ablations

  • Ablations on 2D Square and 3D Approach Ball favor a 20-frame history, while Approach Ball favors a five-step training rollout; the paper therefore uses H = 20 and K = 5 across datasets.Source: §D.2 · p. 27

D.2.1 2D square: history length and predictor depth

  • At H = 20, L = 4 gives position, orientation, and signed-turn errors of 0.745, 13.91°, and 0.153 turns; L = 3 is within 2%–4% on all three metrics with fewer parameters.The L = 3 predictor has 14.57M parameters versus 19.30M for L = 4, motivating the H = 20, L = 3 performance-cost balance. Real-latent controls show H = 30 degrades orientation error to about 11.1° and rotation error to 0.446 turns, indicating representation degradation before autoregressive error accumulates.Source: §D.2.1 · p. 27, Figure 12 · p. 27

D.2.2 Approach Ball: history and training rollout horizon … E.1 Gravity-conditioned rollouts

SG-JEPA leads five of six h44 shape–metric comparisons, while controlled studies attribute its joint-training advantage chiefly to the dynamics-relevant representation learned with the GRU predictor.

D.2.2 Approach Ball: history and training rollout horizon

  • For Approach Ball, H = 20, K = 5, and λSIG = 0.18 gives the lowest probe-validation, direct-probe test, and 20-step rollout NMSE.Its 32-step rollout NMSE is 0.185 versus the best 0.182, while peak host RAM is 61.3 rather than 73.5 GiB; the sweep uses one seed.Source: §D.2.2 · pp. 27–28, Table 3 · p. 28, Figure 13 · p. 28

D.3 Rollout discount

  • At horizon 44, γ = 0.95 gives the lowest joint NMSE and velocity error; its position error is 0.122, close to the 0.118 best at γ = 1.00.The study uses the Arm Catcher Ball dataset and selects γ = 0.95 as the default across datasets.Source: §D.3 · p. 29, Table 4 · p. 29

D.4 Counterfactual gravity conditioning

  • For both predictors, rollout error is minimized when supplied gravity matches true gravity; mismatches generally increase error, and larger mismatches generally increase it more.At gtrue = 4, correct conditioning also gives the lowest GRU position error at every reported horizon, showing that the predictor uses supplied gravity.Source: §D.4 · p. 29, Figure 14 · p. 30

D.5.1 Can Transformer conditioning close the joint-training gap?

  • AdaLN gate initialization to 0.1 lowers the Transformer endpoint from 0.458 to 0.428, closing 38.5% of the GRU gap but remaining above GRU’s 0.380.The best Transformer remains worse than GRU at 22 of 25 gravity values; the controlled studies identify the GRU-trained representation as the main source of the joint-training gain.Source: §D.5.1 · p. 30, §D.5 · p. 30, Figure 15 · p. 31

D.5.2 Encoder–predictor crossover: full results

  • The crossover advantage follows the GRU-trained representation: either predictor performs better on it, despite Transformer error reductions of -0.107 and -0.103 across sources.The Transformer-trained representation has lower Euclidean latent MSE, but this does not yield better physical-state prediction, especially for velocity.Source: §D.5.2 · p. 31, Table 6 · p. 32

D.5.3 History dependence of the frozen representations

  • Shortening context from H20 to H4 yields a 0.015 history penalty over five forecasts for the Transformer-trained source versus the GRU-trained source, rising to 0.089 over 44 steps.The gap becomes pronounced after predicted latents feed back recursively, as the GRU-trained representation retains rollout accuracy with substantially less context.Source: §D.5.3 · pp. 32–33, Figure 16 · p. 32

D.6 Performance on conventional control tasks

  • On OGBench-Cube Single and DMC Reacher, GRU/SSM predictors preserve downstream control success and achieve the lowest planning latency and highest effective rank on both datasets.These conventional-task results show no material loss in control success from using GRU/SSM predictors and support GRU as the default.Source: §D.6 · pp. 33–34, Table 7 · p. 33

E.1 Gravity-conditioned rollouts

  • Five of six h44 shape–metric comparisons favor SG-JEPA; on the triangle, DINO-WM retains the cumulative-rotation lead.On the square, GRU and SSM lead every metric at every reported step; across most held-out gravities, the gains are not confined to the training-distribution center.Source: §E.1 · pp. 34–35, Figure 18 · p. 34, Table 8 · p. 35

E.2 Shape generalization · F 3D datasets full results

SG-JEPA transfers some house-shape dynamics from component shapes it never saw together, though direct house training remains stronger near the training distribution. Position transfers most reliably; recursive prediction particularly widens the rotation gap.

E.2 Shape generalization

  • The shape-generalization experiment trained on triangle plus square, with or without pentagons, and evaluated the unseen composite house against a house-trained model.Both Tiny ViT (D = 192) and Small ViT (D = 384) encoders were tested to address capacity limits.Source: §E.2 · p. 35
  • For g ∈[2.5, 5.5], Small T+S has 1.40× the position error and 1.70× the velocity error of the house-trained reference.Adding pentagons brings no consistent improvement over triangle and square components alone.Source: Table 9 · p. 36, §E.2 · p. 36
  • Position transfers most consistently, velocity is less stable, and rotation remains the hardest quantity to predict.Direct house training is still better near the training distribution.Source: §E.2 · p. 36
  • On encoded ground-truth house frames, Small T+S error ratios versus the house-trained reference are 0.90 for position, 1.23 for velocity, and 1.12 for cumulative rotation.This true-latent probe localizes the remaining gap before autoregressive rollout.Source: §E.2 · p. 36
  • After autoregressive rollout, Small T+S ratios rise to 1.40 for position, 1.70 for velocity, and 3.36 for cumulative rotation versus the house-trained reference.Most of the rotation gap develops during recursive prediction, as angular-velocity errors accumulate and new contact geometry matters.Source: §E.2 · p. 36
  • Across much of both gravity tails, T+S closely follows the house-trained reference on position; Small T+S matches or outperforms it at several extremes.The house shape was unseen during world-model training.Source: §E.2 · p. 36, Figure 19 · p. 37

F.1 Approach-Ball results

SG-JEPA improves Approach-Ball position prediction, but its velocity advantage over DINO-WM is not uniform. The gravity breakdown varies by regime, and each model was evaluated with one training seed.

F.1 Approach-Ball results

  • SG-JEPA (GRU) and SG-JEPA (SSM) have the two lowest mean and final-step position errors, reducing mean position error by roughly 30% relative to DINO-WM.The full state readout covers position and velocity over h1–h44 on 1,600 held-out episodes; h44 corresponds to 2.75 s at 16 fps.Source: Table 10 · p. 37, §F.1 · p. 38
  • GRU has slightly lower rollout-average velocity error than DINO-WM, SSM and DINO-WM are nearly tied, and DINO-WM is more accurate at h44.GRU or SSM has lower error over much of the intermediate rollout, but DINO-WM closes the gap near the end.Source: Figure 20 · p. 38, §F.1 · p. 38
  • GRU is strongest across most central and moderate gravities, DINO-WM at the high-gravity tail, and Original LeWM at the few lowest values.Each world model uses one training seed, so episode averages do not measure variation across training runs.Source: Figure 20 · p. 38, §F.1 · p. 38

F.2 Arm-Catcher-Ball control results

In Arm-Catcher-Ball, SG-JEPA improves capture over DINO-WM across gravity conditions. The discussion attributes its advantage to retaining velocity and contact-phase information, while noting degradation at negative and extreme gravity values.

F.2 Arm-Catcher-Ball control results

  • Evaluation uses diffusion policies trained per checkpoint and tests 23 gravity values with five paired simulator-rollout seeds.At inference, the policy conditions on a 20-frame latent history, predicts 16 Cartesian actions, executes eight, and replans; success requires the ball to enter and latch in the catcher before the episode ends.Source: §F.2 · p. 38
  • 13.8%: SG-JEPA more than doubles overall capture rate versus DINO-WM, with gains across all five paired rollouts.The gain ranged from 12.9% to 14.3%; at the training gravity, SG-JEPA achieved 42.3% versus 16.0% for DINO-WM, and its mean advantage was positive at every gravity.Source: §F.2 · p. 38, Figure 21 · p. 39
  • Both SG-JEPA predictor variants remain ahead of the baselines overall, although Original LeWM uses a one-step Transformer objective while GRU and SSM use the five-step SG-JEPA objective.Original LeWM uses λSIG = 0.09; the GRU and SSM variants use λSIG = 0.18.Source: §F.2 · p. 39
  • SG-JEPA’s advantage is consistent with a representation that preserves ball velocity and contact phase, which matter for phase-sensitive interception as bounce intervals shorten.Negative and extreme gravity values still degrade performance because their motion departs qualitatively from the training regime.Source: §F.2 · p. 39

F.3 Franka paddle-to-basket results · F.4 Arm-Paddle-Ball results

In Franka paddle-to-basket, SG-JEPA leads in strict basket entry overall, though method rankings reverse at two gravity values; in Arm-Paddle-Ball, it predicts held-out ball dynamics more accurately, especially velocity.

F.3 Franka paddle-to-basket results

  • Both SG-JEPA variants achieve the highest strict-entry rates in the four-method Franka comparison, while DINO-WM leads in relaxed success and hits; GRU and SSM are effectively tied.The comparison adds Original LeWM and SG-JEPA (SSM) to the main-text comparison of SG-JEPA (GRU) and DINO-WM.Source: Table 11 · p. 39, §F.3 · pp. 39–40
  • The two gravity-specific reversals persist across all five rollout seeds, but the authors caution that two gravity values cannot support a broader mechanistic conclusion.The runs match episode IDs, action scaling, and checkpoint hashes, varying only the rollout seed.Source: §F.3 · p. 40
  • At g = 13.9 m/s2, DINO-WM averages 61.2% strict entry versus SG-JEPA (GRU) at 1.8%; at g = 16.0 m/s2, SG-JEPA averages 58.6% versus DINO-WM at 0.8%.The entry ordering changes sharply outside the center of the training distribution, while paddle-hit rates remain high beyond the two weakest gravity values.Source: Figure 22 · p. 40, §F.3 · p. 40
  • Franka success depends more on controlling the ball’s post-impact state than on interception, which requires managing contact timing, paddle orientation, incoming velocity, and gravity-dependent trajectories.The authors suggest SG-JEPA better preserves velocity and contact-phase information needed for gravity-dependent impact control.Source: §F.3 · p. 40

F.4 Arm-Paddle-Ball results

  • Arm-Paddle-Ball requires repeated contact control to keep the ball off the floor and produce at least one qualified bounce; the held-out probe tests the model’s understanding of ball dynamics.The larger velocity-error gap is relevant to sustained bouncing because small phase errors carry into the next contact.Source: §F.4 · pp. 40–41
  • At h44, SG-JEPA lowers Arm-Paddle-Ball position error from 0.119 to 0.114 m and velocity error from 0.635 to 0.466 m/s.Across the 44 forecast steps, position error is lower at 36 steps and velocity error at all 44; the probe is evaluated on 5,000 held-out episodes.Source: Figure 23 · p. 41, §F.4 · p. 41

G Sparse-Gravity Post-Training Ablation … H.2 Local prediction error under a fixed gravity

Mixed sparse-gravity post-training outperforms target-only adaptation across the tested combinations, while the theory explains prediction through representation-dependent local error, rollout propagation, and training-horizon effects. The analysis uses an action- and gravity-conditioned autoregressive latent model and distinguishes predictor error from information discarded by the representation.

G Sparse-Gravity Post-Training Ablation

  • Post-training uses 3,200 target episodes at four support gravities, with a mixed arm adding 5,000 source episodes; rollout errors are evaluated at horizon 44 on held-out gravities.Evaluation excludes post-training data, and aggregate improvement is measured over 13 held-out interpolation gravity values relative to each model’s pretrained baseline.Source: G Sparse-Gravity Post-Training Ablation · p. 41
  • Mixed adaptation improves many unobserved interpolation gravities, but gains vary by metric: four mixed configurations pass source-range retention, while none of the target-only configurations do.Target-only DINO-WM worsens triangle rotation by 0.22%, and Mixed DINO-WM worsens square velocity by 4.13%.Source: G Sparse-Gravity Post-Training Ablation · p. 42
  • Mixed post-training reduces aggregate error in all eight shape–model combinations and beats target-only post-training in every combination, raising mean reduction from 6.85% to 14.05% (7.21 percentage points).The largest mixed improvements are 20.84% for triangle SG-JEPA (GRU) and 20.81% for square Original LeWM.Source: G Sparse-Gravity Post-Training Ablation · p. 42, Table 12 · p. 43

H Theory of prediction under changing physical laws

  • The linear feature analysis decomposes one-step latent error into predictor error on learned features, discarded information affecting the next state, and a conditionally zero-mean residual.When true and learned dynamics share gravity dependence, average training error bounds systematic one-step test error scaled by gravity coverage; for affine free flight, the factor is 1 + (g⋆−µtr)^2/σ^2.Source: H Theory of prediction under changing… · p. 46
  • Rollout recursion propagates local error and transition residual through later predicted dynamics, while multi-step loss changes the error geometry and can reverse representation rankings across training horizons.The physical-state bound also retains readout error and transition residual.Source: H Theory of prediction under changing… · p. 46

H.1 Setup: the model

  • The implementation-aligned model conditions the predictor on actions and gravity and recursively rolls out latent predictions from a true context for multiple training steps.The theoretical setup fixes gravity within an episode and models each gravity as indexing a pair of discrete evolution semigroups; the linear model is not an identification of the neural encoder or GRU.Source: §H.1 · pp. 46–47, H Theory of prediction under changing… · p. 46

H.2 Local prediction error under a fixed gravity

  • At fixed gravity, local conditional-mean error separates predictor error on the encoded state from closure error caused by discarded features that still affect the next latent.Predictive closure means the conditional mean of the next latent depends only on the current latent, although residual variation can still make individual transitions unpredictable.Source: §H.2 · p. 48
  • The teacher-forced same-probe diagnostic cancels the true-state probe floor, but its remaining error combines both local-error terms with transition residual and cannot attribute their separate contributions.The cross term vanishes in expectation after a fixed linear probe under the conditional-zero-mean assumption, but no such guarantee applies to nonlinear probes.Source: §H.2 · p. 49

H.3 How gravity coverage controls test error

Gravity generalization is distinct from composing dynamics over time: unseen-gravity error depends on shared law structure and how well training covers its directions. The section gives sufficient transfer results under explicit assumptions and illustrates them for affine free-flight dynamics.

H.3 How gravity coverage controls test error

  • Gravity generalization is separate from semigroup composition: low error at training gravities alone gives no unseen-gravity guarantee, so shared law structure must be stated.Source: §H.3 · p. 49
  • A large law-coverage factor means the test law lies in a weakly covered training direction; the assumed positive-definite moment matrix covers every direction in the chosen basis.Source: §H.3 · p. 49
  • Under shared-basis assumptions and bounded z_t and ϕ_t, Theorem H.4 gives a sufficient transfer bound for local error at an unseen gravity.Lemma H.3 treats transfer for a matrix-valued law family.Source: §H.3 · pp. 49–50
  • For collision-free ballistic motion, gravity enters the transition affinely, yielding the two-dimensional law basis ψ(g) = (1, g)⊤ and a coverage factor determined by training-gravity mean and variance.In the planar experiments, g = max{G, 0.1} with G ∼ N(4, 0.52), so the factor uses moments of that clipped distribution.Source: §H.3 · pp. 50–51
  • Free-flight transitions remain affine in gravity at every horizon, but this applies only during free flight and to physical-coordinate features; learned features need not inherit it, and contact can change the transition branch.In multiple dimensions, gravity forces the aligned position and velocity components, while tangential components have no gravity forcing in this idealized model.Source: §H.3 · p. 51

H.3.1 When the coverage result does not apply … H.8 Supplementary calculation: SIGReg and its limits

The appendix examines the scope and limitations of the linear theory for predicting dynamics under unseen gravities, including the effects of rollout errors, readout constraints, and prediction history. It also characterizes a feature for a simple contact boundary and the local covariance effect and limits of SIGReg.

H.3.1 When the coverage result does not apply

  • If training misses a law direction used at test gravity, residuals can vanish on training data yet remain nonzero at test; analyticity and gravity conditioning alone do not guarantee extrapolation.A singular coverage matrix is not fatal when the test feature lies in its covered range.Source: §H.3.1 · p. 52
  • For a finite gravity design, full column rank allows residual coefficients to be recovered from observed residuals, with the unseen-gravity bound depending on feature leverage.The stated leverage factor is ψ(g⋆)⊤(Ψ⊤Ψ)−1ψ(g⋆).Source: §H.3.1 · p. 52
  • Training error cannot bound a remainder that appears only where training has little or no mass, and basis and remainder terms may cancel at observed gravities.The bound for operator and closure residuals therefore retains their test-gravity remainder terms.Source: §H.3.1 · pp. 52–53

H.4 Local errors under repeated composition

  • The exact latent error recursion carries each one-step defect through remaining learned semigroup iterates, whose contributions can reinforce or cancel rather than grow monotonically.Teacher forcing measures local error before reuse; free rollout feeds each predicted latent back into subsequent predictions.Source: §H.4 · p. 53

H.4.1 Physical readout and the central bound

  • The physical rollout bound separates readout residual, local unseen-law error, recursive propagation, and transition residuals, but is not a stand-alone guarantee from law coverage alone.It also depends on trajectory bounds and rollout gain, and its worst-case factor does not capture cancellation between propagated residuals.Source: §H.4.1 · pp. 54–55

H.4.2 Cross-step second moments

  • Cross-step second moments capture error orientation and dependence; reset comparisons change both these moments and the propagation path, so they do not isolate one spectral gain.A Lipschitz nonlinear transition also yields an iterated error bound, offered as intuition rather than the exact linear result used above.Source: §H.4.2 · p. 55

H.5 Why multi-step composition can change the learned representation

  • Multi-step composition weights one-step intertwining defects by their later consequences, so equal-sized one-step defects can have different multi-step costs.The analysis isolates conditional-mean, action-free dynamics; its quadratic costs measure transition bias rather than irreducible noise.Source: §H.5 · pp. 56–57

H.5.1 A stable ranking reversal

  • One-step profiling prefers WP, but horizon-two loss prefers WF, and full profiling prefers WF when γ > 3600/9007 ≈ 0.39969, including γ = 0.95.This stable construction demonstrates a horizon-dependent reversal, not that longer rollouts always help or that it explains SG-JEPA’s learned mechanism.Source: §H.5.1 · pp. 58–59

H.6 Relating the theory to the history-based predictor

  • The companion-state analysis gives an exact newest-latent rollout-error identity with context, action, and transition-residual contributions, while retaining history shifts and recursive feedback.Its additive-map assumptions do not capture nonlinear interactions in the implemented predictor, and the construction does not claim the nonlinear GRU is globally linear.Source: §H.6 · pp. 60–61, The quadratic form for newest-block… · p. 62

The quadratic form for newest-block loss uses

  • The newest-block loss projects the augmented companion defect, yielding a positive quadratic form on admissible defects when w1 > 0 and ΣX ≻ 0.Matched history shifts contribute zero rows; action-map differences remain explicit terms unless the true and learned augmented action maps agree.Source: The quadratic form for newest-block… · pp. 61–62

H.7 Regime changes and contact

  • Across a fixed shared branch sequence, rollout error telescopes through branch-local defects, but this analysis excludes predictions that change event time or branch.The smooth law-transfer analysis applies within one regime; contact changes the active branch along a trajectory.Source: §H.7 · p. 62
  • A continuous piecewise-affine event needs a hinge beyond affine features; floor-clipped ballistic height therefore has a gravity hinge, without establishing broader contact-modeling gains.The result is limited to one continuous piecewise-affine boundary and does not cover discontinuous, frictional, restitution-dependent, or multi-contact dynamics.Source: §H.7 · pp. 63–64

H.8 Supplementary calculation: SIGReg and its limits

  • Near identity covariance, SIGReg has a strict local quadratic preference for identity, but finite-batch trace bias means identity need not minimize and the calculation does not prove noncollapse.For B > 2, the finite-batch quadratic coefficient is positive, yet the linear trace term remains.Source: §H.8 · pp. 64–65

Related papers

Loading 2609.10464v1…