Multi-task latent world model — checkpoints

Five arms of a latent world model trained on 281 robot manipulation datasets pooled into a single task-tagged mixture. All arms share the same encoder and predictor; they differ only in the action head and its training objective.

Every arm finished cleanly (rc=0). Each checkpoint was verified after training: 558 tensors, no empty or non-finite values, 498.88M parameters.

Files

Every checkpoint here is post-trained, not trained from scratch. The epoch number in each filename counts only that stage; cumulative training is much larger. Read the lineage below before reading the epoch counts.

a 237-dataset pretraining run                    22 epochs  (branch point)
  ├── baseline_237          +25  →  cumulative ~47
  └── baseline_281          +25  →  cumulative ~47
        └── flow_281        +7.3 →  cumulative ~54
              ├── nolang_281      +10  →  cumulative ~64
              └── sigregdz_281    +10  →  cumulative ~64

file	arm	tasks	this stage	cumulative	what differs
baseline_237_epoch25.pt	237-dataset baseline	237	25	~47	the dataset pool before the object-centric sets were added
baseline_281_epoch25.pt	281-dataset baseline	281	25	~47	+ 44 object-centric datasets
flow_281_epoch7.pt	flow action head	281	7.3	~54	regression head → rectified-flow head, head re-initialised
nolang_281_epoch10.pt	flow, language removed	281	10	~64	branches from the flow arm, language table zeroed
sigregdz_281_epoch10.pt	flow + anti-collapse on Δz	281	10	~64	branches from the flow arm, adds a regulariser on latent differences

So flow_281_epoch7.pt is not a model trained for 7 epochs — it is a model with roughly 54 epochs behind it, the last 7.3 of which swapped the action head. It was stopped there deliberately: its curve had flattened, and stopping saved about 13 hours of eight-GPU time. The two arms below it branch from that checkpoint and inherit everything except the single thing each ablates.

Data and architecture

281 datasets in Lance format: 130 + 65 + 44 + 26 + 9 + 5 + 2 across seven manipulation suites. Frame skip 5; actions zero-padded to width 120 (24 native dimensions × 5 sub-actions) with a per-task mask, so a padded slot is never a regression target.

Language conditioning is a frozen pretrained embedding table (nn.Embedding.from_pretrained(..., freeze=True)), 281 × 512. It does not train in any arm — which is why zeroing it is a clean input-feature ablation.

Predictor: 12-layer mixture-of-experts, 8 experts, top-2 routing. Effective expert count rose from 4.9 to 7.1 of 8 over training; no collapse.

Measured planning results

Scored on the largest suite, complete at 130/130 for every arm — the only suite complete across all five:

arm	success
237 baseline	22.2%
281 baseline	22.4%
flow head	21.4%
flow − language	19.0%
flow + Δz regulariser	16.0%

All three flow-family arms score below both baselines, and the Δz term is worst. On the mobile-manipulation suite the gap is far larger (−17.8 / −21.6 / −29.2 paired against the 237 baseline over 125 paired tasks, 121 negative), but that suite is only partially covered (41 / 49 / 35 of 65).

Two things to know before reading those numbers

1. The criterion is goal reaching, not task success. The planner receives a demonstration state 25 environment steps ahead and must drive the scene to match that frame — every object within 0.05 m, end-effector within 0.05 m, joints within 0.1 rad. This is far stricter than "the task was completed", and these numbers are not comparable to published task-success rates on the same suites.

2. The flow arms carry two confounds that are not the objective. Their policy head was reset and trained fresh for 7 epochs while the baseline's head was inherited with a long lineage plus 25 more; and at planning time the flow head draws a single stochastic sample (8-step Euler from noise, no K-sample averaging) where the regression head returns a deterministic conditional mean. So the measured deficit belongs to the flow head as deployed, not yet to rectified flow as an objective.

Loading

import torch
sd = torch.load('flow_281_epoch7.pt', map_location='cpu', weights_only=True)
sd = sd.get('state_dict', sd)   # some files wrap it, some do not