NEEDLEWORK: Offline Rewriting of Robot Data with Verified Local Stitches

2026-10-01

AI Summary

This study addresses the challenge of connection identification and feasibility verification in trajectory stitching within high-dimensional robotic data. To this end, it proposes an offline data augmentation strategy optimized via action bridging, introducing a novel local stitching technique in high-dimensional space that enables sample-efficient reuse without intermediate image generation. Furthermore, the method integrates RGB visual and proprioceptive modalities to reconstruct behaviors, coupled with a sampling ensemble strategy to enhance training efficacy, all achieved without requiring new environment interactions. Evaluated on real-world robotic tasks, the proposed approach yields an average improvement of 21 percentage points in success rate over the strongest baseline, demonstrating its practical effectiveness for robot learning from offline datasets.

Abstract

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/.

Problem

• trajectory stitching
• offline dataset augmentation
• high-dimensional robot data
• robot learning

Innovation

• Offline Data Augmentation
• Trajectory Stitching
• High-dimensional Robot Data
• Action Bridges
• Policy Training