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

Is Forward Prediction Enough? Physical State Grounding for JEPA World Models
Aug 7, 2026Haodong Yan, Jiaguan Zhu, Mingyuan Jia +12Physical Reasoning World Models for Robotics
1The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · 2COCO Matrix
PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Representation Learning World Model Learning
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
Semigroup-JEPA: Latent Dynamics Consistency for Zero-Shot Physics Generalization
Sep 9, 2026Andy Zeyi Liu, Haoran Sun, Lucas Baker +2World Models for Robotics Latent Dynamics Modeling
Yale University · Jump Trading · Brown University

Aug 7, 2026J/K move · Enter open · S save

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

Haodong Yan, Jiaguan Zhu, Mingyuan Jia, Ruiqing Yin+11

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

Open paper AskPDF arXiv

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.

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