eb jepa

The Problem Researchers and engineers need a compact reference implementation for Joint‑Embedding Predictive Architectures (JEPAs) that can be run on a single GPU and that demonstrates image, video, and action‑conditioned video learning as well as planning. Existing codebases are either fragmented or tied to large internal pipelines, making rapid experimentation difficult.

What This Does eb_jepa ships a self‑contained library (eb_jepa/*.py) that implements the core JEPA components: encoders, projectors, loss functions, and a planning module. The three example entry points—examples/image_jepa/main.py, examples/video_jepa/main.py, and examples/ac_video_jepa/main.py—each assemble a model, load a config (*.yaml under examples/*/cfgs), train or evaluate, and write checkpoints/logs to a user‑specified folder.

Key files:

How It Is Wired Execution always starts at one of the three main.py files. For the AC‑video example (examples/ac_video_jepa/main.py:48 → run) the call chain is:

1. run → setup_device (eb_jepa/training_utils.py) → selects GPU/CPU.
2. run → create_video_jepa_model (eb_jepa/architectures.py) → builds the encoder/predictor stack.
3. run → env_creator (examples/ac_video_jepa/main.py:146) → constructs the Two‑Rooms environment (eb_jepa/datasets/two_rooms/env.py).
4. run → unroll (eb_jepa/planning.py) → calls generate_wall_bump_state_and_actions, which repeatedly invokes sample_uniformly_between (used 11 times across the repo).
5. unroll → save_checkpoint (eb_jepa/training_utils.py) → writes checkpoint files (state_dict, save).

The most‑connected internal symbols are unroll (5 callers) and sample_uniformly_between (6 callers). No circular imports were detected, but deep nesting (up to 6 levels) appears in eb_jepa/jepa.py, eb_jepa/datasets/two_rooms/wall_dataset.py, and eb_jepa/planning.py, making those files high‑impact change points.

File‑level responsibilities (sorted by call‑graph reach):

• eb_jepa/architectures.py – 26 functions, 13 classes, central to model construction.
• eb_jepa/losses.py – 23 functions, 9 classes, defines the training objective.
• eb_jepa/planning.py – 17 functions, 6 classes, handles environment rollout and checkpointing.
• examples/image_jepa/dataset.py – 15 functions, 6 classes, provides the CIFAR‑10 data pipeline.
• eb_jepa/jepa.py – 11 functions, 3 classes, the high‑level model wrapper used by all examples.

How To Use It

# Clone the repo
git clone https://github.com/moses-y/eb_jepa.git
cd eb_jepa

# Install dependencies with uv (as documented)
uv sync          # creates .venv and installs from pyproject.toml
source .venv/bin/activate   # optional, uv can run without activation

# Run an example (default config files are under examples/*/cfgs)
python -m examples.image_jepa.main   # image JEPA
python -m examples.video_jepa.main   # video JEPA
python -m examples.ac_video_jepa.main   # action‑conditioned video + planning

Configuration: each example reads a YAML file (examples/*/cfgs/*.yaml). Adjust batch size, learning rate, or checkpoint root by editing those files or by setting the environment variables referenced in the README (EBJEPA_DSETS, EBJEPA_CKPTS). No additional secrets are required.

Real‑World Use A robotics team could replace a custom world‑model training script with examples/ac_video_jepa/main.py. By pointing EBJEPA_DSETS at their own simulator logs and tweaking planning_cem.yaml, the same planning pipeline (unroll → checkpoint → inference) can be reused without rewriting data loaders or loss functions.

Code Health & Issues

• High – GitHub Actions use a mutable tag (astral-sh/setup-uv@v6). Pin to a commit SHA to prevent supply‑chain drift.
• High – No lockfile beside pyproject.toml. Generate and commit a lockfile (e.g., uv lock) to guarantee reproducible builds.
• Medium – Workflows lack explicit permissions for GITHUB_TOKEN. Add permissions: { contents: read } at the top of each workflow.
• Medium – No Dependabot/Renovate configuration; add .github/dependabot.yml.
• Medium – No dependency‑vulnerability scan in CI; integrate dependency-review-action or osv-scanner.
• Low – Workflow jobs omit timeout-minutes; set a reasonable bound (e.g., timeout-minutes: 60).

Other observations: tests exist (tests/*.py) and CI runs them, but the repository has no Dockerfile and no pre‑published lockfile, which limits out‑of‑the‑box reproducibility.

The Bottom Line eb_jepa provides a functional, example‑driven reference for JEPA research, with clear entry points and modular architecture. The code is usable after a simple uv sync, but deep nesting and duplicated blocks raise maintenance overhead, and the lack of a lockfile plus mutable CI actions pose reproducibility and security risks. It is best suited for researchers or engineers who need a quick prototype and are comfortable refactoring the larger modules.