H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

Wancong Zhang*, Basile Terver*, Mike Rabbat, Yann LeCun†, Randall Balestriero†

AMI Labs, NYU, INRIA Paris, Brown University
_{* equal contribution, † equal advising}

H-JEPA learns a hierarchy of action-conditioned JEPA world models end-to-end, with each level predicting farther ahead in its own learned latent space. Hierarchical planning with H-JEPA proceeds top-down: the top level optimizes progress toward the goal, and each level's predictions become subgoals for the level below. On Visual AntMaze, a three-level H-JEPA achieves 73.3% planning success against 18.0% using the flat LeWorldModel, with a fraction of planner compute.

This repository contains the training and planning code for H-JEPA. It supports Visual AntMaze, FourRoom Distractors, OGBench Cube, Push-T, and offline planning on DROID. It contains scripts to reproduce Figure 6 (bottom row), Tables 1 and 9, and Figure 9(b) of the H-JEPA paper. It is built on LeWM and stable-worldmodel.

git clone https://github.com/kevinghst/H-JEPA.git && cd H-JEPA
conda create -n hjepa python=3.11.13 && conda activate hjepa
pip install -e ".[train,env]" --extra-index-url https://download.pytorch.org/whl/cu128
export HJEPA_HOME=/path/to/data   # datasets, expert policies and checkpoints live here

Choose your workflow:

• Train and evaluate on Visual AntMaze — start with a concrete example.
• Reproduce the paper's simulation results
• Reproduce the paper's DROID results

Run commands from h_jepa/ unless noted (cd h_jepa after installation).

Train three-level H-JEPA on Visual AntMaze with seed 42. First prepare the AntMaze datasets, including the probing files, and evaluation tasks.

python main_hjepa.py --config-name ant_hjepa_l3 seed=42

Training configs follow <env>_<model>; see the model/config tables for other choices.

Training automatically runs planning evaluation at the end.

Runs write to $HJEPA_HOME/ckpts/<env>/<env>_<model>/seed<seed>/:

• <env>_<model>_object.ckpt, config.yaml, logs and the stable-pretraining cache spt/;
• planning_eval/epoch_XXXX/metrics.yaml: final planning eval with the matching planner and model seed;
• final_probing_decoding_eval/: frozen-model probes and decoders (config/probing/<env>.yaml).

Use this to evaluate a pretrained checkpoint or rerun evaluation with different planner settings.

Here we evaluate the three-level H-JEPA model from above using three-level hierarchical planning on Visual AntMaze. Prepare the evaluation tasks and either train the model or download its checkpoint first.

python eval.py --config-name ant_l3 seed=42 output.dir=eval_l3 \
  policy="$HJEPA_HOME/ckpts/ant/ant_hjepa_l3/seed42/ant_hjepa_l3_object.ckpt"

Results are saved beside the checkpoint in eval_l3/metrics.yaml; success_rate is reported as a percentage. Completed tasks are skipped on rerun. Use a new output.dir when changing planner settings.

For cluster training and evaluation, see Running on SLURM.

Reproduces the bottom row of Figure 6 and the level-1 columns of Tables 1 and 9. Results report paper success rates (%, mean ± SE over seeds 42, 43, 44; planner seed = model seed). For evaluation only, download tasks and checkpoints; datasets are needed for training and probing.

The datasets are HDF5 files under $HJEPA_HOME.

Push-T and Cube. The training data is LeWM's release (quentinll/lewm-pusht, quentinll/lewm-cube, MIT). The Push-T validation set is in our release. From the repository root:

hf download quentinll/lewm-pusht --repo-type dataset --local-dir $HJEPA_HOME --include "*.zst"
hf download quentinll/lewm-cube --repo-type dataset --local-dir $HJEPA_HOME --include "*.zst"
hf download jepa-world-models/h-jepa pusht_expert_val.h5 --repo-type dataset --local-dir $HJEPA_HOME
unzstd $HJEPA_HOME/pusht_expert_train.h5.zst
tar -I unzstd -xf $HJEPA_HOME/cube_single_expert.tar.zst -C $HJEPA_HOME
python scripts/data/add_pusht_block_ori.py "$HJEPA_HOME/pusht_expert_train.h5"  # add block_ori for probing
python scripts/data/split_cube.py  # first 9,800 episodes for training; last 200 for validation

Visual AntMaze and FourRoom Distractors. Download the eight files from jepa-world-models/h-jepa:

hf download jepa-world-models/h-jepa --repo-type dataset --local-dir $HJEPA_HOME \
  --include "visual_antmaze_medium_*" --include "fourroom_tp35_d1*" --include SHA256SUMS
(cd "$HJEPA_HOME" && sha256sum -c --ignore-missing SHA256SUMS)

bash scripts/data/collect_datasets.sh

mkdir -p $HJEPA_HOME/ogbench_experts
wget -qO- https://rail.eecs.berkeley.edu/datasets/ogbench/experts.tar.gz \
  | tar xz -C $HJEPA_HOME/ogbench_experts --strip-components=1 experts/ant

Each environment has a fixed set of 50 start/goal tasks under h_jepa/assets/eval_trajs/. From h_jepa/, download them from jepa-world-models/h-jepa:

hf download jepa-world-models/h-jepa --repo-type dataset --local-dir assets --include "eval_trajs/*"
(cd assets/eval_trajs && sha256sum -c SHA256SUMS)

# Visual AntMaze: start/goal cells 3 grid cells apart, reached by the expert policy
# (needs the OGBench ant expert in $HJEPA_HOME/ogbench_experts/ant/, see 3.1)
python scripts/generate_maze_expert_grid_eval_tasks.py --config-name ant_flat \
  --output-path assets/eval_trajs/ant/expert_grid_d3_n50.pt \
  --d-low 3 --d-high 3 --num-episodes 50 --rollout-budget 125 --seed 42

# FourRoom Distractors: writes one file per active-distractor count; the evals use _d1
# (no dataset needed)
python scripts/generate_fourroom_eval_tasks.py \
  --data-config-path ../scripts/data/config/fourroom_tp35_d0to5.yaml \
  --output-dir assets/eval_trajs/fourroom --output-stem fourroom_tp35 \
  --cross-n-rooms 2 --max-steps 75 --num-episodes 50

# OGBench Cube: 20-step windows centred on the grasp, from the val split
# (needs cube_single_expert_val.h5, see 3.1)
python scripts/generate_dataset_eval_trajs.py --config-name cube_flat \
  --dataset-name cube_single_expert_val --traj-sampling-mode cube_pickup_centered \
  --goal-offset-steps 20 --eval-budget 50 --seed 42 \
  --output-path assets/eval_trajs/ogbench/goal_offset_20_pickup_val.pt

# Push-T: 75-step windows from the val split, stratified over episodes
# (needs pusht_expert_val.h5, see 3.1)
python eval.py --config-name pusht_flat policy=random load_eval_trajs_path=null \
  eval.goal_offset_steps=75 eval.dataset_name=pusht_expert_val seed=42 \
  dump_eval_trajs_path=assets/eval_trajs/pusht/goal_offset_75_val.pt

Train all 4 environments × 7 models × 3 seeds (84 runs):

scripts/train_all.sh

Runs are sequential. Filter with ENVS, MODELS, SEEDS, e.g. ENVS=cube MODELS="lewm hjepa_l3" SEEDS=42 scripts/train_all.sh, or submit each training command as a separate cluster job. Outputs follow §2.1.

To skip training (§3.3), download all 84 trained models (4 environments × 7 models × 3 seeds, 9.7 GB) from jepa-world-models/h-jepa (the DROID models are in §4.3):

hf download jepa-world-models/h-jepa --local-dir $HJEPA_HOME/ckpts --exclude "droid/*"
# one environment only: --include "cube/*" instead of --exclude; one seed: add --include "*/seed42/*"

Checkpoints and normalizer.pt use the training layout (§2.1), ready for the eval scripts below.

First complete training (§3.3) or download pretrained checkpoints (§3.4), and prepare the evaluation tasks (§3.2). Trained and downloaded checkpoints use the same directory layout expected by this script.

Training runs this eval automatically; to evaluate saved checkpoints:

scripts/eval_depth.sh

Results are saved next to each checkpoint in eval_<planner>/metrics.yaml. Both this script and the cost-ladder script accept ENVS and SEEDS filters. Each requires all relevant model checkpoints for the selected environments and seeds; neither accepts a MODELS filter. For a single checkpoint, use standalone evaluation (§2.2).

See the expected results reported in the paper.

First complete training (§3.3) or download pretrained checkpoints (§3.4), and prepare the evaluation tasks (§3.2).

Evaluate each H-JEPA model with level-1 planning and either native or projected cost:

scripts/eval_cost_ladder.sh

Results go beside each checkpoint in eval_flat/metrics.yaml (native L1) and eval_l<k>_project/metrics.yaml (level-k cost, k = 2 … model depth).

See the expected results reported in the paper.

Evaluate open-loop planning on real-robot videos from DROID. Plans are scored against recorded expert trajectories; this evaluation measures trajectory fidelity, not closed-loop task success.

Reproduces Figure 9(b). For evaluation only, download the evaluation clips (252 MB) and checkpoints; the 91 GB training set is needed for training only.

Download DROID data to $HJEPA_HOME/droid/ (dataset details and license).

Evaluation clips. The 16 evaluation episodes, raw 1280×720 in droid_raw/1.0.1/:

hf download jepa-world-models/h-jepa --repo-type dataset --local-dir $HJEPA_HOME \
  --include droid/droid_raw_eval16.tar --include droid/SHA256SUMS
(cd $HJEPA_HOME/droid && sha256sum -c --ignore-missing SHA256SUMS && tar -xf droid_raw_eval16.tar)

Training data. 91 GB of tar shards, re-encoded at 256×256 in droid_256p/1.0.1/:

hf download jepa-world-models/h-jepa --repo-type dataset --local-dir $HJEPA_HOME --include "droid/*"
bash $HJEPA_HOME/droid/extract.sh    # checks SHA256SUMS, untars all shards; --delete drops the tars

Models train at 5 fps on 74,896 episodes, excluding the 16 evaluation clips, with 64 validation episodes.

Train 3 models × 3 seeds (1, 1000, 10000), for 9 runs:

ENVS=droid scripts/train_all.sh

Each epoch is 292 steps. Run inside a SLURM allocation with 2 GPUs and 2 tasks per node: the script uses srun --ntasks-per-node=2 python main_hjepa.py --config-name droid_<model> seed=<seed>. Alternatively, submit each command separately. MODELS and SEEDS filter runs.

Outputs follow §2.1, with snapshots droid_<model>_epoch_<N>_object.ckpt every save_every_n_epochs.

To skip training (§4.2), download the 9 trained models (3 models × 3 seeds, 2.7 GB) from jepa-world-models/h-jepa:

hf download jepa-world-models/h-jepa --local-dir $HJEPA_HOME/ckpts --include "droid/*"
# one model only: --include "droid/droid_hjepa_l2/*"; one seed: --include "droid/*/seed1/*"

droid_<model>_object.ckpt (the epoch-100 model) and normalizer.pt use the training layout (§4.2), without the epoch snapshots, ready for the eval script below.

First complete training (§4.2) or download pretrained checkpoints (§4.3), and prepare the evaluation clips (§4.1). Trained and downloaded checkpoints use the same directory layout expected by this script.

Plan each clip start → goal in one open-loop call and score against ground-truth actions. Evaluate all models with planner seeds 1, 2, 3, writing eval_{flat,l2}/plan_seed<ps>/eval.csv beside each checkpoint and printing mean ± SE per model (filter with MODELS, SEEDS, PLAN_SEEDS):

scripts/eval_droid.sh

For standalone evaluation, adapt the command in §2.2 to your DROID checkpoint, using droid_flat for LeWM or droid_l2 for HWM and H-JEPA.

Fréchet fidelity measures how closely the planned end-effector path follows the expert trajectory: 0% corresponds to zero motion, and 100% to an exact match. The frechet/skill_mean column in eval.csv reports mean fidelity across the 16 clips; multiply by 100 to express it as a percentage.

See the expected results reported in the paper.

See ARCHITECTURE.md for model and planner internals. Configs use <env> = ant, fourroom, cube, pusht or droid. The simulation environments support the options below. DROID supports models lewm, hjepa_l2 and hwm_l2, with planners flat for LeWM and l2 for H-JEPA/HWM.

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

MIT (see LICENSE).