Accelerating Diffusion Models in Offline RL via Reward-Aware Consistency Trajectory Distillation

2025-06-09 · Xintong Duan, Yutong He, Fahim Tajwar, Ruslan Salakhutdinov, J. Zico Kolter, Jeff Schneider

Although diffusion models have achieved strong results in decision-making tasks, their slow inference speed remains a key limitation. While the consistency model offers a potential solution, its applications to decision-making often struggle with suboptimal demonstrations or rely on complex concurrent training of multiple networks. In this work, we propose a novel approach to consistency distillation for offline reinforcement learning that directly incorporates reward optimization into the distillation process. Our method enables single-step generation while maintaining higher performance and simpler training. Empirical evaluations on the Gym MuJoCo benchmarks and long horizon planning demonstrate that our approach can achieve an 8.7% improvement over previous state-of-the-art while offering up to 142x speedup over diffusion counterparts in inference time.

📄 PDF Abstract BibTeX arXiv:2506.07822

Code (0)

등록된 구현이 없습니다.

Tasks

Decision MakingMuJoCoOffline RL

Methods 이 논문이 사용한 방법론

SPEED The monocular depth estimation (MDE) is the task of estimating depth from a single frame. This information is an essential knowledge in many computer vision tasks such as scene…

Diffusion Diffusion models generate samples by gradually removing noise from a signal, and their training objective can be expressed as a reweighted variational lower-bound…

ConsistencyTTA: Accelerating Diffusion-Based Text-to-Audio Generation with Consistency Distillation

2023-09-19 · Yatong Bai, Trung Dang, Dung Tran, Kazuhito Koishida 외

Diffusion models are instrumental in text-to-audio (TTA) generation. Unfortunately, they suffer from slow inference due to an excessive number of queries to the underlying denoising network per generation. To address thi…

AudioCaps Audio Generation Denoising Diversity

Diffusion Modulation via Environment Mechanism Modeling for Planning

2026-02-23 · Hanping Zhang, Yuhong Guo arxiv

Diffusion models have shown promising capabilities in trajectory generation for planning in offline reinforcement learning (RL). However, conventional diffusion-based planning methods often fail to account for the fact t…

Reinforcement Learning

MetaDiffuser: Diffusion Model as Conditional Planner for Offline Meta-RL

2023-05-31 · Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan 외

Recently, diffusion model shines as a promising backbone for the sequence modeling paradigm in offline reinforcement learning(RL). However, these works mostly lack the generalization ability across tasks with reward or d…

MuJoCo Reinforcement Learning (RL)

Consistency Models as a Rich and Efficient Policy Class for Reinforcement Learning

2023-09-29 · Zihan Ding, Chi Jin

Score-based generative models like the diffusion model have been testified to be effective in modeling multi-modal data from image generation to reinforcement learning (RL). However, the inference process of diffusion mo…

Image Generation Offline RL Reinforcement Learning (RL)

Safe Flow Q-Learning: Offline Safe Reinforcement Learning with Reachability-Based Flow Policies

2026-03-16 · Mumuksh Tayal, Manan Tayal, Ravi Prakash arxiv

Offline safe reinforcement learning (RL) seeks reward-maximizing policies from static datasets under strict safety constraints. Existing methods often rely on soft expected-cost objectives or iterative generative inferen…

Reinforcement Learning Offline RL