Revisiting Diffusion Q-Learning: From Iterative Denoising to One-Step Action Generation

TL;DR

OFQL addresses the computational bottlenecks of Diffusion Q-Learning by replacing multi-step denoising with one-step action generation via Flow Matching. It introduces an average velocity field and the MeanFlow Identity to enable direct one-step transport from noise to action, eliminating backpropagation-through-time through the diffusion chain. Empirically, OFQL achieves state-of-the-art performance on the D4RL suite and substantially faster training and inference (e.g., training time ≈6.3 hours and inference ≈846 Hz) than prior multi-step and distillation-based approaches. This work suggests that one-step flow policies can deliver both high performance and practical efficiency, with potential for real-time control and broader deployment in latency-sensitive domains.

Abstract

Diffusion Q-Learning (DQL) has established diffusion policies as a high-performing paradigm for offline reinforcement learning, but its reliance on multi-step denoising for action generation renders both training and inference slow and fragile. Existing efforts to accelerate DQL toward one-step denoising typically rely on auxiliary modules or policy distillation, sacrificing either simplicity or performance. It remains unclear whether a one-step policy can be trained directly without such trade-offs. To this end, we introduce One-Step Flow Q-Learning (OFQL), a novel framework that enables effective one-step action generation during both training and inference, without auxiliary modules or distillation. OFQL reformulates the DQL policy within the Flow Matching (FM) paradigm but departs from conventional FM by learning an average velocity field that directly supports accurate one-step action generation. This design removes the need for multi-step denoising and backpropagation-through-time updates, resulting in substantially faster and more robust learning. Extensive experiments on the D4RL benchmark show that OFQL, despite generating actions in a single step, not only significantly reduces computation during both training and inference but also outperforms multi-step DQL by a large margin. Furthermore, OFQL surpasses all other baselines, achieving state-of-the-art performance in D4RL.

Figures (7)

• Figure 1: Performance and decision frequency. Performance (i.e., normalized score) and decision frequency are measured on an A100 GPU and averaged across MuJoCo tasks from D4RL. OFQL achieves both high inference speed and strong performance, clearly outperforming prior baselines.
• Figure 2: Comparison between diffusion and flow matching. (a) Conditional flows arise from different $(\epsilon, x)$ pairs, resulting in varying conditional velocities. (b) Marginal velocity is obtained by averaging over these conditional velocities. (c) Flow paths are inherently curved, but average velocity fields enable direct one-step transport from noise to data. (d) Diffusion paths are also curved but noisy, making one-step denoising challenging. Note that all the velocities exhibit symmetry under time reversal. As the model is trained to parameterize the forward flow (from data to noise), inference inverts this direction to generate samples. Accordingly, for clarity, we plot the negative velocity vector to represent the reverse generation trajectory.
• Figure 3: Training Time ($\downarrow$) and Decision Frequency ($\uparrow$) over one million steps, averaged on MuJoCo tasks. NFE (Number of Function Evaluations) denotes the denoising steps required by a flow/diffusion model to generate one action from pure noise. During training and inference, OFQL uses only one NFE, while DQL requires multiple ones. It is worth noting that for inference, FQL runs with a one-step policy, but training still relies on a multi-step flow policy to construct distillation targets.
• Figure 4: Comparison of distribution modeling capabilities between FM with marginal velocity parameterization (left; evaluated at 1,2,5,10 steps generation) and average velocity parameterization (right; evaluated with one-step generation) on a toy dataset with complex multi-modal structure.
• Figure 5: D4RL scores across HalfCheetah datasets under varying flow ratios.