LeWorldModel++ (LeWM++)
LeWorldModel++ (LeWM++) is a planning-interface framework for long-horizon image-goal control with frozen latent world models, such as LeWorldModel and DINO-WM. Short model rollouts can struggle with distant goals because the planning target is too far away, action search has no useful initialization, and a terminal-only cost can miss a good intermediate approach.
LeWM++ combines three components:
- LatentPathFlow generates reachable local targets while retaining the final image goal.
- Action Chunk Prior initializes CEM from goal-directed offline behavior.
- Min-over-Horizon (MoH) scores the closest predicted approach to the current target.
The controller executes one optimized action chunk and replans within a locally reliable prediction horizon. See the code repository and full README for the implementation, training instructions, and evaluation launchers.
What this repository contains
This repository stores the pretrained checkpoints selected for the release evaluation. It has two artifact groups:
| Group | Tasks or datasets | Checkpoint components |
|---|---|---|
lewm-control-suite/ |
cube, pusht, reacher, tworoom |
LeWM, Action Chunk Prior, LatentPathFlow (h25 and longh) |
visual-ogbench/ |
cs_play, cd_play, ct_play, scene_play, cs_noisy, cd_noisy, ct_noisy, scene_noisy |
LeWM, Action Chunk Prior, LatentPathFlow |
The directory structure is:
lewm-control-suite/checkpoints/
├── lewm/{cube,pusht,reacher,tworoom}/
├── action-prior/{cube,pusht,reacher,tworoom}/
└── latent-path-flow/{h25,longh}/{cube,pusht,reacher,tworoom}/
visual-ogbench/checkpoints/
├── lewm/<dataset-tag>/
├── action-prior/<dataset-tag>/
└── latent-path-flow/<dataset-tag>/
LeWM weights are named weights_epoch_10.msgpack; LatentPathFlow weights are checkpoint_200000.msgpack. Action Prior weights are params_100000.pkl for the Control Suite and params_500000.pkl for Visual OGBench. Keep each config.json or flags.json beside its corresponding weights; these files are needed to restore the models.
Datasets and real-robot checkpoints are not included here. The DINO-WM transfer experiments described in the paper use the original DINO-WM implementation and are not included in the LeWM++ code repository.
Download and install
Clone the code and install its training and evaluation dependencies:
git clone https://github.com/pickxiguapi/LeWorldModel-pp.git
cd LeWorldModel-pp
uv sync --extra train --extra dev
source .venv/bin/activate
Download the release checkpoints into artifacts/ at the repository root:
uvx --from huggingface_hub hf download IffYuan/LeWorldModelplusplus \
--include "*/checkpoints/**" --local-dir artifacts
The download preserves the two directory trees above under artifacts/. Evaluation with these pretrained checkpoints does not require retraining or latent-cache precomputation.
Datasets
The datasets are distributed separately. For the LeWM Control Suite, download the four official HDF5 archives from the LeWM Hugging Face collection:
| Dataset | Archive |
|---|---|
quentinll/lewm-cube |
cube_single_expert.tar.zst |
quentinll/lewm-pusht |
pusht_expert_train.h5.zst |
quentinll/lewm-reacher |
reacher.tar.zst |
quentinll/lewm-tworooms |
tworoom.tar.zst |
Extract the four HDF5 files and create JPEG-backed Lance tables with scripts/convert_lewm_hdf5_to_lance.py. The project README gives the expected filenames and conversion command.
For Visual OGBench, provide the eight training .npz datasets and matching -val.npz files in a local data directory. See the Visual OGBench preparation instructions for naming and training details. The training .npz files are also used to recover action normalization during evaluation.
Evaluation
All commands below run from the cloned code repository. Edit the path assignments inside each launcher to point to your local datasets, artifacts/ checkpoint roots, output directory, and GPU IDs.
LeWM Control Suite
Use the LeWM++ launchers for horizons 25, 50, 75, and 100. Set these common roots in each eval_lewmpp_h{25,50,75,100}_4tasks.sh:
LEWM_DATA_ROOT="/absolute/path/to/lewm-control-suite"
EXPERIMENT_ROOT="outputs"
LEWM_CHECKPOINT_ROOT="artifacts/lewm-control-suite/checkpoints/lewm"
ACTION_PRIOR_CHECKPOINT_ROOT="artifacts/lewm-control-suite/checkpoints/action-prior"
For H25, set LATENT_PATH_FLOW_CHECKPOINT_ROOT to artifacts/lewm-control-suite/checkpoints/latent-path-flow/h25. For H50, H75, and H100, use artifacts/lewm-control-suite/checkpoints/latent-path-flow/longh. These roots contain the four task directories; do not append a task name or checkpoint filename. The baseline launchers need only the data, output, and LeWM checkpoint roots.
for horizon in 25 50 75 100; do
bash "experiments/eval/eval_lewmpp_h${horizon}_4tasks.sh"
bash "experiments/eval/eval_lewm_baseline_h${horizon}_4tasks.sh"
done
python impls/aggregate_lewm_control_results.py \
--results-root outputs/eval \
--output outputs/eval/lewm_control_suite_summary.csv
The release protocol evaluates 50 episodes per task. The LeWM baseline uses CEM300x30 with H5/RH5 and an action block of 5; LeWM++ uses CEM300x5 with H2/RH1 and the same action block.
Visual OGBench
Set these roots in experiments/eval/eval_lewmpp_visual_ogbench8.sh:
OGBENCH_DATA_ROOT="/absolute/path/to/visual-ogbench-data"
EXPERIMENT_ROOT="outputs"
LEWM_CHECKPOINT_ROOT="artifacts/visual-ogbench/checkpoints/lewm"
ACTION_PRIOR_CHECKPOINT_ROOT="artifacts/visual-ogbench/checkpoints/action-prior"
LATENT_PATH_FLOW_CHECKPOINT_ROOT="artifacts/visual-ogbench/checkpoints/latent-path-flow"
The baseline launcher needs only the data, output, and LeWM checkpoint roots. Run both launchers after editing their paths:
bash experiments/eval/eval_lewmpp_visual_ogbench8.sh
bash experiments/eval/eval_lewm_baseline_visual_ogbench8.sh
The release protocol uses 50 episodes per official task and three evaluation seeds. The launchers write aggregate summaries; the project README also shows how to regenerate the LeWM++ summary from completed results.
Training and real-robot use
The training section describes how to train LeWM, encode offline data, train the Action Chunk Prior, and train LatentPathFlow in dependency order. It includes launchers for the four Control Suite tasks, eight Visual OGBench datasets, and the example LeRobotDataset v3.0 yaoxianze/push_multi_red_cube.
The real-robot reference adapter loads matching model artifacts and training action statistics. Deployment must reproduce the training camera, image, temporal, and action conventions; the robot-side program must provide its own collision checking, emergency stops, limits, and watchdogs. See the real-robot inference instructions before connecting the policy to a robot.
Acknowledgments and license
LeWM++ builds on LeWorldModel and OGBench. Thanks to their authors for releasing the latent world-model implementation, benchmark environments, datasets, and evaluation APIs. The code repository retains the MIT license. A LeWM++ citation block will be added when the paper receives a public identifier.