Diffusion Models as Optimizers for Efficient Planning in Offline RL

TL;DR

The paper addresses inefficiency in diffusion-model–based planning for offline RL by decomposing trajectory sampling into a fast autoregressive feasible-trajectory generator and a diffusion-based trajectory optimizer. The Trajectory Diffuser framework uses a Transformer to initialize trajectories and retains diffusion refinement to ensure quality, guided by returns-to-go and inverse dynamics for execution. On D4RL benchmarks, it achieves substantial speedups in inference while outperforming prior sequence-modeling methods, especially on data-intensive Adroit and Kitchen tasks. This makes diffusion-based planning practical for offline decision-making, enabling quicker planning without sacrificing performance. The work demonstrates a principled trade-off between efficiency and accuracy through a two-stage inference pipeline and targeted conditional guidance.

Abstract

Diffusion models have shown strong competitiveness in offline reinforcement learning tasks by formulating decision-making as sequential generation. However, the practicality of these methods is limited due to the lengthy inference processes they require. In this paper, we address this problem by decomposing the sampling process of diffusion models into two decoupled subprocesses: 1) generating a feasible trajectory, which is a time-consuming process, and 2) optimizing the trajectory. With this decomposition approach, we are able to partially separate efficiency and quality factors, enabling us to simultaneously gain efficiency advantages and ensure quality assurance. We propose the Trajectory Diffuser, which utilizes a faster autoregressive model to handle the generation of feasible trajectories while retaining the trajectory optimization process of diffusion models. This allows us to achieve more efficient planning without sacrificing capability. To evaluate the effectiveness and efficiency of the Trajectory Diffuser, we conduct experiments on the D4RL benchmarks. The results demonstrate that our method achieves $\it 3$-$\it 10 \times$ faster inference speed compared to previous sequence modeling methods, while also outperforming them in terms of overall performance. https://github.com/RenMing-Huang/TrajectoryDiffuser Keywords: Reinforcement Learning and Efficient Planning and Diffusion Model

Figures (8)

• Figure 1: Illustrative Example and Motivation. We visualize the trajectory planning process of a maze task, where the agent needs to navigate to the higher-reward red target. In the first row, we elucidate that the denoising process can be decomposed into two components: Generate Feasible Trajectory and Trajectory Optimization. In the second row, we illustrate how an autoregressive process is employed to expedite the "Generate Feasible Trajectory" phase.
• Figure 2: Decompose the Trajectory Generation Process. The denoising process is divided into two steps. To expedite the inference process, a more efficient autoregressive model is utilized to generate the feasible trajectory. A small portion of the denoising process is retained for trajectory optimization. This decomposition allows us to achieve improvements in both efficiency and performance simultaneously.
• Figure 3: Planning with Trajectory Diffuser. Trajectory Diffuser first generates the future trajectory $\tau$ autoregressively, denoted as $\tau=( S_{t+1}, \dots, S_{t+k})$, by considering the current state $S_t$ and the history states $( S_{t-h}, \dots, S_{t-1})$ combine with returns-to-go $(\hat{R}_{t-h}, \dots, \hat{R}_{t})$. Subsequently, it utilizes a diffusion model to optimize$\tau \rightarrow \tau^*$, and employs inverse dynamics to extract and execute the action $A_t$ that leads to the immediate future state $S_{t+1}$.
• Figure 4: Tasks and Results Overview. (a) We study the following tasks: Gym-locomotion, Franka Kitchen and Adroit-Pen. (b) Compared to the previous naive diffusion model, Trajectory Diffuser achieves remarkable improvements in terms of inference speed. (c) Our method outperforms prior sequence modeling methods such as Decision Diffuser and Behavioral Cloning (BC) method. For performance metrics, we use normalized average returns fu2020d4rl for D4RL tasks.
• Figure A1: Bandit case. (a) Unfolding of the diffusion process. (b) Diffusion models as optimizers to enhance BC-CVAE.