Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FlowRL: A Taxonomy and Modular Framework for Reinforcement Learning with Diffusion Policies

Chenxiao Gao, Edward Chen, Tianyi Chen, Bo Dai

Abstract

Thanks to their remarkable flexibility, diffusion models and flow models have emerged as promising candidates for policy representation. However, efficient reinforcement learning (RL) upon these policies remains a challenge due to the lack of explicit log-probabilities for vanilla policy gradient estimators. While numerous attempts have been proposed to address this, the field lacks a unified perspective to reconcile these seemingly disparate methods, thus hampering ongoing development. In this paper, we bridge this gap by introducing a comprehensive taxonomy for RL algorithms with diffusion/flow policies. To support reproducibility and agile prototyping, we introduce a modular, JAX-based open-source codebase that leverages JIT-compilation for high-throughput training. Finally, we provide systematic and standardized benchmarks across Gym-Locomotion, DeepMind Control Suite, and IsaacLab, offering a rigorous side-by-side comparison of diffusion-based methods and guidance for practitioners to choose proper algorithms based on the application. Our work establishes a clear foundation for understanding and algorithm design, a high-efficiency toolkit for future research in the field, and an algorithmic guideline for practitioners in generative models and robotics. Our code is available at https://github.com/typoverflow/flow-rl.

FlowRL: A Taxonomy and Modular Framework for Reinforcement Learning with Diffusion Policies

Abstract

Thanks to their remarkable flexibility, diffusion models and flow models have emerged as promising candidates for policy representation. However, efficient reinforcement learning (RL) upon these policies remains a challenge due to the lack of explicit log-probabilities for vanilla policy gradient estimators. While numerous attempts have been proposed to address this, the field lacks a unified perspective to reconcile these seemingly disparate methods, thus hampering ongoing development. In this paper, we bridge this gap by introducing a comprehensive taxonomy for RL algorithms with diffusion/flow policies. To support reproducibility and agile prototyping, we introduce a modular, JAX-based open-source codebase that leverages JIT-compilation for high-throughput training. Finally, we provide systematic and standardized benchmarks across Gym-Locomotion, DeepMind Control Suite, and IsaacLab, offering a rigorous side-by-side comparison of diffusion-based methods and guidance for practitioners to choose proper algorithms based on the application. Our work establishes a clear foundation for understanding and algorithm design, a high-efficiency toolkit for future research in the field, and an algorithmic guideline for practitioners in generative models and robotics. Our code is available at https://github.com/typoverflow/flow-rl.

Paper Structure

This paper contains 40 sections, 19 equations, 9 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Taxonomy of RL with Diffusion Policies
  4. Best-of-N (BoN) Sampling
  5. $Q$-value Guidance
  6. Reparameterization
  7. Weighted Matching
  8. Policy Gradient
  9. Design of FlowRL
  10. Design Principles
  11. Architecture Overview
  12. (1) Network Architecture Layer.
  13. (2) Model Layer.
  14. (3) Algorithm Layer.
  15. (4) Workflow Layer.
  16. ...and 25 more sections

Figures (9)

  • Figure 1: The overview of FlowRL.
  • Figure 2: Representative tasks across three continuous-control suites. Left: MuJoCo Gymnasium. Middle: DeepMind Control Suite. Right: IsaacLab.
  • Figure 3: Performance profile on Gym-Locomotion tasks.
  • Figure 4: Training curves of several off-policy DPRL algorithms across Gym-Locomotion tasks.
  • Figure 5: Training curves of on-policy DPRL algorithms on IsaacLab tasks.
  • ...and 4 more figures