---
title: Drive-JEPA Framework for Visuo-Motor Control
url: https://www.emergentmind.com/topics/drive-jepa-framework
type: topic
---

# Drive-JEPA Framework for Visuo-Motor Control

Drive-JEPA is a class of world-model-plus-planner frameworks in which goal-directed planning is performed in the joint embedding space of observations and actions. It combines a Joint-Embedding Predictive Architecture (JEPA) with lightweight, sampling-based planners to produce highly data-efficient solutions for visuo-motor control and long-horizon policy optimization. Drive-JEPA architectures are characterized by their use of self-supervised (often fully unsupervised) predictive objectives, transformer-based sequence modeling, and empirical model selection via systematic ablation of training and planning components. The result is a unified latent representation space where optimal trajectories can be found using lightweight optimization, enabling robust generalization across simulated and real-world robotics scenarios [2512.24497].

## 1. Drive-JEPA Architecture

Drive-JEPA models are structured around four principal components: encoders, a latent-dynamics predictor, diagnostic decoders, and a latent-space planner.

- **Encoders:**  
  Visual observations $o_t\in\mathbb{R}^{H\times W\times3}$ and, optionally, proprioceptive state vectors $s_t\in\mathbb{R}^{d_\mathrm{prop}}$ are passed through a frozen ViT-based image encoder ($E^{\mathrm{vis}}_\phi$) and a trainable proprioceptive encoder ($E^{\mathrm{prop}}_\theta$), respectively. Their outputs are concatenated to yield a combined latent state embedding. Actions $a_t\in\mathbb{R}^{A}$ are encoded via a trainable action encoder $A_\theta$.

- **Latent-Dynamics Predictor:**  
  A depth-$L$ transformer $P_\theta$, with rotary positional embeddings (RoPE) and Adaptive LayerNorm (AdaLN) conditioning on the action tokens, models temporal evolution in the latent space. The transformer is unrolled over a context window of past $K$ steps, using predicted or teacher-forced state tokens with corresponding actions.

- **Diagnostic Decoders:**  
  Optional visual and proprioceptive decoders enable LPIPS reconstruction or prediction error monitoring during diagnostic analysis, but do not directly impact planning performance.

- **Planner:**  
  Planning is formulated as trajectory optimization in the latent space over a finite horizon $H$, seeking to minimize the distance between the predicted latent state at horizon $z_H$ and the embedding of the goal observation $E_{\phi,\theta}(o_g)$. The main optimization method is the Cross-Entropy Method (CEM), supplemented in real-world contexts by CMA-ES via NeverGrad.

## 2. Training Objectives

Drive-JEPA is trained using a multi-step predictive mean squared error (MSE) computed within the embedding space of frozen encoders. For context length $w\leq W$, the main loss for batch element $b$ with predicted latent state $\hat{z}$ at each step is:

$$
L^{(k)} = \frac{1}{B} \sum_{b=1}^{B} \|P_\theta(\hat{z}_{t-w:t+k-1}^b,\,A_\theta(a_{t-w:t+k-1}^b)) - E_{\phi,\theta}(o_{t+k}^b)\|_2^2
$$

The total loss is the sum over rollout steps up to $K$ (e.g., $K=2$ for simulation, $K=6$ for real-world robotics):

$$
L_\mathrm{train} = \sum_{k=1}^K L^{(k)}
$$

Proprioceptive and visual branches are weighted equally. No contrastive or negative sampling is used. Truncated Backpropagation Through Time (TBPTT) is employed, detaching gradients from past predictions to aid stability.

## 3. Latent-Space Planning Algorithm

Given an initial observation $o_0$ and a goal $o_g$, a candidate action sequence $\{a_0, ..., a_{H-1}\}$ is sampled:

$$
z_H = F_{\phi,\theta}(o_0, a_{0:H-1}) = P_\theta(\ldots P_\theta(E(o_0), A(a_0)), \ldots, A(a_{H-1}))
$$

The planning cost to the goal is:

$$
L^p(o_0, a_{0:H-1}, o_g) = \|z_H - E_{\phi,\theta}(o_g)\|_2^2 + \alpha\|s_H - s_g\|_2^2
$$

if proprioceptive targets are used. Optimization employs CEM: $N=300$ trajectories are sampled, top-$K=10$ are kept, mean and diagonal covariance are updated, and this is repeated for $J=15$ iterations before applying the best $m$ actions in the environment. Empirically, $L_2$-norm cost outperforms $L_1$, and CEM/L2 is superior to gradient-based solvers for contact-rich or highly multi-modal tasks.

## 4. Empirical Results and Ablation Insights

Final Drive-JEPA performance surpasses prior art (DINO-WM, V-JEPA-2-AC) in both navigation and manipulation. Representative success rates (mean $\pm$ std):

| Task               | DINO-WM      | V-JEPA-2-AC | Drive-JEPA     |
|--------------------|--------------|-------------|----------------|
| Maze               | 81.6 ± 3.4   |   —         | 83.9 ± 2.3     |
| Wall               | 64.1 ± 4.6   |   —         | 78.8 ± 3.9     |
| Push-T             | 66.0 ± 4.7   |   —         | 70.2 ± 2.8     |
| MW-Reach           | 44.8 ± 8.9   |   —         | 58.2 ± 9.3     |
| MW-Reach-Wall      | 35.1 ± 9.4   |   —         | 41.6 ± 10.0    |
| Robocasa-Reach     | 19.1 ± 13.4  | 16.2 ± 8.3  | 25.4 ± 16.6    |
| Robocasa-Place     | 21.7 ± 7.2   | 33.1 ± 7.2  | 30.7 ± 8.0     |
| DROID              | 39.4 ± 2.1   | 42.9 ± 2.5  | 48.2 ± 1.8     |

Boosts of +14% on Metaworld Reach and +23% on Wall are observed. Ablations highlight the following:

- **Multistep rollout loss:** 2-step rollout yields a 10–20% success boost in simulation; >2 steps degrades with short context, but up to 6 steps are needed in real robotics.
- **Context length:** Shorter training contexts (W=3 train, $W^p=2$ plan) are optimal in simulation, longer in photorealistic real tasks.
- **Proprioception:** Including the proprio branch increases push/navigation success by 10–20%, with reaching benefitting from additional proprioceptive input (weight $\alpha = 0.1$).
- **Visual encoder:** DINOv2/v3 image encoders outperform video encoders (V-JEPA) for manipulation; DINOv3 further improves in photorealistic environments.
- **Predictor conditioning:** AdaLN+RoPE is optimal on average, though feature- or sequence-based conditioning is task-dependent.
- **Model scaling:** ViT-S/depth-6 is sufficient in simulation; scaling to ViT-L/depth-12 brings an 8–10% improvement in real environments.
- **Planner:** CEM/L2 is robust across a range of tasks, while gradient descent/Adam solvers are competitive only in smooth, single-modal domains.

## 5. Distinguishing Features of Drive-JEPA

Several innovations distinguish Drive-JEPA within the JEPA-WM family [2512.24497]:

- Systematic co-tuning of encoder type, context length, rollout horizon, predictor conditioning, model size, and planner.
- Alignment of training and planning regimes: matching context and rollout for sharp, optimizable latent dynamics.
- Dynamic model scaling: lightweight models suffice in simulators, while larger architectures reap gains in the real world.
- Full zero-reward, goal-conditioned, self-supervised training—no extrinsic reward is required for model learning.
- The resultant latent space yields feasible, optimizable trajectories via lightweight latent-space CEM, with robust generalization to real robot navigation and manipulation domains.

## 6. Relationship to Related Frameworks

Drive-JEPA establishes a new optimum within JEPA-World Model (JEPA-WM) methods by integrating insights from architecture design, training objectives, and planning mechanics. Unlike pixel-reconstruction or contrastive learning world models, Drive-JEPA does not reconstruct raw pixels or generate negative sample pairs; instead, it optimizes readily-computable prediction losses in a structured representation space. In direct empirical comparison, Drive-JEPA's principled ablation-driven methodology leads to superior planning success and transfer, providing a reproducible, scalable solution for both simulated and real-world continuous control challenges [2512.24497].

## 7. Implementation Practices and Open Challenges

Drive-JEPA's implementation best practices emphasize:

- Matching rollout length ($K$ steps), context window, and planning horizon to task domain (simulation vs. real).
- Selection of visual encoders and predictor architectures appropriate to task visual complexity.
- Minimal reward engineering: full self-supervised loss regimes allow for rapid domain transfer.
- Lightweight planning (CEM or CMA-ES) in latent space for computational efficiency.

Open challenges include the extension to long-horizon temporal memory, multi-modal input fusion (e.g., camera, LiDAR), and further investigation of hybrid value–prediction shaping or hierarchical latent structures to close remaining gaps in generalization and sample-efficiency [2512.24497].

---

**Key Reference:**  
- "What Drives Success in Physical Planning with Joint-Embedding Predictive World Models?" [2512.24497]

Source: https://www.emergentmind.com/topics/drive-jepa-framework