ByDM/salt-tworoom

SALT: Two-Room

State-Affine Latent Transition (SALT) world model trained on Two-Room.

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): 98.0 +- 2.0

Usage

Install the SALT code, then place the files under $STABLEWM_HOME:

hf download ByDM/salt-tworoom --local-dir $STABLEWM_HOME/checkpoints/salt/tworoom

and evaluate with checkpoint=salt/tworoom/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}
}