Revisiting Diffusion Q-Learning: From Iterative Denoising to One-Step Action Generation
This addresses the computational inefficiency in offline reinforcement learning for practitioners, though it is incremental as it builds on existing diffusion policy paradigms.
The paper tackles the slow and fragile action generation in Diffusion Q-Learning by introducing One-Step Flow Q-Learning, which enables one-step action generation without auxiliary modules, resulting in significantly faster training and inference while outperforming multi-step DQL and achieving state-of-the-art performance on the D4RL benchmark.
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.