La-JEPA: Compact Robot Policies with Latent World Modeling

📌 Highlights

La-JEPA's central claim is that a robot policy does not need to inherit the architecture of a large video or language generator to obtain strong control performance. The model trains compact World and Action DiTs from scratch in a pretrained V-JEPA 2.1 feature space, adds future-feature prediction as training-time supervision, and generates only actions at inference. In the reported comparisons without embodied pretraining, La-JEPA leads the listed baselines on clean and randomized RoboTwin 2.0 and reaches the strongest corrected LIBERO-Plus results. The policy sizes range from 0.10 to 1.14 billion trainable parameters, and the paper reports lower latency than Fast-WAM at comparable or better benchmark performance.

The most practically important evidence is distributed across representation, robustness, efficiency, and physical evaluation. V-JEPA features outperform Wan2.2 VAE latents, particularly on randomized settings; camera and proprioception dropout improve robustness; and a La-JEPA B policy reaches an average of 17.50 successes out of 20 across four real-robot tasks with a reported 25.16 ms benchmark policy-call latency. The main caveat is that the paper deliberately stops short of exhaustive recipe optimization and tests real hardware on only one robot setup. The results support a strong architecture hypothesis, but not yet a universal replacement for embodied-pretrained policies.

• Compact policy transformers are trained from scratch while visual and language encoders provide pretrained inputs.
• Future prediction is used to shape representations during training but is removed from the inference path.
• La-JEPA B reports a RoboTwin 2.0 clean/randomized mean of 60.8 and corrected LIBERO-Plus success of 82.5%.
• The real-robot average is 17.50/20, comparable to π0.5's 17.25/20 in the listed four-task evaluation.
• The paper identifies non-monotonic RoboTwin scaling and limited real-world scope as unresolved concerns.

🎯 Introduction

The paper starts from a deployment-oriented tension in robot learning. World-action models commonly adapt pretrained video-generation transformers, while vision-language-action policies commonly extend pretrained vision-language models. These routes provide useful visual, temporal, and linguistic priors, but they also transfer representation choices, tokenization, transformer capacity, and inference costs from generation problems into control. The paper asks whether robot policies must inherit those generative architectures to remain competitive, and whether a compact world-action transformer trained from scratch can use a pretrained representation effectively.

The second motivation concerns the latent space itself. Video VAEs preserve information needed for pixel-level reconstruction, whereas control may benefit more from features organized around motion, semantics, depth, and segmentation. La-JEPA therefore uses dense features from V-JEPA 2.1, a video JEPA encoder whose representation is intended for downstream understanding rather than image synthesis. The paper positions this as distinct from earlier JEPA-based robot policies because its World and Action DiTs are initialized from scratch and their capacity can be selected for control independently of the visual encoder.

The paper's objective is consequently both architectural and empirical: construct a compact policy with joint world and action prediction, test whether JEPA features are preferable to video VAE latents, and measure whether the resulting design generalizes beyond the training distribution. The evaluation spans clean and randomized RoboTwin 2.0, LIBERO and LIBERO-Plus, comparisons with latent world-action models, computational cost, and four real-robot tasks. A further objective is to make LIBERO-Plus comparisons interpretable by correcting an evaluator issue that appends perturbation parameters to most released instructions.

• Task-level question: can a small control-oriented transformer compete with much larger generation-oriented policies?
• Representation-level question: do JEPA features provide a better control substrate than video VAE latents?
• Evaluation-level question: how much do corrected LIBERO-Plus instructions change conclusions about robustness?

🔬 Methodology

La-JEPA maps multi-camera observations, proprioception, and a language instruction to an action chunk. The pretrained V-JEPA 2.1 ViT-L encoder processes each camera independently. The current observation becomes context features, while features representing a 32-frame future horizon become prediction targets. The paper normalizes and spatially reduces these grids before projecting them into the World DiT embedding space. The visual encoder is the only visual pathway: the policy does not separately consume RGB and does not reconstruct pixels. A frozen FLAN-T5-large encoder embeds language, and learned projections convert proprioception into conditioning tokens read by both DiTs through cross-attention.

The trainable architecture contains a World DiT over observation and future-feature tokens and an Action DiT over the action chunk. Both are initialized from scratch, with separate attention projections and feed-forward layers. Rotary visual positional encodings retain spatial and temporal information. The attention mask is the key structural constraint: observation tokens attend to observations, future tokens attend to observations and futures, and action tokens attend to observations and actions. Thus actions cannot directly access future tokens, although future prediction still shapes the observation representation through shared World DiT parameters. This design lets the training objective act as a predictive regularizer without requiring the deployed policy to generate a future trajectory.

Training uses independent flow matching for future features and actions. For sampled \tau_f, \tau_a and Gaussian noises \epsilon_f, \epsilon_a, the noisy inputs are defined by Equation (1). The model predicts future and action velocity fields, and Equation (2) sums their mean-squared errors. At inference, future tokens are dropped, the action chunk starts from Gaussian noise, and the action flow is integrated from \tau_a=0 to \tau_a=1 with Euler steps. The World DiT conditioning is fixed because future prediction is disabled. Observation keys and values are computed once and cached, while only the Action DiT runs at each denoising step. The action tokens are converted to robot actions, with the first 10 actions executed on LIBERO and the first 24 on RoboTwin before replanning.

The final recipe is not only the base architecture. The design study selects tube-motion encoding, camera dropout, proprioception dropout, image augmentation, and observation-encoder LoRA. Tube-motion uses a tube observation and jointly encodes 16 future frames spaced two frames apart, allowing each future feature grid to attend across the predicted horizon. The resulting S, B, and L policies have 0.10, 0.34, and 1.14 billion trainable parameters, respectively, although the main text notes that the relative World DiT and Action DiT capacities were not systematically optimized.

📊 Experiments

The paper evaluates on LIBERO, LIBERO-Plus, and RoboTwin 2.0, then adds latent world-action comparisons, latency measurements, and four physical Franka tasks. LIBERO policies are trained on the four standard suites and evaluated both in-distribution and on LIBERO-Plus variants perturbed along camera viewpoint, lighting, sensor noise, and other axes. RoboTwin policies are trained only on clean demonstrations and evaluated in clean and randomized scenes; the randomized setting adds clutter, distractor objects, and new textures. Success percentage is the primary metric, so higher is better. Parameter counts distinguish trainable parameters from deployed parameters, and latency is reported in milliseconds, where lower is better.

On Table 1, corrected-instruction LIBERO-Plus success is 86.3% for La-JEPA S, 87.3% for B, and 88.0% for L, compared with 86.2% for NowWAM. On the released instructions, the corresponding values are 80.0%, 82.5%, and 84.1%, while NowWAM reports 86.2% under the table's common step limits. LIBERO success is 97.9%, 98.0%, and 98.1% for S, B, and L. These results use 0.10, 0.34, and 1.14 billion trainable parameters and 0.74, 0.98, and 1.78 billion deployed parameters. The paper attributes the protocol difference to 8,493 of 10,030 released LIBERO-Plus variants receiving corrupted instructions; corrected evaluation changes Fast-WAM by 19.3 points, LeapBot-WA by 8.1, Cosmos Policy by 4.1, and La-JEPA by 3.9 to 6.3 points.

On Table 2, La-JEPA S, B, and L obtain clean RoboTwin success of 85.5%, 86.2%, and 86.0%, randomized success of 33.9%, 35.4%, and 32.1%, and equally weighted means of 59.7%, 60.8%, and 59.0%. The listed non-embodied-pretrained baselines are substantially lower in the randomized condition: Fast-WAM reaches 1.9% and Faster-WAM 2.2%. The paper emphasizes that RoboTwin scaling is not monotonic: B is best among the three sizes, and the authors hypothesize that a larger Action DiT may overfit the simpler action distribution. Table 3 extends the comparison to latent world-action models. La-JEPA B reports 98.0% LIBERO IID, 82.5% LIBERO-Plus, 95.9% RoboTwin clean-to-clean, and 49.9% clean-to-randomized, exceeding JEPA-WAM's 96.7%, 77.1%, 79.9%, and 36.9% in the corresponding listed columns.

The design analysis provides mechanism-level evidence. Future prediction raises RoboTwin clean success by 7.0 points and randomized success by 3.7 points; on LIBERO-Plus it improves each of three seeds by 2.1 points on average, with paired comparisons reported as significant at p < 0.01. Relative to Wan2.2 VAE latents, V-JEPA improves RoboTwin clean success by 15.9 points and randomized success by 16.0 points, plus 6.0 points on LIBERO-Plus and 2.0 points on LIBERO. Tube encoding adds 1.1 clean and 3.6 randomized RoboTwin points relative to static encoding, while tube-motion gives 3.4 clean and 3.7 randomized points. Camera dropout adds 7.1 randomized RoboTwin points and 10.2 LIBERO-Plus points; proprioception dropout adds 3.8 clean and 4.2 randomized RoboTwin points. On LIBERO-Plus, LoRA alone adds 1.4 points, augmentation alone 5.8, and both 7.2 relative to a plain policy. The paper also reports that complete policy calls with compiled builds run two to five times faster than Fast-WAM and Faster-WAM, while three denoising steps retain success within about one point.

The physical evaluation trains La-JEPA B from scratch on demonstrations for cube placement, fruit picking, Push T, and dish manipulation. It records 19, 20, 15, and 20 successes out of 20, respectively, for an average of 17.50. The listed benchmark policy-call latency is 25.16 ms. This average exceeds Fast-WAM's 16.00 and Faster-WAM's 16.25 among the non-embodied-pretrained policies and is close to π0.5's 17.25, although π0.5 leads on the dish task. These results are encouraging for contact-heavy and multistep control, but the task count and single-robot setup are too narrow to establish broad physical generalization.

Figure 2

Caption: Figure 2: La-JEPA. Left: general architecture. Right: context and future encoding with V-JEPA 2.1.

Why It Matters: This is the central architecture diagram: it shows the context, future, conditioning, World DiT, Action DiT, and tube-motion encoding arrangement used to connect predictive supervision with action generation.

Figure 1

Caption: Figure 1: Success against deployed parameters, including frozen encoders, on LIBERO-Plus with corrected instructions (Section 5.2) and on RoboTwin 2.0, where the panel reports the equally weighted mean of clean and randomized success of clean-trained policies. S, B, and L denote the La-JEPA sizes. The LIBERO-Plus panel shows the policies without embodied pretraining that we re-evaluated with corrected instructions, and the embodied-pretrained π0.5; Table 1 lists all scores, including those with the released instructions. ∗: success estimated from JEPA-WAM’s 20-task results (Appendix B).

Why It Matters: The figure summarizes the paper's efficiency claim by placing success against deployed parameters across corrected LIBERO-Plus and RoboTwin 2.0.

Figure 3

Caption: Figure 3: Gain in success (percentage points) from each design choice over its baseline: A, action-only training; B, Wan2.2 VAE latents; C, static encoding; D, no camera dropout (LIBERO: wrist-only dropout); E, no proprioception dropout; F, a plain policy without augmentation or encoder LoRA. Green: LIBERO; orange: RoboTwin 2.0; solid: in-distribution; hatched: out-of-distribution. Block A averages LIBERO-Plus gains over three training seeds. Values and evaluation budgets in Table 6.

Why It Matters: The ablation figure isolates the contribution of predictive supervision, representation choice, temporal encoding, dropout, augmentation, and LoRA across in-distribution and out-of-distribution evaluations.

🔮 Conclusion

La-JEPA answers the paper's main question affirmatively within the evaluated scope: a compact transformer trained from scratch in a pretrained JEPA feature space can compete with, and on several benchmarks exceed, larger policies built around video or language generation. The method reaches the strongest listed clean and randomized RoboTwin 2.0 results among policies without embodied pretraining, leads corrected LIBERO-Plus comparisons, and reports up to five-times faster inference than Fast-WAM in the paper's latency analysis.

The conclusion is strongest as an architectural result rather than a claim of universal superiority. Future-feature prediction, JEPA representations, and targeted robustness mechanisms each contribute measurable gains, while action-only inference keeps deployment efficient. Physical results are comparable to π0.5 on average across four tasks, but broader embodied pretraining, more robots, and more diverse manipulation demands remain open tests.

🛠️ Future Research Improvements

The clearest next step is to study embodied pretraining. La-JEPA currently learns from target demonstrations alone, so it is unknown whether pretrained robot data would improve generalization, allow transfer across embodiments and environments, or reveal that data rather than policy capacity is the dominant bottleneck. The paper specifically identifies this as its principal future direction. A useful study would hold the V-JEPA representation and policy architecture fixed while varying the amount, diversity, and embodiment coverage of robot data.

The method also needs more systematic recipe optimization. The authors did not search learning rates, weight decay, dropout rates, LoRA configurations, or the relative capacities of the World and Action DiTs. The non-monotonic RoboTwin scaling, where B exceeds L on the reported mean, could therefore reflect capacity allocation or optimization rather than an intrinsic limit. Controlled sweeps should separate World DiT size from Action DiT size and test whether the action expert overfits because its target distribution is simpler than the future-feature distribution.

Finally, evaluation should broaden beyond the current benchmark and hardware scope. Corrected instruction protocols should become standard for LIBERO-Plus comparisons, while future studies should report both clean and randomized success, latency on deployment hardware, and physical performance across more embodiments, contact regimes, and long-horizon tasks. Combining La-JEPA with richer supervision, synthetic data, or robot-specific objectives is plausible, but the paper does not yet measure whether those additions stack with its gains.

🏭 Potential Industry Use Scenarios

The most credible near-term use is efficient closed-loop manipulation on hardware where policy-call latency, memory, or GPU cost matters. La-JEPA's inference path caches observation representations and denoises only action tokens, making the architecture attractive for tabletop pick-and-place, contact-rich pushing, and multistep manipulation where replanning is frequent. The real-robot table reports 25.16 ms benchmark policy-call latency for La-JEPA, alongside an average of 17.50 successes out of 20 across cube placement, fruit picking, Push T, and dish manipulation. These results justify prototyping, not unattended production deployment.

A second scenario is robustness-sensitive automation in visually changing workcells. RoboTwin randomized scenes and LIBERO-Plus perturbations test clutter, distractors, textures, camera changes, lighting, and sensor noise. The reported gains from camera dropout, proprioception dropout, augmentation, LoRA, and JEPA features suggest a recipe for policies that should not rely on one viewpoint or a pixel-reconstruction-oriented latent. Such systems could be evaluated for bin picking, kitting, inspection-assisted manipulation, or workstation reconfiguration, provided each deployment receives task-specific demonstrations and safety validation.

A third scenario is a tiered policy stack in which a compact La-JEPA controller handles fast local action generation while larger models provide occasional language interpretation, task decomposition, or supervision. This division is consistent with the paper's separation of language conditioning, visual representation, world prediction, and action inference. However, the paper does not evaluate industrial throughput, safety incidents, multi-robot transfer, or long-duration reliability. Any product claim should therefore remain conditional until those dimensions are measured.

💬 Critical Analysis

The paper's strongest methodological choice is the clean separation between pretrained perception and trainable control. It tests a meaningful architectural hypothesis rather than merely shrinking an existing generator: V-JEPA 2.1 supplies dense visual features, while the World and Action DiTs are sized and initialized for the control task. The attention mask and cached-observation inference path make the efficiency argument technically coherent. The design analysis is also unusually useful for builders because it connects the headline result to specific choices. Future prediction improves action performance even though future tokens disappear at inference; JEPA features beat Wan2.2 VAE latents most strongly under randomized evaluation; and camera dropout produces large robustness gains.

The benchmark evidence is substantial but requires careful interpretation. Corrected and released LIBERO-Plus instructions are different protocols, and the paper demonstrates that their effect is highly model-dependent. The reported RoboTwin mean is an equally weighted clean/randomized average, not a deployment success rate, and Table 3 includes a JEPA-WAM comparison on a 20-task subset while La-JEPA is trained on all 50 tasks. The paper discusses this mismatch, but builders should still reproduce task overlap and difficulty before treating the comparison as definitive. Similarly, trainable parameters, deployed parameters, denoising steps, policy-call latency, and closed-loop control timing are distinct quantities and should not be collapsed into a single efficiency number.

Reproducibility is helped by explicit architecture and loss descriptions, three La-JEPA sizes, three-seed benchmark means in the principal tables, and a stated code release covering feature precomputation, training, and evaluation. Nevertheless, the paper's own limitations materially constrain causal conclusions. The recipe is not exhaustively tuned, larger models do not improve RoboTwin monotonically, and the strongest embodied-pretrained baselines remain better on randomized RoboTwin. The four-task real-robot study is encouraging but small, and one competitor's dish result uses 10 trials rather than 20. The central claim is therefore well supported as evidence that control-oriented latent world-action policies deserve serious consideration, but broader data, embodiments, safety conditions, and systematic scaling studies are needed before concluding that La-JEPA is generally superior to larger pretrained policy families.