README.md

---
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`:

```bash
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

```bibtex
@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}
}
```