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

TL;DR

This paper introduces a comprehensive taxonomy and a modular JAX-based framework for reinforcement learning with diffusion policies, along with benchmarks to guide future research and practical applications.

Contribution

It provides a unified taxonomy, an open-source codebase for efficient training, and systematic benchmarks for diffusion-based RL methods.

Findings

Unified taxonomy for RL with diffusion policies

Open-source, high-throughput JAX-based codebase

Benchmark results across multiple simulation environments

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.