lewm-pusht-pymunk

A LeWM: ViT-tiny encoder (224 px, patch 14, CLS embedding of 192) and a 6-layer action-conditioned predictor, trained end to end from pixels with next-embedding MSE and SIGReg (weight 0.09). Trained on le-wm's Push-T expert dataset (stable-worldmodel's pusht_expert_train.h5, 18,685 episodes) as released, rendered by the reference pymunk env, at 224 px, with windows of 4 frames 5 steps apart.

• Model: torch_world.models.LeWM of torch-world: LeWM(CLSEncoder(img_size=224, patch_size=14, embed_dim=192, depth=12, num_heads=3), action_dim=10, history=3).
• Inputs: frames normalized with the ImageNet mean and std.
• Actions: blocks of 5 Push-T actions (10 values), z-scored with normalization.action of config.json, in the reference's pixel axes (y down); torch_world's PushT takes y up, so flip y.
• Recipe: AdamW (lr 5e-5, weight decay 1e-3), 5% warmup then cosine decay, batch 128, bf16, gradient clipping at 1, 10 epochs.
• Paper: Maes et al., LeWorldModel, 2026 (https://arxiv.org/abs/2603.19312).
• Intermediate checkpoints: history/ (epoch_01, epoch_02, epoch_03, epoch_04, epoch_05, epoch_06, epoch_07, epoch_08, epoch_09).

Usage

from torch_world import hub

model = hub.from_pretrained("latenthorizon/lewm-pusht-pymunk")  # pip install "torch-world[hub]"

config.json holds the model's spec (what from_pretrained builds) and the normalization of actions.

Results

Push-T goal reaching under stable-worldmodel's protocol: goal 25 expert steps ahead, CEM (300 samples, 30 elites, 30 iterations) over 5 blocks of 5 actions, whole plan executed before replanning, 50 steps, success when the goal is reached at any step.

evaluation	env	episodes	success
stable-worldmodel eval_wm.py (its CEM)	stable-worldmodel swm/PushT-v1 (pymunk)	500	83.8%
same; official checkpoint quentinll/lewm-pusht, for reference	stable-worldmodel swm/PushT-v1 (pymunk)	500	84.2%