Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Basile Terver, Tsung-Yen Yang, Jean Ponce, Adrien Bardes, Yann LeCun

TL;DR

The paper tackles enabling robust planning for physical tasks by examining how joint-embedding predictive world models (JEPA-WMs) can be learned and planned over in latent space. It systematically studies design choices across training and planning, including multistep rollout, proprioception, encoder and predictor architectures, model scaling, and planning optimizers, to identify an optimum JEPA-WM configuration. Across simulated and real-world robotic datasets, the authors demonstrate that carefully combining AdaLN-conditioned predictors, RoPE positional encoding, proprioceptive input, and a 2-step rollout within a latent-space planning framework yields superior performance over baselines such as DINO-WM and V-JEPA-2-AC, with planning robustness depending on task complexity. The work provides practical guidance on when to scale models, which encoders to prefer, and how to select planning strategies (e.g., CEM $L_2$ vs.gradient-based planners) for reliable zero-shot transfer to new robotic tasks. These findings offer actionable insights for building data-efficient, generalizable planning systems in robotics and related domains.

Abstract

A long-standing challenge in AI is to develop agents capable of solving a wide range of physical tasks and generalizing to new, unseen tasks and environments. A popular recent approach involves training a world model from state-action trajectories and subsequently use it with a planning algorithm to solve new tasks. Planning is commonly performed in the input space, but a recent family of methods has introduced planning algorithms that optimize in the learned representation space of the world model, with the promise that abstracting irrelevant details yields more efficient planning. In this work, we characterize models from this family as JEPA-WMs and investigate the technical choices that make algorithms from this class work. We propose a comprehensive study of several key components with the objective of finding the optimal approach within the family. We conducted experiments using both simulated environments and real-world robotic data, and studied how the model architecture, the training objective, and the planning algorithm affect planning success. We combine our findings to propose a model that outperforms two established baselines, DINO-WM and V-JEPA-2-AC, in both navigation and manipulation tasks. Code, data and checkpoints are available at https://github.com/facebookresearch/jepa-wms.

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

TL;DR

The paper tackles enabling robust planning for physical tasks by examining how joint-embedding predictive world models (JEPA-WMs) can be learned and planned over in latent space. It systematically studies design choices across training and planning, including multistep rollout, proprioception, encoder and predictor architectures, model scaling, and planning optimizers, to identify an optimum JEPA-WM configuration. Across simulated and real-world robotic datasets, the authors demonstrate that carefully combining AdaLN-conditioned predictors, RoPE positional encoding, proprioceptive input, and a 2-step rollout within a latent-space planning framework yields superior performance over baselines such as DINO-WM and V-JEPA-2-AC, with planning robustness depending on task complexity. The work provides practical guidance on when to scale models, which encoders to prefer, and how to select planning strategies (e.g., CEM vs.gradient-based planners) for reliable zero-shot transfer to new robotic tasks. These findings offer actionable insights for building data-efficient, generalizable planning systems in robotics and related domains.

Abstract

A long-standing challenge in AI is to develop agents capable of solving a wide range of physical tasks and generalizing to new, unseen tasks and environments. A popular recent approach involves training a world model from state-action trajectories and subsequently use it with a planning algorithm to solve new tasks. Planning is commonly performed in the input space, but a recent family of methods has introduced planning algorithms that optimize in the learned representation space of the world model, with the promise that abstracting irrelevant details yields more efficient planning. In this work, we characterize models from this family as JEPA-WMs and investigate the technical choices that make algorithms from this class work. We propose a comprehensive study of several key components with the objective of finding the optimal approach within the family. We conducted experiments using both simulated environments and real-world robotic data, and studied how the model architecture, the training objective, and the planning algorithm affect planning success. We combine our findings to propose a model that outperforms two established baselines, DINO-WM and V-JEPA-2-AC, in both navigation and manipulation tasks. Code, data and checkpoints are available at https://github.com/facebookresearch/jepa-wms.
Paper Structure (74 sections, 4 equations, 20 figures, 13 tables, 2 algorithms)

This paper contains 74 sections, 4 equations, 20 figures, 13 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related work
  3. World modeling and planning.
  4. Goal-conditioned RL.
  5. Robotics.
  6. Background
  7. Training method.
  8. Planning.
  9. Studied design choices
  10. Planner.
  11. Multistep rollout training.
  12. Proprioception.
  13. Training context size.
  14. Encoder type.
  15. Predictor architecture.
  16. ...and 59 more sections

Figures (20)

  • Figure 1: Left: Training of JEPA-WM: the encoder $E_{\phi,\theta}$ embeds video and optionally proprioceptive observation, which is fed to the predictor $P_{\theta}$, along with actions, to predict (in parallel across timesteps) the next state embedding. Right: Planning with JEPA-WM: sample action sequences, unroll the predictor on them, compute a planning cost $L^p$ for each trajectory, and use this cost to iteratively refine the action sampling. The action encoder $A_{\theta}$ and proprioceptive encoder $E_{\theta}^{prop}$ are not explicitly displayed in this figure for readability.
  • Figure 2: Comparison of different methods on the counterfactual Franka arm lift cup task, where we hardcode 2 actions, either "open and move up" or "close and move up". Each shows 5 model actions in open-loop rollout. Left: "open and move up" action. Right: "close and move up". First row: V-JEPA-2-AC. Second row: DINO-WM. Third row: our best model, described in \ref{['subsec:optimum']}.
  • Figure 3: Left: Comparison of planning optimizers: NG is the Nevergrad-based interface for trajectory optimization that we introduce, compared to the Cross-Entropy Method (CEM), with $L_1$ or $L_2$ distance. Right: Effect of adding multistep rollout loss terms: models are trained with total loss $\mathcal{L}_1 + \dots + \mathcal{L}_k$. Rc-Pl and Rc-R denote the Place and Reach tasks of Robocasa.
  • Figure 4: Left: Models trained with proprioceptive input are denoted "prop", while pure visual world models are named "no-prop". Right: Comparison of JEPA-WMs trained on top of various pretrained visual encoders, all of size ViT-L for fair comparison. Rc-Pl and Rc-R denote the Place and Reach tasks of Robocasa.
  • Figure 5: Left: Maximum number of timesteps of state embedding seen by the predictor at train time in equation \ref{['eq:frame-teacher-forcing-loss']}, the predictor takes up to $\left( E_{\phi,\theta}(o_{t-W+1:t}), A_{\theta}(a_{t-W+1:t}) \right)$ as context. Right: Comparison of model size: we vary from ViT-S to ViT-L the visual encoder size, as well as the predictor embedding dimension, keeping predictor depth constant at 6. Rc-Pl and Rc-R denote the Place and Reach tasks of Robocasa.
  • ...and 15 more figures