README.md

---
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.