Research map / Probabilistic and scientific modeling / Differential equations and PDEs
AdaFlow: Imitation Learning with Variance-Adaptive Flow-Based Policies
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AdaFlow, a flow-based imitation learning policy whose ODE step size adapts to a learned per-state variance estimate, yielding near one-step action generation with strong success rates.
Problem
Diffusion-based imitation learning captures multi-modal actions but needs many recursive denoising steps at inference, which is too slow for real-time robot control.
Approach
Train a state-conditioned rectified flow policy, then learn a variance estimation network whose predicted conditional variance sets the ODE step size at inference: deterministic states get one Euler step, multimodal states get more steps. The method uses only the initially learned ODE, with no reflow or distillation.
Key results
- On RoboMimic, AdaFlow reaches an average NFE of 1.17 with competitive success rates, while Diffusion Policy with NFE=1 collapses to near 0% on most tasks.
- On LIBERO six-task average, AdaFlow achieves 0.91 success with NFE 1.27, versus 0.88 for Diffusion Policy at NFE 20.
- On the Maze2D tasks, AdaFlow achieves an average NFE of 1.56 and success rates of 0.98/1.00/0.96/0.86, outperforming Diffusion Policy at NFE 20.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add variance estimation network to AdaFlow (ablation) | AdaFlow · Maze 1 | success rate: 0.78 -> 0.98 | Helped |
| add variance estimation network to AdaFlow (ablation) | AdaFlow · Maze 2 | success rate: 1.00 -> 1.00 | No clear effect |
| add variance estimation network to AdaFlow (ablation) | AdaFlow · Maze 3 | success rate: 0.92 -> 0.96 | Helped |
| add variance estimation network to AdaFlow (ablation) | AdaFlow · Maze 4 | success rate: 0.80 -> 0.86 | Helped |
| use AdaFlow instead of BC | Maze navigation policies · Maze 1 | success rate: 1.00 -> 0.98 | Hurt |
| use AdaFlow instead of Diffusion Policy with 20 steps | Maze navigation policies · Maze 4 | success rate: 0.82 -> 0.86 | Helped |
| use AdaFlow instead of Diffusion Policy with 20 steps | LIBERO policies · LIBERO average | success rate: 0.88 -> 0.91 | Helped |
| train variance estimator jointly rather than separately (ablation) | AdaFlow · Maze tasks | success rate: Maze 1: 0.98 -> 1.00; Maze 2: 1.00 -> 1.00; Maze 3: 0.96 -> 0.96; Maze 4: 0.86 -> 0.88 | No clear effect |
Surprising findings
- Diffusion Policy with a single step performs near zero success on many RoboMimic and LIBERO tasks, showing that a 1-step diffusion sampler fails without adaptive steps.
- Reflow-based Rectified Flow with NFE 1 underperforms AdaFlow on maze tasks, and the RoboMimic and LIBERO results for the baseline were taken from the paper's reported setup.
Limitations
- Training Diffusion Policy on RoboMimic is very resource-intensive, so the authors fine-tuned a pretrained open-source checkpoint for 500 epochs rather than training from scratch.
- Results for the reflow-based Rectified Flow baseline are reported at NFE 1 only and require an additional costly reflow stage.
- Applied only to offline imitation learning; extension to offline RL is left for future work.
Experimental setup
Datasets: D4RL Maze2D RoboMimic LIBERO Push-T
Benchmarks: D4RL Maze2D RoboMimic LIBERO Push-T
Metrics: success rate number of function evaluations (NFE)
Baselines: Behavioral Cloning LSTM-GMM IBC BET Diffusion Policy Rectified Flow BC-GMM
Builds on: rectified flow Diffusion Policy behavioral cloning D4RL
Closest papers by meaning
- Diffusion Model-Augmented Behavioral Cloning (ICML 2024 · similarity 0.87)
- Ergodic Generative Flows (ICML 2025 · similarity 0.87)
- Learning When to Stop: Selective Imitation Learning Under Arbitrary Dynamics Shift (NeurIPS 2026 · similarity 0.87)
- FlowPolicy: Enabling Fast and Robust 3D Flow-Based Policy via Consistency Flow Matching for Robot Manipulation (AAAI 2025 · similarity 0.87)
- Diffusion Imitation from Observation (NeurIPS 2024 · similarity 0.86)
- Causal Imitation for Markov Decision Processes: a Partial Identification Approach (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.