maskjp/mtdit_plate_croissant_jointbase_fullframe

MTDiT — plate/croissant, joint + base, full frame

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

This is the full-field-of-view variant: the whole 480×640 frame is resized to 224×224 and no crop is applied, so the policy sees 100% of each camera image.

⚠️ This checkpoint has the worst held-out loss of the ten runs

Published deliberately as the field-of-view counterpart to maskjp/mtdit_plate_croissant_jointbase, not because it scored well. Read both numbers before choosing:

	this model (fullimg)	mtdit_plate_croissant_jointbase (noresize)
sees at inference	100% of frame	centre 16% of frame
crop jitter in training	none	±256 v / ±416 h
best eval_loss	0.0081 @ 5K	0.0037 @ 45K
final eval_loss @ 50K	0.0242	0.0038
trend	overfit 3.0×	still improving

The two are a deliberate trade. This model can see objects anywhere in frame but generalizes far worse; the other generalizes far better but is blind outside the centre crop. Denoising loss is a weak proxy for rollout success in diffusion policies, so which one actually works on the robot is an open question — that is the point of publishing both.

If the plate or croissant sits near the edge of frame in your setup, this model is the one that can see it.

Why it overfit

With CPU augmentation disabled, random-cropping out of the full frame is the only regularizer in the pipeline. Removing the crop removes it entirely. Across four image pipelines, held-out loss ordered monotonically with crop jitter:

pipeline	jitter	sees	final eval_loss
no resize, crop 224	±256 v / ±416 h	16%	0.0038
resize [336,448], crop 224	±112 v / ±224 h	33%	0.0087
resize [240,320], crop 224	±16 v / ±96 h	65%	0.0121
resize [224,224], no crop	none	100%	0.0242

With only 93 training episodes, that regularizer matters a great deal.

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)

observation.eef_state (13) and action.eef (13) exist in the config but are not used — LeRobot binds state/action to the exact keys observation.state / action. Language conditioning uses a frozen CLIP ViT-B/16 text encoder with a trained linear projection.

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	[224, 224]
image_crop_shape	None (no crop)
optimizer	AdamW, lr 3e-4, weight_decay 0, no warmup
vision_encoder_lr_multiplier	0.1
steps / batch_size	50,000 / 64 (single H100)
image augmentation	disabled
seed	1000

443.0M params total, 379.9M trainable. Held-out split: 5 eval / 93 train episodes, split at episode level stratified by task.

Note the aspect distortion: 480×640 is 3:4, squashed to 1:1. It is applied identically in training and inference, so the policy simply learns that geometry.

Usage

from lerobot.policies.multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy

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