Skip to content

What Drives Success in Physical Planning with Joint-Embedding Predictive World Models?

🕒 Published (v1): 2025-12-30 22:50 UTC · Source: Arxiv · Venue: TMLR · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

This paper systematically ablates the design space of Joint-Embedding Predictive World Models (JEPA-WMs) — world models that plan in learned representation space rather than pixel space — across architecture, training, and planning algorithm axes. The authors identify the best recipe within this family and combine their findings into a single model that outperforms both DINO-WM and V-JEPA-2-AC on navigation and real-world robotic manipulation benchmarks.

Problem

JEPA-WMs (planning via optimization in latent embedding space, trained with a predictive JEPA loss, no reconstruction or reward head) have shown promise over pixel-space and generative world models, but the literature lacks a systematic study of which design choices actually drive performance. Prior works (DINO-WM, V-JEPA-2-AC) make different architectural and training decisions without cross-comparison, leaving practitioners without principled guidance on encoder type, predictor architecture, multistep rollout training, context length, planning optimizer, and scaling.

Method

The authors define JEPA-WMs as a unified family: a frozen visual encoder \(E_\phi^{vis}\) produces patch embeddings; a jointly-trained predictor \(P_\theta\) autoregressively predicts next-step embeddings conditioned on action embeddings; at plan time, an optimizer minimizes embedding-space distance to the goal embedding over a horizon-H action sequence. They hold the encoder frozen and vary one component at a time from a DINO-WM baseline (ViT-S encoder, depth-6 predictor), testing: - Planners: CEM, NeverGrad (NGOpt), Adam, gradient descent × {L1, L2 cost} - Multistep rollout: 1–6 step loss summation with truncated BPTT - Predictor architecture: sincos+feature-cond, RoPE+sequence-cond, RoPE+feature-cond, AdaLN+RoPE, AdaLN-zero+RoPE - Training context length W: 1–14 frames - Proprioception: joint visual+proprioceptive loss vs. visual-only - Encoder backbone: DINOv2, DINOv3, V-JEPA, V-JEPA-2 (all ViT-L for fair comparison) - Model/data scaling: ViT-S/B/L encoders, predictor depth 3–12, 2–100% of training data

Environments span 2D simulated navigation (Maze, Wall, Push-T), Metaworld arm tasks, and real-world robotics (DROID dataset, zero-shot Robocasa transfer).

Key Contributions

  • Comprehensive ablation study of 7 independent design axes within the JEPA-WM family across both simulated and real-world robotic tasks
  • Identification of CEM+L2 as the best planning optimizer overall; NeverGrad as a hyperparameter-free practical alternative matching CEM on real-world data
  • Finding that 2-step multistep rollout training is optimal for simulation (accuracy-robustness tradeoff), while K=6 is optimal for real-world DROID
  • Evidence that AdaLN+RoPE predictor conditioning is the best architecture (per-layer action modulation avoids vanishing action signal)
  • Demonstration that DINO encoders outperform video encoders (V-JEPA) due to finer object segmentation enabling more localized, learnable token dynamics
  • Quantified benefit of proprioception (resolves metric proximity failures near goal)
  • Practical scaling guideline: model scaling helps only for complex real-world data; simulation tasks saturate at ViT-S / depth-6
  • A combined optimal JEPA-WM that outperforms DINO-WM and V-JEPA-2-AC

Results

  • Maze: Ours 83.9 vs. DINO-WM 81.6 (success %)
  • Wall: Ours 78.8 vs. DINO-WM 64.1
  • Push-T: Ours 70.2 vs. DINO-WM 66.0
  • MW-Reach: Ours 58.2 vs. DINO-WM 44.8
  • MW-ReachWall: Ours 41.6 vs. DINO-WM 35.1
  • Robocasa-Reach: Ours 25.4 vs. V-JEPA-2-AC 16.2 vs. DINO-WM 19.1
  • Robocasa-Place: Ours 30.7 vs. V-JEPA-2-AC 33.1 vs. DINO-WM 21.7
  • DROID (action score): Ours 48.2 vs. V-JEPA-2-AC 42.9 vs. DINO-WM 39.4
  • Gradient-based planners (Adam/GD) fail on non-greedy navigation tasks (stuck in local minima) but are competitive on smooth Metaworld cost landscapes
  • Data scaling consistently improves all methods; performance gaps between ours and baselines are largest at lower data regimes (DROID, Wall)

Limitations

  • Encoder is frozen throughout; encoder–predictor co-training (as in PLDM, EB-JEPA) is explicitly out of scope
  • Planning provides no feasibility guarantees; the approach is trajectory optimization in disguise
  • Zero-shot transfer to Robocasa from DROID suffers from domain gap; proprioception cannot be used across datasets with misaligned proprioceptive spaces
  • Success rates on real manipulation tasks remain low (25–48%), with high variance across seeds
  • Robocasa Place task is not clearly won by the proposed model vs. V-JEPA-2-AC (30.7 vs. 33.1)
  • Study restricted to short-horizon, goal-conditioned tasks; long-horizon or reward-shaped scenarios not evaluated

Relevance to Agentic AI / LLM Agents

JEPA-WMs represent a non-language, perception-action paradigm for building physical agents that plan via internal world models rather than reactive policies — directly relevant to the model-based planning branch of agentic AI. The finding that latent-space planning (no pixel reconstruction, no reward) generalizes to real-world robotic data with careful architecture choices informs how future agentic systems might couple vision foundation models with learned dynamics for embodied task execution. The ablation methodology — systematically isolating planner, architecture, training objective, and data scale — is a template for rigorous evaluation of any world-model-based agent. This work also benchmarks the gap between simulation and real-world transfer, a key open challenge for deploying agentic systems in physical environments.