Offline Reinforcement Learning with Discrete Diffusion Skills

TL;DR

This work addresses long-horizon, sparse-reward offline RL by introducing Discrete Diffusion Skills (DDS), a hierarchical framework that learns a compact discrete skill space with a Transformer encoder and a diffusion-based decoder. A high-level policy trained via offline Q-learning methods selects skills, while the diffusion decoder generates actions conditioned on the chosen skill and current state. DDS demonstrates strong performance on long-horizon tasks like AntMaze and Kitchen, offers improved interpretability through discrete skills, and enhances online exploration when used with pre-trained skills. The approach leverages relabeled offline datasets for high-level learning and shows robustness to various skill-space configurations, though data quality and excessive skill counts can hinder performance. Overall, DDS provides a scalable pathway to combine discrete skill abstractions with diffusion-based action generation in offline RL, with potential applications in robotics and autonomous systems.

Abstract

Skills have been introduced to offline reinforcement learning (RL) as temporal abstractions to tackle complex, long-horizon tasks, promoting consistent behavior and enabling meaningful exploration. While skills in offline RL are predominantly modeled within a continuous latent space, the potential of discrete skill spaces remains largely underexplored. In this paper, we propose a compact discrete skill space for offline RL tasks supported by state-of-the-art transformer-based encoder and diffusion-based decoder. Coupled with a high-level policy trained via offline RL techniques, our method establishes a hierarchical RL framework where the trained diffusion decoder plays a pivotal role. Empirical evaluations show that the proposed algorithm, Discrete Diffusion Skill (DDS), is a powerful offline RL method. DDS performs competitively on Locomotion and Kitchen tasks and excels on long-horizon tasks, achieving at least a 12 percent improvement on AntMaze-v2 benchmarks compared to existing offline RL approaches. Furthermore, DDS offers improved interpretability, training stability, and online exploration compared to previous skill-based methods.

Figures (5)

• Figure 1: General framework of Discrete Diffusion Skill. The offline pre-training for discrete skills is illustrated on the left, where the transformer encoder, discrete skill codebook, and diffusion decoder are jointly trained to minimize the codebook loss, the commitment loss, and the reconstruction loss. The online inference is shown on the right. Every $H$ steps, the high-level policy, trained using IQL and AWR, is called to select a skill index. This index is then used to retrieve the corresponding skill vector from the codebook. Subsequently, the action is generated at each time-step using the pre-trained diffusion decoder.
• Figure 2: Antmaze experiment results: sub-figure (a), (b) and (c) show the trajectories of different skills in Antmaze-Large starting from different positions and sub-figure (d) shows the skill selected in several successful episodes. Sub-figure (e), (f) and (g) show the trajectories of different skills in Antmaze-Medium starting from different positions and sub-figure (h) shows the skill selected in several successful episodes. Different skills are represented with different colors and starting points are denoted with orange stars.
• Figure 3: The the movement of the hopper using two different skills: the upper graph shows smoother motion, while the lower one exhibits more aggressive movement. Joints are connected between frames to illustrate moving pattern.
• Figure 4: Online training results with and without pre-trained discret and continuous skills in AntMaze-Large-Diverse-v2 and AntMaze-Medium-Diverse-v2
• Figure 5: Ablation study on the skill horizon $H$.