cawm — IWS latent world-model checkpoints

Pretrained checkpoints for the interactive world simulator (umi-wm / iws-model): an action-conditioned latent world model with a joint vision + tactile branch. The matching data lives in the dataset repo of the same name, aabyaneh/cawm.

Tasks

Task	Data	Stage 1 (ae.ckpt)	Stage 2 (wm.ckpt)
allen_insert	real	epoch 21, step 70,000	epoch 72, step 60,000
erase_whiteboard	real	epoch 26, step 70,000	epoch 90, step 60,000
pen_in_hole	real	epoch 12, step 90,000	epoch 70, step 95,000
usb_insertion	real	epoch 43, step 80,000	epoch 175, step 80,000
rubik_easy	sim	epoch 11, step 70,000	epoch 41, step 60,000
rubik_medium	sim	epoch 11, step 70,000	epoch 44, step 65,000
rubik_hard	sim	epoch 11, step 70,000	epoch 38, step 50,000
rubik	real	epoch 30, step 90,000	epoch 88, step 65,000
rubik_easy_2cam	sim, 2 cameras	epoch 10, step 60,000	— not trained yet

Baselines

Checkpoints for other world models, trained on the same data for comparison. They are not IWS models and do not follow the <task>/checkpoints/ layout above.

Model	Folder	Task	Code
ContactWorld	baselines/contactworld	rubik_easy, wrist + tactile, 100k steps	umi-wm/baselines/contactworld

Layout

<task>/
  README.md              provenance + loading snippet
  checkpoints/ae.ckpt    stage 1 — encoder + decoder
  checkpoints/ae.yaml     └─ its Hydra config
  checkpoints/wm.ckpt    stage 2 — dynamics
  checkpoints/wm.yaml     └─ its Hydra config

Most tasks have all four files. rubik and rubik_easy_2cam currently ship stage 1 only — their wm.ckpt does not exist yet, because stage 2 has not been run for them.

Every *.yaml was verified to rebuild the model and load its checkpoint with strict=True: 1,522 tensors, no missing or unexpected keys, for all 14 checkpoints.

Do not call OmegaConf.resolve() on these configs. algorithm.dtype is ${torch:float}, which resolves to a torch.dtype — a value OmegaConf can read lazily (as Hydra does) but cannot store back into the container. Load and pass cfg.algorithm straight through.

Architecture is identical across tasks

The 15 single-camera checkpoints all have the same shape: action_dim=4, one camera view, latent 4 channels, tactile 1→2 channels, 24x32 grid, 1,522 tensors. Their only genuine per-task differences are dataset.dataset_dir and the run name. This was read out of the weights, not assumed.

rubik_easy_2cam is the one exception and is deliberately not interchangeable with the rest: two camera views put x_shape at [6,128,128] and num_latent_channel at 8, giving 1,528 tensors. Its camera order is load-bearing — obs_keys fixes the channel order of the concatenated image tensor, so it must stay [wrist, agentview] in every stage that loads it.

One inherited wart is corrected here: pen_in_hole reused sim_dexsuite_dataset.yaml, whose shape_meta.action.shape says [8], while the model it trained is 4-dimensional (algorithm.action_dim=4, and the normalizer and action_emd weights are both 4-wide). shape_meta.action is never read by the data path — actions come straight from the HDF5 — so the mismatch was harmless, but the configs published here say [4] throughout.

Training stages

1. Stage 1 (ae.ckpt) — trains the image encoder/decoder and the tactile encoder/decoder.
2. Stage 2 (wm.ckpt) — freezes the encoders and trains the image dynamics plus the tactile dynamics, cross-conditioned through a vision ControlNet. Loads stage 1 via algorithm.load_ae, so wm.ckpt already contains the stage-1 weights.

Every checkpoint is a full Lightning checkpoint for LatentWorldModel (1,522 tensors): encoder, decoder, dynamics, encoder_tactile, decoder_tactile, dynamics_tactile, vision_controlnet, both noise schedulers and the normalizer. All have tactile enabled with a 24x32 force grid.

Note on the sim tasks

rubik_{easy,medium,hard} supersede the earlier cube_play_* checkpoints, which have been removed from this repo. Use the rubik ones.

rubik (no suffix) is the real-robot Rubik's cube task and is unrelated to the simulated rubik_{easy,medium,hard}. rubik_easy_2cam is the two-camera variant of rubik_easy — same episodes, plus the front view recovered from the dataset repo's LeRobot copy. That merged dataset is not stored as HDF5 on the hub; rebuild it with scripts/merge_agentview_from_lerobot.py in the umi-wm repo.