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

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.

• The paper introduces a joint-embedding predictive model that leverages a frozen visual encoder and a trainable predictor to improve planning using latent space optimization.
• The paper demonstrates that employing multistep rollout loss and AdaLN-based action conditioning significantly enhances performance on both simulated and real-world manipulation and navigation tasks.
• The paper finds that sampling-based optimizers like CEM and NG outperform gradient-based methods in non-convex cost landscapes, ensuring robust planning in complex environments.

Success Determinants in Physical Planning with Joint-Embedding Predictive World Models

Introduction and Motivation

This paper provides a comprehensive empirical analysis of design choices in Joint-Embedding Predictive World Models (JEPA-WMs) for visual goal-conditioned planning and robotics. JEPA-WMs leverage a frozen, pretrained encoder to map visual observations (optionally combined with proprioception) to latent spaces, where a trainable predictor models the dynamics, providing representations for planning without reliance on explicit reward supervision. This paradigm has shown performance benefits over conventional state/input-based planning, particularly for manipulation and navigation where semantic abstraction is crucial.

The authors systematically study architectural and training components—including multistep rollout losses, action and state encoding, planning optimizers, context length, and model scaling—using both synthetic and real-world robotic data to determine what factors most directly drive planning performance. They introduce improved variants of JEPA-WMs, ultimately outperforming strong baselines (DINO-WM, V-JEPA-2-AC) across standard manipulation and navigation tasks.

Formulation of JEPA-WMs

JEPA-WMs are composed of: (i) a frozen visual encoder ($E^{vis}_\phi$), possibly augmented with a proprioceptive encoder ($E^{prop}_\theta$); (ii) a jointly trained action encoder ($A_\theta$); (iii) a predictor ($P_\theta$), generally a transformer, receiving sequences of embedded observations and actions, trained to predict future latent states using a mean squared error objective.

At planning time, an action sequence is optimized in the latent space to minimize a planning cost—defined as a combination of visual and (optionally scaled) proprioceptive embedding distances to a goal observation's embedding. Open-loop and MPC-style replanning are used, with planning horizon and context window as key hyperparameters.

Figure 1: Training and planning with JEPA-WMs: the encoder and predictor operate on embedded observations and actions, with planning optimizing in latent space.

Empirical Study: Design Choices

Planning Optimizer Selection

A variety of planning optimizers are evaluated, including population-based (CEM, Nevergrad/NG) and gradient-based (GD, Adam) methods. CEM ($L_2$) is found to deliver consistently strong results, especially in multimodal or contact-rich planning landscapes. GD/Adam are only competitive in environments with smooth cost surfaces (e.g., Metaworld reach tasks) but perform poorly in tasks with local minima or discontinuities (e.g., navigation with obstacles, dexterous manipulation). NG provides practical performance nearly matching CEM on real-world data, with less hyperparameter tuning required.

Figure 2: Left: Comparison of planning optimizers (NG, CEM) across environments. Right: Impact of multistep rollout loss.

Multistep Rollout Loss

Models incorporating multi-horizon rollout losses during training outperform those relying solely on frame-wise (teacher-forcing) prediction. The benefit peaks at moderate rollout lengths; excessive rollout can degrade performance, presumably due to reduced specialization to the test-time rollout regime, and worse utilization of limited unique trajectory slices per batch. For real-world datasets (DROID), longer rollout depth (up to 6 steps) proves optimal, reflecting more complex dynamics.

Proprioceptive Input

Proprioceptive input confers a clear performance advantage, especially in environments where precise control of the agent's configuration is critical to success. On reach and navigation, proprioceptive signals help reduce goal oscillations and plan more accurate action sequences.

Figure 3: Left: Impact of proprioception. Right: Comparison of encoder types (DINO vs. V-JEPA) for Robocasa tasks.

Context Window and Model Scaling

Performance benefits from increasing the context window up to the point where velocity or acceleration can be inferred (typically $W=2$ or $3$), but longer context can reduce performance due to data sparsity and increased prediction difficulty. Model scaling shows task-dependent benefits: larger encoders and deeper predictors yield marked improvements on real-world data but offer little gain on synthetic/simulated tasks, sometimes even degrading optimization in the larger latent space.

Figure 4: Maximum context length effects on training and evaluation.

Action Conditioning and Predictor Architecture

AdaLN-based action conditioning (adaptive LayerNorm) outperforms feature or sequence concatenation, enabling stronger and more persistent propagation of action signal through the predictor. Differences in conditioning scheme prove task-dependent—spatially simple 2D navigation prefers direct feature concatenation; complex 3D manipulation benefits from per-layer, sequence-based routing.

Figure 5: Left: Action conditioning architectures; Right: Influence of predictor depth.

Encoder Choice

DINOv2/v3 encoders, trained for fine-grained spatial segmentation, outperform video-based encoders (V-JEPA) for both manipulation and navigation, attributable to superior local object-centric features. DINOv3 is superior on photorealistic/real-world imagery, while on stylized/simulated tasks, DINOv2 suffices.

Experimental Results

The proposed model configuration—carefully incorporating each positive finding—demonstrates consistent gains in success rate over strong baselines (DINO-WM, V-JEPA-2-AC) across all evaluated environments. The ablations provide detailed insight into which factors are principal drivers for various task types.

Figure 6: Qualitative comparison of model predictions in manipulation for three methods on distinct action sequences.

Planning Failure Modes

Gradient-based planners frequently get trapped in local minima, illustrated on navigation tasks where the agent becomes stuck at walls or in corners rather than achieving visually-goal relevant transitions. In manipulation, rollouts may overoptimistically hallucinate object grasping without accurate contact consequences.

Figure 7: Typical failure modes for gradient-based planners (Wall task).

Implications and Future Directions

This analysis underscores several key findings for practitioners:

• Multistep rollout objectives and nontrivial action conditioning are critical for robust, sample-efficient planning in complex environments.
• Transferability to real-world domains mandates scaling model capacity and utilizing vision encoders explicitly optimized for fine spatial discrimination.
• Sampling-based optimizers (CEM, NG) remain necessary for tasks with non-convex cost landscapes, while gradient-based planners remain limited to smooth, near-expert domains.

In practical robotics and embodied AI, these insights inform model and optimizer selection as domains scale in complexity, and strong zero-shot transfer is required.

Future research should address the limitations of current planning objectives and explore more expressive latent cost functions, as well as improved alignment between training and roll-out regimes. Attention to initial trajectory diversity and sample complexity remains essential, especially when scaling into the broader real-world manipulation spectrum.

Conclusion

This work delivers a rigorous empirical dissection of component choices in JEPA-World Models for goal-conditioned visual planning. The establishment of best practices—multistep rollout, DINO-vision encoding, AdaLN conditioning, and robust sampling-based planning—sets a new baseline and outperforms contemporary methods in both simulated and real-robot tasks. These findings are highly relevant for researchers and practitioners developing generalist, transferable agents leveraging world modeling and planning in complex, visual domains.