JEPA-x: Cross-Predictive Physics Grounding for Forecastable Latent Dynamics

Kehan Wen; Ziming Li; Siyuan Luo; Fan Shi

cs.LGAug 25, 2026

JEPA-x: Cross-Predictive Physics Grounding for Forecastable Latent Dynamics

Authors: Kehan Wen, Ziming Li, Siyuan Luo, Fan Shi

Organizations: Department of Electrical and Computer Engineering, NUS

Abstract

Latent world models plan by predicting how candidate actions advance learned latent dynamics. In self-predictive models, however, the encoder and predictor are optimized jointly and can co-adapt to latent transitions that are easy to predict but weakly constrained by the physical evolution of the scene. We introduce the cross-predictive JEPA (JEPA-x), which grounds latent dynamics in privileged physical trajectories. JEPA-x treats visual observations and physical states as corresponding views of the same action-conditioned trajectory, advances both through a shared predictor, and matches each prediction to the future representations of both modalities. This encourages the action-conditioned predictor to learn a common transition rule across the two views. Privileged physical state is used only during training, leaving a visual-only model at deployment. Empirical results show that JEPA-x reduces the rollout drift of a newly fitted predictor from $0.361$ to $0.104$ and increases mean control success from $53.6\%$ to $78.2\%$ on a multi-task suite spanning six evaluation subfamilies. We additionally show that direct physical-state regression improves decodability without improving forecastability or control, indicating that the benefit comes from shaping latent dynamics rather than merely encoding physical variables.

Figures & tables

Figure 1: Cross-predictive physical grounding constrains representation geometry and predictive dynamics. Visual self-prediction jointly learns a visual encoder and predictor, while state regression makes physical variables decodable from individual visual latents. Neither directly couples a predicted visual transition to the corresponding physical future. JEPA- x instead treats visual and physical streams as paired views of the same action-conditioned trajectory and matches predictions from either history to future representations in both modalities. The physical branch is removed after training, leaving a visual-only model at deployment.

Figure 2: Unified privileged-state interface. A scene is represented by a fixed set of object slots, together with end-effector (EE) and table tokens; unused object slots are masked. Each object token encodes SDF geometry $G_{i}$ , extents $e_{i}$ , rotation $R_{i}$ , and position $t_{i}$ . Masked self-attention followed by learned-query attention pooling maps the resulting token set to the physical latent $z_{t}^{s}$ .

Figure 3: JEPA- x produces more forecastable latent dynamics throughout the rollout. Multi-task suite, three training seeds.

	Relative drift $\downarrow$			Object-position $R^{2}$ $\uparrow$
Environment	JEPA- x	Visual	Regress	JEPA- x	Visual	Regress
Two-Room	$0.221_{\pm 0.021}$	$0.503_{\pm 0.010}$	$0.480_{\pm 0.008}$	$0.998_{\pm 0.000}$	$0.921_{\pm 0.011}$	$0.936_{\pm 0.007}$
OGBench-Block	$0.269_{\pm 0.007}$	$0.507_{\pm 0.026}$	$0.393_{\pm 0.018}$	$0.982_{\pm 0.003}$	$0.879_{\pm 0.009}$	$0.983_{\pm 0.001}$
Push-T	$0.258_{\pm 0.007}$	$0.399_{\pm 0.118}$	$0.343_{\pm 0.039}$	$0.991_{\pm 0.001}$	$0.944_{\pm 0.009}$	$0.968_{\pm 0.007}$
Reacher	$0.233_{\pm 0.009}$	$0.313_{\pm 0.009}$	$0.309_{\pm 0.003}$	$0.999_{\pm 0.000}$	$0.999_{\pm 0.000}$	$0.999_{\pm 0.000}$

Table 1: JEPA- x produces the most forecastable visual dynamics across all four matched single-task environments. Forecastability is measured using a predictor fitted from scratch to each frozen representation. Object-position

R^{2}

measures how well the environment’s object position is decoded from the frozen visual latent. Values are means

\pm

sample SD over three training seeds. Bold marks the best mean in each comparison; exact ties are all bold.

Figure 4: Grounding improves multi-task control. JEPA- x raises overall success from

53.6\%

78.2\%

, with improvements across all six interaction families. Regressing a comprehensive pose-based physical-state target from the visual latent does not recover the gain. Bars show means

\pm

sample SD over three training seeds; dots show individual seeds.

Figure 5: Rolled over identical expert actions, JEPA- x imagines the eraser reaching the block; Visual imagines it beside the block. Middle row: ground truth. Top and bottom: each model’s own rollout over the same action chunks from the same observation, with the decoded object and end-effector positions written back into the simulator and rendered.

Environment	JEPA- x	Visual	Regress
Two-Room	$94.8_{\pm 1.0}$	$74.7_{\pm 4.1}$	$75.8_{\pm 3.8}$
OGBench-Block	$78.7_{\pm 2.7}$	$64.9_{\pm 1.4}$	$70.2_{\pm 0.8}$
Push-T	$94.1_{\pm 0.4}$	$94.9_{\pm 1.1}$	$96.9_{\pm 2.7}$
Reacher	$83.6_{\pm 1.4}$	$82.7_{\pm 2.1}$	$82.6_{\pm 0.4}$

Table 2: Matched single-task control success. Values are mean success rates (%)

\pm

sample SD over three training seeds. Bold marks the best mean in each row; exact ties are all bold.

Variant	Cross-pred.	Shared	Corresp.	Drift $\downarrow$	Control $\uparrow$
JEPA- x	✓	✓	✓	$0.104_{\pm 0.003}$	$78.2_{\pm 1.3}$
Cross-only	✓	–	✓	$0.101_{\pm 0.001}$	$73.0_{\pm 0.9}$
Share-only	–	✓	–	$0.309_{\pm 0.020}$	$56.7_{\pm 3.5}$
Align-only	–	✓	✓	$0.158_{\pm 0.001}$	$72.6_{\pm 3.6}$
Distill	–	–	✓	$0.516_{\pm 0.004}$	$57.7_{\pm 4.4}$
Shuffle	✓	✓	$\times$	$0.540_{\pm 0.023}$	$3.6_{\pm 1.3}$

Table 3: Correct trajectory pairing supports control, while direct cross-modal prediction improves forecastability beyond alignment. The design columns indicate cross-modal prediction, predictor sharing, and correct visual–physical pairing (✓ present,

\times

present but mispaired, – absent). Drift and control are means

\pm

sample SD over three training seeds; bold marks the best mean in each column.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Interaction group	Configurations
Grasp–move–place	battery, block, can, tennis ball
Peg / plug insertion	peg-align-insert, peg-insert-pull, plug-insert
Planar pushing	block, book, eraser
Stacking	battery-on-phone, block-on-book, eraser-on-block
Righting / toppling	stand-up-can, stand-up-thermos, topple-battery, topple-can, topple-thermos
Tossing	battery, block, eraser, tennis ball

Appendix

Table 4: The

22

configurations in the multi-task suite.

Family	Success condition
Grasping	Object at least $5$ cm above its initial table height over the final $1$ s, with in-grasp slip at most $1$ cm.
Stacking	Support contact between the stacked object and its base after a $2$ s settle window, with the base orientation preserved.
Insertion	Insertion depth at least $90\%$ of the target cavity depth, axis tilt at most $3^{\circ}$ , and gripper open at termination.
Placement	Object at rest within the asset-specific positional tolerance of the target, with orientation compared modulo symmetry.
Pushing	Object within the asset-specific tolerance of the planar goal, with orientation compared to the initial pose modulo symmetry.
Tossing	Object inside the container interior at the end of the trajectory.

Appendix

Table 5: Success predicates for the six evaluated interaction families.

Family	JEPA- x	Visual	Regress	Distill	Shuffle
Insertion	$95.3_{\pm 3.1}$	$60.7_{\pm 10.1}$	$54.7_{\pm 9.2}$	$84.0_{\pm 14.0}$	$16.0_{\pm 6.0}$
Grasping	$66.7_{\pm 7.0}$	$34.0_{\pm 10.0}$	$48.0_{\pm 3.5}$	$32.0_{\pm 8.7}$	$1.3_{\pm 1.2}$
Stacking	$90.7_{\pm 3.1}$	$58.7_{\pm 4.2}$	$64.0_{\pm 10.4}$	$62.7_{\pm 7.6}$	$0.7_{\pm 1.2}$
Placement	$72.7_{\pm 1.2}$	$54.0_{\pm 2.0}$	$46.0_{\pm 8.7}$	$48.7_{\pm 4.6}$	$2.0_{\pm 2.0}$
Tossing	$94.7_{\pm 1.2}$	$79.3_{\pm 4.6}$	$76.7_{\pm 4.2}$	$82.7_{\pm 14.5}$	$1.3_{\pm 2.3}$
Pushing	$49.3_{\pm 2.3}$	$34.7_{\pm 2.3}$	$32.7_{\pm 2.3}$	$36.0_{\pm 4.0}$	$0.0_{\pm 0.0}$

Appendix

Table 6: Per-family control success (%), mean

\pm

sample SD over three training seeds. Families are ordered by the JEPA- x – Visual gap, matching Figure 4 .

Arm	Raw $\downarrow$	Copy	Drift $\downarrow$	Manip.-object pos. $R^{2}$ $\uparrow$
JEPA- x	1.18	12.17	$0.104_{\pm 0.003}$	$0.978_{\pm 0.002}$
Visual	4.43	12.42	$0.361_{\pm 0.024}$	$0.978_{\pm 0.000}$
Regress	4.28	12.07	$0.373_{\pm 0.005}$	$0.991_{\pm 0.000}$
Distill	6.11	12.50	$0.516_{\pm 0.004}$	$0.991_{\pm 0.001}$
Shuffle	8.41	15.51	$0.540_{\pm 0.023}$	$0.958_{\pm 0.002}$

Appendix

Table 7: Forecastability and manipulated-object position decodability on the multi-task suite. Both metrics use the same frozen encoders and held-out split, with a separately fitted probe for each metric. Drift and

R^{2}

are means

\pm

sample SD over three training seeds; bold marks the best reported mean. Raw and Copy are averaged over horizons and seeds; Drift is the mean of the per-horizon ratios and therefore differs from their ratio.

Figure 6: Under the same actions, JEPA- x imagines the object coming to rest on its support. Decoded height of the manipulated object during a stacking episode: ground truth against each model’s own rollout over the identical expert action chunks, one panel per planning cycle. Visual ’s imagined object passes through the support and stays there; JEPA- x has a mean terminal error of

1.4

cm across the three cycles.

	Effective rank		Relative drift $\downarrow$
Projection	JEPA- x	Visual	JEPA- x	Visual
None (native)	$56.5$	$95.7$	$0.104_{\pm 0.004}$	$0.359_{\pm 0.018}$
$k=57$	$51.7$	$52.8$	$0.121_{\pm 0.008}$	$0.453_{\pm 0.022}$
$k=32$	$30.7$	$30.4$	$0.175_{\pm 0.011}$	$0.402_{\pm 0.031}$
$k=16$	$15.7$	$15.6$	$0.193_{\pm 0.011}$	$0.379_{\pm 0.025}$
$k=8$	$\phantom{0}7.9$	$\phantom{0}7.9$	$0.214_{\pm 0.016}$	$0.407_{\pm 0.033}$

Appendix

Table 8: Matching effective rank does not close the forecastability gap. Relative drift after projection onto the leading

k

principal components, reported as mean

\pm

sample SD over three training seeds.

Family	Fut-delta	Fut-abs	Visual	Regress
Insertion	$54.0_{\pm 6.9}$	$56.7_{\pm 22.7}$	$60.7_{\pm 10.1}$	$54.7_{\pm 9.2}$
Grasping	$39.3_{\pm 19.4}$	$40.0_{\pm 6.9}$	$34.0_{\pm 10.0}$	$48.0_{\pm 3.5}$
Stacking	$60.0_{\pm 2.0}$	$60.0_{\pm 5.3}$	$58.7_{\pm 4.2}$	$64.0_{\pm 10.4}$
Placement	$52.0_{\pm 4.0}$	$56.0_{\pm 4.0}$	$54.0_{\pm 2.0}$	$46.0_{\pm 8.7}$
Tossing	$72.7_{\pm 4.6}$	$77.3_{\pm 6.4}$	$79.3_{\pm 4.6}$	$76.7_{\pm 4.2}$
Pushing	$30.0_{\pm 3.5}$	$36.7_{\pm 8.1}$	$34.7_{\pm 2.3}$	$32.7_{\pm 2.3}$

Appendix

Table 9: Per-family control success (%) for the two future-state auxiliary variants, with Visual and Regress for reference. Mean

\pm

sample SD over three training seeds, same protocol as Table 6 .

Figure 7: JEPA- x against Visual on the same scenario: grasping, insertion, and placement. Each family contributes a pair of rows executed from the same initial state using the same planner, JEPA- x above and Visual below, cropped to the working volume. The goal is shared by both arms and is therefore shown once per comparison, spanning the family’s two rows. Scenarios are the representative episodes in which the two arms disagree, so each pair is a case JEPA- x completes and Visual does not; per-family success rates are in Table 6 .

Figure 8: JEPA- x against Visual on the same scenario: pushing, stacking, and tossing. Same protocol as Figure 7 : paired rows from one initial state, JEPA- x above and Visual below, with the shared goal shown once per comparison. In stacking, Visual leaves the eraser beside the block; in tossing, it leaves the object outside the carton.

Latent Dynamics Modeling

Latent dynamics modeling learns compact hidden-state representations and transition rules to predict, explain, and control evolving systems.

See all 253 papers

Aug 7, 2026

Is Forward Prediction Enough? Physical State Grounding for JEPA World Models

Learning structured and control-relevant latent representations remains a key challenge for world models. Recent JEPA-based world models learn action-conditioned predictive latent dynamics from observation sequences. However, their forward-prediction objectives do not explicitly enforce reliable identifiability of robot-centric physical state from individual latents or state changes from latent pairs, which can limit downstream planning and policy performance. We propose PSG-JEPA, a physically grounded JEPA world model that shapes its latent space with two complementary grounding objectives beyond forward prediction: grounding individual latents in robot proprioceptive state, and grounding latent pairs in multi-horizon joint-angle changes. Both objectives are applied only during training, leaving the inference architecture and computational cost unchanged. To comprehensively evaluate PSG-JEPA, we conduct experiments at three levels: (1) latent identifiability via probing, (2) goal-conditioned planning on frozen latents, and (3) policy learning in simulation and on a real robot. Experiments demonstrate that our PSG-JEPA consistently outperforms state-of-the-art latent world-model baselines at all three levels.

Haodong Yan, Jiaguan Zhu, Mingyuan Jia +12

1The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · 2COCO Matrix

Physical Reasoning World Models for Robotics World Model-Based Planning Joint-Embedding Predictive Architecture Action-Conditioned World Models State Representation Learning

Ask PDF

Aug 6, 2026

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

We propose PhyLatent, a dynamics-relevant training objective for Joint-Embedding Predictive Architecture (JEPA) world models. Our key observation is that preventing global latent collapse does not necessarily ensure that the learned representation preserves physically meaningful state and action relationships. We identify three failure modes: sensitivity to appearance changes that leave the physical state unchanged, insufficient separation of distinct physical states, and insufficient separation of different action-conditioned futures. We refer to these as Physical Invariance Collapse, Physical Distinguishability Collapse, and Counterfactual Dynamics Collapse, respectively. PhyLatent targets these failures through three coordinated training pathways, implemented with static visual invariance, physical state grounding, future representation alignment, counterfactual branch separation, and latent denoising. On OGBench-Cube, PhyLatent reduces the three collapse rates by 43.9%, 27.9%, and 47.1%, respectively, while improving model predictive control (MPC) success by 12.0 percentage points (17.2% relative). Across four visual-control tasks, average planning success increases by 6.62 percentage points (8.3% relative). These results show that global non-collapse alone is insufficient for learning a reliable JEPA world-model state space, and that explicitly preserving dynamics-relevant structure can improve closed-loop planning.

Xi Zeng, Haojie Ren, Ziying Song +2

School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

Representation Learning World Model Learning World Model-Based Planning Joint-Embedding Predictive Architecture State Representation Learning

Ask PDF

Sep 9, 2026

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

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.

Andy Zeyi Liu, Haoran Sun, Lucas Baker +2

Yale University · Jump Trading · Brown University

World Models for Robotics Latent Dynamics Modeling Joint-Embedding Predictive Architecture

Ask PDF

Abstract

Latent Dynamics Modeling

Explore similar work

Is Forward Prediction Enough? Physical State Grounding for JEPA World Models

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

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