metadata
license: mit
library_name: pytorch
tags:
- world-model
- jepa
- visual-planning
- model-based-planning
- robotics
SALT: OGBench-Cube
State-Affine Latent Transition (SALT) world model trained on OGBench-Cube.
Paper: Beyond One-Step Accuracy: State-Affine Latent Transition for Reliable Visual Planning Code: TODO: add the GitHub URL
SALT keeps the LeWM encoder and replaces the Transformer predictor with an action-conditioned state-affine transition
z' = A(c) z + B c + b, A(c) = A_0 + sum_r (W_g c)_r N_r, sigma_max(A_0) < 1.
Files
| File | Content |
|---|---|
weights.pt |
PyTorch state_dict (fp32) |
config.json |
Hydra instantiation config (targets salt.models.*) |
training_config.yaml |
Training configuration and evaluation record |
SHA256SUMS |
Checksums |
Model
- Encoder: ViT-Tiny (patch 14, 224x224 input) with a 192-2048-192 projector
- Predictor: 703,872 parameters, latent dimension 192, R = 16 modulation modes
- Training: 10 epochs, rollout objective with K = 5 recursive steps, SIGReg weight 0.09, initialization seed 0, data split seed 3072
Evaluation
CEM planning with horizon H = 5, receding horizon 5 action blocks (25 environment steps), goal offset 25, budget 50, 50 episodes for each of 3 evaluation seeds.
Success rate of this checkpoint (mean and std over the 3 evaluation seeds): 94.0 +- 2.0
Usage
Install the SALT code, then place the files under $STABLEWM_HOME:
hf download ByDM/salt-cube --local-dir $STABLEWM_HOME/checkpoints/salt/cube
and evaluate with checkpoint=salt/cube/weights.pt (see the repository README).
Citation
@misc{salt2026,
title = {{Beyond One-Step Accuracy: State-Affine Latent Transition for Reliable Visual Planning}},
author = {Zhang, Boyuan and Du, Yingjun and Zhen, Xiantong and Shao, Ling},
year = {2026},
eprint = {2609.33595},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.33595}
}