Research map / Reinforcement learning / Offline RL
DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DiffStitch, a diffusion-based data augmentation pipeline that generates stitching sub-trajectories connecting low-reward and high-reward offline trajectories, improving IQL, TD3+BC and Decision Transformer on D4RL.
Problem
Offline RL performance depends heavily on dataset quality; when offline datasets contain few high-reward trajectories, agents struggle to reach high-reward regions because trajectories are disjoint and distant.
Approach
Estimate the number of steps needed to connect two trajectories by imagining future states from one trajectory's end and matching to the other's start, then use a masked conditional diffusion model to generate stitching states, and wrap them with inverse dynamics and reward models, filtering with a learned dynamics model.
Key results
- DiffStitch with IQL reaches 83.8 locomotion average vs 79.5 original and 83.1 for SER; TD3+BC reaches 80.2 vs 75.4 original.
- On Adroit pen, DiffStitch improves average pen score for IQL to 75.7 (from 62.5) and TD3+BC to 20.6 (from -3.2).
- With only 20% of original data, DiffStitch raises IQL average on the two small-sample tasks to 79.7 vs 70.2 original.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add trajectory stitching augmentation | IQL · MuJoCo locomotion average | normalized score: 79.5 -> 83.8 | Helped |
| Add trajectory stitching augmentation | TD3+BC · MuJoCo locomotion average | normalized score: 75.4 -> 80.2 | Helped |
| Add trajectory stitching augmentation | DT · MuJoCo locomotion average | normalized score: 73.1 -> 76.1 | Helped |
| Add trajectory stitching augmentation | IQL · Adroit pen average | normalized score: 62.5 -> 75.7 | Helped |
| Add trajectory stitching augmentation | TD3+BC · Adroit pen average | normalized score: -3.2 -> 20.6 | Helped |
| Add trajectory stitching augmentation | DT · Adroit pen average | normalized score: 60.3 -> 64.9 | Helped |
| Add trajectory stitching augmentation | TD3+BC · Hopper-Medium-Replay | normalized score: 64.8$\pm$19.2 -> 79.6$\pm$13.5 | Helped |
| Add trajectory stitching augmentation | IQL · HalfCheetah-Medium-Expert, 20% samples | normalized score: 89.3$\pm$4.3 -> 91.5$\pm$1.5 | Helped |
| Add trajectory stitching augmentation | IQL · Hopper-Medium-Replay, 20% samples | normalized score: 51$\pm$3.9 -> 67.8$\pm$13.5 | Helped |
| Add the step-estimation module rather than fixing stitching length at 40 (ablation) | DiffStitch · Hopper | normalized score | Helped |
| Add qualification filtering (ablation) | DiffStitch · Hopper | normalized score | Helped |
| Add augmented data (ablation) | IQL and TD3+BC · Hopper | normalized score | Helped |
Surprising findings
- Augmenting with 0:1 ratio (only augmented data) collapses performance to near zero for all algorithms, indicating heavy reliance on original data.
Limitations
- Choice of which trajectory pairs to stitch is naive (random low-reward and high-reward pairs), left as future work.
- Performance is sensitive to data ratio and qualification threshold delta, which must be tuned per dataset.
- SER and TATU could not be used to augment DT, so DT comparison is limited to original vs DiffStitch.
Experimental setup
Datasets: D4RL HalfCheetah datasets D4RL Hopper datasets D4RL Walker2d datasets D4RL Adroit pen human D4RL Adroit pen cloned
Benchmarks: D4RL
Metrics: normalized average return
Baselines: SER TATU Original (no augmentation)
Builds on: TD3+BC IQL Decision Transformer SER TATU
Closest papers by meaning
- State-Covering Trajectory Stitching for Diffusion Planners (NeurIPS 2025 · similarity 0.91)
- Stitching Sub-trajectories with Conditional Diffusion Model for Goal-Conditioned Offline RL (AAAI 2024 · similarity 0.91)
- Graph-Assisted Stitching for Offline Hierarchical Reinforcement Learning (ICML 2025 · similarity 0.90)
- DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning (NeurIPS 2024 · similarity 0.89)
- Offline RL with Observation Histories: Analyzing and Improving Sample Complexity (ICLR 2024 · similarity 0.89)
- Reasoning with Latent Diffusion in Offline Reinforcement Learning (ICLR 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.