File size: 2,643 Bytes
3325882 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | ---
license: apache-2.0
tags:
- robotics
- lerobot
- flow-matching
- imitation-learning
- manipulation
---
# 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](https://github.com/huggingface/lerobot) 0.6.1 (pin `8a74e0a`)
- **Dataset:** [`Bigenlight/banana_in_pot_lerobot_v3`](https://huggingface.co/datasets/Bigenlight/banana_in_pot_lerobot_v3) — 51 episodes / 21,524 frames / 30 fps
- **fp16 sibling (Diffusion):** [`Bigenlight/diffusion_banana_in_pot_joint_fp16`](https://huggingface.co/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.
|