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

2026-10-05

AI Summary

This study addresses the challenges of out-of-distribution extrapolation failure and multimodal behavior collapse for unseen goals in goal-conditioned imitation learning. To this end, it proposes a bilinear flow strategy that pioneers bilinear conditional flows to model the joint influence of anchors and residuals. By integrating transductive retrieval, the method maps unseen goal-state pairs to known anchor samples, complemented by an anchor selection algorithm with provable error bounds. Experimental results demonstrate that this approach achieves a success rate 2.63 times that of baselines in simulation tasks and yields a 32% improvement in real-world scenarios. Furthermore, it provides a diagnostic tool for evaluating extrapolation capabilities prior to deployment, thereby enabling robust out-of-distribution generalization.

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.

Problem

• Goal-conditioned imitation learning
• Distributional extrapolation
• Flow matching
• Out-of-distribution generalization
• Visuomotor policy

Innovation

• Bilinear Flow Policy
• Distributional Extrapolation
• Goal-Conditioned Imitation Learning
• Flow Matching
• Anchor Selection