mtdit_croissant_jointbase_h48

Multi-Task DiT policy trained with LeRobot on l5vel-peng/base4-plate-croissant-eef-merged-v30 (98 episodes, ~155k frames, 50 fps).

Final checkpoint: step 100,000 (train loss 0.001, held-out eval loss 0.0141).

Action space

Joint + base, 10-D — [joint1..joint6, gripper, base_x, base_y, base_yaw] under the keys observation.state / action. The dataset also carries 13-D EEF variants (observation.eef_state / action.eef); those columns ride along in the metadata but are not fed to the model, since LeRobot keys the policy off the exact names observation.state and action.

Observations: three 480×640 RGB cameras (left, right, wrist), randomly cropped to 224×224.

Configuration

objective	diffusion (DDPM, 100 train timesteps)
horizon	48 (0.96 s @ 50 Hz)
n_action_steps	32 (0.64 s executed per call)
n_obs_steps	2
hidden_dim / layers / heads	512 / 6 / 8
vision + text encoder	openai/clip-vit-base-patch16
optimizer lr	1e-4, 2000 warmup steps
batch size	32 per GPU × 7 GPUs (DDP) = 224 effective
steps	100,000 (~151 epochs)
precision	bf16
seed	1000
eval split	0.05 (~5 held-out episodes)

Usage

from lerobot.policies.multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy

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

Training

Reproduced with train_mtdit_croissant.sh in the training repo:

torchrun --nproc_per_node=7 -m lerobot.scripts.lerobot_train \
  --dataset.repo_id=l5vel-peng/base4-plate-croissant-eef-merged-v30 \
  --dataset.eval_split=0.05 \
  --policy.type=multi_task_dit --policy.objective=diffusion \
  --policy.horizon=48 --policy.n_action_steps=32 --policy.n_obs_steps=2 \
  --policy.optimizer_lr=1e-4 --policy.scheduler_warmup_steps=2000 \
  --batch_size=32 --steps=100000 --accelerator.mixed_precision=bf16 --seed=1000