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 3, DROID: section 4). The evaluation tasks and datasets are in jepa-world-models/h-jepa (datasets).

License

The checkpoints are released under the MIT License, the same license as the code release. The DROID checkpoints were trained on DROID, which its authors release under CC BY 4.0; please cite DROID when using them.

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}
}