---
license: mit
language: en
tags:
  - robotics
  - world-model
  - planning
  - model-based-control
  - jax
---

# 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:

1. **LatentPathFlow** generates reachable local targets while retaining the final image goal.
2. **Action Chunk Prior** initializes CEM from goal-directed offline behavior.
3. **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](https://github.com/pickxiguapi/LeWorldModel-pp) 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:

```text
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](https://github.com/gaoyuezhou/dino_wm) and are not included in the [LeWM++ code repository](https://github.com/pickxiguapi/LeWorldModel-pp).

## Download and install

Clone the code and install its training and evaluation dependencies:

```bash
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:

```bash
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](https://huggingface.co/collections/quentinll/lewm):

| Dataset | Archive |
| --- | --- |
| [`quentinll/lewm-cube`](https://huggingface.co/datasets/quentinll/lewm-cube) | `cube_single_expert.tar.zst` |
| [`quentinll/lewm-pusht`](https://huggingface.co/datasets/quentinll/lewm-pusht) | `pusht_expert_train.h5.zst` |
| [`quentinll/lewm-reacher`](https://huggingface.co/datasets/quentinll/lewm-reacher) | `reacher.tar.zst` |
| [`quentinll/lewm-tworooms`](https://huggingface.co/datasets/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`](https://github.com/pickxiguapi/LeWorldModel-pp/blob/main/scripts/convert_lewm_hdf5_to_lance.py). The [project README](https://github.com/pickxiguapi/LeWorldModel-pp#lewm-control-suite) 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](https://github.com/pickxiguapi/LeWorldModel-pp#visual-ogbench) 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](https://github.com/pickxiguapi/LeWorldModel-pp/tree/main/experiments/eval) for horizons 25, 50, 75, and 100. Set these common roots in each `eval_lewmpp_h{25,50,75,100}_4tasks.sh`:

```bash
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.

```bash
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`:

```bash
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
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](https://github.com/pickxiguapi/LeWorldModel-pp#more-challenging-tasks-on-visual-ogbench) also shows how to regenerate the LeWM++ summary from completed results.

## Training and real-robot use

The [training section](https://github.com/pickxiguapi/LeWorldModel-pp#training) 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`](https://huggingface.co/datasets/yaoxianze/push_multi_red_cube).

The [real-robot reference adapter](https://github.com/pickxiguapi/LeWorldModel-pp/blob/main/impls/real_robot_lewmpp.py) 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](https://github.com/pickxiguapi/LeWorldModel-pp#real-robot-inference) before connecting the policy to a robot.

## Acknowledgments and license

LeWM++ builds on [LeWorldModel](https://github.com/lucas-maes/le-wm) and [OGBench](https://github.com/seohongpark/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](https://github.com/pickxiguapi/LeWorldModel-pp/blob/main/LICENSE). A LeWM++ citation block will be added when the paper receives a public identifier.