# NEEDLEWORK: Offline Rewriting of Robot Data with Verified Local Stitches

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 47 of its 52 references that are in the map.
- First posted on arXiv 1 Oct 2026.

## Builds on (most cited first, 30 of 47)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [DINOv3](https://mlsanity.com/paper/2508.10104/) (2025)
- [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)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Universal Manipulation Interface: In-The-Wild Robot Teaching Without In-The-Wild Robots](https://mlsanity.com/paper/2402.10329/) (2024)
- [IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies](https://mlsanity.com/paper/2304.10573/) (2023)
- [Zero-Shot Robotic Manipulation with Pretrained Image-Editing Diffusion Models](https://mlsanity.com/paper/2310.10639/) (2023)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [Search on the Replay Buffer: Bridging Planning and Reinforcement Learning](https://mlsanity.com/paper/1906.05253/) (2019)
- [HIQL: Offline Goal-Conditioned RL with Latent States as Actions](https://mlsanity.com/paper/2307.11949/) (2023)
- [Q-Transformer: Scalable Offline Reinforcement Learning via Autoregressive Q-Functions](https://mlsanity.com/paper/2309.10150/) (2023)
- [Zero-Shot Visual Imitation](https://mlsanity.com/paper/1804.08606/) (2018)
- [DexMimicGen: Automated Data Generation for Bimanual Dexterous Manipulation via Imitation Learning](https://mlsanity.com/paper/2410.24185/) (2024)
- [All Roads Lead to Likelihood: The Value of Reinforcement Learning in Fine-Tuning](https://mlsanity.com/paper/2503.01067/) (2025)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [From Foresight to Forethought: VLM-In-the-Loop Policy Steering via Latent Alignment](https://mlsanity.com/paper/2502.01828/) (2025)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)
- [SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation](https://mlsanity.com/paper/2509.25358/) (2025)
- [Generative Trajectory Stitching through Diffusion Composition](https://mlsanity.com/paper/2503.05153/) (2025)
- [Robot Data Curation with Mutual Information Estimators](https://mlsanity.com/paper/2502.08623/) (2025)
- [DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching](https://mlsanity.com/paper/2402.02439/) (2024)
- [Diffusion Meets DAgger: Supercharging Eye-in-hand Imitation Learning](https://mlsanity.com/paper/2402.17768/) (2024)
- [Hybrid Inverse Reinforcement Learning](https://mlsanity.com/paper/2402.08848/) (2024)
- [Sparse Graphical Memory for Robust Planning](https://mlsanity.com/paper/2003.06417/) (2020)
- [CCIL: Continuity-based Data Augmentation for Corrective Imitation Learning](https://mlsanity.com/paper/2310.12972/) (2023)
- [DataMIL: Selecting Data for Robot Imitation Learning with Datamodels](https://mlsanity.com/paper/2505.09603/) (2025)

## Abstract (arXiv)

Robot demonstrations may contain useful behavior even when individual episodes are inefficient or unsuccessful. Trajectory stitching offers a way to compose these behaviors into improved training data, but identifying useful connections and verifying their feasibility is difficult in high-dimensional robot data, where many prior methods rely on low-dimensional state representations. We introduce NEEDLE, an offline dataset-augmentation algorithm that addresses these challenges by adding short, verified action bridges between recorded observations in high-dimensional robot demonstrations. First, NEEDLE identifies and creates connections that bypass suboptimal detours, broaden action coverage, and augment the original dataset with failed trajectories, using only RGB images, proprioception, and episode-level outcomes, without new environment interaction or privileged object state. Next, we present a sampling technique that incorporates accepted bridges into policy training without synthesizing intermediate images or discarding the original demonstrations, allowing policies to learn alternative actions while retaining the original dataset's coverage. On real-robot tasks, NEEDLE improves success rate over the strongest baseline on each task by an average of 21 percentage points. Videos and supplementary materials are on https://needle-work.github.io/.