mtdit-spatial-multitask-spatial-mask

MTDiT-Spatial — multi-task, spatial_mask

A 450M-parameter multi-task diffusion transformer trained on 949 real-robot episodes. Vision path spatial_softmax, state dropout 0.15, absolute joint actions.

Read this before using the weights

These checkpoints are published as a negative result. They fit their data well and largely ignore their scene cameras, predicting joint targets from proprioception instead. Arm camera sensitivity is 0.272 against a null-test floor of 0.000 and a grounded reference of ~1.0.

Do not deploy this as a vision-conditioned policy. It is useful as a reproduction target, a baseline for work on visual grounding, or a test case for sensitivity metrics.

Measurements

Camera sensitivity holds the proprioceptive state fixed, swaps all three camera streams for matched-phase frames from a different held-out episode, and reports the mean change in predicted joints over the natural cross-episode spread of the ground truth. A policy that tracks the scene scores ~1.0; one that ignores it scores ~0.

Two controls make the number trustworthy. Every variant is sampled from an identical noise sequence (common random numbers), without which diffusion sampling noise alone contributes 0.077. And the null test — feed the anchor its own cameras — returns 0.000 here, so the reported value is signal, not floor.

For reference under the same evaluator: stock multi_task_dit on LIBERO scores 0.92–1.10, and π₀.₅ on LIBERO scores 1.19–1.31.

The sweep

Ranking by loss is the inverse of ranking by grounding. The best-fitting arms are the blindest, and the two arms that withhold proprioception 15% of the time recover 2.5–3.5× the camera sensitivity for a ~12% loss penalty. That inversion is the result, so all five arms are published together.

All five arms reached 100,000 steps. The three unmasked arms converge to an identical 0.0017 and the two masked ones to 0.0019-0.0020; the floor was already reached at 30K, so 3.3x more compute moved it by 0.0001.

Training

The 30K checkpoint was extended to 100K by resuming, which re-stretched the cosine schedule; loss bumps at 35K and recovers by ~65K. The floor was already reached at 30K — 3.3× more compute moved it by 0.0001.

Usage

from lerobot.policies.factory import make_policy

policy = make_policy("maskjp/mtdit-spatial-multitask-spatial-mask")

Requires the mtdit_spatial plugin from maskjp/lerobot_policy_mtdit_spatial, which also contains scripts/vision_sensitivity.py used for every number above.

Known limitations

• Ignores the scene cameras for arm motion; the left and right cameras contribute least.
• Trained on absolute joint targets, which correlate 0.96 with the current state — the shortcut this failure runs on. Datasets with per-step deltas (LIBERO, −0.14) do not show it.
• Single robot, single embodiment, 3 fixed camera poses. No environment-rollout success rate is reported; only offline metrics.