Stitching Sub-Trajectories with Conditional Diffusion Model for Goal-Conditioned Offline RL

TL;DR

This paper tackles offline goal-conditioned RL with sparse rewards and long-horizon planning. It introduces SSD, a value-conditioned diffusion framework that generates sub-trajectories conditioned on the goal and a return-to-go, and stitches them via a multi-step goal chaining strategy. The method employs a Condition-Prompted-Unet architecture to produce realistic trajectories and trains the diffusion model jointly with a relabeled-value function, enabling end-to-end planning without explicit subgoal hierarchies. Empirical results on Maze2D, Multi2D, and Fetch demonstrate state-of-the-art performance and effective stitching of suboptimal sub-trajectories into high-quality plans, with strong practical implications for long-horizon offline RL tasks.

Abstract

Offline Goal-Conditioned Reinforcement Learning (Offline GCRL) is an important problem in RL that focuses on acquiring diverse goal-oriented skills solely from pre-collected behavior datasets. In this setting, the reward feedback is typically absent except when the goal is achieved, which makes it difficult to learn policies especially from a finite dataset of suboptimal behaviors. In addition, realistic scenarios involve long-horizon planning, which necessitates the extraction of useful skills within sub-trajectories. Recently, the conditional diffusion model has been shown to be a promising approach to generate high-quality long-horizon plans for RL. However, their practicality for the goal-conditioned setting is still limited due to a number of technical assumptions made by the methods. In this paper, we propose SSD (Sub-trajectory Stitching with Diffusion), a model-based offline GCRL method that leverages the conditional diffusion model to address these limitations. In summary, we use the diffusion model that generates future plans conditioned on the target goal and value, with the target value estimated from the goal-relabeled offline dataset. We report state-of-the-art performance in the standard benchmark set of GCRL tasks, and demonstrate the capability to successfully stitch the segments of suboptimal trajectories in the offline data to generate high-quality plans.

Figures (13)

• Figure 1: Our approach uses a diffusion model conditioned on the goal and return-to-go that generates sub-trajectories. These sub-trajectories are then be stitched together to achieve the goal.
• Figure 2: Plot of sampled trajectory from Decision Diffuser (DD) ajay2022conditional. It generates unrealistic trajectories that pass through the wall. The blue point is the initial point and the red point is the goal. left: Maze2D-Medium. right: Maze2D-Large.
• Figure 3: Architecture of the proposed Condition-Prompted-Unet model. Trajectory input is tokenized into state-action pairs. The goal and value are included as prompt tokens into the transformer layers, which constitute the Unet structure. Downsampling and Upsampling layers are integrated with each transformer block. The process is repeated for each reverse step, using the given goal and the trained action value as conditions.
• Figure 4: Layouts of Maze2d environment. left: Umaze, middle: Medium, and right: Large
• Figure 5: Comparision of behaviors of diffusion based methods. (a), (b) and (c) are result of Diffuser janner2022planning, DD ajay2022conditional, and SSD (Ours) respectively. The blue point is the initial point and the red point is the goal. All of them share the same initial point and the same goal.