Flow-Matching Policy — Put the right banana in the pot (UR7e, JOINT action space)

• Policy: LeRobot multi_task_dit with --policy.objective=flow_matching — a conditional / rectified flow-matching action head (linear-interpolation probability path, velocity target, Euler ODE integration at inference) on a small DiT transformer denoiser (hidden_dim = 512, 6 layers, 8 heads, RoPE), conditioned on CLIP ViT-B/16 vision and text features (shared encoder, text encoder frozen). Receding-horizon action generation: horizon = 32, n_obs_steps = 2, n_action_steps = 24.
• Flow-matching process: linear-interpolation path x_t = t·data + (1 − (1 − σ)t)·noise, velocity target data − (1 − σ)·noise, MSE on the predicted velocity field, sigma_min = 0.0, timestep_sampling = beta. Inference integrates the learned velocity field with Euler ODE steps (num_integration_steps, default 100; evals here used 10).
• Trained on: Bigenlight/banana_in_pot_lerobot_v3 — 51 teleoperated episodes / 21,524 frames, UR7e follower + GELLO leader, 2 RGB cameras.
• This checkpoint: step 70,000 (best by open-loop held-out MAE; see Results).
• Framework: LeRobot v0.6.1 (pin 8a74e0a).

Why this policy? multi_task_dit + flow_matching is the only non-VLA, train-from-scratch flow-matching policy in LeRobot 0.6.1. The other flow-matching / action-expert policies (pi0, pi0_fast, pi05, smolvla, evo1) all require heavy pretrained VLM backbones that do not fit a single RTX 3060 12 GB and/or need a language stack. This policy trains its DiT denoiser from scratch and only borrows a frozen CLIP ViT-B/16 for vision + text conditioning, so it fits the same 12 GB budget as the ACT and Diffusion siblings — making it a fair third data point in the same overfit-diagnostic study.

Headline finding (read this first): on a held-out split the flow-matching eval_loss ROSE ~2.6× (min ≈ 0.089 near 6k → 0.233 @ 100k) over training, which naively screams "severe overfitting". But the deployment-relevant open-loop rollout MAE kept IMPROVING then plateaued, reaching its minimum at step 70k (0.0816 → 0.0735 rad). For a flow-matching policy too, the held-out flow/denoising loss is a misleading overfit/early-stop signal — select checkpoints by open-loop MAE, not by eval_loss. This now reproduces the diffusion sibling's lesson for a third policy family.

Task & data

• Dataset: Bigenlight/banana_in_pot_lerobot_v3 (LeRobot v3.0 format).
• Scale: 51 episodes / 21,524 frames / 30 fps / ~12 min.
• Action / state space: 7-D absolute joint (6 UR joints in radians + gripper), i.e. [cmd1..cmd6, grip_cmd]. The gripper channel is effectively binary (open/close).
• Cameras: two RGB viewpoints (Intel RealSense D435 + D435if), captured at 1280×720 (720p) @ 30 fps, RGB only (no depth / IR). cam1 ↔ cam2 order is fixed and must be preserved at deploy time.
• Text conditioning: unlike the ACT and Diffusion siblings (which ignore the task string), this policy is text-conditioned through CLIP. The task string is fixed to "put the right banana in the pot" for both training and inference.

Train / held-out split

Model architecture

Note: config.json records the raw dataset image feature shape as [3, 720, 1280], but the policy resizes internally to [224, 224] before the CLIP encoder. Unlike the diffusion sibling (which expected an external 360×640 resize), this policy does its own 224×224 resize — feed it the native/full-res decoded frame and let it resize. See why.

Why images must be 224×224 (and must NOT be pre-resized)

Training setup

Repro note — crash & resume

Results & the headline finding

The misleading eval_loss (the lesson)

Comparison — Flow matching vs Diffusion vs ACT (JOINT)

Caveat: eps 45–50 are the chronological tail (not an i.i.d. split); poseMAE is teacher-forced open-loop, not a closed-loop success rate; Euler-10 is a fast sampler (more integration steps could shift absolute numbers but not the ranking). Raw per-checkpoint metrics: results/eval_fm_final_gpu.csv.

Usage / inference

Load the policy (LeRobot 0.6.1)

Prerequisite: openai/clip-vit-base-patch16 must be in your Hugging Face cache (the CLIP vision + text towers), or construction fails.

Run the control loop

Reproduce the offline evaluation

Deployment on a real UR7e

• A py3.10 ROS node (policy_leader_node) acts as a synthetic GELLO leader, publishing /gello/joint_states at 30 Hz and reusing the safety-tuned gello_ur_bridge unmodified.
• The node talks over a localhost ZMQ REQ/REP split to a py3.12 policy server that runs the torch / LeRobot inference. The split exists because Humble's rclpy is py3.10 but LeRobot needs py3.12.

1. it must load the policy generically (get_policy_class) rather than hardcoding DiffusionPolicy;
2. FM uses num_integration_steps (Euler) — not a DDIM scheduler; and
3. FM images must NOT be externally resized to 360×640 — the policy does its own 224×224 resize, so the server forwards the native frame; and
4. FM is text-conditioned (CLIP), so the server must send the real task string ("put the right banana in the pot") into the obs each tick — unlike the single-task diffusion/ACT servers, which send an empty task.

1. reads the UR7e measured joints + gripper → observation.state (7,);
2. grabs both camera frames, BGR→RGB, CHW [0, 1] at native res (no external resize) → observation.images.cam1 / cam2, plus the fixed task string;
3. preprocessor → policy.select_action → postprocessor → q_target (7,);
4. streams q_target[:6] to the arm and drives the gripper from grip_cmd.

1. Start near the dataset initial pose before enabling the policy, or the first absolute command is a large jump.
2. First-command jump guard: if max(|q_target − getActualQ()|) exceeds a small threshold (~0.15 rad), abort.
3. Clamp per-tick joint change and clamp to UR software joint limits; run at reduced speed for first trials with a hand on the E-stop.
4. cam1/cam2 mapping is fixed — swap the two views and the policy fails silently. Verify wiring every session.

Limitations & intended use

• Small, single-task lab dataset: 51 demonstrations, one scene layout, one operator. Expect limited generalization to novel object arrangements, lighting, or camera placement.
• Success-only demonstrations: no failure/recovery data; not suited as-is for methods that need negative examples.
• Offline metrics only: the best checkpoint (70k) reaches held-out poseMAE ≈ 0.0735 rad and gripper accuracy ≈ 0.95–0.96 in open-loop rollout. These are not closed-loop task success rates — real closed-loop success on hardware has not been measured here and must be validated on the arm.
• Absolute-joint action space demands the safety guards above; the policy was only ever conditioned on states near the data-collection start pose.
• Not for production. Intended for research in imitation learning / flow-matching policies for robot manipulation. Workspace-, robot-, and camera-specific.
• Encoder provenance: the CLIP ViT-B/16 vision + text towers are web-pretrained (not robotics-pretrained; text tower frozen, vision fine-tuned at 0.1× LR); the DiT denoiser is trained from scratch on this task.

Links

• This model: Bigenlight/flow_matching_banana_in_pot_joint
• Dataset: Bigenlight/banana_in_pot_lerobot_v3
• Experiments repo: github.com/Bigenlight/banana-in-pot-experiments
• Diffusion sibling model: Bigenlight/diffusion_banana_in_pot_joint
• ACT sibling model: Bigenlight/act_banana_in_pot
• Deployment stack (ROS 2 Humble): github.com/Bigenlight/gello_software
• Framework: LeRobot v0.6.1 (pin 8a74e0a)

Citation