Bilinear Flow Policy: Distributional Extrapolation for Goal-Conditioned Visuomotor Imitation

Wonsuhk Jung; Sundhar Vinodh Sangeetha; Chen Xu; Abhishek Gupta; Masha Itkina; Shreyas Kousik; Haruki Nishimura

cs.ROOct 5, 2026

Bilinear Flow Policy: Distributional Extrapolation for Goal-Conditioned Visuomotor Imitation

Authors: Wonsuhk Jung, Sundhar Vinodh Sangeetha, Chen Xu, Abhishek Gupta, Masha Itkina, Shreyas Kousik, Haruki Nishimura

Organizations: Georgia Institute of Technology · Toyota Research Institute · University of Washington

Abstract

Goal-conditioned imitation learning (GCIL) with flow matching is a promising framework that can represent multimodal behaviors while adapting to diverse, user-specified goals, yet often fails when goals lie outside the demonstration support. To extrapolate to such unseen goals without collapsing multimodality - a problem we call distributional extrapolation - we introduce Bilinear Flow Policy (BFP), a generative visuomotor policy that combines transductive retrieval with a bilinear conditional flow. Given an unseen observation-goal pair, BFP retrieves an "anchor" training example and transductively reformulates the unseen pair as this familiar anchor plus a residual term. For this decomposition to guide action prediction, the residual must compactly encode how the current observation-goal pair differs from the anchor, and the anchor must be chosen so that this difference is predictive of the corresponding action distribution. BFP achieves this with pretrained visual features and a novel learned anchor-selection algorithm. The novel bilinear flow then models how the anchor and the residual jointly determine the multimodal action distribution. We prove that, for bilinear flow under suitable assumptions, action distribution error at unseen goals is bounded by the in-distribution flow-matching error up to problem-dependent factors. Across five manipulation tasks in simulation, BFP achieves 2.63x the out-of distribution success rate of a GCIL policy and 1.36x that of the strongest extrapolation-targeted baseline. On two real-world tasks, BFP improves over GCIL by 32%. Finally, our theory yields practical, pre deployment diagnostics for predicting which trained policies will extrapolate well and to which unseen goal.

Figures & tables

Figure 1: BFP Overview. BFP learns how changes across observation-goal conditions affect behavior and reuses this relation to extrapolate from a demonstrated anchor to an unseen goal. A bilinear conditional flow generates the resulting multimodal action distribution, while the explicit extrapolation rule enables theoretical guarantees and practical diagnostics.

Figure 2: Distributional extrapolation. The expert conditional distribution $q^{\star}(y\mid x)$ is a mixture of two modes (grey), evaluated outside $\operatorname{supp}(P_{\mathrm{tr}}^{X})$ (shaded red). BT averages the modes (orange), while CFM captures multimodality but not extrapolation (purple). BFP achieves both (red).

Figure 3: Overview of BFP . Given an unseen observation–goal condition ${\bm{x}}$ , BFP retrieves an anchor ${\bm{x}}^{\mathsf{\dagger}}$ and predicts from the resulting anchor–residual pair. The novel retrieval policy ( Section 4.2 ) uses a phase estimator to form a phase-compatible candidate set and a learned scorer $s_{\psi}$ , trained from held-out BFP loss, to select ${\bm{x}}^{\mathsf{\dagger}}$ . A frozen visual encoder computes ${\bm{z}}=\phi({\bm{x}})$ and ${\bm{z}}^{\mathsf{\dagger}}=\phi({\bm{x}}^{\mathsf{\dagger}})$ , with residual $\bm{\delta}={\bm{z}}-{\bm{z}}^{\mathsf{\dagger}}$ . The novel bilinear conditional-flow predictor (Sec. 4.1) predicts the CFM velocity through the low rank- $r$ interaction; flow integration then generates the action chunk.

Method	$\operatorname{NW}_{2}$ ID $\downarrow$	$\operatorname{NW}_{2}$ OOD $\downarrow$	Balance OOD $\downarrow$
Oracle	$0.078$	$0.086$	$0.019$
Flow	$\mathbf{0.221}$	$0.506$	$0.160$
BT	$0.726$	$0.763$	$0.500$
BFP	$0.224$	$\mathbf{0.321}$	$\mathbf{0.103}$

Table 1: BFP distributionally extrapolates. Results averaged over three problems and three seeds. Lower is better for all metrics. Oracle samples directly from the expert conditional and gives the Monte Carlo floor.

Panel of Figure 4 — Figure 4: BFP generalizes in visuomotor manipulation. (a) Success rate / Normalized reward averaged across seven simulation settings (b) Same across hardware tasks. Success rate violins show Bayesian posterior estimates under a uniform Beta prior. Normalized reward violins illustrate empirical distributions of the results. Dots represent sample means and horizontal lines show posterior means (success-rate only). Letters on top of the plots show statistical comparison, where columns not sharing the same alphabet are significantly different under 5% FWER control.

Panel of Figure 5 — Figure 5: Pre-rollout diagnostics and BFP ablations. (a) Held-out ID flow error ranks OOD error for BFP but not for Flow ; within a fixed BFP checkpoint, better residual coverage is associated with lower extrapolation error. (b) Ablations of the visual encoder, anchor policy, backbone, and bilinear rank. Frozen pretrained features and phase-conditioned retrieval both lead to better performance; BFP can exploit stronger sequence backbones and performs similarly across the ranks.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Task Setup. (Left) InsertPipette (Right) PnPMug

Method	Metric	PnPMug	InsertPipette	Avg.
(a) In-distribution (ID)
Flow	SR	$86{\scriptstyle\,\pm\,5}$	$83{\scriptstyle\,\pm\,7}$	$85$
Flow	Reward	$95{\scriptstyle\,\pm\,2}$	$\underline{92{\scriptstyle\,\pm\,3}}$	$93$
BT	SR	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\mathbf{100}$
BT	Reward	$\underline{97{\scriptstyle\,\pm\,1}}$	$\mathbf{99{\scriptstyle\,\pm\,1}}$	$\underline{98}$
BFP	SR	$\underline{98{\scriptstyle\,\pm\,2}}$	$\underline{97{\scriptstyle\,\pm\,3}}$	$\underline{97}$

Appendix

Table 2: Performance in- and out-of-distribution on real-world robot hardware. Success rate (SR, %), our primary metric, is highlighted in blue. Reward denotes normalized dense reward (

\times 100

). Higher is better for both metrics. Values are means

\pm

standard error over trials: PnPMug 50 ID / 100 OOD trials, InsertPipette 30 ID / 110 OOD trials per method. Bold and underlined values indicate the best and second-best means, respectively, including ties.

Figure 7: OOD Rollout Strips. Successful BFP rollouts in OOD labeled with the progress phase used to compute dense reward. BT and Flow failure rollouts in OOD exhibit common failure modes. In PnPMug and InsertPipette , Flow often fails early in the task in grasping the manipuland despite its unchanging pose in the scene between ID training demonstrations and OOD rollouts. BT ’s time-synchronized anchor cannot re-synchronize once the robot lags behind it, so the residual drifts off support and leads to a freezing failure mode where the policy fails to make further progress. In PnPMug this often occurs in the grasping phase, and in InsertPipette often occurs during the alignment phase.

Figure 8: PnPMug spatial distribution of successes and failures . Each dot represents a goal object (plate) IC used in evaluation. Failures increase as the goal object is moved further from the training data distribution (ID). BFP achieves successful generalization to these OOD conditions more often and further away from the training data goal pose distribution compared to Flow .

Figure 9: InsertPipette spatial distribution of successes and failures . Each circle represents a test-tube IC used for evaluation in a

4\times 8

test-tube tray. Generalization to unseen, OOD goals is substantially improved with BFP .

Figure 10: Success rate and Dense reward under OOD evaluation on real-world robot hardware.

Task (Benchmark)	Goal	Success criterion	Dense reward
PushCube (MetaWorld)	2D target position	target dist. $<5$ cm	reach, grasp, and target progress
Relocate (Adroit)	3D target position	target dist. $<10$ cm	reach, lift, and target progress
RotatePen (Adroit)	target orientation	ori. cos. $>0.95$ , pos. err. $<7.5$ cm	orientation and position progress
StackCube (ManiSkill3)	target cube position	stacked, released, static	reach, grasp, stack, and release progress
PickCube (ManiSkill3)	3D target position	target dist. $<2.5$ cm	reach, grasp, and place progress

Appendix

Table 3: Simulation task definitions. Goal denotes the compact representation used for policy conditioning. Success is evaluated at the final step of the episode. Dense reward is the environment’s task-progress reward, normalized by the maximum episode return for each task.

Setting	Demos	ID goal support	OOD goal support	Object initial condition
PushCube-D0	200	$x\in[0,0.3]$ $y\in[0.5,0.7]$	$x\in[-0.3,0]$ $y\in[0.5,0.7]$	fixed
PushCube-D1	200	$x\in[0,0.3]$ $y\in[0.5,0.7]$	$x\in[-0.3,0]$ $y\in[0.5,0.7]$	puck $x\in[-0.06,0.06]$
Relocate-D0	100	$x,y\in[0,0.3]$ $z\in[0.15,0.35]$	$x,y\in[-0.3,0]$ $z\in[0.15,0.35]$	fixed
Relocate-D1	250	$x,y\in[0,0.3]$ $z\in[0.15,0.35]$	$x,y\in[-0.3,0]$ $z\in[0.15,0.35]$	ball $x\in[-0.1,0.1]$
RotatePen	187	roll $\in[0.2,1.0]$ pitch $\in[-0.2,0.6]$	roll $\in[-0.3,0]$ pitch $\in[-0.2,0.6]$	fixed
StackCube	200	$x\in[0,0.2]$ $y\in[0,0.27]$	$x\in[-0.19,0]$ $y\in[-0.25,0]$	$x\in[-0.19,0.19]$ $y\in[-0.29,0.30]$

Appendix

Table 4: Simulation data and evaluation splits. ID goals are sampled from the training goal support while OOD goals are held out. D0 fixes the object initial-state variation. Positions and angles are reported in meters and radians, respectively.

Panel of Figure 11 — Figure 11: Evaluation Tasks. (a) PushCube , (b) Relocate , (c) RotatePen , (d) StackCube , (e) PickCube . The suite spans pushing, dexterous object relocation and rotation, pick-and-place, stacking, and insertion. For definition of ID/OOD per task, please see Figure 12 .

Panel of Figure 12 — Figure 12: Distribution Shift. Each task panel visualizes the support of the training conditions (ID, left) and the held-out evaluation conditions (OOD, right). The paired views use the same camera and scene configuration so that the corresponding distribution shift is directly visible.

Method	Metric	PushCube-D0	PushCube-D1	Relocate-D0	Relocate-D1	RotatePen	StackCube	PickCube	Avg.
(a) State
Flow	SR	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$98{\scriptstyle\,\pm\,1}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\underline{97{\scriptstyle\,\pm\,1}}$	$62{\scriptstyle\,\pm\,3}$	$\underline{96{\scriptstyle\,\pm\,1}}$	$93$
Flow	Reward	$\mathbf{89{\scriptstyle\,\pm\,0}}$	$\mathbf{88{\scriptstyle\,\pm\,0}}$	$81{\scriptstyle\,\pm\,1}$	$80{\scriptstyle\,\pm\,1}$	$\underline{87{\scriptstyle\,\pm\,0}}$	$64{\scriptstyle\,\pm\,0}$	$63{\scriptstyle\,\pm\,0}$	$\underline{79}$
DARP	SR	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\underline{97{\scriptstyle\,\pm\,0}}$	$\underline{75{\scriptstyle\,\pm\,0}}$	$87{\scriptstyle\,\pm\,1}$	$94$
DARP	Reward	$\mathbf{89{\scriptstyle\,\pm\,0}}$	$\mathbf{88{\scriptstyle\,\pm\,0}}$	$\mathbf{84{\scriptstyle\,\pm\,1}}$	$\underline{82{\scriptstyle\,\pm\,1}}$	$\underline{87{\scriptstyle\,\pm\,0}}$	$\underline{67{\scriptstyle\,\pm\,0}}$	$64{\scriptstyle\,\pm\,0}$	$\mathbf{80}$
CTA	SR	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\underline{96{\scriptstyle\,\pm\,3}}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\mathbf{100{\scriptstyle\,\pm\,0}}$	$\underline{97{\scriptstyle\,\pm\,0}}$	$52{\scriptstyle\,\pm\,2}$	$82{\scriptstyle\,\pm\,2}$	$89$

Appendix

Table 5: In-distribution performance with state and visuomotor observations. Success rate (SR, %), our primary metric, is highlighted in blue. Reward denotes normalized dense reward (

\times 100

). Higher is better for both metrics. Bold and underlined values indicate the best and second-best means, respectively, including ties. Mean

\pm

standard error over three training seeds.

Method	Metric	PushCube-D0	PushCube-D1	Relocate-D0	Relocate-D1	RotatePen	StackCube	PickCube	Avg.
(a) State
Flow	SR	$31{\scriptstyle\,\pm\,4}$	$49{\scriptstyle\,\pm\,2}$	$26{\scriptstyle\,\pm\,7}$	$24{\scriptstyle\,\pm\,2}$	$52{\scriptstyle\,\pm\,5}$	$41{\scriptstyle\,\pm\,1}$	$80{\scriptstyle\,\pm\,2}$	$43$
Flow	Reward	$27{\scriptstyle\,\pm\,4}$	$42{\scriptstyle\,\pm\,1}$	$18{\scriptstyle\,\pm\,5}$	$22{\scriptstyle\,\pm\,2}$	$49{\scriptstyle\,\pm\,5}$	$55{\scriptstyle\,\pm\,0}$	$50{\scriptstyle\,\pm\,3}$	$38$
DARP	SR	$\underline{36{\scriptstyle\,\pm\,1}}$	$45{\scriptstyle\,\pm\,1}$	$43{\scriptstyle\,\pm\,11}$	$47{\scriptstyle\,\pm\,8}$	$75{\scriptstyle\,\pm\,10}$	$2{\scriptstyle\,\pm\,1}$	$5{\scriptstyle\,\pm\,1}$	$36$
DARP	Reward	$31{\scriptstyle\,\pm\,1}$	$40{\scriptstyle\,\pm\,1}$	$29{\scriptstyle\,\pm\,7}$	$36{\scriptstyle\,\pm\,6}$	$65{\scriptstyle\,\pm\,8}$	$38{\scriptstyle\,\pm\,1}$	$40{\scriptstyle\,\pm\,0}$	$40$
CTA	SR	$\underline{36{\scriptstyle\,\pm\,8}}$	$48{\scriptstyle\,\pm\,9}$	$15{\scriptstyle\,\pm\,4}$	$24{\scriptstyle\,\pm\,4}$	$77{\scriptstyle\,\pm\,12}$	$12{\scriptstyle\,\pm\,3}$	$25{\scriptstyle\,\pm\,3}$	$34$

Appendix

Table 6: OOD performance with state and visuomotor observations. Success rate (SR, %), our primary metric, is highlighted in blue. Reward denotes normalized dense reward (

\times 100

). Higher is better for both metrics. Bold and underlined values indicate the best and second-best means, respectively, including ties.

Figure 13: Success rate under OOD evaluation.

Figure 14: Dense reward under OOD evaluation.

Figure 15: Success rate under ID evaluation.

Figure 16: Dense reward under ID evaluation.

Figure 23

Hyperparameter	Value
Optimizer	Adam, $\beta=(0.9,0.999)$
Learning rate (BFP predictor/scorer)	$10^{-4}$
Batch size	32 queries
Gradient clipping	Global norm 1.0
EMA power / maximum decay	$0.75/0.9999$
Training steps	120k (Relocate-D0), 420k (Relocate-D1)

Appendix

Table 7: Key hyperparameters for simulation experiments. Anchor parameters apply only to BFP.

Figure 19: Distributional extrapolation sample visualization. Rows correspond to different tasks: mm_pg_sm , mm_pg_pm , and mm_widen , while columns show the samples of training data (gray), BFP (red), Flow (purple), and BT (orange).

	$\mathrm{NW}_{2}\downarrow$		Balance $\downarrow$
Method	ID	OOD	ID	OOD
mm_pg_sm
Oracle	$0.10{\pm}0.01$	$0.11{\pm}0.01$	$0.02{\pm}0.00$	$0.02{\pm}0.00$
Flow	$0.23{\pm}0.01$	$0.30{\pm}0.01$	$0.09{\pm}0.00$	$0.07{\pm}0.01$
BT	$0.51{\pm}0.00$	$0.53{\pm}0.00$	$0.50{\pm}0.00$	$0.50{\pm}0.00$
BFP	$\mathbf{0.15{\pm}0.01}$	$\mathbf{0.17{\pm}0.00}$	$\mathbf{0.03{\pm}0.00}$	$\mathbf{0.03{\pm}0.01}$

Appendix

Table 8: Controlled distributional extrapolation. Mean

\pm

standard deviation over three seeds. Bold denotes the best learned method; Oracle is the finite-sampling floor.

Problem	$f_{1}$	$f_{2}$
mm_pg_sm	$f_{\mathrm{pg}}$	$f_{\mathrm{sm}}$
mm_pg_pm	$f_{\mathrm{pg}}$	$f_{\mathrm{pm}}$
mm_widen	$f_{\mathrm{sm}}$	$f_{\mathrm{sbl}}$

Appendix

Table 9: Distributional extrapolation problem setting. Each problem combines two mode functions from equation 92 .

Panel of Figure 20 — Figure 20: Sensitivity to representation and prescribed extrapolation rules. (a) ID/OOD success rate under visuomotor, averaged over five tasks. Scratch uses a jointly trained ResNet-34 for Flow and BFP and IMPALA encoder for CTA. Bold denotes the best-performing method within each encoder setting and split. (b) OOD sensitivity of EquiDP to the prescribed symmetry and camera viewpoint.

Method	Representation	Extrapolation Rule	Prediction Model
BFP	PVR	learned anchor–residual bilinear	conditional flow
Flow	PVR	none	conditional flow
BT	PVR	synchronized anchor–residual bilinear	deterministic regression
DARP	PVR	local retrieved residual	conditional flow
EquiDP	Unfrozen	prescribed geometric equivariance	conditional flow
CTA	Unfrozen	learned temporal-distance analogy	hierarchical policy

Appendix

Table 10: Design axes of the main baselines. The table separates three components of each method: the observation representation, the extrapolation rule used beyond demonstrated goals, and the prediction model used to generate actions. This decomposition highlights which components are matched to BFP and which define each baseline’s extrapolation mechanism.

Explore similar work

Potential-Guided Flow Matching for Vision-Language-Action Policy Improvement
Jun 3, 2026Yunpeng Mei, Jiakai He, Hongjie Cao +12Flow Matching Visuomotor Policy Learning
1Beijing Institute of Technology · 2Tsinghua University · 3The Hong Kong University of Science and Technology (Guangzhou)
SeFA-Policy: Fast and Accurate Visuomotor Policy Learning with Selective Flow Alignment
Nov 11, 2025Rong Xue, Jiageng Mao, Mingtong Zhang +1Rectified Flow Visuomotor Policy Learning
University of Southern California
Source-Lifted Flow Matching for Intervenable Multimodal Imitation
Jul 11, 2026He Zhang, Ying Sun, Ziyang Chen +6Robotic Control Flow Matching
The Hong Kong University of Science and Technology (Guangzhou) · AI2 Robotics

Jun 3, 2026, cs.ROJ/K move · Enter open · S save

Potential-Guided Flow Matching for Vision-Language-Action Policy Improvement

Yunpeng Mei, Jiakai He, Hongjie Cao, Chenyu Wang+11

1Beijing Institute of Technology · 2Tsinghua University · 3The Hong Kong University of Science and Technology (Guangzhou)

Flow Matching Visuomotor Policy Learning Vision-Language-Action Models Flow-Based Generative Modeling

Open paper AskPDF arXiv

Large vision-language-action (VLA) policies are increasingly trained as conditional generative models over action chunks. Yet deployment produces mixed-quality experience-successful demonstrations, partial completions, recoverable mistakes, and failures-that is difficult to use with standard imitation. Full behavior cloning (BC) imitates failures, filtered BC discards useful sub-trajectories, and offline reinforcement learning adds a large critic. We introduce ForesightFlow, a self-guided flow-matching policy that augments each generated action chunk with a learned success-potential trajectory. The same flow proposes and scores candidate actions, enabling best-of-

K

inference without an external critic. The key issue is that policy improvement and value calibration require different supervision: advantage weighting should emphasize high-quality actions, but applying the same weights to potential coordinates suppresses failure gradients and creates overconfident scores. We address this with decoupled advantage-weighted flow matching, applying exponentiated advantage weights only to action velocities while training potential velocities uniformly. We further derive a one-step boundary estimator for conditional flow matching, allowing advantage computation with a single stop-gradient forward pass. Across five BEHAVIOR-1K simulation tasks and five real-world bimanual tasks, ForesightFlow improves over imitation baselines, matches the strongest separate-critic baseline in simulation success, improves real-world success, and reduces training compute by

38\%

. Ablations show that decoupling prevents value hallucination, the one-step estimator preserves candidate-ranking fidelity, and self-guided sampling improves long-horizon execution.

Abstract

Explore similar work

Potential-Guided Flow Matching for Vision-Language-Action Policy Improvement

SeFA-Policy: Fast and Accurate Visuomotor Policy Learning with Selective Flow Alignment

Source-Lifted Flow Matching for Intervenable Multimodal Imitation