maskjp/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

Metric	Value
Validation loss @ 100,000	0.0019
Arm camera sensitivity	0.272
Gripper camera sensitivity	0.930
Arm prediction error (held-out)	0.102
Null test (own cameras)	0.000

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.

Arm	Vision path	State dropout	Val loss	Arm sensitivity	Gripper sensitivity
cls_nodrop	cls	0.00	0.0017	0.077	0.706
spatial_nodrop	spatial_softmax	0.00	0.0017	0.089	0.693
patch_tokens	patch_tokens	0.00	0.0017	0.075	0.567
spatial_mask	spatial_softmax	0.15	0.0019	0.272	0.930
cls_mask	cls	0.15	0.0020	0.198	0.922

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

Dataset	l5vel-peng/base4-multitask-eef-merged-v30 — 7 tasks, 949 episodes, 1,524,905 frames, 50 Hz
Held out	eval_split=0.05, applied per task
Action space	absolute joints, 10-dim
Steps	100,000, batch 64, seed 1000
Architecture	4 layers, hidden 512, 8 heads, dropout 0.1, RoPE
Objective	DDPM, 100 train timesteps, horizon 32, 24 executed steps, 2 observation steps
Encoders	openai/clip-vit-base-patch16 for vision and text, vision LR × 0.1
Images	3 cameras, 480×640 → resize 240×320 → random crop 224×224
Optimizer	3e-4, weight decay 0, no warmup, cosine decay
Normalization	visual MEAN_STD, state MIN_MAX, action MIN_MAX

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.