Bigenlight/flow_matching_banana_in_pot_joint_bf16

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.