EfficientFlow: Efficient Equivariant Flow Policy Learning for Embodied AI

TL;DR

EfficientFlow introduces an equivariant flow-based policy learning framework that improves data efficiency and inference speed for embodied AI tasks, demonstrating superior performance with limited data across robotic benchmarks.

Contribution

The paper presents a novel equivariant flow matching approach with acceleration regularization, enhancing data efficiency and inference speed in embodied AI policy learning.

Findings

Achieves competitive or superior performance with limited data.

Dramatically faster inference compared to existing methods.

Theoretically proven equivariance leads to better generalization.

Abstract

Generative modeling has recently shown remarkable promise for visuomotor policy learning, enabling flexible and expressive control across diverse embodied AI tasks. However, existing generative policies often struggle with data inefficiency, requiring large-scale demonstrations, and sampling inefficiency, incurring slow action generation during inference. We introduce EfficientFlow, a unified framework for efficient embodied AI with flow-based policy learning. To enhance data efficiency, we bring equivariance into flow matching. We theoretically prove that when using an isotropic Gaussian prior and an equivariant velocity prediction network, the resulting action distribution remains equivariant, leading to improved generalization and substantially reduced data demands. To accelerate sampling, we propose a novel acceleration regularization strategy. As direct computation of acceleration is intractable for marginal flow trajectories, we derive a novel surrogate loss that enables stable and scalable training using only conditional trajectories. Across a wide range of robotic manipulation benchmarks, the proposed algorithm achieves competitive or superior performance under limited data while offering dramatically faster inference. These results highlight EfficientFlow as a powerful and efficient paradigm for high-performance embodied AI.