Flow-Matching Policy — banana-in-pot (JOINT, bf16)
Flow-Matching policy (multi_task_dit, objective=flow_matching) trained on the
"put the right banana in the pot" task (UR7e + GELLO teleoperation, 2 RGB cameras),
in JOINT action space (6 joints + gripper), using bf16 mixed-precision training.
- Checkpoint: step 60,000 (best open-loop MAE)
- Base library: LeRobot 0.6.1 (pin
8a74e0a)
- Dataset:
Bigenlight/banana_in_pot_lerobot_v3 — 51 episodes / 21,524 frames / 30 fps
- fp16 sibling (Diffusion):
Bigenlight/diffusion_banana_in_pot_joint_fp16
Architecture
CLIP ViT-B/16 vision-language backbone → DiT (diffusion transformer) velocity field,
trained with the flow-matching objective (Euler integration at inference). ~186M
learnable / ~249M total params. Images resized/cropped to 224×224.
Training
- Precision: bf16 via HF Accelerate
mixed_precision=bf16 (no GradScaler needed;
bf16 preferred over fp16 for the CLIP+DiT stack for numerical headroom).
- Requires a
dtype field on MultiTaskDiTConfig (absent upstream at this pin); launched
with --policy.dtype=bfloat16.
- Batch 8, 80k steps, AdamW, seed 1000, 45 train / 6 held-out episodes.
- Hardware: single RTX A4000. ~4.14 step/s, wall-clock 5:22:07. No NaN/instability.
Open-loop evaluation (Euler-10, held-out episodes 45–50)
| step | poseMAE (rad) | gripAcc | overallL1 |
|---|
| 20k | 0.08048 | 0.954 | 0.07629 |
| 40k | 0.07713 | 0.959 | 0.07273 |
| 60k ⭐ | 0.07605 | 0.961 | 0.07135 |
| 80k | 0.07648 | 0.959 | 0.07185 |
fp32 FM baseline: poseMAE 0.0735 @70k. bf16 lands at 0.07605 (60k) with slightly
higher gripper accuracy (0.961 vs fp32) — within run-to-run noise, no quality regression,
and bf16 removes fp16's overflow risk on the CLIP+DiT stack while cutting VRAM/wall-clock.
Select the deploy checkpoint by open-loop MAE, not eval_loss (which rises during
training for generative policies — here 0.0722@5k → 0.1707@80k — without indicating
overfitting).
Intended use & limitations
Research artifact. Small single-task, single-scene, real-world (noisy) dataset of 51
success-only demonstrations; offline metrics only — no closed-loop hardware success rate
measured yet. Not safety-validated for autonomous operation.
