Flow Q-Learning

TL;DR

Flow Q-Learning (FQL) introduces a flow-matching policy for offline RL that models complex action distributions efficiently, avoiding recursive backpropagation and iterative action generation, leading to strong performance on diverse tasks.

Contribution

FQL proposes a novel approach that trains an expressive one-step policy with RL to improve offline RL performance without recursive backpropagation.

Findings

Achieves strong results on 73 challenging offline RL tasks.

Effectively models complex action distributions with flow-matching.

Avoids unstable recursive backpropagation during training.

Abstract

We present flow Q-learning (FQL), a simple and performant offline reinforcement learning (RL) method that leverages an expressive flow-matching policy to model arbitrarily complex action distributions in data. Training a flow policy with RL is a tricky problem, due to the iterative nature of the action generation process. We address this challenge by training an expressive one-step policy with RL, rather than directly guiding an iterative flow policy to maximize values. This way, we can completely avoid unstable recursive backpropagation, eliminate costly iterative action generation at test time, yet still mostly maintain expressivity. We experimentally show that FQL leads to strong performance across 73 challenging state- and pixel-based OGBench and D4RL tasks in offline RL and offline-to-online RL. Project page: https://seohong.me/projects/fql/