RoboJEPA: Scaling Robotic Latent World Models

What it claims

RoboJEPA is a family of action-conditioned latent world models, ranging from 22M to 8B parameters. Each model is trained as a predictor on top of a frozen V-JEPA 2.1-G encoder. The paper fits the forward-prediction loss on out-of-distribution held-out data (DROID real-robot, RoboCasa sim) to L(C) = A·C^(α−γ ln C) + E, where C is training compute and E is the irreducible loss. The authors report that this fit extrapolates beyond the scale at which it was fit, which they describe as the first JEPA scaling law with verified extrapolation.

Capability shows a compute threshold. Only models trained with more than 10²² FLOPs (the 2B, 4B and 8B models) solve RoboCasa object-interaction tasks via non-greedy planning. On the real Franka, only models of 2B or larger reliably keep the gripper closed after grasping; smaller ones often drop the object.

How it works

• Input. Up to V camera views per timestep are encoded by the frozen V-JEPA 2.1-G into 16×16-patch grids of 1664-dim features. Each patch is linearly projected to the model width. Actions and proprioceptive state get linear encoders specific to each robot id.
• Token sequence. Tokens are interleaved per timestep: all visual tokens, then one action token, then one state token, as a single flat sequence.
• Backbone. A ViT with a timestep-causal mask and joint space-time-view attention (not factorized). It uses pre-norm, RMSNorm, GELU and QK-norm. The authors call QK-norm particularly important for stability when scaling data and model size.
• Position and view information. A 4-axis RoPE covers time, height, width and view. A learnable per-layer view bias tells the model which camera a view is, so view order can be randomized.
• Prediction. The model is deterministic: one forward pass predicts the next latent, and each spatial token predicts its counterpart at t+1. Actions are either state deltas or absolute commands.
• Training losses:
• Teacher-forcing ℓ1 loss in the encoder's latent space.
• Autoregressive rollout loss: the model gets the full action sequence but only a prefix of encoded frames, then predicts the rest with a KV-cache.
• Planning. A goal image is encoded once. At each step, the cross-entropy method searches for the H-step action sequence whose imagined rollout minimizes ℓ1 distance to the goal latent. Candidate actions are sampled uniformly and then ranked by the world model.

flowchart LR
  X[Multiview frames] --> E[Frozen V-JEPA 2.1-G]
  E --> P[RoboJEPA predictor]
  A[Actions + proprio, robot id] --> P
  P --> TF[Teacher-forcing L1]
  P --> AR[K-step rollout L1, KV-cache]
  P -. at test time .-> CEM[CEM vs goal-image latent]

Evidence

Training data. The mixture combines 23 public datasets with 26 action spaces and 12 embodiments: about 2.87M episodes, 15,022 video hours and 6,692 action hours.

Evaluation scale. There are over 50,000 planning episodes across two robot platforms. Planning hyperparameters are fixed per robot, and only the horizon is varied for harder tasks.

• Real Franka: Grasp, Object Lift, Pick and Place.
• RoboCasa simulation: four tiers, each short or long horizon and with or without object interaction.

Results.

• Success rises with model size on the real Franka and with compute in RoboCasa.
• Forward-prediction error correlates strongly with planning success, so the authors propose it as a proxy for real-robot evaluation.
• The fetched text does not give exact success rates; they appear only in Figure 1 plots.

How to use it

Checkpoints at all sizes, plus training and robot-deployment code, are released at https://github.com/facebookresearch/robo_jepa (project site: https://robojepa.github.io). The page does not state a licence or install commands.

Running the model requires the V-JEPA 2.1-G encoder. Actions and state need a known robot id, because the action and state encoders are specific to each robot id.

Limits

• Predictor-only scaling. The encoder is held fixed, so the laws describe scaling of the predictor on a shared representation, not end-to-end scaling.
• Encoder dependence. With the V-JEPA 2 encoder, scaling was not predictable. The authors only hypothesize why (Appendix K).
• Dataset sizes. They are as collected and verified by the authors and may be smaller than the original papers report.
• Tasks tested. Real-robot tasks are limited to grasp, lift and pick-and-place on a Franka, with a single goal image as the task specification.
• Truncated text. The fetched text cuts off partway through the method section. Appendices, ablations and exact numbers were not available here.