Skip to content
Preprint

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

Aug 2026 · 2 citations · 41 references
Computer Science

TL;DR

PSG-JEPA is proposed, 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.

Abstract

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.

View source

Similar papers

#machine learning Preprint Aug 2026

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

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.

Kehan Wen, Ziming Li, Siyuan Luo et al. · 0 citations
Preprint Aug 2026

Predicting Consequences and Reinforcing Navigation Policies with Latent World Models

This work proposes a compatibility prediction Latent World Model for robot navigation that predicts action-conditioned latent feature compatibility rather than reconstructing observations and demonstrates how the learned world model can supervise policy learning from unlabeled video data and improve policies through reinforcement learning entirely within the world model.

Zengmao Wang, Wei Gao, Shuhan Shen · 0 citations
Preprint Aug 2026

Orthogonal JEPA: Factorized Predictive States for Latent World Models

Method, a latent world-modeling framework based on orthogonal predictive factorization, is introduced, a latent world-modeling framework based on orthogonal predictive factorization that can be used by a readout, decoder, planner, or autoregressive rollout of an underlying system.

Taoyong Cui, Pheng-Ann Heng, Wanli Ouyang · 0 citations
Preprint Jul 2026

On the Identifiability of Controlled World Models

A joint identifiability condition for controlled world models with Gaussian latent states with Gaussian latent states is presented, which consists of two coupled components: representation identifiability and transition identifiability, and it is proved that when this condition holds, minimizing the LeJEPA-style predictive objective can recover both latent states and controlled dynamics in the sense of orthogonal transformation.

Xiangteng Zhang, Yang Guan, Bo Zhang et al. · 0 citations
Preprint Jul 2026

ActSWM: Action-Sensitive World Models for Long-Horizon Planning in Open-World Games

This work proposes ActSWM, an action-sensitive latent world model grounded in a transition-separation principle, which preserves larger action-dependent rollout gaps than existing baselines, improves task success in long-horizon interactive settings, and enables world-model-based action recovery from offline gameplay videos.

Zhe Gan, Zitong Zeng, Jiajun Cheng et al. · 1 citation
Preprint Jul 2026

SAGE: Subgoal-Conditioned Action Generation for Latent World Model Planning

Latent world models have emerged as a powerful planning paradigm by learning action-conditioned predictive dynamics and using them as internal simulators to imagine and evaluate candidate action sequences. However, as the planning horizon grows, performance becomes increasingly constrained by proposal quality: a fixed candidate budget must search an exponentially larger action space, making it difficult to expose the world model to high-quality candidate futures for evaluation. In this paper, we introduce a prior-conditioned planner that replaces random proposal initialization with structured guidance. At each planning stage, a goal-conditioned generator predicts the next reachable latent subgoal for a specified duration, which is then used to condition the generation of candidate action sequences. To capture semantic information across temporal scales, we use subgoals of varying durations as priors, balancing fine-grained local control with higher-level long-horizon progress. Then the frozen world model evaluates and refines these subgoal-conditioned proposals before execution. Experiments on PushT and OGBench Cube show that coupling latent subgoal decomposition with prior-conditioned action generation substantially improves long-horizon planning while preserving strong short-horizon performance. To be specific, when the target offset is $150$, it raises PushT success from $12.7\%$ to $64.7\%$ and OGBench Cube success from $26.7\%$ to $67.3\%$.

Le-Tian Cheng, Qi Zhang, Yisen Wang · 2 citations