Wancong Zhang, Basile Terver, Artem Zholus, Soham Chitnis, Harsh Sutaria,
Mido Assran, Amir Bar, Randall Balestriero, Adrien Bardes, Yann LeCun*, Nicolas Ballas*
- Implements Hierarchical Planning with Latent World Models (HWM)
- Demonstrates long-horizon planning in Diverse Maze (PLDM)
- Achieves higher success and lower planning cost vs flat planners
Disclaimer: While HWM is evaluated across multiple world models (VJEPA2, DINO-WM, and PLDM), this repository provides a minimal implementation on PLDM (Diverse Maze). For full results across additional world models and tasks, see the project page and paper.
Figure 1a: Hierarchical planning in latent space. A high-level planner optimizes macro-actions using a long-horizon world model to reach the goal; the first predicted latent state serves as a subgoal for a low-level planner, which optimizes primitive actions with a short-horizon world model.
Figure 1b: Hierarchical planning improves success on non-greedy, long-horizon tasks across multiple latent world models.
Tested on python 3.9, CUDA 13.0
git clone git@github.com:kevinghst/HWM_PLDM.git
cd HWM_PLDM
conda create -n pldm python=3.9 -y
conda activate pldm
pip install -r requirements.txt
pip install -e .
mkdir -p "$HOME/.mujoco" cd "$HOME/.mujoco" wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz tar -xzf mujoco210-linux-x86_64.tar.gz
export MUJOCO_GL=egl export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco210/bin" export D4RL_SUPPRESS_IMPORT_ERROR=1
- Go to
pldm_envs/, follow instructions to set up dataset for the environment of your choice - Go to
pldm/, follow instruction to run training or evaluation
To see the datasets we used to train our models, see folders inside pldm_envs/. The readmes there will guide you on how to download and set up the datasets
