Modelpublic

maskjp/mtdit-spatial-multitask-patch-tokens

sourceHugging Faceapache-2.0updated 1mo agoView on Hugging Face

0likes8downloads

Model Card

MTDiT-Spatial — multi-task, `patch_tokens`

A 450M-parameter multi-task diffusion transformer trained on 949 real-robot episodes. Vision path patch_tokens, state dropout 0.00, 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.075 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.0017
Arm camera sensitivity	0.075
Gripper camera sensitivity	0.567
Arm prediction error (held-out)	0.095
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 MEANSTD, state MINMAX, 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

python

from lerobot.policies.factory import make_policy

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

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.