AdaFlow: Imitation Learning with Variance-Adaptive Flow-Based Policies

TL;DR

AdaFlow introduces a variance-adaptive flow-based policy model for imitation learning that achieves fast inference and maintains diverse, multi-modal action generation by adaptively solving ODEs during decision-making.

Contribution

It proposes a novel variance-adaptive ODE solver for flow-based policies, enabling rapid inference while preserving diversity in multi-modal decision-making.

Findings

AdaFlow outperforms diffusion-based methods in inference speed.

It maintains high diversity in generated actions.

Automatically reduces to a one-step generator for uni-modal actions.

Abstract

Diffusion-based imitation learning improves Behavioral Cloning (BC) on multi-modal decision-making, but comes at the cost of significantly slower inference due to the recursion in the diffusion process. It urges us to design efficient policy generators while keeping the ability to generate diverse actions. To address this challenge, we propose AdaFlow, an imitation learning framework based on flow-based generative modeling. AdaFlow represents the policy with state-conditioned ordinary differential equations (ODEs), which are known as probability flows. We reveal an intriguing connection between the conditional variance of their training loss and the discretization error of the ODEs. With this insight, we propose a variance-adaptive ODE solver that can adjust its step size in the inference stage, making AdaFlow an adaptive decision-maker, offering rapid inference without sacrificing diversity. Interestingly, it automatically reduces to a one-step generator when the action distribution is uni-modal. Our comprehensive empirical evaluation shows that AdaFlow achieves high performance with fast inference speed.