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:
| file | purpose |
|---|---|
lewm_weights.ckpt |
Lightning training checkpoint |
lewm_epoch_*_object.ckpt |
Object checkpoint used by native CEM eval |
config.yaml |
Training configuration snapshot |
hydra/.hydra/* |
Hydra-resolved training config |
eval/native_results/*.json |
Native PushT CEM evaluation outputs |
Training Setup
Common settings:
| item | value |
|---|---|
| environment | PushT |
| source repo | third_party/le-wm at commit ca231f9f |
| encoder | ViT-Tiny, trained from scratch |
| epoch target | anchor epoch 10 |
| effective batch | 128 |
| precision | FP32 |
| hardware | RLL02, RTX 2080 Ti DDP |
Predictor settings:
| family | predictor | trainable params | batch |
|---|---|---|---|
vanilla_vits |
ViT-S-style, hidden dim 384, depth 12, heads 6 | about 42M total | 32 x DDP4 |
vanilla_vitb |
ViT-B-style, hidden dim 768, depth 12, heads 12 | about 140.6M total | 32 x DDP4 |
Frame-skip step budgets:
| fs | max steps |
|---|---|
| 5 | 154,822 |
| 10 | 125,627 |
| 20 | 68,997 |
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:
| item | value |
|---|---|
| eval episodes | 50 per eval seed |
| eval seeds | 42, 123, 7 |
| goal offset | 25 |
| horizon | 5 |
| action block | frame skip |
| receding horizon | 5 |
| eval budget | 50 for fs=5 and fs=10, 100 for fs=20 |
| CEM samples | 300 |
| CEM iterations | 30 |
| CEM top-k | 30 |
| success criterion | PushT pose tolerance, position < 20 and angle < pi/9 |
Paper-train results:
| predictor | fs | seed 42 | seed 123 | seed 7 | mean ± std |
|---|---|---|---|---|---|
| ViT-S | 5 | 96.0% | 92.0% | 96.0% | 94.67 ± 2.31% |
| ViT-S | 10 | 72.0% | 76.0% | 64.0% | 70.67 ± 6.11% |
| ViT-S | 20 | 50.0% | 48.0% | 54.0% | 50.67 ± 3.06% |
| ViT-B | 5 | 86.0% | 90.0% | 88.0% | 88.00 ± 2.00% |
| ViT-B | 10 | 70.0% | 78.0% | 74.0% | 74.00 ± 4.00% |
| ViT-B | 20 | 48.0% | 38.0% | 40.0% | 42.00 ± 5.29% |
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:
| predictor | fs | success rate |
|---|---|---|
| ViT-S | 5 | 80.0% |
| ViT-S | 10 | 62.0% |
| ViT-S | 20 | 40.0% |
| ViT-B | 5 | 80.0% |
| ViT-B | 10 | 70.0% |
| ViT-B | 20 | 42.0% |
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.