💻 computer science

Depth-Regularized JEPA World Models Learn More Transferable Representations from Real Outdoor Robot Data

This paper introduces a depth-regularized JEPA world model that leverages geometric priors and latent regularization to learn more transferable and robust representations from real-world outdoor robot data, significantly outperforming baselines in visual odometry, surprise detection, and multi-step prediction without increasing inference overhead.

Original authors: Usman M. Khan

Published 2026-07-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Usman M. Khan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot how to drive a tractor through a field. The robot has a camera, but the world it sees is messy: the sun glares, shadows stretch and shrink, leaves rustle in the wind, and the ground looks different every time it turns a corner. If you ask the robot to memorize every single pixel of every image, it will get overwhelmed and confused. Instead, scientists are trying to teach robots "World Models." Think of a World Model as the robot's internal daydream. Instead of trying to redraw the whole picture, it learns to predict what will happen next in a simplified, abstract way. It's like a chess player who doesn't memorize the exact color of every square on the board but understands the rules of how pieces move and how the game evolves.

The specific type of "daydream" this paper focuses on is called a JEPA (Joint Embedding Predictive Architecture). Imagine a student taking a test where they have to guess the next sentence in a story based on the previous ones, but they aren't allowed to write the whole sentence—they just have to guess the vibe or the essence of the next part. This is great because it ignores the boring details (like the font color) and focuses on the important stuff (the plot). However, when these robots try to learn from real, outdoor videos, their internal "vibe" guesses often get messy or collapse into nonsense because the real world is so chaotic. The big question is: How do we teach these robots to understand the physics of the real world—like how things are arranged in 3D space—without needing a super-complex brain that slows them down?

This paper introduces a clever trick to solve that problem. The researchers took a compact robot brain (an 18-million-parameter model called LeWorldModel) and gave it a special "training-only" tutor: depth information. Imagine you are learning to draw a landscape. You have a regular photo (RGB), but you also have a 3D map (depth) that shows exactly how far away every tree and rock is. The researchers let the robot look at both the photo and the 3D map while it was studying (training). They told the robot, "Your guess about the next scene should match the 3D structure you see in the map." But here's the magic: once the robot finished its homework, they took the 3D map away. When the robot goes out to work in the real world, it only uses the camera photo, just like before. However, because it learned with the 3D map, its internal understanding of the world is now much sharper and more grounded in reality.

The results show that this simple trick works wonders. By using depth only during training, the robot's internal "daydream" became much better at understanding how the world moves. When the researchers tested it, the robot's ability to guess its own movement (visual odometry) improved by 33%, making it much more accurate. It also got much better at spotting "weird" things. If you suddenly teleported the robot to a different spot or flipped the video backward, the robot's internal alarm went off much louder and clearer than before. Surprisingly, this even helped the robot understand things that aren't strictly 3D, like sudden changes in lighting or shadows, suggesting that understanding the shape of the world helps the robot make sense of everything else, too.

The paper also found that this "depth-trained" robot was much more flexible when it moved to a completely new environment (a different robot in a different field). While the standard robot struggled to keep its predictions accurate over time in the new place, the depth-trained robot kept its cool, and its predictions stayed accurate for longer stretches. The researchers suggest that by grounding the robot's learning in the physical geometry of the world, they created a more robust and transferable brain. It's a bit like teaching a child to ride a bike on a smooth, marked track with a balance beam nearby; once they learn the balance, they can ride on any bumpy path without the beam, because they truly understood how to stay upright. This approach offers a practical way to make small, efficient robot brains smarter and more adaptable without needing expensive sensors or super-computers when they are actually working.

Technical Summary: Depth-Regularized JEPA World Models for Real Outdoor Robot Data

Problem Statement

While Joint Embedding Predictive Architectures (JEPA) have demonstrated success in learning robust environmental dynamics, their application to visually complex, real-world outdoor robotics remains challenging. Existing approaches often rely on pixel-level reconstruction, which is vulnerable to nuisance variations (e.g., lighting shifts, foliage clutter) in unstructured environments. Conversely, latent-space JEPA models, such as LeWorldModel (LeWM), offer a more stable abstraction by predicting future latent representations rather than pixels. However, LeWM and similar compact models have primarily been validated in simulated or controlled indoor settings. When applied to real outdoor data, these models risk representational collapse or failure to capture physically meaningful dynamics without additional structural constraints. The core challenge is to enhance the robustness and transferability of compact JEPA models trained on real-world data without increasing inference-time computational costs or requiring auxiliary sensors (like depth cameras) during deployment.

Methodology

The authors propose a training-time intervention that augments the standard LeWorldModel (LeWM) architecture with a depth regularization objective and predictor-side overparameterization.

1. Base Architecture

The model is a compact, end-to-end JEPA variant (18M parameters) based on LeWM. It takes a history of RGB frames and actions as input, encodes them via a shared encoder-projector, and uses a predictor to forecast the next latent embedding. The training objective typically consists of a prediction loss ( $L_{pred}$ ) and a latent regularizer (SIGReg) that encourages isotropic Gaussian distributions to prevent mode collapse.

2. Training-Only Depth Regularization

To inject geometric structure into the latent space without requiring depth at inference, the authors introduce a depth alignment loss ( $L_{depth}$ ):

Dual-Modal Training: During training, the model processes paired RGB and stereo depth observations through the same encoder-projector stack.
Alignment Objective: The RGB embedding is optimized to align with the corresponding depth embedding (computed with a stop-gradient) using a cosine similarity loss.
Complementary Roles: The SIGReg term maximizes entropy (promoting task-agnostic diversity), while the depth term acts as a geometric constraint, organizing that diversity along physically meaningful directions. The combined objective seeks the highest-entropy representation consistent with scene geometry.
Inference: At test time, the depth branch is discarded; the model operates solely on RGB inputs.

3. Predictor Overparameterization

To accommodate the increased complexity of satisfying both isotropic and geometric constraints within a fixed capacity, the authors introduce training-only overparameterization.

Mechanism: During training, low-rank (rank-16) RepLinear branches are added to the predictor's MLP up-projection and down-projection layers.
Merging: After training, these branches are algebraically merged into the base predictor weights. This ensures the deployed model remains architecturally identical to the baseline with no inference overhead.

4. Data and Evaluation

The model was trained on ~31,000 frames from a single real agricultural robot (Aigen Element) and evaluated on:

In-domain: Aigen validation set.
Out-of-Domain (OOD): TartanGround benchmark (different robot, environment, and sim-to-real gap).
Metrics: Frozen-embedding Visual Odometry (VO) probes, predictor-based surprise detection (under synthetic perturbations), and multi-step latent rollout fidelity.

Key Contributions

Real-World Training: Demonstrated that a compact end-to-end JEPA world model can be successfully trained directly on real outdoor agricultural robot data, a setting previously unreported for this architecture family.
Training-Only Geometric Prior: Introduced a depth alignment objective that injects geometric structure into JEPA representations. This improves physical understanding and robustness without requiring depth sensors, task labels, or pretrained encoders at inference.
Comprehensive Evaluation: Provided a consistent evaluation protocol showing that depth regularization improves representation quality, predictor sharpness, and OOD robustness across visual odometry, surprise detection, and multi-step rollouts.

Results

The depth-regularized model consistently outperformed the vanilla LeWM baseline, particularly in transfer scenarios:

Representation Quality: The depth-regularized model reduced Visual Odometry (VO) probe error by 33% (MSE dropped from 0.0022 to 0.0015) compared to the baseline, surpassing a frozen DINOv2 reference.
Surprise Detection: The model showed significantly improved separation between plausible and implausible transitions.
- In-domain: Surprise separation increased from 0.067 to 0.120.
- Out-of-domain (TartanGround): Separation increased nearly 3-fold, from 0.035 to 0.099.
- Non-Geometric Gains: Notably, the model also improved on non-geometric perturbations (brightness jumps, color swaps), suggesting the geometric prior induces a scaffold for general physical consistency.
Out-of-Domain Generalization: While the baseline slightly outperformed the depth model in-domain, the depth-regularized model demonstrated superior robustness under domain shift. On TartanGround, it achieved higher goal similarity (0.745 vs. 0.716 at horizon $h=10$ ) and trajectory similarity, with the performance gap widening at longer horizons ( $h=20$ ).
Latent Geometry: Analysis of principal subspaces revealed that the depth-regularized model produced latent spaces with higher cross-domain overlap in the mid-rank regime ( $k \approx 5-50$ ), indicating a more transferable geometry for task-relevant structure.

Significance and Claims

The paper argues that a lightweight, training-time geometric prior can significantly enhance the utility and transferability of compact JEPA world models for real-world deployment.

Practicality: The approach is "cheap at training time and free at test time," as the depth signal and overparameterization are removed or merged before deployment.
Generalization: The results suggest that moving the training objective from pure isotropy toward a geometry-informed anisotropy creates a more robust latent space. This holds even for tasks not directly tied to 3D geometry (e.g., lighting consistency), implying that grounding representations in physical structure improves overall predictive expectations.
Limitations: The authors acknowledge that the study is a proof of concept using offline metrics on a single robot platform and a relatively small dataset. They note that closed-loop planning performance and scalability to larger models/data remain untested.

In conclusion, the work demonstrates that incorporating physically grounded priors like depth during training can bridge the gap between controlled simulation benchmarks and the messy reality of outdoor robotics, enabling compact models to learn more transferable and physically aware dynamics.