pyjepa 0.2.0

A unified Python library for JEPA, V-JEPA, LeJEPA world models, and physical intelligence research.

pyjepa

A unified Python library for training, finetuning, evaluating, and deploying Joint Embedding Predictive Architecture (JEPA) models, end-to-end world models (LeWM), and physical intelligence systems.

One consistent API. Any hardware. From research to production.

pip install pyjepa

---

Why pyjepa

Most JEPA implementations are single-paper reference code -- tightly coupled to one dataset, one cluster, one set of hyperparameters. Extending them requires reverse-engineering thousands of lines of undocumented training infrastructure.

pyjepa provides a modular, production-grade toolkit where every component -- encoder, predictor, loss, mask strategy, optimizer schedule, planner -- is independently usable, swappable, and well-documented. You can use the full training pipeline end-to-end, or drop a single module into your existing codebase.

Designed for two audiences:

• Easy for practitioners. Train an I-JEPA with 10 lines. One-call functions for common workflows. Automatic device selection, dtype, and hardware optimization.
• Customizable for researchers. Every layer, loss function, and schedule is exposed. Register custom encoders, predictors, and masking strategies. Hook into any point in the training loop. Replace any component without touching the rest.

---

Table of Contents

• Installation
• Supported Models
• Quick Start
• Architecture
• Training
• Finetuning
• Planning and Control
• Inference
• Custom Models and Extensions
• Hardware and Optimization
• CLI Reference
• Configuration
• API Reference
• Citation
• License

---

Installation

From PyPI

pip install pyjepa            # core library (torch, numpy, einops)

Optional extras for specific workflows:

pip install pyjepa[vision]    # torchvision + Pillow for image datasets
pip install pyjepa[video]     # av + decord for video loading
pip install pyjepa[data]      # h5py + HuggingFace datasets for trajectory data
pip install pyjepa[train]     # wandb + tensorboard for experiment tracking
pip install pyjepa[planning]  # scipy for optimization-based planners
pip install pyjepa[all]       # everything above
pip install pyjepa[dev]       # pytest + ruff + black for development

From Source

git clone https://github.com/aneesaslam/py-jepa.git
cd py-jepa
pip install -e ".[dev]"

Apple Silicon (macOS M1/M2/M3/M4)

No extra configuration. pyjepa auto-detects the MPS backend. Install PyTorch 2.1+ from pytorch.org -- the MPS backend is included in the standard macOS wheel.

pip install torch torchvision
pip install pyjepa[vision]

Windows (CUDA)

Install the CUDA-enabled PyTorch wheel first:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install pyjepa[all]

Verify Installation

pyjepa info       # prints device, memory, dtype support, recommendations
pyjepa list       # lists all available encoders, models, losses, planners

import pyjepa
print(pyjepa.__version__)
print(pyjepa.get_device())       # auto-detects best available device
print(pyjepa.model_summary(pyjepa.vit_small()))

---

Supported Models

Model	Paper	What it does
I-JEPA	Assran et al., CVPR 2023	Predict representations of masked image regions in latent space. EMA target encoder prevents collapse.
V-JEPA	Bardes et al., 2024	Extends I-JEPA to video with 3D spatiotemporal masking (tube + multi-block).
LeWM	Maes, Le Lidec et al., 2026	End-to-end world model from raw pixels. Two losses only: prediction MSE + SIGReg. No EMA encoder needed. ~15M parameters. Trains on a single GPU in hours.

All three share the same VisionTransformer encoder backbone, optimizer utilities, and checkpoint format.

---

Quick Start

I-JEPA: Image Pre-training (10 lines)

import pyjepa

# Build model -- encoder + EMA target + predictor, fully wired
model = pyjepa.create_jepa("ijepa", encoder="vit_small", img_size=224)

# Print architecture summary
print(pyjepa.model_summary(model))

# Train (one-call convenience function)
pyjepa.train_ijepa(model, your_loader, epochs=300, lr=1e-3, output_dir="./runs/ijepa")

I-JEPA: Full Control

import torch
from pyjepa import (
    IJEPA, IJEPATrainer, IJEPAConfig, TrainConfig,
    vit_small, MultiBlockMaskCollator,
    MultiMaskWrapper, PredictorMultiMaskWrapper,
)
from pyjepa.models.predictor import vit_predictor

# Build encoder and predictor separately
enc = vit_small(img_size=224, patch_size=16)
pred = vit_predictor(
    img_size=224, patch_size=16, embed_dim=enc.embed_dim,
    predictor_embed_dim=192, depth=6, num_heads=enc.num_heads,
)
model = IJEPA(
    encoder=MultiMaskWrapper(enc),
    predictor=PredictorMultiMaskWrapper(pred),
)

# Configure mask collator
cfgs_mask = [
    {"spatial_scale": (0.85, 1.0), "aspect_ratio": (0.75, 1.5), "num_blocks": 1},
    {"spatial_scale": (0.15, 0.2), "aspect_ratio": (0.75, 1.5), "num_blocks": 4},
]
collator = MultiBlockMaskCollator(cfgs_mask, crop_size=224, patch_size=16)

# Configure training with full control over every hyperparameter
train_cfg = TrainConfig(
    output_dir="./runs/ijepa",
    device="auto",              # auto-selects CUDA > MPS > CPU
    epochs=300,
    amp_dtype="auto",           # auto-selects bfloat16/float16/float32
    grad_clip=10.0,
    save_every_epochs=10,
    keep_last_ckpts=3,
)
ijepa_cfg = IJEPAConfig(
    lr=1e-3,
    start_lr=2e-4,
    final_lr=1e-6,
    weight_decay=0.04,
    warmup_epochs=40,
    momentum_start=0.996,       # EMA momentum schedule
    momentum_end=1.0,
    loss_exp=1.0,               # 1.0 = L2 loss, 2.0 = L4 loss
    reg_coeff=0.0,              # variance regularizer weight
)
trainer = IJEPATrainer(model, train_cfg, ijepa_cfg, iterations_per_epoch=len(loader))

# Hook into the training loop
@trainer.on_step
def log_metrics(trainer, metrics):
    if trainer.state.step % 100 == 0:
        print(f"step {trainer.state.step}: loss={metrics['loss']:.4f}")

trainer.fit(loader)

LeWM: World Model Training

from pyjepa import vit_small, build_lewm, LeWMTrainer, LeWMConfig, TrainConfig
from pyjepa.data import TrajectoryDataset

# Build world model
encoder = vit_small(img_size=84, patch_size=12)
model = build_lewm(
    encoder,
    action_dim=4,
    history_size=3,
    emb_dim=384,
    action_emb_dim=64,
    pred_depth=6,
    pred_heads=6,
    pred_mlp_dim=1024,
)

# Load trajectory data
dataset = TrajectoryDataset(episodes, window=8)  # history_size + num_preds
loader = DataLoader(dataset, batch_size=128)

# Train
train_cfg = TrainConfig(output_dir="./runs/lewm", epochs=100, device="auto")
lewm_cfg = LeWMConfig(
    lr=3e-4,
    history_size=3,
    num_preds=5,
    sigreg_weight=0.1,          # SIGReg collapse prevention
    sigreg_knots=17,
    sigreg_num_proj=1024,
)
trainer = LeWMTrainer(model, train_cfg, lewm_cfg)
trainer.fit(loader)

LeWM: Planning with CEM

from pyjepa import CEMPlanner, PlanConfig

model.eval()
info = {
    "pixels": obs_window,    # (B, T_ctx, C, H, W)
    "goal": goal_image,      # (B, 1, C, H, W)
}

planner = CEMPlanner(PlanConfig(
    horizon=10,
    action_dim=4,
    num_samples=512,
    num_elites=64,
    num_iters=5,
    action_min=-1.0,
    action_max=1.0,
))

with torch.no_grad():
    best_actions = planner(model, info)   # (horizon, action_dim)
    action = best_actions[0]              # execute first action

---

Architecture

I-JEPA

Trains a VisionTransformer encoder with a masked prediction objective in latent space. A context block is encoded by the online encoder; target blocks are encoded by an EMA copy. A lightweight predictor maps context embeddings to target embeddings.

L_JEPA = (1/M) * sum_m ||predictor(enc(x_ctx), mask_m) - sg(enc_ema(x, mask_m))||^2

V-JEPA

Extends I-JEPA to video with 3D spatiotemporal masking. Masks are volumetric blocks (tubes) across time. The VisionTransformer switches to 3D tubelet patch embeddings when num_frames > 1.

LeWM (LeWorldModel)

End-to-end world model from pixels with only two losses:

L = MSE(AR_predictor(enc(x_t), a_t), enc(x_{t+1})) + lambda * SIGReg(enc(x))

SIGReg (Sketched Isotropic Gaussian Regularizer) prevents collapse by matching the empirical characteristic function of embeddings to a standard Gaussian via random 1D projections.

The ARPredictor is a causal Transformer conditioned on action embeddings via AdaLN-zero blocks.

---

Training

One-Call Functions

For simple use cases:

pyjepa.train_ijepa(model, loader, epochs=300, lr=1e-3, output_dir="./out")
pyjepa.train_vjepa(model, loader, epochs=200, lr=6e-4, output_dir="./out")
pyjepa.train_lewm(model, loader, epochs=100, lr=3e-4, sigreg_weight=0.1,
                  history_size=3, output_dir="./out")

Trainer Classes

For full control:

from pyjepa.training import IJEPATrainer, LeWMTrainer, TrainConfig

trainer = IJEPATrainer(model, train_cfg, ijepa_cfg, iterations_per_epoch=len(loader))
trainer.fit(train_loader, val_loader)

Training Hooks

Register callbacks at any point in the loop:

@trainer.on_step
def on_step(trainer, metrics):
    # Called after every training step with loss, lr, grad_norm, etc.
    wandb.log(metrics, step=trainer.state.step)

@trainer.on_epoch_end
def on_epoch_end(trainer):
    # Called after every epoch
    trainer.save_checkpoint()

Checkpointing

# Save
trainer.save_checkpoint()

# Resume automatically (looks for latest checkpoint in output_dir)
trainer = IJEPATrainer(model, TrainConfig(output_dir="./runs", resume=True), cfg)
trainer.fit(loader)  # resumes from where it left off

Distributed Training

torchrun --nproc_per_node=8 train.py --config configs/ijepa_base.yaml

from pyjepa.utils.distributed import init_distributed
init_distributed()  # reads RANK/WORLD_SIZE from environment

Config-Driven Training

pyjepa train ijepa --config configs/ijepa_base.yaml
pyjepa train lewm  --config configs/lewm_base.yaml

---

Finetuning

Linear Probe

from pyjepa.training.finetune import LinearProbeTrainer, FinetuneConfig

ft_cfg = FinetuneConfig(
    num_classes=1000,
    head="linear",
    freeze_encoder=True,
    lr=1e-3,
    label_smoothing=0.1,
)
trainer = LinearProbeTrainer(model.encoder, train_cfg, ft_cfg)
trainer.fit(train_loader, val_loader)

Attentive Pooler

Learned query tokens + cross-attention over patch tokens -- stronger than CLS pooling:

ft_cfg = FinetuneConfig(num_classes=1000, head="attentive", freeze_encoder=True)

One-Call Finetuning

pyjepa.finetune_classifier(
    encoder=model.encoder,
    train_loader=train_loader,
    val_loader=val_loader,
    num_classes=1000,
    head="attentive",
    freeze_encoder=True,
    epochs=100,
    output_dir="./runs/finetune",
)

Standalone Modules

from pyjepa import AttentivePooler, AttentiveClassifier

pooler = AttentivePooler(embed_dim=384, num_heads=6, num_query_tokens=1)
classifier = AttentiveClassifier(embed_dim=384, num_heads=6, num_classes=1000)

---

Planning and Control

All planners implement __call__(model, info) -> Tensor(H, A):

from pyjepa import PlanConfig, CEMPlanner, MPPIPlanner, RandomShootingPlanner

cfg = PlanConfig(
    horizon=10,
    action_dim=4,
    num_samples=512,
    num_elites=64,        # CEM only
    num_iters=5,
    action_min=-1.0,
    action_max=1.0,
    temperature=0.1,      # MPPI only
)

cem  = CEMPlanner(cfg)    # Cross-Entropy Method: iterative Gaussian refit to elite trajectories
mppi = MPPIPlanner(cfg)   # MPPI: soft-min reweighting, smoother gradient landscape
rs   = RandomShootingPlanner(cfg)  # single-shot baseline

WorldModelPolicy

Stateful policy wrapper for environment interaction:

from pyjepa.planning import WorldModelPolicy

policy = WorldModelPolicy(model=lewm, solver=CEMPlanner(cfg), config=cfg)
action = policy.act({"pixels": obs, "goal": goal})

---

Inference

Embedding

from pyjepa import Encoder, embed_image, embed_video

# Stateful wrapper with pooling and normalization
enc = Encoder(model.encoder, device="mps", pooling="mean", normalize=True)
feats = enc.embed_image(images)   # (B, D)
feats = enc.embed_video(video)    # (B, D)

# Functional (stateless)
feats = embed_image(model.encoder, images)

Latent Rollout

from pyjepa import rollout_latent
from pyjepa.inference import Rollout

# Batch rollout
predicted = rollout_latent(model, initial_info, actions, history_size=3)

# Step-by-step
roll = Rollout(model, history_size=3)
roll.reset(initial_obs)
emb = roll.step(action)

Surprise Scoring

Measures prediction error as a proxy for physical plausibility:

from pyjepa import surprise_score

scores = surprise_score(model, info, history_size=3, reduction="per_step")  # (B, T)
mean_err = surprise_score(model, info, reduction="mean")                    # (B,)

---

Custom Models and Extensions

Register a Custom Encoder

from pyjepa import register_encoder, create_encoder
from pyjepa.models.vit import VisionTransformer

@register_encoder("my_vit_128")
def my_vit(**kwargs):
    return VisionTransformer(
        img_size=128, patch_size=8,
        embed_dim=512, depth=10, num_heads=8,
        **kwargs
    )

enc = create_encoder("my_vit_128")  # now available everywhere by name

Register a Custom Predictor

from pyjepa import register_predictor
import torch.nn as nn

@register_predictor("mlp_predictor")
class MLPPredictor(nn.Module):
    def __init__(self, embed_dim, **kwargs):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.GELU(),
            nn.Linear(embed_dim * 4, embed_dim),
        )

    def forward(self, x, masks_x=None, masks=None):
        return [self.net(xi) for xi in x]

Register a Custom JEPA

@pyjepa.register_jepa("my_jepa")
def build_my_jepa(encoder, predictor, **kwargs):
    return MyCustomJEPA(encoder=encoder, predictor=predictor)

model = pyjepa.create_jepa("my_jepa", encoder="my_vit_128")

Listing Registered Components

print(pyjepa.available_encoders())    # ['my_vit_128', 'vit_base', 'vit_small', ...]
print(pyjepa.available_predictors())  # ['ar_predictor', 'mlp_predictor', ...]
print(pyjepa.available_jepas())       # ['my_jepa']

---

Hardware and Optimization

Automatic Device Selection

import pyjepa

device = pyjepa.get_device()           # auto: CUDA > MPS > CPU
device = pyjepa.get_device("mps")      # explicit
device = pyjepa.get_device("cuda:1")   # specific GPU

Override globally via environment variable:

export PYJEPA_DEVICE=mps

Hardware Report

info = pyjepa.DeviceInfo.from_current()
print(info)
# Prints device, backend, memory, dtype support, torch.compile support,
# and device-specific recommendations

Supported Hardware

Backend	AMP dtype	torch.compile	Notes
NVIDIA CUDA	float16 or bfloat16 (SM80+)	Supported	Flash-Attention 2 via SDPA
Apple Silicon MPS	float32	Not recommended	Math SDPA kernel; autocast experimental
CPU	float32	Supported (slow)	Fallback for testing
AMD ROCm	float16	Supported	Detected via torch.version.hip
Intel XPU	--	--	Basic support via torch.xpu

Automatic Mixed Precision

pyjepa auto-selects AMP dtype:

• CUDA SM80+ (A100, H100): bfloat16
• CUDA SM70- (V100, T4): float16 with gradient scaler
• MPS (Apple Silicon): float32 (MPS autocast is experimental)
• CPU: float32

Override: TrainConfig(amp_dtype="bfloat16")

torch.compile

Enabled on CUDA only. Disabled by default on MPS due to graph-break issues.

from pyjepa.device import supports_compile
print(supports_compile())  # True on CUDA, False on MPS

DataLoader on Apple Silicon

from pyjepa.device import pin_memory_supported, get_device

loader = DataLoader(
    dataset,
    batch_size=256,
    num_workers=4,
    pin_memory=pin_memory_supported(get_device()),  # True only on CUDA
)

Model Summary

import pyjepa

model = pyjepa.vit_small()
print(pyjepa.model_summary(model))

Output:

================================================================================
 VisionTransformer
--------------------------------------------------------------------------------
 Layer                                    Output Shape                 Params
--------------------------------------------------------------------------------
 patch_embed (PatchEmbed)                 --                           295.3K
 blocks (ModuleList)                      --                           21.29M
 norm (LayerNorm)                         --                              768
================================================================================
 Total params:         21,664,896  (21.66M)
 Trainable params:     21,589,632  (21.59M)
 Frozen params:            75,264  (75.3K)
 Est. size:               82.6 MB
--------------------------------------------------------------------------------

---

CLI Reference

pyjepa --help                          Show all available commands
pyjepa version                         Print version string
pyjepa info                            Hardware report with recommendations
pyjepa info --device mps               Inspect a specific device
pyjepa list                            List all registered components
pyjepa summary vit_small               Model architecture and parameter counts
pyjepa summary vit_base --depth 3      Deeper layer breakdown
pyjepa benchmark                       Quick throughput benchmark
pyjepa benchmark -e vit_base -b 16     Benchmark specific model and batch size
pyjepa train ijepa -c config.yaml      Train I-JEPA from YAML config
pyjepa train vjepa -c config.yaml      Train V-JEPA
pyjepa train lewm  -c config.yaml      Train LeWM world model

Environment Variables

Variable	Default	Description
PYJEPA_DEVICE	auto	Override device selection (cuda, mps, cpu, cuda:N)
PYJEPA_LOG	INFO	Log level (DEBUG, INFO, WARNING, ERROR)
PYJEPA_SEED	0	Global random seed

---

Configuration

YAML Configs

# configs/ijepa_base.yaml
model:
  type: ijepa
  encoder: vit_base
  img_size: 224
  patch_size: 16
  pred_depth: 6

mask:
  type: multiblock
  cfgs_mask:
    - spatial_scale: [0.85, 1.0]
      aspect_ratio: [0.75, 1.5]
      num_blocks: 1
    - spatial_scale: [0.15, 0.2]
      aspect_ratio: [0.75, 1.5]
      num_blocks: 4

training:
  output_dir: ./runs/ijepa_base
  device: auto
  epochs: 300
  lr: 0.001
  weight_decay: 0.04
  warmup_epochs: 40
  amp_dtype: auto
  grad_clip: 10.0

Programmatic Config

from pyjepa import load_config, save_config

cfg = load_config("configs/ijepa_base.yaml", overrides={"training.lr": 5e-4})
print(cfg.training.lr)  # 0.0005
save_config(cfg, "my_config.yaml")

---

API Reference

Top-Level Functions

Function	Description
create_jepa(name, ...)	Build a fully wired JEPA model by name
build_lewm(encoder, ...)	Build a LeWM world model around any encoder
train_ijepa(model, loader, ...)	One-call I-JEPA training
train_vjepa(model, loader, ...)	One-call V-JEPA training
train_lewm(model, loader, ...)	One-call LeWM training
finetune_classifier(encoder, ...)	One-call linear/attentive finetuning
embed_image(encoder, images)	Extract image embeddings
embed_video(encoder, clips)	Extract video embeddings
rollout_latent(model, info, actions)	Latent-space trajectory prediction
surprise_score(model, info)	Per-step prediction error scoring
get_device(prefer)	Auto-select best device
model_summary(model)	Print parameter count and architecture
load_config(path)	Load YAML config as nested namespace

Encoder Presets

Name	Embed dim	Depth	Heads	Params
vit_tiny	192	12	3	~5.5M
vit_small	384	12	6	~21.7M
vit_base	768	12	12	~86.6M
vit_large	1024	24	16	~304M
vit_huge	1280	32	16	~632M
vit_giant	1408	40	16	~1.01B

Loss Functions

Loss	Description
JEPALoss	L2/L4 prediction loss with optional variance regularizer
SIGReg	Sketched Isotropic Gaussian Regularizer (collapse prevention)
VICRegLoss	Variance-Invariance-Covariance regularizer
VarianceRegularizer	Standard deviation hinge penalty

Masking Strategies

Collator	Description
MultiBlockMaskCollator	I-JEPA style multi-block spatial masking
MultiBlock3DMaskCollator	V-JEPA style 3D spatiotemporal masking
RandomTubeMaskCollator	Same spatial mask across all frames
RandomPatchMaskCollator	Bernoulli per-patch masking (MAE-style baseline)

---

Citation

If you use pyjepa in your research, please cite the relevant works:

@inproceedings{assran2023ijepa,
  title     = {Self-Supervised Learning from Images with a Joint-Embedding
               Predictive Architecture},
  author    = {Mahmoud Assran and Quentin Duval and Ishan Misra and
               Piotr Bojanowski and Pascal Vincent and Michael Rabbat
               and Yann LeCun and Nicolas Ballas},
  booktitle = {CVPR},
  year      = {2023},
}

@article{bardes2024vjepa,
  title   = {V-JEPA: Latent Video Prediction for Visual Representation Learning},
  author  = {Adrien Bardes and Quentin Garrido and Jean Ponce and Xinlei Chen
             and Michael Rabbat and Yann LeCun and Mido Assran and Nicolas Ballas},
  journal = {arXiv preprint arXiv:2404.08471},
  year    = {2024},
}

@article{maes2026lewm,
  title   = {LeWorldModel: Stable End-to-End Joint-Embedding Predictive
             Architecture from Pixels},
  author  = {Quentin Maes and Etienne Le Lidec and others},
  journal = {arXiv preprint arXiv:2603.19312},
  year    = {2026},
  url     = {https://le-wm.github.io/},
}

---

License

Apache 2.0. See LICENSE for details.