H-JEPA checkpoints

Trained models of H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning (Wancong Zhang*, Basile Terver*, Mike Rabbat, Yann LeCun†, Randall Balestriero†). The code release trains, evaluates and reproduces the paper's planning results with them.

Contents

Folder	Models	Seeds	Checkpoints	Size
ant/ (Visual AntMaze)	LeWM, H-JEPA 2/3/4 levels, HWM 2/3/4 levels	42, 43, 44	21	2.6 GB
fourroom/ (FourRoom Distractors)	same 7 models	42, 43, 44	21	1.7 GB
cube/ (OGBench Cube)	same 7 models	42, 43, 44	21	2.7 GB
pusht/ (Push-T)	same 7 models	42, 43, 44	21	2.7 GB
droid/ (DROID, 5 fps)	LeWM + IDM, HWM 2 levels, H-JEPA 2 levels	1, 1000, 10000	9	2.7 GB

Each run is a folder <env>/<env>_<model>/seed<seed>/ with <env>_<model>_object.ckpt (the trained model) and normalizer.pt (action/proprio normalization), the layout the training code writes. <model> is lewm, hjepa_l<n> or hwm_l<n>, with <n> the number of levels. The DROID checkpoints are the epoch-100 models.

Download

Into $HJEPA_HOME/ckpts, where the code reads checkpoints:

# the 4 simulation environments (84 checkpoints)
hf download jepa-world-models/h-jepa --local-dir $HJEPA_HOME/ckpts --exclude "droid/*"
# DROID (9 checkpoints)
hf download jepa-world-models/h-jepa --local-dir $HJEPA_HOME/ckpts --include "droid/*"
# one environment or one seed: --include "cube/*", --include "*/seed42/*"

Usage

The _object.ckpt files are pickled HJEPA modules, so loading them needs the code release on the Python path (run from its h_jepa/ directory):

import torch

model = torch.load("cube/cube_hjepa_l3/seed42/cube_hjepa_l3_object.ckpt", weights_only=False)

To reproduce the paper's planning results, pass a checkpoint as policy= to eval.py or run the evaluation scripts, as described in the code release README (simulation: section 4, DROID: section 5). The evaluation tasks and datasets are in jepa-world-models/h-jepa (datasets).

Citation

@article{zhang2026hjepa,
  title   = {{H-JEPA}: End-to-End Learning of Hierarchical World Models for Visual Planning},
  author  = {Zhang, Wancong and Terver, Basile and Rabbat, Mike and LeCun, Yann and Balestriero, Randall},
  journal = {arXiv preprint arXiv:2610.06805},
  year    = {2026}
}