Exploring Temporal Context-Aware Physical Grounding in Latent World Models

Yilei Man; Yang Chen; Jingcai Guo

cs.AIOct 20, 2026

Exploring Temporal Context-Aware Physical Grounding in Latent World Models

Authors: Yilei Man, Yang Chen, Jingcai Guo

Organizations: Department of Computing, The Hong Kong Polytechnic University · Department of Land Surveying and Geospatial Science, The Hong Kong Polytechnic University

Abstract

Learning latent dynamics is fundamental to building action-conditioned world models for planning. Among existing families, Joint Embedding Predictive Architectures (JEPAs) offer a natural framework by modeling transitions directly in latent space rather than at the pixel level. However, two issues remain in JEPA-based world models: (i) physical grounding lacks explicit aggregation of temporal context and (ii) action planning always overly relies on generic sampling distributions. To this end, we propose $\mathcal{G}$ -JEPA, which incorporates temporal context into physical grounding for latent world models and couples the learned models with prior-guided action planning. Specifically, during learning, we incorporate a shared causal temporal module that aggregates temporal context within observation and prediction sequences under physical state supervision, grounding both in task-relevant physical states. When planning, we further introduce action inference through stepwise conditioning of a generative action distribution on the current and goal latent representations together with their evolving discrepancy, yielding anchored proposals that guide model-based online refinement. Together, physical state supervision guides temporal representation learning, while the action prior guides search through the learned world model. Experiments demonstrate improved planning performance on manipulation tasks, including tasks with increasing temporal separation between initial and goal observations. Code is publicly available at https://github.com/myl11045/G-JEPA.

Figures & tables

Figure 1: Temporal physical grounding and prior-guided planning. (A) PhyLatent ( Zeng et al., 2026 ) grounds individual latents in physical states, while PSG-JEPA ( Yan et al., 2026 ) additionally supervises state changes between latent pairs. (B) Gaussian CEM generates action candidates from a proposal distribution adapted through online elite updates. (C)

\mathcal{G}

-JEPA uses a shared causal grounding module to aggregate temporal context within encoded and predicted latent sequences under physical state supervision. (D) Learned action proposals supplement Gaussian candidates, with joint world-model evaluation and shared elite selection guiding subsequent Gaussian updates. Curves illustrate schematic candidate motion paths.

Figure 2: Overview of $\mathcal{G}$ -JEPA. World-model learning combines latent matching and SIGReg with (a) nuisance invariance under photometric augmentation and (b) temporal physical grounding. A shared S4D module aggregates temporal context separately within observed latents and aligned one-step predictions, supervising physical state estimates from both sequences. (c) During planning, the learned action prior is instantiated as a latent-conditioned diffusion policy that generates action proposals alongside Gaussian candidates. The frozen world model scores both sources using latent goal-matching costs, and a shared elite set updates the Gaussian search distribution. The policy remains fixed, while action execution provides new observations for subsequent planning rounds.

Method	PushT-Ada	OGBench Cube	TwoRooms
$\mathcal{G}$ -JEPA w/o AD	$79.43_{\pm 1.58}$	$78.60_{\pm 4.12}$	$97.00_{\pm 1.00}$
$\mathcal{G}$ -JEPA	$\mathbf{87.11}_{\pm 1.17}$	$\mathbf{100.00}_{\pm 0.00}$	—
LeWM	$56.90_{\pm 2.35}$	$72.80_{\pm 2.53}$	$\mathbf{99.00}_{\pm 0.00}$
TD-JEPA	$56.64_{\pm 2.71}$	$78.00_{\pm 3.13}$	—
PSG-JEPA	$17.32_{\pm 1.58}$	$75.40_{\pm 2.50}$	—

Table 1: Cross-method planning comparison. — denotes not evaluated.

	PushT-Distance			PushT-Challenging
Method	25 steps	50 steps	75 steps	25 steps	50 steps	75 steps
LeWM	$90.49_{\pm 0.81}$	$52.99_{\pm 0.60}$	$16.54_{\pm 0.60}$	$80.69_{\pm 1.65}$	$42.42_{\pm 0.73}$	$10.97_{\pm 0.64}$
$\mathcal{G}$ -JEPA w/o AD	$\mathbf{96.09}_{\pm 1.17}$	$80.34_{\pm 0.81}$	$47.53_{\pm 1.63}$	$\mathbf{92.06}_{\pm 2.38}$	$75.92_{\pm 1.00}$	$44.03_{\pm 1.73}$
$\mathcal{G}$ -JEPA	$91.80_{\pm 1.17}$	$\mathbf{85.55}_{\pm 0.39}$	$\mathbf{55.21}_{\pm 0.90}$	$83.33_{\pm 2.38}$	$\mathbf{82.30}_{\pm 0.48}$	$\mathbf{52.22}_{\pm 0.96}$

Table 2: Planning with increasing goal distances. We evaluate goal-reaching success with start–goal separations of 25, 50, and 75 environment steps along recorded trajectories.

State ( $S$ )	Invariance ( $I$ )	Predicted state ( $P$ )	PushT-Ada	OGBench Cube
✗	✗	✗	$56.90_{\pm 2.35}$	$72.80_{\pm 2.53}$
✓	✗	✗	$55.73_{\pm 1.58}$	$75.40_{\pm 1.65}$
✓	✓	✗	$\mathbf{79.43}_{\pm 0.90}$	$76.20_{\pm 2.57}$
✓	✓	✓	$\mathbf{79.43}_{\pm 1.58}$	$\mathbf{78.60}_{\pm 4.12}$

Table 3: Grounding-objective ablations.

S

I

, and

P

denote observed state grounding, nuisance invariance, and predicted physical state grounding.

Architecture	Grounding params. (K)	Success (%)
S4D	55.431	$\mathbf{78.12}_{\pm 3.34}$
GRU	54.279	$69.14_{\pm 14.89}$
Causal Transformer	55.607	$61.46_{\pm 10.61}$
Attention-only	55.555	$53.39_{\pm 34.02}$

Table 4: Temporal grounding architectures. Mean success on PushT-Ada across three training replicates with a fixed planning seed.

Figure 3: Qualitative comparison of action selection. Execution examples from the same initial state and goal within each task. (a) PushT shows recorded states, with goal poses in green and pusher paths in blue. (b) Cube shows simulator observations. Labels report success and final position/orientation errors (pixels/degrees) for PushT or cube–goal distance (cm) for Cube. $t$ counts environment steps, and Final marks each method’s stopping state in PushT.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Setting	PushT-Ada	OGBench Cube	TwoRooms
CEM budget	100 / 10 / 10	300 / 30 / 10	300 / 30 / 30
Executed environment steps / replan	5	25	25
Total environment-step budget	100	50	50
Position success threshold	20 px	4 cm	16 px
Additional angle threshold	$20^{\circ}$	—	—
DDIM inference steps	100	10	—

Appendix

Table 6: Cross-dataset planning settings. CEM budgets denote candidates / elites / iterations.

World Models

World models represent and predict how environments change, helping agents plan, act, and reason about possible futures.

See all 369 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
LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
Sep 28, 2026Luzhe Huang, Lei Chu, Jingyi Liang +1Disentangled Representation Learning World Models for Robotics
Independent researcher

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

World Models

Explore similar work

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

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models