Source-linked AI summary
Toward Physically Grounded JEPA World Models for Goal-Conditioned Robotic Planning
Muyuan Liu, Yue Huang, Zheng Liang, Xiang Gao
TL;DR
The paper addresses the risk that latent prediction in action-conditioned JEPA world models fails to retain control-relevant physical information. It introduces end-to-end inverse dynamics and state alignment, achieving the highest reported success on TwoRoom, PushT, and OGBench-Cube while improving planning over IDM alone across four tasks.
Problem
Latent prediction alone can collapse and does not explicitly ensure that representations retain information relevant to robotic control.
Method
The model augments action-conditioned latent prediction with inverse dynamics and pair-based state alignment using measured physical states.
Results
The model attains the highest success rates on TwoRoom, PushT, and OGBench-Cube, performs comparably to LeWorldModel on Reacher, and improves planning over IDM alone across all four tasks.
Takeaways & Limitations
State alignment is supported as a complement to inverse dynamics, while transition-subspace analysis shows that higher straightening can coincide with lower effective transition dimension.
Abstract
from arXiv · showhide
Action-conditioned JEPA world models enable planning toward visually specified goals without reconstructing future pixels, yet latent prediction alone does not explicitly encourage the learned representations to retain information relevant to robotic control. We introduce an end-to-end JEPA world model that augments latent prediction with inverse dynamics (IDM) and state alignment (SA). While inverse dynamics discourages latent collapse and makes latent transitions informative of the actions that produced them, state alignment grounds consecutive representations in their associated physical configuration and motion. Across four benchmark tasks, our model attains the highest success rates on TwoRoom (100%), PushT (98%), and OGBench-Cube (87%), while performing comparably to LeWorldModel on Reacher. Our ablation further shows that adding state alignment consistently improves planning success over IDM alone across all four tasks. Although LeWorldModel, our primary baseline, attains higher average straightening on OGBench-Cube, transition-subspace analysis shows that its transition energy is concentrated in a substantially lower-dimensional subspace. Our state-aligned model exhibits a higher effective transition dimension than LeWorldModel and improves planning over IDM alone, supporting state alignment as an effective complement to inverse dynamics for robotic planning.
I. INTRODUCTION
JEPA world models avoid future-pixel reconstruction, but latent prediction alone can collapse or fail to preserve control-relevant physical information. The proposed approach complements inverse dynamics with state alignment and evaluates both planning success and latent transition structure.
- Latent prediction can map all observations to one representation, while existing approaches rely on visual pretraining, distributional regularization, or composite objectives to stabilize learning.
- State alignment complements inverse dynamics by grounding consecutive latent representations in measured physical states and motion.
- The model is evaluated on four goal-conditioned planning tasks using planning success, temporal straightening, and transition-subspace analyses.
- The contributions include an end-to-end state-aligned JEPA model and evidence that higher straightening can coincide with lower transition dimension and weaker planning.
II. METHOD
The method uses a fixed offline dataset collected by a behavior policy, with trajectories containing observations, physical measurements, and action chunks.
- The dataset is fixed and collected offline by a behavior policy.
- Each trajectory records visual observations, corresponding physical measurements, and action chunks between successive observations.Action chunks contain a fixed number of low-level actions executed between observations.
B. Model Architecture
The model jointly trains an action-conditioned JEPA predictor with inverse dynamics and state alignment. State alignment uses consecutive latent representations to connect latent transitions with physical configuration and velocity.
- The backbone uses a visual encoder and action-conditioned latent predictor trained end-to-end without future-pixel reconstruction.
- Latent prediction alone admits collapse, so inverse dynamics predicts executed actions from consecutive representations and encourages action-informative variation.
- State alignment adds physical supervision because inverse dynamics does not directly anchor latent pairs to measured physical states.
- Consecutive latent pairs provide temporal context for aligning both instantaneous physical configuration and velocity measurements.
- Physical measurements supervise training only; deployment plans entirely in latent space using the encoder and predictor.
C. Implementation
The implementation uses a lightweight LeWorldModel-style backbone and evaluates qualitative goal-conditioned control on OGBench-Cube. The state-alignment and inverse-dynamics heads are jointly optimized with the backbone.
- The encoder is a randomly initialized ViT-Tiny/14 producing 192-dimensional representations, and the predictor is a six-layer causal Transformer.The predictor uses up to three recent latent–action pairs as context.
- Action and Euclidean state quantities are standardized per dimension, while angular quantities use continuous coordinates to avoid wrap-around discontinuities.
- In the OGBench-Cube visualization, the top row gives the reference trajectory and the bottom row shows simulator observations produced by CEM actions.
- The state-alignment and inverse-dynamics heads are two-layer MLPs, and all components are optimized jointly with AdamW.
D. Latent Planning
At deployment, the fixed encoder and predictor evaluate candidate action sequences by recursively predicting terminal latent representations and minimizing their distance to the goal representation.
- Candidate action sequences are recursively rolled out through the predictor to obtain a terminal latent representation.
- Planning selects actions by minimizing terminal latent distance to the encoded goal representation.
- The Cross-Entropy Method samples action sequences, evaluates latent rollouts, updates toward elite candidates, and repeats execution until success or budget exhaustion.
A. Experimental Setup
The evaluation follows LeWorldModel’s four-task benchmark with fixed start–goal problems, a 25-step horizon, and a 50-step interaction budget; auxiliary-loss weights are selected on PushT validation.
- The benchmark covers TwoRoom, Reacher, PushT, and OGBench-Cube, spanning navigation, reaching, pushing, and robot-arm manipulation.
- Each variant is evaluated on 50 fixed start–goal problems with the goal 25 environment steps ahead and an interaction budget of 50 steps.
- Table I reports goal-conditioned planning success rates, using LeWorldModel baseline means and three independently trained seeds for the proposed results.
- The tied auxiliary-loss weights α = β = 1.0 are chosen from {0.01, 0.1, 1.0} using PushT validation success and applied across all four tasks.
B. Planning Performance in Latent Space
State alignment improves planning success over IDM alone across all four tasks, while the proposed method is strongest on three tasks and comparable to LeWorldModel on Reacher.
- The proposed method attains the highest success rates on TwoRoom, PushT, and OGBench-Cube, while performing comparably to LeWorldModel on Reacher.
- Adding state alignment consistently improves planning success over the IDM-only ablation across all four benchmark tasks.
- A representative CEM-controlled OGBench-Cube trial reaches the image-defined goal within the evaluation budget.
- Figure 3 compares time-colored PCA projections and local straightening scores for a representative OGBench-Cube trajectory.
- Table II reports OGBench-Cube temporal-straightening means and standard deviations over 100 fixed trajectories.
C. Temporal Straightening
Temporal straightening measures local alignment of consecutive latent displacements, but higher average straightening does not necessarily indicate better planning or richer transition structure.
- Temporal straightening quantifies how consistently consecutive latent displacements align along temporally coherent physical trajectories.
- Adding state alignment lowers the IDM-only model’s mean straightening score on OGBench-Cube.
- LeWorldModel has the highest average straightening despite lower OGBench-Cube planning success, with isolated near-reversals unlike the more distributed changes of IDM and SA+IDM.
- Transition-subspace analysis applies uncentered SVD to latent displacements and measures r95, the components needed to retain 95% of transition energy.
- Mean r95 is 17 for LeWorldModel, 30.2 for IDM, and 32.8 for SA+IDM, linking higher straightening with a smaller effective transition dimension.
IV. CONCLUSION
The proposed JEPA world model combines inverse dynamics with pair-based state alignment, improving planning success over IDM alone across four goal-conditioned tasks. Transition-subspace analysis indicates that higher temporal straightening can coincide with lower effective transition dimension.
- State alignment consistently improves planning success over the IDM-only ablation across four goal-conditioned tasks.
- Higher temporal straightening can coincide with a substantially lower effective transition dimension and weaker planning performance.