RoboJEPA: Scaling Robotic Latent World Models

AI summary

Overview

Research area: Robot learning / world models — specifically latent, action-conditioned world models for planning, and empirical scaling laws for them.

Technical level: Advanced. The paper assumes familiarity with joint-embedding predictive architectures (JEPA), vision transformers, model-predictive control with the cross-entropy method, and neural scaling-law fitting.

Scope (one sentence): The paper introduces RoboJEPA, a family of action-conditioned latent world models trained on 23 public manipulation datasets across 12 embodiments, and shows that its latent rollout error follows a predictable second-order power law in training compute that extrapolates to 4B and 8B models and correlates with real-robot planning success.

What This Paper Is About

Scaling laws — empirical formulas relating a model's loss to the compute, data, and parameters used to train it — are well established for language models, image generation, and image classification, but essentially uncharacterized for robotic world models. Prior latent world models were trained at a single scale, so nobody knew how their quality changes as you spend more compute, and prior scaling studies in robotics scaled policies rather than world models, or studied video generation outside of robot control. This paper's goal is to build a scalable latent world model and establish the concrete functional form of its scaling law, so practitioners can forecast what a given compute budget will buy before spending it.

Key Contributions

1. RoboJEPA, a scalable multi-embodiment latent world model. A deterministic, feed-forward transformer that predicts the next latent state in a single forward pass, operating in the frozen representation space of a V-JEPA 2.1-G encoder and trained on a unified mixture of 23 public manipulation datasets spanning 12 robot platforms (real and simulated), 26 action spaces, 15,022 hours of video, and 6,692 hours synchronized with actions.

2. A second-order power-law scaling law for imagination error. The paper compares four candidate functional forms (first-order power law, second-order power law, Broken Neural Scaling Law, Unified Neural Scaling Law) and shows that the second-order form, L(C) = E + A·C^(α − γ ln C), best extrapolates from models fit at 22M–2B parameters to held-out 4B and 8B models, on both a real-robot DROID holdout and simulated RoboCasa scenes.

3. Evidence that offline imagination error predicts downstream planning. Downstream planning success in RoboCasa improves predictably with training compute, and imagination error is strongly correlated with it, making it a usable proxy for expensive real-robot evaluation. The authors also show that longer-horizon training rollout (K=10 instead of K=2) is needed for predictable scaling on non-greedy, long-horizon tasks.

4. Zero-shot deployment on real hardware, plus open release. RoboJEPA plans toward a single goal image without any task-specific policy, evaluated over 50,000 episodes across two robot platforms, and the authors release all checkpoints together with training and robot-deployment code.

Main Findings

• Imagination error follows a second-order power law in compute. Fit on 22M–2B models and extrapolated to 4B and 8B models, the second-order law achieved held-out extrapolation error of 0.6 × 10⁻³ on DROID and 1.4 × 10⁻³ on RoboCasa, versus 2.0 and 3.7 × 10⁻³ for the standard first-order power law. BNSL (n=1) gave 1.0 and 1.8; UNSL (n=1) gave 1.4 and 5.7.
• The fitted law's parameters differ by domain. On DROID: E = 0.218, A = 7.96 × 10⁻¹⁹, α = 1.91, γ = 0.02308. On RoboCasa: E = 0.1709, A = 1.66 × 10⁻¹⁷, α = 1.761, γ = 0.02105.
• The models are close to saturation under the current data budget. The second-order law and BNSL agree on nearly the same irreducible error: ≈0.2 on DROID and ≈0.17 on RoboCasa.
• No single model size sits on the compute-optimal frontier. The lower envelope across model sizes forms the frontier, and different sizes occupy it at different compute budgets.
• Greedy planning capabilities emerge at identifiable compute thresholds. Basic end-effector 3D movement (pose reaching) emerges around 10²⁰ FLOPs, while reaching while holding an object emerges around 3 × 10²⁰ FLOPs — both corresponding to training 50M–100M-parameter models on the full multi-embodiment dataset. These tasks then saturate early.
• Non-greedy object interaction requires much more compute. Only models trained with more than 10²² FLOPs (at sizes 2B, 4B, and 8B) can solve object-interaction tasks in RoboCasa via non-greedy planning; smaller models trained with less compute cannot.
• Long-horizon tasks do not scale predictably with the short-horizon recipe. On the obstacle-reach task, using the recipe with K=2 rollout steps, reward correlated with compute but no clean frontier emerged — for example, the 300M planner reached a higher final reward than the 1B one.
• Aligning training with deployment fixes this. Increasing the training rollout to K=10 and mixing in higher-resolution inputs (240×320 and 480×640) during the cooldown makes training more compute-efficient beyond about 10²¹ FLOPs of total training compute (around 300M–1B parameters at the compute-optimal frontier).
• Real-robot success improves with model size. On a Franka robot across Grasp, Object Lift, and Pick and Place, success rate shows an improvement trend as model size increases. Qualitatively, only models with 2B parameters or more keep the gripper closed after grasping; smaller models often accidentally open it and drop the object.
• Larger models imagine more faithfully. Decoded rollouts on an AgiBot episode show the 0.3B, 2B, and 8B models tracking true dynamics progressively better, and in two-view RoboCasa rollouts larger models keep the two camera views mutually consistent deeper into the rollout. All world models share the same V-JEPA 2.1-G encoder and decoder.

Methodology in Plain English

The model. RoboJEPA takes in a history of what the robot has seen (encoded as latent features by a frozen, pretrained V-JEPA 2.1-G vision encoder), the robot's actions, and its proprioceptive state, and predicts — in one forward pass — the latent features of the next frame. It is a plain vision transformer with a causal mask (a token at time t can only see tokens at time t or earlier) and local attention (no token attends more than 8 timesteps back, so the model never sees a longer history than it was trained on, even during long rollouts). Position information comes from a factorized rotary embedding with four axes — time, height, width, and camera view — combined with a per-layer learnable view bias so the model knows both that two tokens come from different cameras and what content each camera holds. QK-normalization was found particularly important for stability when scaling dataset and model size.

Handling heterogeneity. Datasets differ in number of cameras, state size, controller type, and frame rate. Actions and states are projected by robot-id-specific linear encoders, then visual, action, and state tokens are interleaved by timestep into one flat sequence. To handle variable views and resolutions, GPUs are partitioned into four groups: single-view low-res (240×320), dual-view low-res, single-view high-res (480×640), and dual-view high-res, with batch sizes tuned to balance step times across groups.

The training loss. Two ℓ1 losses are averaged: a teacher-forcing loss (predict the next encoded frame given the real past) and an autoregressive rollout loss (given only a short prefix of encoded frames but the full action sequence, predict the remaining frames using the model's own outputs, accelerated with a KV-cache). The second loss is what teaches the model to survive its own errors over long horizons.

How planning works. At test time the goal is a single image, encoded once into goal features. At each step, the cross-entropy method searches for the H-step action sequence whose imagined latent rollout lands closest (in ℓ1 distance) to the goal features. It is receding-horizon: the robot executes, then re-plans. Only a single goal image is used — no subgoals — and the same planning hyperparameters are used across embodiments and tasks, varying only the planning horizon for harder tasks. A rolling KV-cache, precompiled inference saved to disk, and multi-GPU planning make this practical: a multi-step plan takes only several seconds even for the 8B model.

Measuring the law. For every model at every training duration the authors count the FLOPs of every operation in a training step and multiply by steps taken, giving total training compute. They plot the per-token ℓ1 error over a 9-step latent rollout on held-out, unseen scenes (two views: one side camera and one wrist camera), keep the best error at each compute budget to form the compute-optimal frontier, fit each candidate curve to the 22M–2B points only, and score each candidate by how well it predicts the unseen 4B and 8B models.

Training details. Model sizes range from 22M to 8B parameters, training compute from 2 × 10¹⁹ to 9.5 × 10²² FLOPs, with AdamW, weight decay 0.1, gradient-norm clipping at 1.0, and a learning rate warmed up then held constant at a peak of 5 × 10⁻⁵ before a short linear cooldown. Because the flat stage is run once and annealing is done from intermediate checkpoints, compute budgets can be overlapped across models.

A diffusion decoder for visualization. To see what the world model imagines, a diffusion decoder denoises Cosmos tokenizer tokens (Cosmos-CV4x8x8) spatially conditioned on V-JEPA 2.1 features. Since V-JEPA features are computed with a frame skip of 4 while Cosmos tokens have no skip, the decoder must generate the 3 intermediate RGB frames — expanding, for example, 5 Hz latent predictions into a 20 FPS video.

Why This Matters

For research: This is, to the authors' knowledge, the first work to establish scaling laws for multi-embodiment robotic world models trained on real robot data, and the first scaling law for JEPA with verified extrapolation. It also produces the largest JEPA predictor trained to date at 8B parameters. It moves world-model research from "train one model at one scale" to a predictable engineering discipline, and it offers a concrete methodological finding: whether scaling laws appear at all depends on matching the training rollout horizon to the deployment horizon.

Real-world applications:

• Budget planning for robot learning teams — deciding in advance whether an additional order of magnitude of compute is worth spending, or whether the data budget has already saturated.
• Cheap evaluation of world models — using offline imagination error as a proxy instead of running tens of thousands of expensive real-robot episodes.
• Cross-embodiment robot foundation models — the architecture and tokenization handle arbitrary numbers of cameras, state sizes, and action spaces, so a single model can serve single-arm, dual-arm, and humanoid platforms.
• Goal-image-conditioned manipulation — specifying a task by showing a single picture of the desired end state, with no task-specific policy or subgoal engineering.

Industry relevance: The finding that irreducible error is near ≈0.2 and ≈0.17 on both evaluations suggests that, at this data budget, simply adding compute has limited headroom — a signal that data collection, not compute, may be the binding constraint. The engineering contributions (precompiled inference, KV-cache rollout, multi-GPU CEM) are also what make an 8B world model fast enough to be usable on physical hardware.

Future Directions

• Breaking the saturation wall. Both best-fitting laws estimate nearly the same irreducible error, indicating RoboJEPA is close to saturation under its current training and data budget. The natural next question is whether more or better data — rather than more compute — moves that floor, and whether a data-aware scaling law is needed.
• Scaling the long-horizon recipe. Long-horizon training (K=10, mixed resolution) only becomes compute-efficient beyond about 10²¹ FLOPs and was applied as a cooldown from pretrained checkpoints rather than by retraining from scratch. Whether an end-to-end long-horizon recipe produces an even better frontier is left open.
• Extending verified extrapolation further. The law was validated by holding out 4B and 8B models; a key open question is how many orders of magnitude beyond 8B and 9.5 × 10²² FLOPs the second-order form continues to hold before breaking.
• Evaluation-first model design. The authors explicitly recommend building downstream evaluations that reflect the desired capability first, then estimating the minimal training compute needed to reach it. Turning that recommendation into a general practice — including for capabilities beyond those tested here (only grasping, lifting, pick-and-place, reaching, obstacle avoidance, and pushing were studied) — remains open.

Target Audience

Researchers and engineers working on robot learning, world models, and model-based control; scaling-law researchers interested in whether language-model-style predictability transfers to physical-world models; and practitioners at robotics or foundation-model companies who need to justify compute and data budgets for embodied AI. It is also relevant to anyone building on JEPA architectures, since it documents architecture choices (QK-norm for stability, factorized 4D RoPE with view bias, local causal attention, V-JEPA 2.1 rather than V-JEPA 2 as the encoder) that the authors found necessary for predictable scaling. Readers without background in latent world models, vision transformers, or CEM planning will find the paper dense.

Authors’ abstract

Latent world models have shown a remarkable ability to predict future states and to plan in the real world. In practice, however, we lack a principled way to estimate how their capabilities scale with model size, data, and compute, an open problem that slows progress in the field. In this work we present RoboJEPA, a world model based on the Joint Embedding Predictive Architecture (JEPA) and trained on a large-scale dataset spanning 12 robotic embodiments. We show that RoboJEPA's imagination error, the error of its latent rollouts, follows a second-order power law in compute, allowing us to predict model quality well beyond the scale at which the law is fit. We further show that downstream robotic planning performance improves predictably with compute, and that imagination error is strongly correlated with it, making it a reliable proxy for real-robot evaluation. Finally, we demonstrate that latent world models can be deployed zero-shot as robotic agents, planning toward a single goal image to solve tasks requiring long-horizon planning on real hardware. We release all model checkpoints together with our training and robot deployment code. To our knowledge, this is the first work to establish scaling laws for multi-embodiment robotic world models trained on real robot data, and RoboJEPA, at 8B parameters, is the largest JEPA predictor model trained to date.