marsmute/pusht-leworldmodel-horizon3

PushT LeWorldModel — three-action horizon

This repository contains a locally trained LeWorldModel checkpoint for PushT. It predicts the latent visual state after a group of three raw environment actions. The release is intended for world-model research, planning experiments, and comparisons of temporal abstraction in robotic manipulation.

This is a predictive world model, not a standalone control policy. It does not select actions, interact with PushT by itself, or decode its prediction back into an image. A planner such as CEM/MPC must propose and score actions around the model.

Model description

The checkpoint contains four learned components:

1. A 12-layer Vision Transformer that encodes each 224 × 224 RGB observation.
2. An action encoder that converts three 2-D PushT actions into one action token.
3. A six-layer autoregressive conditional transformer that predicts future latent states from visual history and grouped actions.
4. Input and prediction projectors mapping to a 192-dimensional latent space.

The predictor uses three observation frames. Consecutive observations are separated by three raw environment actions. A complete input contains two historical action groups connecting the observed frames and one proposed future action group.

Property	Value
Environment	PushT
History observations	3
Raw actions per transition	3
Raw action dimension	2
Image resolution	224 × 224 RGB
Latent dimension	192
Predictor layers	6
Predictor attention heads	16
Parameters	18,034,438
Training episodes	2,000
Training epochs	10
Training seed	3072
Training precision	BF16

The complete architecture and preprocessing specification is recorded in config.json.

Inputs and outputs

Images

Images have shape [batch, 3, 3, 224, 224]: batch, history, RGB channels, height, and width. Images are converted to RGB, resized using bicubic interpolation, rescaled to [0, 1], and normalized with ImageNet mean and standard deviation.

Actions

Raw actions have shape [batch, 3, 3, 2]: batch, three temporal groups, three raw actions per group, and the 2-D PushT action. Actions are standardized using the values in action_stats.json and flattened to grouped tokens with shape [batch, 3, 6].

Prediction

predict_next returns one tensor with shape [batch, 192]. It represents the predicted latent visual state after the final proposed group of three actions. The model has no pixel decoder.

Installation

Python 3.10 or newer is recommended.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements-mac.txt

Despite the filename, requirements-mac.txt contains ordinary Python dependencies and inference also works on CPU. The loader selects Apple MPS when available and otherwise falls back to CPU.

Loading the model

from model import load_model

model = load_model("pusht_lewm_fs3_epoch10.safetensors")

The loader reads config.json, reconstructs the architecture, loads only tensor data from SafeTensors, switches to evaluation mode, and disables gradients. Pass normalized images to model.encode_images, then pass the resulting history and normalized grouped actions to model.predict_next. The exact tensor shapes and normalization values are documented above and in config.json and action_stats.json.

Validation

The release checkpoint was converted from the epoch-10 training checkpoint and validated against a fixed synthetic input:

• SafeTensors output exactly matches the legacy object checkpoint in FP32.
• The CPU golden-output regression passes.
• Apple MPS agrees with CPU with a maximum absolute error of 1.67e-06 on the fixed regression input.

SafeTensors SHA-256:

eb51b5ba8e4e9102eadda739590056790f14aee2aca3096db8e45dd703ecde5a

The legacy pickled checkpoint is not required for normal inference and is intentionally excluded from this release.

Planning results and scope

The model was also used as the low-level world model in a separately trained Hi-LeWM constrained hierarchical planning system. Under each system's natural configuration, the following aggregate PushT results were observed across three seeds with 50 episodes per seed:

Goal distance	Released five-action system	Three-action system
75 steps	23.3%	34.0%
100 steps	9.3%	20.7%

These are system-level results, not scores produced by this checkpoint alone. They additionally require a separately trained high-level checkpoint, the Hi-LeWM constrained planner, a PushT environment, and their associated planning configurations. The three-action system also replans more frequently and performs more planning work under its natural configuration. These results therefore do not establish that a three-action horizon is universally better.

Intended uses

• Studying temporal granularity in learned world models.
• Building CEM, MPC, or hierarchical planners for PushT.
• Profiling and optimizing small visual world models on CPU, Apple MPS, or GPU.
• Reproducing the low-level component of the three-action Hi-LeWM experiments.
• Fine-tuning or comparing predictors using the same observation/action contract.

Limitations

• The model was trained only for PushT and should not be expected to transfer directly to another robot or environment.
• It predicts a deterministic latent vector and does not represent multiple possible futures explicitly.
• It has no image decoder, reward model, value model, action selector, or success classifier.
• Prediction errors can compound when it is unrolled repeatedly.
• The model card does not claim that three-action transitions are generally superior to five-action transitions.
• The hierarchical benchmark results require components that are not included in this low-level checkpoint release.

Training and provenance

• Hi-LeWM source commit: 4bb21a2888e8f22b8d084762c80361e398968775
• LeWorldModel source commit: 83f97d72ad067855bc89a1b74b4aff11d4dfdf0c
• Training configuration: 2,000 PushT episodes, 10 epochs, frame skip 3, three-frame history, seed 3072, BF16.

The inference implementation is a dependency-reduced copy of the relevant LeWorldModel path vendored by Hi-LeWM. The source revisions above identify the exact upstream implementations.

Related work

• LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
• Mind the Gap: Promises and Pitfalls of Hierarchical Planning in LeWorldModel
• Hierarchical Planning with Latent World Models
• DINO-WM: World Models on Pre-trained Visual Features Enable Zero-shot Planning
• Implicit Behavioral Cloning

License

The released model weights and accompanying code are available under the MIT License. The included implementation derives from MIT-licensed LeWorldModel code; the upstream and model-release copyright notices are preserved in LICENSE.

Contact and citation

This is an experimental checkpoint. Until an accompanying paper or archived release receives a permanent citation, cite the upstream LeWorldModel and Hi-LeWM papers and link directly to this model repository with the revision or commit used in your experiment.