A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures

Basile Terver; Randall Balestriero; Megi Dervishi; David Fan; Quentin Garrido; Tushar Nagarajan; Koustuv Sinha; Wancong Zhang; Mike Rabbat; Yann LeCun; Amir Bar

Source-linked AI summary

A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures

Basile Terver, Randall Balestriero, Megi Dervishi, David Fan, Quentin Garrido, Tushar Nagarajan, Koustuv Sinha, Wancong Zhang, Mike Rabbat, Yann LeCun, Amir Bar

arXiv:2602.03604v3Feb 3, 2026cs.CVcs.AI

TL;DR

JEPA research has promising applications but lacks accessible implementations that connect its principles to practical image, video, and world-model settings. EB-JEPA provides modular, small-scale implementations across these settings, with approximately 90–91% CIFAR-10 probing accuracy and 97% planning success on Two Rooms; its experiments also show that regularization remains important while its theoretical interactions are incomplete.

Problem
Accessible implementations that bridge JEPA theory and practical application remain scarce, creating a high barrier to entry for researchers and learners.
Method
EB-JEPA provides modular, documented JEPA implementations for image representation learning, video prediction, and action-conditioned planning, trainable on a single GPU in a few hours.
Results
The library achieves approximately 90–91% linear probing accuracy on CIFAR-10, improves multi-step video prediction with k-step training, and reaches 97% planning success on Two Rooms.
Takeaways & Limitations
EB-JEPA serves as a practical toolkit and educational resource for studying JEPA-based representations, video prediction, planning, and regularization.
Takeaways & Limitations
The theoretical understanding of why particular regularization combinations work remains incomplete.

Abstract

from arXiv · show

We present EB-JEPA, an open-source library for learning representations and world models using Joint-Embedding Predictive Architectures (JEPAs). JEPAs learn to predict in representation space rather than pixel space, avoiding the pitfalls of generative modeling while capturing semantically meaningful features suitable for downstream tasks. Our library provides modular, self-contained implementations that illustrate how representation learning techniques developed for image-level self-supervised learning can transfer to video, where temporal dynamics add complexity, and ultimately to action-conditioned world models, where the model must additionally learn to predict the effects of control inputs. Each example is designed for single-GPU training within a few hours, making energy-based self-supervised learning accessible for research and education. We provide ablations of JEA components on CIFAR-10. Probing these representations yields 91% accuracy, indicating that the model learns useful features. Extending to video, we include a multi-step prediction example on Moving MNIST that demonstrates how the same principles scale to temporal modeling. Finally, we show how these representations can drive action-conditioned world models, achieving a 97% planning success rate on the Two Rooms navigation task. Comprehensive ablations reveal the critical importance of each regularization component for preventing representation collapse. Code is available at https://github.com/facebookresearch/eb_jepa.

1 INTRODUCTION

EB-JEPA addresses the accessibility gap in JEPA research with modular, documented implementations spanning image learning, video prediction, and action-conditioned planning. The examples are designed for small-scale training and practical experimentation.

Scope: The codebase covers self-supervised image learning, latent-space video prediction, and action-conditioned world models for goal-directed planning.These progressively complex settings are presented as self-contained implementations.
Motivation and contribution: EB-JEPA provides modular, well-documented JEPA implementations designed for educational use and rapid experimentation.The library addresses the high barrier to entry created by production-scale implementations and specialized world-model setups.
Accessible implementations: Accessible implementations span image representation learning, video prediction, and action-conditioned planning, each trainable on a single GPU in a few hours.The library is designed as a practical and educational entry point to JEPA-based models.
Modular architecture: The library includes reusable encoders, predictors, regularizers, and planners that can be recombined for new applications.Its modular architecture is paired with systematic experiments and ablations.

2 RELATED WORK

EB-JEPA builds on joint-embedding methods and latent world models while emphasizing accessible educational implementations rather than state-of-the-art performance on complex benchmarks.

Joint-Embedding methods: Joint-embedding methods use regularization-based collapse prevention in EB-JEPA rather than stop-gradient techniques.The paper situates this design within prior work on VICReg, SIGReg, and related non-contrastive methods.
World models for planning: Latent world models have been studied extensively for model-based reinforcement learning and planning.The paper relates EB-JEPA to PLDM, IWM, DINO-WM, Navigation World Models, and JEPA-WMs.
World models for planning: Unlike related planning systems, EB-JEPA focuses on accessible, educational implementations rather than state-of-the-art performance on complex benchmarks.This distinction defines the paper’s positioning within prior world-model work.

3 PRELIMINARIES: A UNIFIED JEPA FRAMEWORK

The unified JEPA framework predicts target representations from encoded inputs, optionally conditioned on actions, while regularization prevents collapse. Its instances progress from view-invariant images to temporal video prediction and action-conditioned dynamics.

Unified formulation: The unified objective combines representation prediction with regularization: L = Lpred(gϕ(z, u), z′) + λR(z).Here z encodes the input, u optionally encodes conditioning such as robotic controls, z′ is the target representation, and λ balances the terms.
Image-JEPA: Image-JEPA learns view-invariant representations by predicting representations of differently augmented views of the same image.Low energy corresponds to similarity between the two view representations.
Video-JEPA: Video-JEPA predicts the next representation from a temporal context, capturing dynamics without access to future frames during prediction.The predictor operates over encoded frame representations and a temporal receptive field.
Action-conditioned video-JEPA: Action-conditioned video-JEPA conditions prediction on encoded actions to model how control inputs affect latent dynamics.The action encoder maps action sequences to control representations used by the predictor.
Regularization: SIGReg is more stable across hyperparameter configurations than VICReg, while VICReg reaches similar peak performance with more careful tuning.This comparison is reported on CIFAR-10.
Regularization: Regularization prevents representation collapse by maintaining feature spread, decorrelating dimensions, or enforcing Gaussianity in projected embeddings.EB-JEPA includes VICReg and SIGReg; SIGReg tests Gaussianity along random one-dimensional projections and uses a single hyperparameter with linear time and memory complexity.

4 TRAINING AND PLANNING WITH WORLD MODELS

The paper trains video and action-conditioned JEPAs with recursive multistep prediction, regularization, and trajectory-level energy objectives. These components support temporal rollouts and low-energy goal-directed planning.

Video Prediction: Video-JEPA rollouts maintain coherent digit motion over extended horizons, capturing trajectories and dynamics in Moving MNIST.The visualization compares input frames, one-step predictions, and full autoregressive rollouts.
Multistep Rollout Training: Multistep rollout training computes losses Lk for k = 1, . . . , K by recursively unrolling the predictor.L1 recovers the single-step loss, while higher-order predictions require repeated predictor calls.
Training Objectives: Video and world-model objectives combine prediction loss with regularization, with action-conditioned models adding control inputs and additional trajectory constraints.The listed losses are Lvideo = Lpred + λR(z1:T) and Lworld = Lpred + λR(z1:T, u1:T).
Multistep Rollout Training: Training with k-step rollouts aligns optimization with autoregressive inference, reducing exposure bias and improving long-horizon prediction quality.This alignment is presented as the reason recursive training improves rollout behavior.
Goal-Conditioned Planning: Goal-conditioned planning accumulates prediction energy over imagined trajectories toward a goal, and MPPI searches for minimum-energy action sequences.MPPI samples trajectories, weights them by exponentiated negative energy, and iteratively refines the proposal distribution.

5 EXPERIMENTS

Experiments evaluate modular JEPA components across image representation learning, video prediction, and action-conditioned planning. Results show useful CIFAR-10 representations, improved temporal prediction from multistep training, and strong Two Rooms planning with regularization ablations.

Experimental Setup: Three tasks test increasing complexity: CIFAR-10 image representation learning, Moving MNIST video prediction, and Two Rooms goal-conditioned planning.The experiments use modular encoders, predictors, regularizers, and planners across these settings.
Image Representation Learning: Approximately 90–91% linear probing accuracy is achieved on CIFAR-10 with both VICReg and SIGReg.A learned projector improves performance by around 3 points over directly regularizing encoder outputs.
Video Prediction: Multistep autoregressive rollouts maintain Moving MNIST prediction quality over extended horizons, while k-step training significantly improves downstream Average Precision.The improvement is attributed to reduced exposure bias between teacher-forced training and autoregressive inference.
Action-Conditioned Video-JEPA: 97% success rate is achieved on the randomized Two Rooms planning task using MPPI with cumulative cost over the planning horizon.The task includes randomized initial and goal states and may require temporarily moving farther from the goal.
Regularization Ablations: Removing inverse dynamics causes success to collapse to 1%, while variance, covariance, and temporal-similarity terms provide substantial absolute improvements.The reported gains are approximately 50% for each variance or covariance term and approximately 35% for temporal similarity.
Planning Cost Ablation: 50%?

6 FUTURE DIRECTIONS

EB-JEPA is designed as a small-scale testbed for rapid algorithmic exploration, with future directions spanning regularization theory, hierarchical prediction, learned costs, and transition to larger codebases.

Single-GPU training, simple datasets, and controlled simulated environments support rapid prototyping of JEPA algorithms.The library is intended for fast iteration before scaling to more complex settings.
Advancing Regularization Theory: Regularization theory remains incomplete, motivating systematic study of variance, covariance, temporal similarity, inverse dynamics, and hyperparameter selection.The controlled setting enables iteration without confounding factors from large-scale distributed training.
Hierarchical World Models: Hierarchical world models could combine fine-grained dynamics for local control with coarse abstractions for long-horizon planning.Current JEPA models predict at a single temporal resolution.
Learned Cost and Value Functions: Learned task-specific cost or value functions could improve goal-directed behavior beyond the current simple distance-based representation-space costs.The proposed direction includes learning from demonstrations or sparse rewards.
Complementary to Large-Scale Codebases: After small-scale validation, researchers can transition promising approaches to distributed codebases, pretrained backbones, and more complex environments.EB-JEPA is positioned as complementary to larger-scale systems rather than a replacement for them.

7 CONCLUSION

EB-JEPA is an open-source library for JEPA-based representations and world models, designed for reproducible single-GPU experiments and educational use. Its scope is limited to public benchmarks and procedurally generated environments without personally identifiable information.

EB-JEPA spans image representation learning, video prediction, and action-conditioned planning, with examples trainable on a single GPU within a few hours.Experiments and ablations examine component importance for preventing collapse and enabling planning.
The library is intended as both a practical toolkit for JEPA research and an educational resource for energy-based self-supervised learning.
All experiments use CIFAR-10, Moving MNIST, or procedurally generated Two Rooms environments, none of which contain personally identifiable information.The paper reports no direct ethical concerns with this work.
The full codebase, training scripts, model implementations, evaluation code, hyperparameters, and documented Two Rooms seeds support reproducibility.Each example is self-contained and trains on a single GPU in a few hours.

A HYPERPARAMETERS

The appendix documents training and evaluation hyperparameters for the library’s examples, including rollout and trajectory settings and the image-JEPA CIFAR-10 configurations.

Tables 5 and 6 summarize training hyperparameters, including multistep rollout count K and temporal trajectory slice length T.These parameters cover the temporal examples and multistep prediction training.
Table 7 specifies the MPPI planning configuration for goal-conditioned navigation in the action-conditioned video-JEPA example.
The ViT Image-JEPA CIFAR-10 example achieves 87% linear probing accuracy, while the main-paper image results use a ResNet-18 backbone.

B PLANNING ALGORITHM

The planning algorithm uses MPPI to optimize action sequences under a learned planning energy, combining soft importance weighting with elite selection in the implementation.

MPPI iteratively refines action sequences through sampling-based importance weighting of trajectory costs.Unlike CEM, it weights all samples by exponentiated costs rather than fitting a Gaussian only to elite samples.
The planner minimizes the planning energy E_plan over action sequences using the trained encoder, predictor, and action encoder.
The implementation selects the top Q trajectories before applying exponential weighting among those elites.This combines MPPI’s soft weighting with elite selection.
MPPI’s temperature τ controls weight sharpness, while its importance-sampling formulation accommodates multi-modal cost landscapes.

C EXTENDED RELATED WORK

The section contrasts diffusion-based trajectory generation with JEPA-based latent dynamics and outlines MPPI planning for action-conditioned video-JEPA models. It also identifies the associated video-training and planning hyperparameter tables.

Diffusion-Based Planning: Diffusion planning treats trajectory optimization as iterative denoising, with later work extending the paradigm to predictive control and visuomotor policy learning.Diffuser pioneered diffusion-based planning; Diffusion MPC and Diffusion Policy broadened its application areas.
Experimental Configuration: The accompanying tables specify training hyperparameters for video-JEPA examples and planning hyperparameters for the action-conditioned video-JEPA MPPI example.The video-training table defines K as rollout steps and T as trajectory-slice length; the planning table uses Algorithm 1 notation and defines M as total replanning steps.
Diffusion-Based Planning: JEPAs provide efficient latent dynamics for fast online planning, whereas diffusion models excel at generating diverse, multimodal trajectories.
MPPI Planning: MPPI encodes initial and goal observations, samples noisy action-sequence perturbations, unrolls the predictor, and selects low-cost elite trajectories.The algorithm iterates sampling and evaluation before weighting elite samples and updating the action distribution.
MPPI Planning: MPPI executes the first m actions of the selected mean sequence, then replans from the new observation until the episode reaches M steps.