README.md

---
license: other
tags:
  - world-model
  - robot-manipulation
  - multi-task
---

# 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

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