# DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching

arXiv: https://arxiv.org/abs/2402.02439 · HTML: https://mlsanity.com/paper/2402.02439/

## Key facts

- Cited by 35 papers in the map (data through 7 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 27 of its 31 references that are in the map.
- First posted on arXiv 4 Feb 2024.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [A Survey on Offline Reinforcement Learning: Taxonomy, Review, and Open Problems](https://mlsanity.com/paper/2203.01387/) (2022)
- [Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learning](https://mlsanity.com/paper/2002.08396/) (2020)
- [Offline RL Without Off-Policy Evaluation](https://mlsanity.com/paper/2106.08909/) (2021)
- [AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners](https://mlsanity.com/paper/2302.01877/) (2023)
- [Deployment-Efficient Reinforcement Learning via Model-Based Offline Optimization](https://mlsanity.com/paper/2006.03647/) (2020)
- [CORL: Research-oriented Deep Offline Reinforcement Learning Library](https://mlsanity.com/paper/2210.07105/) (2022)
- [Synthetic Experience Replay](https://mlsanity.com/paper/2303.06614/) (2023)
- [Human-centric Dialog Training via Offline Reinforcement Learning](https://mlsanity.com/paper/2010.05848/) (2020)
- [Offline Reinforcement Learning with Reverse Model-based Imagination](https://mlsanity.com/paper/2110.00188/) (2021)
- [Bootstrapped Transformer for Offline Reinforcement Learning](https://mlsanity.com/paper/2206.08569/) (2022)
- [Offline Reinforcement Learning for Autonomous Driving with Safety and Exploration Enhancement](https://mlsanity.com/paper/2110.07067/) (2021)
- [Double Check Your State Before Trusting It: Confidence-Aware Bidirectional Offline Model-Based Imagination](https://mlsanity.com/paper/2206.07989/) (2022)
- [Semi-Markov Offline Reinforcement Learning for Healthcare](https://mlsanity.com/paper/2203.09365/) (2022)
- [Uncertainty-driven Trajectory Truncation for Data Augmentation in Offline Reinforcement Learning](https://mlsanity.com/paper/2304.04660/) (2023)

## Built on by (most cited first, 30 of 35)

- [Generative Trajectory Stitching through Diffusion Composition](https://mlsanity.com/paper/2503.05153/) (2025)
- [Graph-Assisted Stitching for Offline Hierarchical Reinforcement Learning](https://mlsanity.com/paper/2506.07744/) (2025)
- [State-Covering Trajectory Stitching for Diffusion Planners](https://mlsanity.com/paper/2506.00895/) (2025)
- [Local Manifold Approximation and Projection for Manifold-Aware Diffusion Planning](https://mlsanity.com/paper/2506.00867/) (2025)
- [DmC: Nearest Neighbor Guidance Diffusion Model for Offline Cross-domain Reinforcement Learning](https://mlsanity.com/paper/2507.20499/) (2025)
- [Zero-Shot Trajectory Planning for Signal Temporal Logic Tasks](https://mlsanity.com/paper/2501.13457/) (2025)
- [Goal-Conditioned Data Augmentation for Offline Reinforcement Learning](https://mlsanity.com/paper/2412.20519/) (2024)
- [BiTrajDiff: Bidirectional Trajectory Generation with Diffusion Models for Offline Reinforcement Learning](https://mlsanity.com/paper/2506.05762/) (2025)
- [Return-Aligned Decision Transformer](https://mlsanity.com/paper/2402.03923/) (2024)
- [Is Temporal Difference Learning the Gold Standard for Stitching in RL?](https://mlsanity.com/paper/2510.21995/) (2025)
- [RedFlow: Redirect Failure into Action-Level Corrections for Flow-matching VLA Policy](https://mlsanity.com/paper/2607.27782/) (2026)
- [Enhancing Decision Transformer with Diffusion-Based Trajectory Branch Generation](https://mlsanity.com/paper/2411.11327/) (2024)
- [Offline Reinforcement Learning via Inverse Optimization](https://mlsanity.com/paper/2502.20030/) (2025)
- [STITCH-OPE: Trajectory Stitching with Guided Diffusion for Off-Policy Evaluation](https://mlsanity.com/paper/2505.20781/) (2025)
- [Compositional Monte Carlo Tree Diffusion for Extendable Planning](https://mlsanity.com/paper/2510.21361/) (2025)
- [PROF: An LLM-based Reward Code Preference Optimization Framework for Offline Imitation Learning](https://mlsanity.com/paper/2511.13765/) (2025)
- [Refining Compositional Diffusion for Reliable Long-Horizon Planning](https://mlsanity.com/paper/2605.03075/) (2026)
- [Offline Trajectory Optimization for Offline Reinforcement Learning](https://mlsanity.com/paper/2404.10393/) (2024)
- [DRDT3: Diffusion-Refined Decision Test-Time Training Model](https://mlsanity.com/paper/2501.06718/) (2025)
- [Extendable Planning via Multiscale Diffusion](https://mlsanity.com/paper/2503.20102/) (2025)
- [Treatment Stitching with Schr"odinger Bridge for Enhancing Offline Reinforcement Learning in Adaptive Treatment Strategies](https://mlsanity.com/paper/2511.12075/) (2025)
- [Compositional Transduction with Latent Analogies for Offline Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2605.20609/) (2026)
- [From Incomplete Coarse-Grained to Complete Fine-Grained: A Two-Stage Framework for Spatiotemporal Data Reconstruction](https://mlsanity.com/paper/2410.05323/) (2024)
- [Model-Based Offline Reinforcement Learning with Reliability-Guaranteed Sequence Modeling](https://mlsanity.com/paper/2502.06491/) (2025)
- [Improving Trajectory Stitching with Flow Models](https://mlsanity.com/paper/2505.07802/) (2025)
- [What Do You Need for Compositional Generalization in Diffusion Planning?](https://mlsanity.com/paper/2505.18083/) (2025)
- [Decision Flow Policy Optimization](https://mlsanity.com/paper/2505.20350/) (2025)
- [RAD: Retrieval High-quality Demonstrations to Enhance Decision-making](https://mlsanity.com/paper/2507.15356/) (2025)
- [Mildly Conservative Regularized Evaluation for Offline Reinforcement Learning](https://mlsanity.com/paper/2508.05960/) (2025)
- [Target-Aligned Fusion for Decision-Sequence Learning under Dynamics Shift](https://mlsanity.com/paper/2511.09173/) (2025)

## Abstract (arXiv)

In offline reinforcement learning (RL), the performance of the learned policy highly depends on the quality of offline datasets. However, in many cases, the offline dataset contains very limited optimal trajectories, which poses a challenge for offline RL algorithms as agents must acquire the ability to transit to high-reward regions. To address this issue, we introduce Diffusion-based Trajectory Stitching (DiffStitch), a novel diffusion-based data augmentation pipeline that systematically generates stitching transitions between trajectories. DiffStitch effectively connects low-reward trajectories with high-reward trajectories, forming globally optimal trajectories to address the challenges faced by offline RL algorithms. Empirical experiments conducted on D4RL datasets demonstrate the effectiveness of DiffStitch across RL methodologies. Notably, DiffStitch demonstrates substantial enhancements in the performance of one-step methods (IQL), imitation learning methods (TD3+BC), and trajectory optimization methods (DT).