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")