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

A flow-matching visuomotor policy (multi_task_dit, objective=flow_matching) trained by imitation learning to perform the manipulation task "put the right banana in the pot" on a Universal Robots UR7e arm with two RGB cameras. Actions are 7-D absolute joint targets (6 UR joints in radians + gripper).

• 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

"put the right banana in the pot." The tabletop holds several distractor objects — two bananas, an apple, carrots/peppers, and a slice of watermelon — plus a silver pot. The operator must grasp the RIGHT banana (the target) and place it inside the pot. Success = the right banana ends up inside the pot. Every demonstration is a success.

• 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

Training used --dataset.eval_split=0.117, which holds out the LAST ceil(51 × 0.117) = 6 episodes (indices 45–50) as a true validation split and trains on the other 45 episodes (0–44), i.e. 18,807 frames. The held-out episodes 45–50 are used both for the in-training flow-matching eval_loss probe and for all offline open-loop evaluation below.

Model architecture

LeRobot multi_task_dit with objective = flow_matching. All values below are quoted directly from the checkpoint's config.json / train_config.json.

Observation encoder (vision + text, CLIP ViT-B/16):

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.

Denoiser (DiT transformer):

Flow-matching process:

Normalization (normalization_mapping):

Normalizer statistics are baked into the pre/post-processor pipelines saved alongside the checkpoint (policy_preprocessor.json / policy_postprocessor.json), not into forward(). select_action returns a normalized action; the post-processor converts it back to radians.

I/O summary:

Parameter count: 186 M learnable parameters (249 M total including the frozen CLIP text tower).

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

CLIP ViT-B/16 is built for a fixed 224 × 224 patch grid. This policy therefore sets image_resize_shape = [224, 224] and performs the resize inside the policy's own image transform (image_crop_shape = [224, 224] equals the resize, so the crop is a no-op). This differs from the diffusion sibling, whose SpatialSoftmax encoder expected an external 360×640 resize.

Practical consequence for deployment: feed the decoded RGB frame at its native / full resolution and let the policy resize to 224×224. Do NOT pre-resize to 360×640 (that is a diffusion-specific step and would double-resize / distort the input here). The same internal 224×224 path is reproduced in eval_offline.py, so offline eval, training, and deploy all feed CLIP identically.

Training setup

Trained with lerobot-train (LeRobot 0.6.1, pin 8a74e0a). Exact invocation: train_fm_joint_valdiag.sh. Values below are from that script and the saved train_config.json.

Prerequisite — CLIP weights must be cached. The policy loads openai/clip-vit-base-patch16 from the Hugging Face cache. Make sure those weights are present (download once with network access) before training or offline eval, or construction will fail.

Repro note — crash & resume

The run crashed at step 20k when the SSD filled during a checkpoint save (ENOSPC). It was resumed from the 10k checkpoint and completed 100k cleanly. To reproduce a resume:

# train_fm_joint_resume.sh
lerobot-train \
  --config_path=outputs/train/fm_joint_val_diag/checkpoints/010000/pretrained_model/train_config.json \
  --resume=true

--config_path=<last_ckpt>/pretrained_model/train_config.json plus --resume=true restores the optimizer / scheduler / step counter. (Old runs were archived to an HDD to free space; see EXPERIMENT_LOG.md.)

Results & the headline finding

Offline open-loop evaluation on the held-out episodes 45–50 with eval_offline.py (each logged observation is fed to select_action; the predicted action is compared to the dataset ground truth). Sampling used Euler with 10 integration steps (--num-inference-steps 10, which maps onto num_integration_steps) for ~10× faster rollouts — the flow-matching analogue of the diffusion sibling's DDIM-10. poseMAE is the mean absolute error over the 6 joint dims (radians); gripAcc is the binary gripper-open/close accuracy (threshold 0.5); overall L1 averages all 7 dims.

Best checkpoint = 70k (min poseMAE and min overall L1). Open-loop poseMAE improves early (already ~0.076 by 30k) then settles onto a ~0.074 rad plateau from ~40k onward (70k/80k/90k/100k = 0.0735/0.0745/0.0748/0.0746, within eval noise); gripper accuracy holds steady at ~0.95–0.96. There is no destructive open-loop overfitting through 100k.

The misleading eval_loss (the lesson)

During training the held-out flow-matching eval_loss (LeRobot's in-training validation probe, computed under policy.eval() on eps 45–50) did the opposite of the rollout metric:

Read naively, the held-out eval_loss bottoms near step 6k and then rises ~2.6×, so an early-stop rule would pick ~step 6k and declare "severe overfit". That recommendation is wrong for deployment: the same held-out episodes, evaluated by open-loop rollout, get better out to 70k and then hold flat.

Why: the flow-matching loss is a per-sample MSE on the predicted velocity field, at a randomly re-sampled noise vector and flow timestep at every forward pass — it is (a) high-variance/stochastic by construction and (b) only loosely coupled to closed-loop action quality. As the model sharpens its learned action distribution, the average velocity-MSE on unseen frames can rise even while the sampled (Euler-integrated) action trajectories become more accurate. Takeaway: for a flow-matching policy, select checkpoints and early-stop by open-loop rollout MAE, not by held-out eval_loss. This is now the third policy family (after Diffusion) to show the divergence — the ACT sibling did not, so this is a pitfall of denoising/flow-style generative action heads specifically.

Comparison — Flow matching vs Diffusion vs ACT (JOINT)

Held-out eps 45–50, open-loop poseMAE (radians), best checkpoint of each family:

Flow matching is the best of the three on JOINT open-loop pose accuracy — about 13% lower poseMAE than Diffusion and ~25% lower than ACT, at comparable (excellent) gripper accuracy. FM also beats Diffusion at every checkpoint (e.g. 10k: 0.0816 vs 0.1193; 30k: 0.0763 vs 0.0921) and converges fast (already on-plateau by ~30k).

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)

Use the generic loader (get_policy_class) — it works for any policy type and resolves to MultiTaskDiTPolicy here. Normalization is not baked into forward() in LeRobot 0.6.1; it lives in the pre/post-processor pipelines saved with the checkpoint.

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

import torch
from lerobot.configs import PreTrainedConfig
from lerobot.policies.factory import get_policy_class, make_pre_post_processors

CKPT = "Bigenlight/flow_matching_banana_in_pot_joint"
device = "cuda"

cfg = PreTrainedConfig.from_pretrained(CKPT)
cfg.pretrained_path = CKPT
cfg.device = device
cfg.num_integration_steps = 10        # optional: fewer Euler steps = faster (default 100)

policy = get_policy_class(cfg.type).from_pretrained(CKPT, config=cfg)  # -> MultiTaskDiTPolicy
policy.to(device)
policy.eval()

preprocessor, postprocessor = make_pre_post_processors(
    policy_cfg=cfg,
    pretrained_path=CKPT,
    preprocessor_overrides={"device_processor": {"device": device}},
)

Run the control loop

Build the observation dict exactly as training did: joint state (7,) plus both cameras as RGB CHW tensors in [0, 1] at native / full resolution (the policy resizes to 224×224 internally — do NOT pre-resize to 360×640), plus the fixed task string. cam1/cam2 must map to the same physical viewpoints as at collection.

policy.reset()          # once at the start of each episode/rollout
preprocessor.reset()
postprocessor.reset()

# obs = {
#   "observation.state":        state_7,          # (7,) float32, radians + gripper
#   "observation.images.cam1":  img1_chw,         # (3, H, W) float32 in [0,1], NATIVE res
#   "observation.images.cam2":  img2_chw,         # (3, H, W) float32 in [0,1], NATIVE res
#   "task": "put the right banana in the pot",    # text conditioning (CLIP) — required
# }

with torch.inference_mode():
    proc   = preprocessor(obs)             # rename -> add batch dim -> device -> normalize
    action = policy.select_action(proc)    # (1, 7) NORMALIZED
    action = postprocessor(action)         # (1, 7) radians, on cpu
q_target = action.squeeze(0).numpy()       # (7,) -> [cmd1..cmd6, grip_cmd]

select_action returns one action per call from an internal queue. Because n_action_steps = 24, the policy integrates a fresh action sequence, executes 24 actions from it, then replans (with n_obs_steps = 2 frames of observation context) — i.e. it replans ~every 24 ticks (~0.8 s @ 30 Hz). Call policy.reset() at the start of every episode to clear that queue. The gripper channel grip_cmd is ~binary — threshold at > 0.5 → close and map to your gripper driver.

Reproduce the offline evaluation

The repo's eval_offline.py runs the exact open-loop protocol used for the results table (same internal 224×224 resize, same normalization via the saved processors). The --num-inference-steps flag maps onto the flow-matching num_integration_steps:

python eval_offline.py \
  --checkpoint outputs/train/fm_joint_val_diag/checkpoints/070000/pretrained_model \
  --episodes 45,46,47,48,49,50 \
  --device cuda \
  --num-inference-steps 10 \
  --out eval_out_fm_70k

--num-inference-steps 10 gives the ~10× rollout speedup; raise it (up to 100) to integrate the velocity field more finely.

Deployment on a real UR7e

Closed-loop deployment targets a real UR7e through the ROS 2 Humble stack in Bigenlight/gello_software, package gello_policy — exactly like the ACT and Diffusion JOINT models. The 7-D joint action + two-camera observation contract is byte-identical to the diffusion JOINT deploy, so the entire ROS / safety side is unchanged:

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

What's new for flow matching — a dedicated py3.12 server fm_server.py (added in gello_policy alongside act_server.py / diffusion_server.py), because:

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.

The run script defaults to Euler-10 integration (the sampler used for the results above) to keep refill latency under the leader's fault timeout. Point fm_server.py at this checkpoint (70k). See the deploy runbook gello_software/docs/ros2/GELLO_UR7E_FM_DEPLOY.md.

Each control tick (target 30 Hz) the pipeline:

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.

Safety — actions are ABSOLUTE joint positions (same guards as the diffusion card):

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

@misc{theo2026bananainpotflowmatching,
  title        = {Flow-Matching Policy for "put the right banana in the pot"
                  (UR7e, joint action space)},
  author       = {Theo and {Bigenlight}},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/Bigenlight/flow_matching_banana_in_pot_joint}},
  note         = {LeRobot 0.6.1 multi_task_dit, objective=flow_matching,
                  trained on banana_in_pot_lerobot_v3}
}

License: Apache-2.0.