PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow Q-Learning

Seohong Park, Qiyang Li, Sergey Levine

arXiv:2502.02538v2cs.LGcs.AI

TL;DR

Offline RL needs expressive policies for complex dataset action distributions, but RL training for iterative flow policies is difficult. FQL instead applies value maximization to a distilled one-step policy while retaining a BC-trained flow policy, and reports strong performance across 73 OGBench and D4RL tasks in offline and offline-to-online RL.

  • Problem

    Iterative flow and diffusion policies complicate value maximization because policy extraction lacks a straightforward training procedure and can be unstable or costly.

  • Method

    FQL trains a flow policy with behavioral cloning and a separate expressive one-step policy with value maximization and distillation from the flow policy.

  • Results

    FQL shows strong performance across 73 diverse state- and pixel-based OGBench and D4RL tasks in offline and offline-to-online RL.

  • Takeaways & Limitations

    The one-step output avoids backpropagation through time and costly iterative flow generation at evaluation while leveraging the flow policy’s expressivity.

  • Takeaways & Limitations

    Weighted behavioral cloning can perform subpar, especially on complex tasks, because it uses few effective samples and has limited expressivity.

Abstract

from arXiv · show

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/

1. Introduction

Offline RL must optimize returns while staying close to increasingly complex, multimodal dataset behavior, motivating expressive policies. FQL addresses the difficulty of applying RL to iterative flow policies by training a separate one-step policy for value maximization and distillation.

  • Offline RL: Offline RL learns from fixed datasets while maximizing returns under the constraint of remaining within the dataset’s state-action distribution.The setting avoids costly environment interactions and also includes offline-to-online fine-tuning with limited online data.
  • Motivation: Growing and more diverse datasets create complex, multimodal behavioral distributions that require expressive policy classes for accurate behavioral constraints.Flow matching is proposed as a way to model arbitrarily complex action distributions.
  • Challenge: Iterative flow and diffusion policies are difficult to train with learned values because their action-generation procedures complicate policy extraction.Prior extraction schemes can suffer from unstable backpropagation through time, limited effective samples, or high computational cost.
  • Flow Q-learning: FQL trains the iterative flow policy only with behavioral cloning, while a separate expressive one-step policy maximizes values and distills from the flow policy.The resulting one-step policy avoids steering the iterative process and is deployed at evaluation time.
  • Results: FQL is reported to outperform Gaussian- and diffusion-policy offline RL methods on complex tasks and to support direct offline-to-online fine-tuning without iterative flow steps at test time.The evaluation covers 73 state- and pixel-based tasks from OGBench and D4RL.
  • Flow policies: Flow matching uses deterministic ODE-based generation, with a state-dependent vector field mapping noise to actions through numerical integration.The paper uses linear paths, uniform time sampling, and the Euler method for inference.

3. Flow Q-Learning

FQL separates behavioral modeling from value maximization: a flow policy is trained with behavioral cloning, while a one-step policy is trained to maximize values and distill the flow policy. This avoids backpropagation through time and iterative flow generation at evaluation while retaining expressive action modeling.

  • Flow Q-Learning: FQL uses an expressive flow-matching policy to model complex and multimodal behavioral action distributions.The flow policy is trained only with behavioral cloning to provide the behavioral reference.
  • Flow Q-Learning: Naïvely maximizing values with a flow policy requires backpropagation through time because numerical ODE solvers are recursive.This procedure is described as unstable and costly in practice.
  • Flow Q-Learning: One-step guidance trains a separate policy to maximize the value function while distilling outputs from the full behavioral-cloning flow policy.The one-step policy learns a direct mapping from noise z to the flow policy’s output action.
  • Flow Q-Learning: The one-step actor objective combines a Q loss with a distillation-based behavioral regularizer weighted by α.The critic is trained using the one-step policy, while the flow policy remains behavior-cloned.
  • Flow Q-Learning: The deployed one-step policy avoids iterative flow steps at test time while maintaining most of the full flow model’s expressivity.FQL also avoids unstable recursive gradients and can be implemented on top of a standard actor-critic framework with one major hyperparameter α.
  • Flow Q-Learning: The distillation loss upper-bounds the squared 2-Wasserstein distance between the one-step policy and the behavioral flow policy.This gives the behavioral regularizer awareness of the Euclidean metric in action space.

4. Prior Work

Prior work extracts RL policies from iterative diffusion or flow models through weighted behavioral cloning, reparameterized gradients, or rejection sampling. FQL instead trains the flow policy with behavioral cloning and guides a distilled one-step policy with value maximization, avoiding recursive backpropagation.

  • Policy extraction strategies: Weighted behavioral cloning selectively clones high-advantage transitions but often performs poorly because few samples receive effective weight and expressivity is limited.The paper reports especially subpar performance on complex tasks.
  • Policy extraction strategies: Reparameterized policy gradients maximize Q with gradients through the generative policy while regularizing it with a flow or diffusion loss.For iterative generative models, prior methods implement this through backpropagation through time, which can create stability issues and suboptimal performance.
  • Policy extraction strategies: Rejection sampling selects the highest-value action among N samples from a fixed behavioral-cloning policy.This approach is simple and stable but requires querying the policy and value function N times per environment step.
  • FQL: FQL trains the iterative flow policy only with behavioral cloning and trains a separate one-step policy to maximize values while distilling from the flow model.The resulting policy avoids recursive backpropagation and costly iterative flow steps during evaluation.
  • FQL: FQL uses reparameterized policy gradients without backpropagation through time and reports better performance than prior distillation-based and other policy-extraction methods.The comparison includes Consistency-AC and SRPO.

5. Experiments

The experiments evaluate FQL against Gaussian, diffusion, and flow-policy baselines on diverse OGBench tasks and selected D4RL tasks. The setup uses controlled training comparisons, while offline-to-online experiments test fine-tuning against both offline RL and online-RL methods.

  • Benchmarks: The evaluation covers OGBench’s diverse robotic locomotion and manipulation tasks with state and pixel observations, alongside selected D4RL tasks.OGBench tasks are generally more challenging than standard D4RL tasks.
  • Evaluation protocol: Table 2 aggregates results across 73 diverse, challenging offline RL tasks, reporting averages over 8 seeds or 4 seeds for pixel-based tasks.Cells without ± signs are taken from prior work.
  • Methods: The offline comparison includes Gaussian baselines BC, IQL, and ReBRAC; diffusion baselines IDQL, SRPO, and CAC; and flow variants such as FAWAC and FBRAC.CAC is the closest diffusion baseline because it also uses behavior-regularized actor-critic distillation, while still using backpropagation through time.
  • Offline-to-online evaluation: Offline-to-online experiments compare FQL with fine-tunable IQL, ReBRAC, and IFQL, plus Cal-QL and RLPD, which are designed for data-driven online RL.FQL can be fine-tuned directly with online rollouts without algorithmic changes.
  • Evaluation protocol: Offline results are evaluated after a fixed number of gradient steps rather than at the best evaluation epoch, with individually tuned baselines and matched network size and discount factor.The study uses 8 seeds for state-based tasks and 4 seeds for pixel-based tasks.

Q: How good is FQL for offline RL?

FQL achieves the best or near-best performance on most evaluated offline RL tasks, with particularly strong results on complex manipulation environments. It consistently outperforms its closest diffusion baseline and often surpasses the closest Gaussian baseline on multimodal manipulation tasks.

  • Overall performance: FQL achieves the best or near-best performance on most of the 73 state- or pixel-based offline RL tasks.The benchmark spans robotic locomotion and manipulation.
  • Baseline comparisons: FQL consistently outperforms its closest diffusion baseline, CAC, across the aggregated offline RL results.The comparison is reported as part of the 73-task benchmark.
  • Baseline comparisons: FQL often significantly outperforms its closest Gaussian baseline, ReBRAC, especially on manipulation tasks with highly multimodal action distributions.These results support the benefit of the method on complex manipulation environments.
  • D4RL result: 84% is FQL’s reported performance on D4RL antmaze-large-play, described as one of the benchmark’s hardest tasks.The value is reported in Table 3.

Q: Can’t I just use existing policy extraction schemes?

Existing extraction schemes can be used, but policy extraction alone substantially affects offline RL performance, and FQL’s one-step guidance performs best among the compared flow-based methods.

  • Previous policy extraction schemes generally lead to much worse performance than FQL.
  • Policy extraction alone can significantly affect performance across 50 state-based OGBench tasks.
  • FQL’s one-step guidance significantly outperforms FAWAC, FBRAC, and IFQL on aggregated offline RL performance.These baselines use different policy extraction mechanisms and share the same codebase, architecture, and per-environment tuning.

Q: Can FQL be fine-tuned with online rollouts?

FQL can be directly fine-tuned with online rollouts by adding new transitions and continuing the same training objective, achieving the best reported fine-tuning performance on the evaluated tasks.

  • FQL can be directly fine-tuned without modifications and often significantly outperforms previous methods.
  • FQL fine-tuning adds online transitions to the dataset while continuing to train all networks with the offline objective.
  • FQL achieves the best fine-tuning performance across 5 representative OGBench tasks and 10 D4RL antmaze and adroit tasks.The comparison includes previous offline RL methods and methods specifically designed for online fine-tuning.
  • The BC coefficient is the most important hyperparameter and needs environment-specific tuning based on dataset suboptimality.

Q: Do I need to tune flow-related hyperparameters?

Flow-related hyperparameters generally require little tuning: uniform time sampling is often sufficient, and performance is robust to flow-step counts that are not too small.

  • Performance is generally robust to flow-related hyperparameters, and uniform time sampling is often sufficient.
  • Performance is generally robust to the number of flow steps as long as the count is not too small.
  • FQL is only slightly slower than Gaussian policy-based methods and faster than most flow-based baselines in training and inference costs.
  • One-step guidance is the main contribution, while applying it to diffusion policies would require appropriate modifications to convert SDEs to ODEs.
  • FQL generally achieves the best performance across challenging robotic, offline, offline-to-online, state-based, and pixel-based settings.
  • The method’s simplicity is presented as important because offline RL is sensitive to implementation details.

A. Limitations

FQL’s main tuning burden is the BC coefficient, while several flow and target-value choices are comparatively robust; implementation details and training-time ODE solving remain relevant scope considerations.

  • Limitations: FQL requires numerically solving ODEs during training to minimize the distillation loss.The authors identify this as a potential limitation, though it was not a significant speed bottleneck in their experiments.
  • Limitations: FQL lacks a built-in exploration mechanism and does not achieve the best online fine-tuning on puzzle-4x4.
  • Implementation details: Pixel-based tasks use frame stacking with three images, which is important for some cube and puzzle tasks.
  • BC coefficient α: The BC coefficient α must be tuned for each task according to dataset suboptimality.
  • Target value aggregation methods: Not using clipped double Q-learning often performs better, although some tasks use minimum target aggregation.
  • Flow steps: Performance is generally robust to flow-step counts as long as the number of steps is not too small.The default is 10 Euler steps.
  • Time distributions for flow matching: Performance is generally robust to the time distribution, and uniform sampling is often enough to achieve the best performance.

D. Additional Results

FQL provides strong performance across the reported offline and offline-to-online results while maintaining favorable training and inference speed. The experiments use standardized evaluation protocols and report results across 73 tasks.

  • Run time comparison: FQL is faster than most other flow-based methods in both training and inference, while being only slightly slower than Gaussian policy-based methods.Runtime measurements use one A5000 GPU, the same machine, and averages over 8 seeds.
  • Full results: The full offline RL evaluation covers 73 OGBench and D4RL tasks, with results averaged over 8 seeds or 4 seeds for pixel-based tasks.Tables mark values at or above 95% of the best performance in bold and report standard deviations after ±.
  • Offline-to-online results: Offline-to-online evaluation begins online fine-tuning at 1M steps, and results are averaged over 8 seeds unless otherwise noted.

E.1. Environments, Tasks, and Datasets

The evaluation spans diverse state- and pixel-based OGBench tasks plus 18 D4RL tasks, covering navigation, locomotion, dexterous manipulation, and long-horizon object control. OGBench uses single-task variants with semi-sparse rewards and multiple evaluation goals.

  • OGBench: OGBench evaluation uses 10 environments, 50 state-based tasks, and 5 pixel-based tasks, with five evaluation goals per environment.The benchmark uses single-task variants to evaluate standard reward-maximizing offline RL methods.
  • OGBench: OGBench single-task variants label dataset transitions with semi-sparse rewards based on the number of remaining subtasks.Locomotion rewards are −1 or 0, while manipulation rewards range from −ntask to 0, with up to 16 subtasks.
  • Task diversity: The selected OGBench environments cover maze navigation, quadrupedal soccer, dexterous manipulation, and long-horizon control of multiple objects.Maze agents have 8 or 21 degrees of freedom, and scene tasks involve up to 8 objects.
  • D4RL: The experiments additionally use 18 relatively hard D4RL tasks, comprising 6 antmaze and 12 adroit tasks.D4RL antmaze uses binary success rates, while adroit uses normalized returns; adroit has a 24-D action space.

E.2. Methods and Hyperparameters

The study compares FQL with 11 prior offline and offline-to-online RL approaches using shared implementations and standardized defaults where possible. Prior methods and FQL receive environment-specific tuning for relevant hyperparameters.

  • Comparison protocol: The evaluation considers 11 previous offline and offline-to-online RL approaches, generally using the same default hyperparameters, architectures, and codebase.Method-specific hyperparameters are individually tuned for each environment, with OGBench tuning performed on each environment’s default task.
  • Baseline methods: BC uses a Gaussian policy with unit standard deviation, while IQL, ReBRAC, and Cal-QL receive searches over their principal method-specific coefficients.IQL fixes the expectile at 0.9 and tunes its AWR inverse temperature per environment; ReBRAC tunes actor and critic BC coefficients.
  • Baseline methods: IDQL, SRPO, and Consistency-AC use their official implementations or reimplementations with selected architectures and tuned value- or Q-learning coefficients.The authors report unstable training for 4-layer IDQL value networks and unsuccessful initial SRPO results with its official implementation.
  • Flow-based baselines: FAWR, FBRAC, and IFQL share FQL’s flow-matching implementation and use individually tuned policy-extraction hyperparameters for each environment.FAWR’s inverse temperature and the flow-based methods’ test-time settings are tuned for comparison.
  • FQL configuration: FQL hyperparameters are documented in Table 5, while task-specific offline and offline-to-online settings are provided in Tables 6 and 7.
Loading 2502.02538v2…