NEEDLEWORK: Offline Rewriting of Robot Data with Verified Local Stitches

Juntao Ren; Yifan Hou; Shuran Song

cs.ROOct 1, 2026

NEEDLEWORK: Offline Rewriting of Robot Data with Verified Local Stitches

Authors: Juntao Ren, Yifan Hou, Shuran Song

Organizations: Stanford University · NVIDIA

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

Figures & tables

Figure 1: Needle produces synthetic bridges within trajectories that bypass inefficient detours, cross-trajectory stitches that discover more successful ways to complete the task, and recover useful behavior from failed episodes.

Figure 2: From candidate connection to verified bridge. Given a source-target pair $(s,g)$ , the inverse dynamics model (IDM) samples $K$ action proposals. The action-conditioned verifier predicts whether and when each proposal reaches $g$ . We select the shortest accepted prefix as a bridge.

Figure 3: Training on a bridge without intermediate observations. Each example pairs a recorded observation with the original actions leading to the bridge (cyan), followed by the generated bridge actions (green). The loss stops when the bridge ends (red). Original demonstration windows remain available as separate training examples (blue).

Figure 4: Real-robot tasks and evaluation. We evaluate Needle on two real-robot bimanual tasks: Sweater Folding and Dish Racking. We show task sequences, reset distributions, success rates, and completion times. Each policy is evaluated 50 times per task at paired starting configurations across all methods. Completion times show mean and standard error over successful trials.

Figure 5: Fewer failures across task stages. Trials completing each successive stage, out of 50 per policy and task, from video review. Endpoint labels give final success counts, and the negative numbers mark each baseline’s largest loss of trials at a single stage.

Figure 6: Policy performance on suboptimal Robomimic demonstrations. (a) Needle achieves higher average success than all baselines. (b) Its successful rollouts take fewer steps on average than DP’s on Can and Square . Bars show the mean and standard error over five training seeds, with 50 held-out initial conditions per seed; completion steps are averaged over successful trials.

Figure 7: Needle finishes more trials within a fixed budget. Labels give final success rates per time limit over the 50 paired trials.

Figure 8: Shorter continuations and higher success within a step budget. For each task, the left panel shows remaining completion steps after executing either the original continuation or an accepted bridge, restricted to pairs where both continuations succeed. On average, bridges lead to shorter remaining continuations. The right panel shows the fraction of matched nominal starts solved within each budget. Needle reaches comparable success sooner and solves more starts by the full horizon. We plot mean and S.E. for five training seeds, each with 50 held-out starts.

Figure 9: Expanded training coverage and recovery from failure states. For each task, the left panel shows successful-demonstration states and starts added by failure funnels, visualized as 90% density regions in a t-SNE embedding of the visual encoder. The right panel measures recovery after resuming from OOD states, where Needle recovers more often than DP and AWR, the non-DP baseline with the highest average success rate across all three tasks on nominal resets. Curves show mean and standard error for five training seeds, each with 50 held-out starts.

Figure 10: Verifier-based selection connects action sensitivity, bridge quality, and policy success. (a) With the source and target fixed, the verifier distinguishes demonstrated actions from logged substitutes on 1,500 held-out pairs per task. (b) Accepted proposals more often land near the target than rejected proposals. Scene gap is the larger of the arm and object position errors, measured only for this diagnostic; Needle itself never observes object positions. (c) Removing verifier filtering reduces downstream success on all three tasks with the IDM and bridge counts per stage held fixed. Bars show success and standard error over five training seeds of 50 evaluations each.

Task	DP	Ours
Coffee	62.0 $\pm 1.2$	63.3 $\pm 0.7$
Three Piece Assembly	51.3 $\pm 0.7$	55.3 $\pm 1.8$
Threading	46.0 $\pm 1.2$	46.0 $\pm 1.2$

Table 1: Policy success with near-optimal DexMimicGen demonstrations. Needle matches or modestly improves on DP across the three tasks. Values are success rates (%), reported as mean and standard error over three training seeds with 50 held-out initial conditions per seed.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 11: Simulation benchmarks. We evaluate dataset rewriting on three Robomimic tasks with suboptimal demonstrations ( Mandlekar et al., 2022 ) and three DexMimicGen tasks with near-optimal demonstrations ( Jiang et al., 2025 ) .

Figure 12: Retraced motion on the real robot. Fraction of gripper travel revisiting a previously exited 2 cm region, over 100 trials per policy (50 per task). Left: pooled over steps. Right: averaged over trials. Bars show the estimate and standard error.

Figure 13: Which connections help the policy? Stage subsets compared with the main-paper DP and Needle references. Bars show mean success and standard error over five training seeds.

Figure 14: How should a policy learn from a bridge? Rows below Needle name alternative ways to construct training examples. Bars show mean success and standard error over five training seeds; DP and Needle are main-paper references.

Figure 15: Do the bridge actions matter? The first row removes every training window that contains bridge actions and keeps the reweighted original windows. Bars show mean success and standard error over five training seeds; DP and Needle are main-paper references.

Task	DP	SBR	AWR	IDQL	SARM	Ours
Coffee	62.0 $\pm 1.2$	60.7 $\pm 2.4$	52.7 $\pm 1.8$	49.3 $\pm 1.8$	39.3 $\pm 1.8$	63.3 $\pm 0.7$
Three Piece Assembly	51.3 $\pm 0.7$	44.7 $\pm 1.8$	47.3 $\pm 1.3$	43.3 $\pm 3.5$	26.7 $\pm 1.8$	55.3 $\pm 1.8$
Threading	46.0 $\pm 1.2$	40.7 $\pm 2.4$	50.0 $\pm 1.2$	40.0 $\pm 1.2$	34.0 $\pm 2.0$	46.0 $\pm 1.2$

Appendix

Table 2: Full baseline comparison on DexMimicGen. Success rates (%) reported as mean and standard error over three training seeds with 50 evaluation initial conditions per seed. Needle has the highest mean success on two tasks; AWR leads on Threading.

Figure 16: Adapting baselines to visual robot control. Hollow bars show released-policy implementations evaluated on cached DINOv3 features and proprioception; filled bars show the adapted and tuned configurations used in Figure 6 . Adaptation improves success across tasks. Error bars show standard errors over three training seeds for original baseline implementations, and five training seeds for our adaptation; zero-height bars indicate measured zero success.

Policy, IDM, and verifier
Setting	Value
Image encoder / resolution	Frozen DINOv3 ViT-B/16 / $224\times 224$
Optimizer / learning rate	AdamW / $10^{-4}$
Weight decay	$10^{-6}$
Adam coefficients / epsilon	$(0.95,0.999)$ / $10^{-8}$
Learning-rate schedule / warmup	Cosine decay / 500 steps

Appendix

Table 3: Shared visual and training settings. The encoder and optimizer settings apply to the policy, IDM, and verifier; the remaining rows specify their scope.

Architecture
Setting	Value
Visual feature pooling	Spatial softmax over a $14\times 14$ patch grid
DDIM sampling steps (model configuration)	50
Maximum executable bridge length	23 actions
Action representation	Absolute pose (simulation); source-relative pose (real robot); 6D rotation and gripper command
U-Net channel widths	512, 1024, 2048

Appendix

Table 4: IDM architecture and hyperparameters. Shared visual and optimizer settings are in Table 3 .

Architecture
Setting	Value
Source and target visual tokens	$7\times 7$ patch grid per camera
Transformer layers / attention heads / hidden width	2 / 4 / 512
Feed-forward width / dropout	4096 / 0
Attention / positional encoding	Causal / rotary
Proprioception / action normalization	Source only / per action dimension

Appendix

Table 5: Verifier architecture and hyperparameters. Sampling and training values describe the Robomimic runs. Shared visual and optimizer settings are in Table 3 .

Setting	Value
Candidate search and acceptance
IDM proposals per pair	8
Bridges retained per source	1
Successful-source preference	Shortest bridge-plus-suffix path; positive steps saved
Failure-source preference	Shortest bridge-plus-suffix path
Relative sampling weights per window

Appendix

Table 6: Stitch generation and training-window construction. Shared settings for the nominal Robomimic runs; task-specific candidate budgets appear in Table 7 . Weights control how often a window is sampled, rather than scaling its action loss.

Observations, data, and evaluation
Setting	Can	Square	Transport
RGB cameras	External + wrist	External + wrist	2 shoulder + 2 wrist
Action dimensions	10	10	20
Successful / unsuccessful episodes	40 / 150	40 / 150	40 / 150
Full evaluation budget	500	500	700
Evaluation starts per seed	50	50	50

Appendix

Table 7: Robomimic task settings. Columns describe the nominal policies used for the main success comparison. Shared training settings appear in Table 3 .

Robot Learning

Robot learning develops methods for robots to acquire skills and adapt from demonstrations, interaction, data, and physical models.

See all 323 papers

RESTA: Time-Scaling Augmentation for Data-Efficient Robot Imitation Learning
Oct 13, 2026Yunpeng Jiang, Jianshu Hu, Han Fang +3Data-Efficient Robot Learning Robot Imitation Learning
Global College, Shanghai Jiao Tong University, Shanghai, China · Shanghai Institute of Physical Intelligence and Robotics (PAIR), Shanghai, China · Duke Kunshan University, Jiangsu, China
ARSTAG: An Agentic Real2Sim2Real System for Task-Specific Robot Data Generation
Sep 21, 2026Bowei Li, Yuner Zhang, Changliu LiuRobotics Simulation Sim-to-Real Transfer
Carnegie Mellon University, Pittsburgh, PA 15213, USA
REDIRECT: A 1% Fix for Bad Robot Habits
Oct 2, 2026Yu Zhang, Jiazhuo Li, Yancong Wei +3Learning from Suboptimal Demonstrations Robot Policy Adaptation
Tsinghua University. · University of Michigan.

Oct 13, 2026J/K move · Enter open · S save

RESTA: Time-Scaling Augmentation for Data-Efficient Robot Imitation Learning

Yunpeng Jiang, Jianshu Hu, Han Fang, Zhiyuan Guan+2

Global College, Shanghai Jiao Tong University, Shanghai, China · Shanghai Institute of Physical Intelligence and Robotics (PAIR), Shanghai, China · Duke Kunshan University, Jiangsu, China

Data-Efficient Robot Learning Robot Imitation Learning Data Augmentation Time Series Data Augmentation Bimanual Robotic Manipulation

Open paper AskPDF arXiv

Robot manipulation policies trained through imitation learning commonly rely on costly demonstrations recorded in one task direction and at one temporal resolution. We propose Reversal and Stride-based Trajectory Augmentation (RESTA), which exploits time-scaling data augmentation to improve data efficiency through trajectory reversal and temporal resampling. Trajectory-reversal augmentation reverses and relabels demonstrations between paired tasks with opposite objectives, reusing recorded rollouts without additional environment interaction. Because contact-rich manipulation is not exactly reversible, we treat reversal as an approximate paired-task prior rather than a physical symmetry. Temporal-stride augmentation varies observation and action sampling strides and conditions the policy on the resulting offsets, enabling fixed or randomized temporal schedules. Experiments across 16 paired bimanual manipulation tasks show that trajectory reversal is particularly effective in the low-data regime and that temporal-stride augmentation improves success under randomized temporal schedules. With 25 and 50 independently collected demonstrations per task, RESTA achieves 78.4% and 85.2% average success, respectively, matching or exceeding baselines trained with twice as many demonstrations (77.3% and 85.4%). These results show that RESTA can improve data efficiency in robot imitation learning.

Abstract

Robot Learning

Explore similar work

RESTA: Time-Scaling Augmentation for Data-Efficient Robot Imitation Learning

ARSTAG: An Agentic Real2Sim2Real System for Task-Specific Robot Data Generation

REDIRECT: A 1% Fix for Bad Robot Habits