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

AI summary

Overview

Research area: Model-based reinforcement learning and robot learning, specifically action-conditioned Joint-Embedding Predictive World Models (JEPA-WMs) used for latent-space planning.

Technical level: Advanced. The paper assumes familiarity with self-supervised visual encoders (DINOv2/DINOv3, V-JEPA), transformer predictors, and trajectory optimization.

Scope: A systematic ablation study of the design space of JEPA-WMs, covering architecture, training objective, and planning algorithm, evaluated on simulated navigation and manipulation plus real-world robotic data.

What This Paper Is About

Recent world-model methods learn to predict the future in the embedding space of a pretrained visual encoder and then plan by optimizing action sequences inside that embedding space. The paper asks a question it describes as a gap in the literature: how do you efficiently learn a dynamics model (the predictor) on top of a pretrained visual encoder for manipulation and navigation planning? Rather than proposing a brand-new algorithm, the authors characterize this family as JEPA-WMs and perform a controlled study of which technical choices actually drive planning success, then combine the best settings into a single recommended model.

Key Contributions

1. A unified formalization of JEPA-WMs as a training and planning recipe — a dynamics model trained solely through a predictive loss in embedding space, with no reconstruction, reward prediction, or value/policy heads — presented explicitly as a unified implementation recipe rather than a novel algorithm.

2. A comprehensive empirical study of eight design axes: multistep rollout training, predictor architecture, training context length, use of proprioception, encoder type, model size, data scaling, and the planning optimizer (including the introduction of a NeverGrad-based planner the authors denote NG).

3. A per-task recommended recipe (simulated navigation vs. real-world manipulation) summarized in the paper's Table 1, across encoder, predictor architecture, predictor depth, rollout steps, context length, proprioception, planning optimizer, and cost function.

4. A combined model that outperforms two established baselines, DINO-WM and V-JEPA-2-AC, on both navigation and manipulation tasks, with code, data, and checkpoints released.

Main Findings

• Planner choice matters most, and the best overall planner is CEM with an L2 cost. The Cross-Entropy Method with L2 distance is either best or competitive on all environments, so the authors fix it for all subsequent experiments. Across all planning setups and models, the L2 cost consistently outperforms L1.

• Gradient-based planners are task-dependent. Adam with L2 achieves the best overall performance on Metaworld, and Gradient Descent is also competitive with CEM there, which the authors attribute to relatively smooth cost landscapes where the goal is greedily reachable. On 2D navigation (Wall, Push-T, Maze) that requires non-greedy planning, gradient-based methods perform very poorly because GD gets stuck in local minima. On DROID they also perform significantly worse, due to multi-modal cost landscapes. Robocasa, being simulated but closer in nature to Metaworld, allows gradient-based methods to perform reasonably well again.

• NG is a practical tuning-free alternative. On 2D navigation, CEM clearly outperforms NG because these tasks need precise action sequences; NG converges more slowly, indicating more exploration. On DROID and Robocasa, CEM and NG perform similarly, and NG requires no hyperparameter tuning — unlike CEM, which needs the top-K trajectories parameter and the initialization of the proposal Gaussian distribution.

• Multistep rollout training helps up to a point. Performance increases from pure teacher-forcing models to 2-step rollout loss models, then decreases for models trained in simulated environments. On DROID, reducing the effective Lipschitz constant more than compensates the increase in one-step error, and the optimal point shifts to 6 steps.

• Proprioception consistently helps. Models trained with proprioceptive input are consistently better than pure-visual models. On Metaworld, without proprioception the arm reaches the vicinity of the goal but oscillates, unable to resolve the remaining distance from vision alone. Robocasa results are not displayed because the proprioceptive space is not aligned between DROID and Robocasa.

• Planning context must not exceed training context. The authors confirm that W_p ≤ W must hold, otherwise predictions degrade rapidly through unrolling. There is a large performance gap between models trained with W = 1 and W = 2, since the predictor needs two frames to infer velocity and three to infer acceleration. Optimal context is 3 on simulated datasets but 5 on DROID, which the authors attribute to more complex dynamics. On DROID, W = 14 retains only 86% of videos.

• DINO encoders beat V-JEPA encoders. The authors posit this is due to DINO's better fine-grained object segmentation, which makes object motion translate into localized, sparse token changes the predictor can learn efficiently. DINOv3 clearly outperforms DINOv2 only on the more photorealistic environments, Robocasa and DROID; on Maze and Wall, DINOv3 models take longer to converge to a lower success rate.

• AdaLN with RoPE wins on average, but narrowly. On Metaworld, sincos with feature conditioning actually performs best, and the authors see no substantial improvement from RoPE over sincos positional embedding. AdaLN-zero achieves higher average performance but underperforms AdaLN on the environments providing the most reliable signal (DROID, PushT, Maze), so AdaLN is chosen for the final model.

• Scaling helps real-world data but not simulated tasks. Increasing encoder size (with predictor width) or predictor depth does not improve performance on simulated environments, where the optimal predictor depth is 6 for most environments and possibly as low as 3 for the simplest 2D navigation tasks (Wall, Maze). On DROID there is a clear and consistent positive correlation between both encoder size and predictor depth and planning performance.

• Data scaling helps across the board. Performance clearly increases when scaling data for all datasets and methods considered, evaluated at 2%, 10%, 50%, and 100% of available training data. (The provided excerpt cuts off mid-sentence at this point, so the remaining insights from this experiment are not reported here.)

• Successful unrolling is not sufficient for successful planning. The authors stress that even models able to faithfully unroll a large number of actions do not automatically succeed at the planning task.

Methodology in Plain English

The authors start from a simple base configuration: DINO-WM without proprioception, using a ViT-S encoder and a depth-6 predictor of the same embedding dimension. They then vary one component at a time from this base to isolate its effect, prioritizing planning-time choices first (because they affect every evaluation) and then fixing the best planner per environment for the remaining experiments.

Training works by embedding observations with a frozen visual encoder plus an optional shallow proprioceptive encoder, embedding actions with an action encoder, and having a transformer predictor output the next state embedding. Only the predictor, action encoder, and proprioceptive encoder are trained; the visual encoder stays frozen. The predictor uses a frame-causal attention mask, so it is trained simultaneously to predict from all context lengths from w = 0 to W − 1, with W = 3 by default (raised to 7 for the longer-horizon models).

Planning is cast as an optimization over a horizon-H product action space. Given an initial and goal observation, candidate action trajectories are sampled, the predictor is unrolled on them, and a planning cost combining visual and proprioceptive embedding distances (with weight α = 0.1, or α = 0 for DROID and Robocasa) is minimized. Four optimizers are compared: CEM, the NeverGrad-based NG planner (using the default NGOpt meta-optimizer with no parameter tuning), Adam, and plain gradient descent.

Evaluation uses Metaworld tasks gathered by training TD-MPC2 online agents (Reach and Reach-Wall), the offline datasets released by Zhou et al. — Push-T, Wall, and PointMaze — with a 90% train split, plus DROID for training with zero-shot evaluation on Robocasa (custom pick-and-place tasks "Place" and "Reach") and on 16 videos of a real Franka arm filmed in the lab. DROID models are not finetuned on Robocasa trajectories. The main metric is success rate; the authors also track embedding space error through unrolling, proprioceptive decoding error, visual decoding of open-loop rollouts, and LPIPS between decodings and groundtruth frames. Final models are trained with 3 seeds, and evaluation uses 96 episodes per epoch for most environments, 64 for DROID, and 32 for Robocasa.

Why This Matters

Impact on research: The paper provides a controlled, component-by-component map of an increasingly popular model family, and reframes the design question away from the encoder and toward the predictor and planner. It also contributes an analytical framing of the accuracy-robustness tradeoff in multistep training: because the predictor is a network with Lipschitz constant Λ ≥ 1 in general, compounding errors grow exponentially with the horizon, so increasing the number of rollout steps K raises the one-step error but reduces the effective Lipschitz constant — making the multistep loss act as data augmentation against compounding error, analogous to scheduled sampling.

Real-world applications:

• Robot manipulation from camera input, including pick-and-place and reach tasks trained on large real-world datasets such as DROID.
• Robotic navigation in 2D and physical environments, including maze and wall-avoidance settings.
• Autonomous driving and egocentric video-game agents, which the introduction identifies as domains where large-scale world models have achieved impressive simulation accuracy.
• Visual servoing and model-predictive control loops for robot arms with exteroceptive cameras.

Industry relevance: The paper is co-authored across Meta FAIR, Inria Paris, ENS/PSL, and NYU, and releases code, data, and checkpoints at a public repository. The finding that the NG planner performs on par with CEM on real-world manipulation data while requiring no hyperparameter tuning is directly relevant to practitioners who would otherwise pay a large tuning cost when moving to new tasks or datasets.

Future Directions

• Extending the study beyond frozen-encoder predictors toward lightweight adaptation of frozen visual foundation model encoders, which the authors identify as a complementary direction explored by concurrent work.

• Understanding why scaling model capacity helps real-world data but saturates on simulated tasks, and whether the proposed explanations (task simplicity, harder optimization landscapes in higher dimensions, fewer gradient updates per parameter at fixed compute) can be disentangled.

• Investigating whether the recommended recipe transfers to new embodiments and datasets, given the observation that the proprioceptive space is not aligned between DROID and Robocasa and therefore models using proprioception were irrelevant for zero-shot transfer in that setting.

• Filling out the data-scaling analysis: the provided excerpt ends mid-discussion of the 2%/10%/50%/100% regimes, so the full conclusions of that experiment and the direct comparison against DINO-WM and V-JEPA-2-AC per data regime are not reported in the available content.

Target Audience

Researchers and graduate students working on model-based reinforcement learning, world models, and robot learning; engineers building latent-space planners on top of pretrained visual encoders; and practitioners deciding between sampling-based and gradient-based trajectory optimizers for real-world manipulation. Readers without a background in JEPA-style self-supervised learning, transformer predictors, or trajectory optimization will find the paper demanding, but the per-component summary table makes the concrete recommendations accessible to applied roboticists.

Authors’ 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.