ameyawagh555/flow_matching_pusht

Model Card for Flow Matching Policy / PushT

A Flow Matching robot policy (continuous-time linear-interpolant forward process + a configurable fixed-step ODE solver, default forward-Euler) trained for the PushT environment from gym-pusht, built as an out-of-tree LeRobot plugin: lerobot_policy_flow_matching.

It reuses Diffusion Policy's vision encoder and conditional U-Net architecture verbatim, replacing only the generative process (DDPM denoising -> flow matching) so the two policies are directly comparable on the same task, at the same training budget, with matched checkpoint-selection methodology.

How to Get Started with the Model

See the LeRobot library and the lerobot_policy_flow_matching plugin for instructions on how to load and evaluate this model, e.g.:

lerobot-eval --policy.path=ameyawagh555/flow_matching_pusht --env.type=pusht \
    --eval.n_episodes=200 --eval.use_async_envs=false

Training Details

Trained with the lerobot_policy_flow_matching plugin (on top of lerobot>=0.6.1,<0.7) and the pusht dataset, using:

scripts/train_flow_matching.sh 200000 64 outputs/train/m4b_full_fixed -- \
    --env_eval_freq=25000 --save_freq=25000 --eval.n_episodes=50

Key config (see config.json in this repo for the full list):

This checkpoint is the 175,000-step checkpoint of that run, picked by periodic in-training rollout eval (50 episodes every 25,000 steps) -- the best of 8 candidates, following the same best-of-periodic- eval selection methodology lerobot/diffusion_pusht itself uses (its own shipped checkpoint is its 175k-step one out of a 200k-step run, not the final step). Training curves: https://wandb.ai/ameya555-ieee/lerobot/runs/hmsg3mdx

use_group_norm=True requires pretrained_backbone_weights=None -- swapping BatchNorm for GroupNorm in an ImageNet-pretrained backbone would ruin its pretrained weights, so the backbone is trained from scratch on the ~206 PushT demonstration episodes, matching `lerobot/diffusion_pusht`'s own recipe.

Evaluation

Evaluated on the PushT environment from gym-pusht and compared against `lerobot/diffusion_pusht`, a DDPM-based Diffusion Policy checkpoint trained on the same dataset at the same training budget (200,000 steps, batch 64). Two eval metrics, per episode:

• —Average max overlap with target (avg_max_reward), in [0, 1].
• —Success: whether max overlap reaches >= 95%.

Read this as parity, with a possible modest edge, not a clear win. A two-proportion test against the baseline's published number gives z ~= 2.07, p ~= 0.038 -- it clears the conventional significance bar as a single comparison, but this is one of several comparisons made during the investigation (see below) and the baseline's 65.4% is their own reported figure rather than a fresh eval run under identical conditions on the same machine, so it should not be read as a confidently-established win.

What's more solidly established: the same success rate is reached with 10 ODE integration steps per action chunk instead of ~100 DDPM steps, and separately, that 50,000 training steps already reaches the same success rate as 200,000 (a 50k-step checkpoint from the same recipe scored 69.0%, n=200, statistically indistinguishable from both the baseline and this 175k checkpoint, p=0.32) -- once a drop_n_last_frames configuration bug (documented below) was fixed. The efficiency gain, not a success-rate win, is the well-supported claim here.

Full run details, ablations, and the investigation that led to this configuration -- including two config bugs that initially made this policy look roughly half as good as the baseline (n_action_steps mismatched against lerobot's current DiffusionConfig defaults, and an incorrectly recomputed drop_n_last_frames that broke terminal-alignment training data) -- are written up in `docs/comparison_pusht.md` of the plugin repo.

Open items, tracked in the repo, none blocking: the plugin's ode_solvers.py module (Heun, RK4) is implemented but not yet benchmarked against the default Euler solver; the num_inference_steps and training-time t-sampling distribution have not been re-swept on this (bug-fixed) configuration; and the handful of hard failure episodes at this checkpoint haven't been individually inspected for a common failure mode.