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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 G\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) G\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 G\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.
MethodPushT-AdaOGBench CubeTwoRooms
G\mathcal{G} -JEPA w/o AD79.43±1.5879.43_{\pm 1.58}78.60±4.1278.60_{\pm 4.12}97.00±1.0097.00_{\pm 1.00}
G\mathcal{G} -JEPA87.11±1.17\mathbf{87.11}_{\pm 1.17}100.00±0.00\mathbf{100.00}_{\pm 0.00}—
LeWM56.90±2.3556.90_{\pm 2.35}72.80±2.5372.80_{\pm 2.53}99.00±0.00\mathbf{99.00}_{\pm 0.00}
TD-JEPA56.64±2.7156.64_{\pm 2.71}78.00±3.1378.00_{\pm 3.13}—
PSG-JEPA17.32±1.5817.32_{\pm 1.58}75.40±2.5075.40_{\pm 2.50}—
Table 1: Cross-method planning comparison. — denotes not evaluated.
 PushT-DistancePushT-Challenging
Method25 steps50 steps75 steps25 steps50 steps75 steps
LeWM90.49±0.8190.49_{\pm 0.81}52.99±0.6052.99_{\pm 0.60}16.54±0.6016.54_{\pm 0.60}80.69±1.6580.69_{\pm 1.65}42.42±0.7342.42_{\pm 0.73}10.97±0.6410.97_{\pm 0.64}
G\mathcal{G} -JEPA w/o AD96.09±1.17\mathbf{96.09}_{\pm 1.17}80.34±0.8180.34_{\pm 0.81}47.53±1.6347.53_{\pm 1.63}92.06±2.38\mathbf{92.06}_{\pm 2.38}75.92±1.0075.92_{\pm 1.00}44.03±1.7344.03_{\pm 1.73}
G\mathcal{G} -JEPA91.80±1.1791.80_{\pm 1.17}85.55±0.39\mathbf{85.55}_{\pm 0.39}55.21±0.90\mathbf{55.21}_{\pm 0.90}83.33±2.3883.33_{\pm 2.38}82.30±0.48\mathbf{82.30}_{\pm 0.48}52.22±0.96\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 ( SS )Invariance ( II )Predicted state ( PP )PushT-AdaOGBench Cube
✗✗✗56.90±2.3556.90_{\pm 2.35}72.80±2.5372.80_{\pm 2.53}
✓✗✗55.73±1.5855.73_{\pm 1.58}75.40±1.6575.40_{\pm 1.65}
✓✓✗79.43±0.90\mathbf{79.43}_{\pm 0.90}76.20±2.5776.20_{\pm 2.57}
✓✓✓79.43±1.58\mathbf{79.43}_{\pm 1.58}78.60±4.12\mathbf{78.60}_{\pm 4.12}
Table 3: Grounding-objective ablations. SS , II , and PP denote observed state grounding, nuisance invariance, and predicted physical state grounding.
ArchitectureGrounding params. (K)Success (%)
S4D55.43178.12±3.34\mathbf{78.12}_{\pm 3.34}
GRU54.27969.14±14.8969.14_{\pm 14.89}
Causal Transformer55.60761.46±10.6161.46_{\pm 10.61}
Attention-only55.55553.39±34.0253.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. tt counts environment steps, and Final marks each method’s stopping state in PushT.
Figure 8
Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

SettingPushT-AdaOGBench CubeTwoRooms
CEM budget100 / 10 / 10300 / 30 / 10300 / 30 / 30
Executed environment steps / replan52525
Total environment-step budget1005050
Position success threshold20 px4 cm16 px
Additional angle threshold20∘20^{\circ}——
DDIM inference steps10010—
Appendix
Table 6: Cross-dataset planning settings. CEM budgets denote candidates / elites / iterations.
More in this topic

World Models

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

See all 369 papers

Explore similar work

CardsList
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 ReasoningWorld Models for RoboticsWorld Model-Based PlanningJoint-Embedding Predictive ArchitectureAction-Conditioned World ModelsState Representation Learning
AskPDF
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 Sheffield, Sheffield, United Kingdom · School of Artificial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

Representation LearningWorld Model LearningWorld Model-Based PlanningJoint-Embedding Predictive ArchitectureState Representation Learning
AskPDF
Sep 28, 2026

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

Compact JEPA world models enable efficient latent-space planning, but low-dimensional representation trained under reward-free self-supervision must encode both action-conditioned dynamics and predictable visual context. This competition can entangle controllable state with high-rank nuisance appearance and degrade planning as scenes become more complex. We introduce LRC-JEPA, a lightweight end-to-end world model that routes information into a compact predictive latent z\mathbf{z} and learned-query residual-context embeddings u\mathbf{u}. Only z\mathbf{z} is propagated by the dynamics model and used for planning, while u\mathbf{u} captures temporally persistent information for cross-attention reconstruction; a differentiable residual connection encourages the latent to retain complementary dynamic content. Under explicit assumptions, we show that the resulting representation is sufficient, minimal, nuisance-invariant, and disentangled. Across four simulated control environments, LRC-JEPA improves average planning success over a parameter-matched JEPA baseline by 9 percentage points and matches or exceeds substantially larger pretrained models. On the real-world Bridge-v2 set, its 5.5M-parameter active encoder outperforms DINO-WM (22.1M) and V-JEPA2 (303.9M) encoders while also enabling faster planning. Physical-state probes, reconstruction interventions, and ablations confirm the effectiveness of LRC-JEPA's representation disentanglement.

Luzhe Huang, Lei Chu, Jingyi Liang +1

Independent researcher

Disentangled Representation LearningWorld Models for RoboticsWorld Model-Based PlanningAction-Conditioned World ModelsModel-Based PlanningLatent World Models+1
AskPDF