maskjp/mtdit_plate_croissant_jointbase

MTDiT — plate/croissant, joint + base control

Multi-Task Diffusion Transformer trained on l5vel-peng/base4-plate-croissant-eef-merged-v30 for 10-D joint + base control.

Best arm of a 7-way hyperparameter sweep, selected on held-out episode loss.

⚠️ Read before deploying

image_resize_shape is None, so the policy center-crops 224×224 out of each 480×640 camera frame at inference — it sees roughly the centre 16% of each image. Anything outside that centre window is invisible to the policy.

Make sure the cameras are framed the same way as during data collection. If the object drifts to the edge of frame, the policy is blind to it, and this will look like a policy failure when it is actually a framing problem.

Companion model: maskjp/mtdit_plate_croissant_jointbase_fullframe is the same architecture trained to see 100% of the frame (resize to 224x224, no crop). It has a much worse held-out loss (0.0242 vs 0.0038) because dropping the crop removes the only regulariser in the pipeline -- but it can see objects the centre crop cannot. Denoising loss is a weak proxy for rollout success, so the two are worth comparing on the real robot rather than on loss alone.

(During training the crop is random across the full frame, which is what makes this the only sweep arm that did not overfit — see below.)

Inputs / outputs

	key	shape
in	observation.state	10 (joint + base)
in	observation.images.left	3×480×640
in	observation.images.right	3×480×640
in	observation.images.wrist	3×480×640
out	action	10 (joint + base)

The dataset also carries observation.eef_state (13) and action.eef (13). These are present in the config but not used — LeRobot binds the state/action features to the exact keys observation.state / action, so the EEF columns are loaded and ignored.

Language conditioning uses a frozen CLIP ViT-B/16 text encoder; only a linear projection on top of it is trained.

Config

objective	diffusion (DDPM, 100 train timesteps)
hidden_dim / layers / heads	768 / 8 / 12
horizon / n_action_steps / n_obs_steps	48 / 40 / 2
image_resize_shape	None
image_crop_shape	224×224 (random in train, centre at eval)
optimizer	AdamW, lr 3e-4, weight_decay 0, no warmup
vision_encoder_lr_multiplier	0.1
dropout / RoPE	0.1 / enabled
steps / batch_size	50,000 / 64 (single H100)
image augmentation	disabled
seed	1000

Params: 443.0M total, 379.9M trainable.

Results

Held-out split: 5 eval episodes / 93 train episodes, split at the episode level stratified by task (no frame leakage). Eval uses policy.eval(), deterministic centre crop, and no augmentation.

step	5K	15K	25K	35K	45K	50K
eval_loss	0.0079	0.0051	0.0045	0.0042	0.0037	0.0038

This was the only arm of seven whose held-out loss kept falling. Every other configuration bottomed out around 10–20K steps and then rose — e.g. the otherwise identical image_resize_shape=[240,320] arm ended at 0.0121, a 12× train/eval gap. The strong implicit augmentation from random-cropping 224×224 out of the full 480×640 frame appears to be doing the regularisation work, which matters with only 93 training episodes.

Caveat: denoising loss is a weak proxy for rollout success in diffusion policies. This checkpoint is the best available by held-out loss; it has not been validated on a real robot.

Usage

from lerobot.policies.multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy

policy = MultiTaskDiTPolicy.from_pretrained("maskjp/mtdit_plate_croissant_jointbase")