lewm-pusht

LeWM PushT Checkpoint Collection

This repository contains reproduced LeWM checkpoints for PushT across frame skips and predictor capacities.

These are reproduction checkpoints, not official checkpoints from the LeWM authors.

Contents

Checkpoints are organized by predictor family and frame skip:

vanilla_vits/fs5/
vanilla_vits/fs10/
vanilla_vits/fs20/
vanilla_vitb/fs5/
vanilla_vitb/fs10/
vanilla_vitb/fs20/

Each subdirectory contains:

Training Setup

Common settings:

Predictor settings:

Frame-skip step budgets:

Native PushT Planning Evaluation

The reported success rates use the LeWM paper-style native PushT planning protocol: sample start and goal states from the offline pusht_expert_train trajectories, then run CEM-MPC in the PushT simulator.

Native planning config:

Paper-train results:

The ViT-S fs=5 reproduction matches the LeWM paper's reported PushT result of 96.0 ± 2.83% within the reported band.

Held-out reference, seed 42 on pusht_common_eval:

The held-out numbers are stricter references and are not the paper's native evaluation protocol.

Loading

Download the repo and point AutoCostModel to the desired subfolder:

from pathlib import Path
from huggingface_hub import snapshot_download
from stable_worldmodel.policy import AutoCostModel

repo_dir = Path(snapshot_download("MasonJK99/lewm-pusht"))
model = AutoCostModel(repo_dir / "vanilla_vits" / "fs10")

For Lightning checkpoint inspection:

from huggingface_hub import hf_hub_download
import torch

ckpt_path = hf_hub_download(
    "MasonJK99/lewm-pusht",
    "vanilla_vitb/fs10/lewm_weights.ckpt",
)
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

Notes

• The native planner requires the LeWM/stable-worldmodel environment used for training and evaluation.
• Evaluation JSONs are included so the reported success rates can be audited.
• Dataset files are not included in this repository.
• The LeWM source code is MIT licensed; dataset usage is subject to the PushT / stable-worldmodel dataset terms.