Efficient Planning with Latent Diffusion

TL;DR

This paper introduces LatentDiffuser, a novel framework using score-based diffusion models for continuous latent action space planning in offline reinforcement learning, improving efficiency and flexibility especially in high-dimensional tasks.

Contribution

It proposes a unified approach for continuous latent action space learning and planning, establishing theoretical equivalence with energy-guided sampling and introducing a sequence-level exact sampling method.

Findings

Competitive performance on low-dimensional tasks

Outperforms existing methods on high-dimensional tasks

Efficient and flexible planning in continuous latent spaces

Abstract

Temporal abstraction and efficient planning pose significant challenges in offline reinforcement learning, mainly when dealing with domains that involve temporally extended tasks and delayed sparse rewards. Existing methods typically plan in the raw action space and can be inefficient and inflexible. Latent action spaces offer a more flexible paradigm, capturing only possible actions within the behavior policy support and decoupling the temporal structure between planning and modeling. However, current latent-action-based methods are limited to discrete spaces and require expensive planning. This paper presents a unified framework for continuous latent action space representation learning and planning by leveraging latent, score-based diffusion models. We establish the theoretical equivalence between planning in the latent action space and energy-guided sampling with a pretrained diffusion model and incorporate a novel sequence-level exact sampling method. Our proposed method, $\texttt{LatentDiffuser}$, demonstrates competitive performance on low-dimensional locomotion control tasks and surpasses existing methods in higher-dimensional tasks.