English

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

Machine Learning 2026-02-09 v3 Artificial Intelligence

Abstract

Although diffusion models have achieved strong results in decision-making tasks, their slow inference speed remains a key limitation. While consistency models offer a potential solution, existing applications to decision-making either struggle with suboptimal demonstrations under behavior cloning or rely on complex concurrent training of multiple networks under the actor-critic framework. 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 achieves single-step sampling while generating higher-reward action trajectories through decoupled training and noise-free reward signals. Empirical evaluations on the Gym MuJoCo, FrankaKitchen, and long horizon planning benchmarks demonstrate that our approach can achieve a 9.7% improvement over previous state-of-the-art while offering up to 142x speedup over diffusion counterparts in inference time.

Keywords

Cite

@article{arxiv.2506.07822,
  title  = {Accelerating Diffusion Planners in Offline RL via Reward-Aware Consistency Trajectory Distillation},
  author = {Xintong Duan and Yutong He and Fahim Tajwar and Ruslan Salakhutdinov and J. Zico Kolter and Jeff Schneider},
  journal= {arXiv preprint arXiv:2506.07822},
  year   = {2026}
}
R2 v1 2026-07-01T03:07:09.057Z