import torch

# Load our best pretrained JEPA-WMs
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'jepa_wm_droid')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'jepa_wm_metaworld')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'jepa_wm_pusht')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'jepa_wm_pointmaze')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'jepa_wm_wall')

# Load reproduced DINO-WM baseline models
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'dino_wm_droid')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'dino_wm_metaworld')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'dino_wm_pusht')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'dino_wm_pointmaze')
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'dino_wm_wall')

# Load fixed V-JEPA-2-AC baseline model
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'vjepa2_ac_droid')
# Load V-JEPA-2-AC official ckpt from https://github.com/facebookresearch/vjepa2
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'vjepa2_ac_oss')

from huggingface_hub import hf_hub_download
import torch

# Download a specific checkpoint
checkpoint_path = hf_hub_download(
    repo_id="facebook/jepa-wms",
    filename="jepa_wm_droid.pth.tar"
)

# Load the checkpoint
checkpoint = torch.load(checkpoint_path, map_location="cpu")

# Or use directly with torch.hub (automatically tries HF Hub first)
model, preprocessor = torch.hub.load('facebookresearch/jepa-wms', 'jepa_wm_droid')

Available model files on HF Hub:

• jepa_wm_droid.pth.tar, jepa_wm_metaworld.pth.tar, jepa_wm_pusht.pth.tar, jepa_wm_pointmaze.pth.tar, jepa_wm_wall.pth.tar
• dino_wm_droid.pth.tar, dino_wm_metaworld.pth.tar, dino_wm_pusht.pth.tar, dino_wm_pointmaze.pth.tar, dino_wm_wall.pth.tar
• vjepa2_ac_droid.pth.tar, vjepa2_ac_oss.pth.tar
• Decoder heads: dinov2_vits_224.pth.tar, dinov2_vits_224_INet.pth.tar, dinov3_vitl_256_INet.pth.tar, vjepa2_vitg_256_INet.pth.tar

$JEPAWM_HOME/
├── jepa-wms/          # This repository
├── dinov3/            # DINOv3 repository (optional)
├── robocasa/          # RoboCasa repository (optional)
└── robosuite/         # RoboSuite repository (optional)

DINOv2 is automatically downloaded via TorchHub when first used. Other encoders require manual setup.

Why manual download for V-JEPA v2? We centralize all model architectures around our own src/models/ for clarity. TorchHub loading can cause import conflicts since both repos share similar file structures.

Organize checkpoints in $JEPAWM_OSSCKPT:

$JEPAWM_OSSCKPT/
├── vjepa1_opensource/     # V-JEPA v1 checkpoints
│   └── vitl16.pth.tar
├── vjepa2_opensource/     # V-JEPA v2 checkpoints
│   ├── vjepa2_vit_large.pth
│   └── vjepa2_vit_giant.pth
└── dinov3/                # DINOv3 checkpoints
    ├── dinov3_vits16_pretrain_lvd1689m.pth
    └── dinov3_vitl16_pretrain_lvd1689m-<hashkey>.pth

Download from:

• V-JEPA v1: facebookresearch/jepa → ViT-L/16
• V-JEPA v2: facebookresearch/vjepa2 → ViT-L/16 or ViT-G/16
• DINOv3: facebookresearch/dinov3 → Download weights and clone repo to $JEPAWM_HOME/dinov3/

Only required for PointMaze (uses d4rl → mujoco-py). Other environments use the modern mujoco package.

# Download MuJoCo 2.1.0
mkdir -p ~/.mujoco && cd ~/.mujoco
wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz
tar -xzvf mujoco210-linux-x86_64.tar.gz

# Add to ~/.bashrc
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco210/bin
source ~/.bashrc  # or ~/.zshrc

# Verify installation
python -c "import mujoco_py; print('mujoco-py works!')"

Required for RoboCasa/RoboSuite environments:

# Install RoboSuite
git clone https://github.com/Basile-Terv/robosuite.git && cd robosuite
uv pip install -e . && cd ..

# Install RoboCasa
git clone https://github.com/Basile-Terv/robocasa.git && cd robocasa
uv pip install -e .
python robocasa/scripts/download_kitchen_assets.py   # Caution: Assets to be downloaded are around 20GB.
python robocasa/scripts/setup_macros.py && cd ..

DROID requires separate download via gsutil:

Download the DROID dataset following the instructions. This requires uv pip install gsutil. We only use the left camera and not the SVO cam files hence you can run the second of the two below commands to obtain the raw dataset of full-HD resolution (720 x 1280) MP4 files.

# Raw DROID dataset in stereo HD, stored as MP4 videos (8.7TB)
gsutil -m cp -r gs://gresearch/robotics/droid_raw <path_to_your_target_dir>
# Raw DROID dataset, non-stereo HD video only (5.6TB, excluding stereo video & raw SVO cam files)
gsutil -m rsync -r -x ".*SVO.*|.*stereo.*\.mp4$" "gs://gresearch/robotics/droid_raw" <path_to_your_target_dir>

After downloading, generate the paths CSV file required by the dataloader:

python src/scripts/generate_droid_paths.py \
    --droid_root <path_to_your_target_dir>/droid_raw/1.0.1 \
    --output_path $JEPAWM_DSET/DROID/droid_paths.csv \
    --num_workers 16 \

This script scans the dataset directory structure in parallel and creates a CSV file listing all valid episode paths.

$JEPAWM_DSET/
├── pusht_noise/           # Push-T dataset
├── point_maze/            # PointMaze dataset
├── wall_single/           # Wall dataset
├── Metaworld/             # Metaworld dataset
│   └── data/
│       └── train-00000-of-00001.parquet
├── robocasa/              # RoboCasa dataset
│   └── combine_all_im256.hdf5
├── franka_custom/         # Franka custom dataset
│   └── data/
│       ├── folding/
│       ├── pick/
│       └── push/
│           ├── brownboxpush_v0/
│           │   └── run_0001/
│           │       ├── episode.h5
│           │       └── trajectory.hdf5
│           └── push_various_objects/
├── DROID/                 # DROID dataset
│   └── droid_paths.csv
├── kinetics400/           # Kinetics-400 dataset (optional)
│   ├── k400_train_paths.csv
│   └── k400_val_paths.csv
├── kinetics710/           # Kinetics-710 dataset (optional)
│   ├── k710_train_paths.csv
│   └── k710_val_paths.csv
├── ssv2/                  # Something-Something-v2 dataset (optional)
│   ├── ssv2_train_paths.csv
│   └── ssv2_val_paths.csv
└── howto100m/             # HowTo100M dataset (optional)
    └── howto100m_paths.csv

All configs under configs/vjepa_wm/.

Decoder heads enable visualization and light evals (rollout decoding via val_rollout() in the training loop). See VM2M Decoder Heads for pretrained weights.

Note: Decoder heads are not required for training world models or running planning evaluations. The training configs in configs/vjepa_wm/*_sweep/ have heads_cfg: null by default.

Two training strategies:

• Cross-environment (recommended if datasets available): Train one decoder on VideoMix2M (HowTo100M + SSv2 + K400) — works across all environments. See configs in configs/vjepa_wm/vm2m/open_source_decs/.
• In-domain: Train one decoder per encoder per environment on environment-specific data

# Cross-environment decoder (recommended)
python -m app.main --fname configs/vjepa_wm/vm2m/open_source_decs/step2_lpips_vm2m_<enc>_<params>.yaml --debug

# State head (environment-specific)
python -m app.main --fname configs/vjepa_wm/<env>/step2_<env>_state_head_<enc>_<params>.yaml --debug

# Image decoder head (environment-specific)
python -m app.main --fname configs/vjepa_wm/<env>/step2_lpips_<env>_<enc>_<params>.yaml --debug

To reproduce the design choice comparison plots from the paper (e.g., encoder comparison, predictor architecture, rollout steps), train models using the configs in configs/vjepa_wm/*_sweep/ and then run the plotting commands in app/plan_common/plot/logs_plan_joint_per_design_choice.py.

Example commands:

# Encoder comparison
python app/plan_common/plot/logs_plan_joint_per_design_choice.py \
    --design_choices_file app/plan_common/plot/local/design_choice_yamls/enc.yaml \
    --output enc_comparison --verbose

# Predictor architecture comparison
python app/plan_common/plot/logs_plan_joint_per_design_choice.py \
    --design_choices_file app/plan_common/plot/local/design_choice_yamls/pred_arch.yaml \
    --output pred_arch_comparison --verbose

# Rollout steps comparison
python app/plan_common/plot/logs_plan_joint_per_design_choice.py \
    --design_choices_file app/plan_common/plot/local/design_choice_yamls/rollout_steps.yaml \
    --output rollout_steps_comparison --plot_line --verbose

# Final baseline comparison (LaTeX table)
python app/plan_common/plot/logs_plan_joint_per_design_choice.py \
    --design_choices_file app/plan_common/plot/local/design_choice_yamls/final_baseline_comp.yaml \
    --output final_baseline_comp --generate_latex --verbose

See the main() docstring in the script for the full list of commands used to generate paper figures.

Counterfactual decoding evaluation that generates predictions with hardcoded custom actions. This is useful for visualizing how the world model responds to specific action scenarios (e.g., "open gripper + move up" vs "close gripper + move up").

Note: This evaluation is designed to work only with DROID or franka_custom data.

To run unroll decode evaluation, set meta.unroll_decode_eval_only_mode: true in your training config and configure unroll_decode_evals:

meta:
  unroll_decode_eval_only_mode: true
unroll_decode_evals:
  specific_video: true  # Use a specific video file
  specific_video_path: /path/to/video.npz  # Optional: path to npz file
  play_in_reverse: false
  repeat_hardcode_act: 5  # Number of times to repeat hardcoded actions
  wrapper_kwargs:  # Same structure as evals.wrapper_kwargs
    ctxt_window: 2

The hardcoded actions can be customized by modifying the create_counterfactual_actions() function in evals/unroll_decode/eval.py.

The SLURM job submission is configured in src/utils/cluster.py. This file may need to be modified depending on your cluster's setup:

• Account/Partition/QoS: The function slurm_account_partition_and_qos() reads SLURM environment variables from the current job. Some clusters don't use all these concepts (account, partition, QoS) — the function handles None values gracefully.

• Low-priority QoS: For evaluation jobs, set the SLURM_QOS_LOW_PRIORITY environment variable to your cluster's low-priority QoS name (e.g., export SLURM_QOS_LOW_PRIORITY="lowest").

If you encounter MuJoCo rendering errors during evaluation (especially on headless servers or clusters), you may need to configure the rendering backend by setting these environment variables before running your scripts:

# For systems with EGL support (e.g., NVIDIA GPUs with recent drivers)
export MUJOCO_GL=egl
export PYOPENGL_PLATFORM=egl

# For systems without EGL (e.g., CPU-only rendering)
export MUJOCO_GL=osmesa
export PYOPENGL_PLATFORM=osmesa

When to use each backend:

• EGL: Preferred for GPU-accelerated rendering on headless servers with NVIDIA GPUs and recent drivers. Provides better performance.
• OSMesa: Fallback option for CPU-based rendering when EGL is not available. Slower but more compatible.

Common error messages:

• "ERROR: GLEW initialization error: Missing GL version" → Try using osmesa backend
• "Cannot initialize EGL" → Try using osmesa backend or check GPU drivers
• Rendering appears blank or corrupted → Verify the correct backend for your system

You cannot launch a main_distributed.py job from a GPU node if you do not clear the env variables, as is done with with submitit.helpers.clean_env(): in app/vjepa_wm/train.py.

If you encounter errors when loading checkpoints from torchhub such as urllib.error.HTTPError: HTTP Error 503: Service Unavailable, you should rm uv.lock, then recreate your uv venv with uv sync, activate this new env and rerun your command.

if running into issues with numba/numpy because of the numba dependency of robocasa, run:

conda install -c numba numba=0.56.4 -y