PaperlayerView exact revision on arXiv

Source-linked AI summary

NEEDLEWORK: Offline Rewriting of Robot Data with Verified Local Stitches

Juntao Ren, Yifan Hou, Shuran Song

arXiv:2610.02339v1cs.ROcs.LG

TL;DR

Robot demonstrations may contain useful behaviors even when individual episodes are inefficient or unsuccessful, but finding useful, feasible stitches in high-dimensional data is difficult. NEEDLE adds locally verified action bridges between recorded observations and incorporates them into policy training while preserving original examples. Across real-robot and simulated manipulation tasks, the paper reports that this formulation can turn inefficient, complementary, and unsuccessful behavior into useful new supervision.

  • Problem

    In high-dimensional robot data, identifying useful trajectory connections and verifying their feasibility is difficult, while many prior stitching methods rely on low-dimensional or privileged state.

    Source: §1 · p. 1

  • Method

    NEEDLE finds useful unconnected segments, proposes short action bridges using a goal-conditioned inverse dynamics model, verifies them with an action-conditioned reachability model, and adds accepted bridges while preserving recorded supervision.

    Source: §1 · p. 2

  • Results

    Across real-robot and simulated manipulation tasks, NEEDLE turns inefficient, complementary, and unsuccessful behavior into useful new supervision.

    Source: §5 · p. 10

  • Takeaways & Limitations

    Improving offline robot data can involve carefully connecting behaviors already present, rather than only selecting which recorded behaviors to learn from.

    Source: §5 · p. 10

    Short bridges may not handle connections requiring sustained contact, substantial object reconfiguration, or feedback during execution.

    Source: §4.6 · pp. 9–10

Abstract

from arXiv · show

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

1 INTRODUCTION

Useful robot behavior may be distributed across imperfect demonstrations, but stitching it is difficult when connections must be both valuable and physically feasible. NEEDLE identifies such connections and adds verified action bridges while retaining recorded supervision.

1 INTRODUCTION

  • Robot demonstrations can contain complementary behaviors whose connection would yield a better demonstration than either episode alone.The laundry-folding example combines reliable sleeve folding from one episode with a clean bottom fold from another.Source: §1 · p. 1
  • Prior stitching methods often assume low-dimensional or privileged state, while mixed-quality-data methods typically filter or reweight behavior without supplying connecting actions.Source: §1 · pp. 1–2
  • NEEDLE identifies useful unconnected segments and proposes verified action bridges, reducing high-dimensional stitching to locally supervised reachability problems.It learns feasibility and duration from recorded action sequences using a goal-conditioned diffusion inverse dynamics model and an action-conditioned verifier.Source: §1 · p. 2
  • NEEDLE pairs accepted bridges with actions leading to their source, exposing synthesized actions from multiple recorded observations without generating intermediate images or dropping original examples.This preserves supervision for situations and behaviors covered by the demonstrations while adding alternative ways to make progress.Source: §1 · p. 2
  • The introduction identifies evaluations on real robots and in simulation covering success, completion efficiency, recovery, and key design decisions.Source: §1 · p. 2

2 RELATED WORK

Prior work stitches datasets through selected connections and learned connectivity, curates behavior by filtering or reweighting, or composes behavior through policies and planning. NEEDLE instead materializes offline action supervision between recorded observations.

Explicit Stitching of Offline Datasets

  • Explicit stitching selects connections using value, return, temporal structure, or coverage, then assesses connectivity with dynamics, successor likelihood, or learned temporal proximity.Some generative approaches also synthesize states or longer trajectories for training.Source: §2 · p. 2

Curation via Filtering, Retrieval, and Reweighting

  • Filtering, retrieval, and reweighting identify useful recorded behavior at different granularities but leave recorded action targets unchanged.NEEDLE instead constructs alternative action sequences connecting recorded observations.Source: §2 · p. 2

Explicit Stitching of Offline Datasets

  • NEEDLE verifies a candidate connection by predicting whether and when proposed actions reach the target, then uses the shortest accepted prefix as the bridge.Source: Figure 2 · p. 3

Implicit Stitching and Compositional Planning

  • Offline RL, goal-conditioned policies, generative planners, graph methods, and policy steering compose behavior through learned values or policies, planning, or deployment-time search.NEEDLE materializes connections offline as action supervision, so its imitation policy needs neither connection search nor verification at deployment.Source: §2 · p. 3

3 HOW NEEDLEWORK WORKS

NEEDLE separates the search for useful stitches from verifying that proposed actions can reach their targets. It trains local reachability models from short recorded segments, then incorporates accepted bridges into policy examples while preserving original supervision.

3 HOW NEEDLEWORK WORKS

  • NEEDLE separates candidate selection for useful connections from action generation and verification of their feasibility.The method uses local reachability supervision rather than relying on long-horizon return estimates or reconstructed intermediate states.Source: §3 · p. 3, §3.3 · p. 4

3.2 WHERE SHOULD WE STITCH?

  • Candidate pairs cover within-trajectory shortcuts, shorter successful continuations from another trajectory, and failure-to-success connections.For successful sources, NEEDLE retains bridges with positive estimated continuation gain; for failed sources, it prefers the shortest combined bridge and successful suffix.Source: §3.2 · pp. 3–4

3.3 TRAINING THE NEEDLE AND THIMBLE

  • Short recorded segments supervise an inverse dynamics model to propose actions and a verifier to predict whether each action prefix reaches its goal.The verifier supports choosing both a connection and its duration, using local reachability rather than long-horizon returns or reconstructed states.Source: §3.3 · p. 4
  • For each candidate pair, the IDM samples action proposals, and the verifier’s prefix scores are thresholded using held-out trajectories with low empirical false-positive rates prioritized.False positives would introduce incorrect action supervision.Source: §3.3 · pp. 4–5
  • Verifier training uses positive same-trajectory reachability examples alongside beyond-horizon and cross-trajectory negative pairs for the same action chunk.A negative cross-trajectory pair does not rule out an alternative action sequence connecting its observations.Source: §3.3 · pp. 4–5

3.4 TRAINING A POLICY FROM STITCHES

  • Bridge training examples begin at preceding recorded observations and combine their original lead-in actions with generated bridge actions, without intermediate images.The action target stops at the bridge endpoint; actions after it are masked rather than concatenated with the target demonstration’s continuation.Source: §3.4 · p. 5, Figure 3 · p. 5
  • Original actions at bridge sources and along bypassed segments remain available, preserving the original supervision alongside alternative bridge actions.The sampling frequency of original and new bridge windows is controlled relative to one another.Source: §3.4 · p. 5

4 EXPERIMENTS

Experiments compare NEEDLE with baselines on real-robot and simulated manipulation tasks, then examine completion budgets, bridge effects, verification, near-optimal demonstrations, and limitations.

4.1 EXPERIMENTAL SETUP

  • Evaluations use real-robot sweater folding and dish racking, suboptimal Robomimic tasks, and near-optimal DexMimicGen tasks, with RGB, proprioception, and episode outcomes.The methods use no per-step rewards, privileged object state, or new environment interaction.Source: §4.1 · p. 6, Figure 4 · p. 6

4.2 DOES NEEDLE IMPROVE POLICY SUCCESS?

  • 24 percentage points on sweater folding and 18 points on dish racking over the strongest baseline; on Robomimic, gains average 8.4 points over standard DP and 6.9 over the strongest competing baseline.Source: §4.2 · pp. 6–7
  • Baseline policies often repeat suboptimal motions cyclically, while NEEDLE has the lowest percentage of cyclic repeats across the two real-robot tasks.Source: §4.2 · pp. 6–7, Figure 5 · p. 6

4.3 WHAT DOES NEEDLE CHANGE ABOUT THE LEARNED BEHAVIOR?

  • On real-robot and simulated tasks, NEEDLE reaches comparable success sooner and completes more trials within fixed budgets; accepted bridges also shift simulated continuations toward fewer remaining steps.The bridge comparison is restricted to pairs where both continuations succeed.Source: §4.3 · p. 7, Figure 7 · p. 7, Figure 8 · p. 8
  • Failure funnels add supervision beyond successful-demonstration coverage, and NEEDLE recovers more often than DP and AWR from sampled failure states across all three evaluated populations.Source: §4.3 · pp. 7–8, Figure 9 · p. 8

4.4 DOES VERIFYING FEASIBILITY MATTER?

  • The verifier distinguishes demonstrated actions from logged substitutes (AUROC 0.85–0.98) and accepted bridges more often land within 5 cm than rejected bridges on all three tasks.The landing-error diagnostic uses arm and object positions in simulation; NEEDLE does not observe object positions.Source: §4.4 · pp. 8–9, Figure 10 · p. 9
  • Without verifier filtering, policy success falls on all three tasks despite holding bridge counts per stage fixed; on Transport it drops from 53.2% to 20.4%.Source: §4.4 · p. 9

4.5 WHAT IF EXISTING DEMONSTRATIONS ARE ALREADY NEAR-OPTIMAL?

  • With near-optimal demonstrations, NEEDLE’s gain is at most four percentage points and it does not reduce success relative to standard BC.The authors identify useful recomposition opportunities as the boundary on what NEEDLE can add.Source: §4.5 · p. 9, Table 1 · p. 9

5 CONCLUSION

NEEDLE composes useful behavior within imperfect demonstrations through offline dataset augmentation. Results across real-robot and simulated manipulation tasks suggest useful supervision can be created by connecting recorded behaviors, not only by selecting among them.

5 CONCLUSION

  • NEEDLE is an offline dataset-augmentation method for composing useful behavior within imperfect robot demonstrations.Source: §5 · p. 10
  • Across real-robot and simulated manipulation tasks, NEEDLE turns inefficient, complementary, and unsuccessful behavior into useful new supervision.Source: §5 · p. 10
  • The results suggest offline robot data can be improved by connecting behaviors already present, rather than only selecting which recorded behaviors to learn from.Source: §5 · p. 10

A DATASETS AND TASK SETUPS

A.1 REAL-ROBOT TASKS

The real-robot evaluation uses bimanual tasks with RGB wrist-camera views and relative proprioception, without a shared coordinate frame. Sweater folding tests coordinated deformable-object manipulation, while dish racking tests precise pickup, handover, reorientation, and insertion.

A.1 REAL-ROBOT TASKS

  • Policies receive RGB wrist-camera views and relative proprioception from each trajectory’s starting position, with no shared-coordinate reference frame or shared absolute proprioception.Source: §A.1 · p. 15
  • Sweater folding requires folding both sleeves toward the center and then folding the bottom upward, testing a longer-horizon sequence on a deformable object.The dataset contains 108 successful and 78 failure trajectories.Source: §A.1 · p. 15
  • Dish racking requires a left-hand pickup, a supported handover and reorientation to the right hand, then precise insertion into the rack.The dataset contains 179 successful and 59 failure trajectories.Source: §A.1 · p. 15
  • All policies are deployed on a bimanual ARX-5 setup using 50 fixed starting states per task, with image overlays to reset objects to precise poses.Source: §A.1 · p. 15

A.2 SIMULATION TASKS

On the real robot, NEEDLE shows less retraced gripper travel than DP and AWR. Simulation analyses indicate that stage contributions vary by task, while datasets, fixed evaluation starts, and baseline tuning provide experimental context.

A.2 SIMULATION TASKS

  • Robomimic uses 40 successful demonstrations per Multi-Human task, split between the most suboptimal and most optimal pools, and adds 150 unsuccessful policy rollouts for methods using failure data.DexMimicGen uses 100 successful demonstrations per task and makes 200 unsuccessful episodes available to methods using failure data; its successful behavior is described as near-optimal.Source: §A.2 · p. 15
  • Simulation checkpoints are evaluated every 25 epochs on the same 50 fixed initial conditions, with mean and standard error of the maximum observed success rate across seeds reported.Source: §A.2 · p. 15
  • The baselines were extensively tuned rather than evaluated only with their released configurations.Source: §A.2 · p. 16

B.1 DOES REWRITING REDUCE RETRACED MOTION ON THE REAL ROBOT?

  • 9.2% of pooled gripper travel revisited a previously exited 2 cm region with NEEDLE, versus 10.8% with DP and 16.7% with AWR.The measure covers 50 evaluation trials per task across two tasks; travel pooling gives longer paths more weight.Source: §B.1 · p. 16, Figure 12 · p. 16

B.2 CAN STITCHING IMPROVE WITH A SUBSET OF STAGES?

  • Adding failure-to-success bridges to S1+S2 raises success from 41.2% to 47.2% on Square and from 45.6% to 53.2% on Transport.Cross-demonstration bridges alone reach 80.4% on Can and 46.4% on Square; no single stage matches the three-stage configuration on all three tasks.Source: §B.2 · p. 16, Figure 13 · p. 16

B.3 WHAT IS THE BEST WAY TO TRAIN ON STITCHED DATA?

Training on bridge actions contributes to downstream success, while the training construction preserves original behavior and limits bridge targets to their endpoints. Ablations compare these choices, including retaining original actions, skipped-segment supervision, and separate learning of post-bridge actions.

B.3 WHAT IS THE BEST WAY TO TRAIN ON STITCHED DATA?

  • A bridge is an alternative action continuation, not a complete synthetic trajectory, because its endpoint may differ from the recorded target or use another coordinate system.Source: §B.3 · p. 17
  • The training construction retains recorded actions alongside bridges and preserves supervision from bypassed segments, rather than removing original examples or simply increasing their sampling weight.Source: §B.3 · p. 17
  • Removing supervision from skipped segments yields 50.0% success on Transport, compared with the 53.2% reference.Source: §B.3 · p. 17
  • Bridge targets stop at the bridge endpoint, with remaining target actions learned separately; appending them immediately reaches 42.0% on Square and 45.6% on Transport, below references of 47.2% and 53.2%.The alternative treats the bridge-to-target connection as exact, a condition that may be unreliable when the bridge lands slightly away from its target.Source: §B.3 · p. 17

B.4 DO THE BRIDGE ACTIONS THEMSELVES MATTER?

  • Removing windows containing bridge actions lowers success from 83.6% to 77.2% on Can, 47.2% to 42.8% on Square, and 53.2% to 45.6% on Transport.The comparison keeps reweighted original windows, indicating that bridge actions account for a large part of the gain over DP rather than reweighting alone.Source: §B.4 · p. 17

B.5 FULL DEXMIMICGEN RESULTS

On DexMimicGen, NEEDLE leads on two of three tasks, while AWR leads on Threading; the appendix also compares adapted baselines and examines bridge-action contributions.

B.5 FULL DEXMIMICGEN RESULTS

  • Figure 15 tests bridge-action contributions by removing training windows containing bridge actions while retaining reweighted original windows.Bars report mean success and standard error over five training seeds; the caption does not state the ablation outcome.Source: Figure 15 · p. 18
  • NEEDLE has the highest mean success on Coffee and Three Piece Assembly, while AWR leads on Threading.The comparison uses 100 successful demonstrations per method, with the same additional 200 failure episodes for methods using unsuccessful data.Source: §B.5 · pp. 17–18, Table 2 · p. 18

C.1 SHARED SETUP

  • Adapting baselines to the shared visual-control setting improves success across tasks over their released-policy implementations.The comparison includes policy-architecture changes and task-specific tuning.Source: §C.1 · p. 18

C.2 BASELINE-SPECIFIC ADAPTATIONS

The appendix adapts several offline-learning baselines to the visual-control setting and reports that adaptation improves success across tasks. The adaptations retain each baseline's core mechanism while changing inputs, architectures, or supervision to fit the experiments.

C.2 BASELINE-SPECIFIC ADAPTATIONS

  • SBR retrieves failure transitions by similarity to successful demonstrations and weights them for training; raw-feature retrieval returns none on Can or Transport.The selected configurations remove the internal validation split, which otherwise excludes four of 40 successful demonstrations; they use autoencoder latents on Can and Transport.Source: §C.2 · p. 18
  • AWR applies advantage weighting from a frozen IQL critic to the shared diffusion actor, with selected actors trained on successful demonstrations only.Selected critics use all episodes on Can and Transport, but successful demonstrations only on Square.Source: §C.2 · pp. 18–19
  • Adapted baseline configurations improve success across tasks over released-policy implementations in the visual-control setting.The comparison includes policy-architecture changes and task-specific tuning.Source: Figure 16 · p. 19
  • IDQL reranks diffusion-policy action samples by critic score; the selected configuration uses 16 candidates, maximum-Q selection, and a success-only actor.The critic checkpoint is selected by highest held-out ranking accuracy; an actor without reranking is also evaluated.Source: §C.2 · p. 19
  • SARM is evaluated without language conditioning, using whole-episode progress or stage-aware progress with proportional temporal boundaries.Selected configurations use stage-aware progress on Can and rewind augmentation on Square and Transport; the boundaries provide temporal structure, not semantic subtask annotations.Source: §C.2 · p. 19
  • MBTS rewrites the dataset by proposing joins between nearby states and accepting them only after dynamics, density, and value checks.For this setting, privileged simulator state is replaced by frozen DINOv3 features and proprioception, and sparse episode success replaces simulator reward.Source: §C.2 · p. 19

D.1 SHARED ARCHITECTURE AND OPTIMIZATION

The appendix specifies shared visual features and training settings, then details the inverse dynamics model, verifier, stitch construction, and compute requirements.

D.1 SHARED ARCHITECTURE AND OPTIMIZATION

  • A frozen DINOv3 encoder supplies visual features to the policy, IDM, and verifier; the verifier retains image-token grids while the policy and IDM pool patches into spatial features.The shared encoder also allows image features to be cached during training.Source: §D.1 · pp. 19–20, Table 3 · p. 20

D.2 INVERSE DYNAMICS MODEL

  • The 24-row action prediction includes one row aligned with the earlier observation, leaving a 23-step bridge from the current source; the policy executes 12 actions before replanning.After an IDM training goal is reached, remaining predicted rows hold the target pose rather than continuing along the demonstration.Source: §D.2 · p. 20, Table 4 · p. 20

D.3 VERIFIER

  • The verifier learns from recorded-action source–goal pairs labeled as connected or not, including negative goals beyond the action horizon or in other trajectories.Sampling controls which examples it sees, while positive and negative loss terms receive equal total weight.Source: §D.3 · p. 20, Table 5 · p. 21

D.4 STITCHING AND TRAINING-WINDOW CONSTRUCTION

  • NEEDLE keeps the shortest bridge-plus-suffix path per source, avoiding source states with multiple competing action targets and, for successful sources, maximizing planned step reduction.Window weights control sampling frequency rather than scaling action loss; task-specific candidate budgets are listed separately.Source: §D.4 · p. 21, Table 6 · p. 21

D.5 COMPUTE RESOURCES

  • Simulation training uses one RTX PRO 6000 GPU unless noted; IDM training takes 2–4 hours, verifier training 6–11 GPU-hours, stitching about 2 hours, and policy training 2–9 hours.IDM and verifier training can run in parallel on separate GPUs; baseline policies cost the same as downstream policy training under the shared recipe.Source: §D.5 · p. 21, Table 7 · p. 22

D.6 EVALUATION AND DIAGNOSTIC PROTOCOLS

The evaluation protocols define how policy success, bridge continuations, and recovery from failure states are measured.

D.6 EVALUATION AND DIAGNOSTIC PROTOCOLS

  • Robomimic policies are evaluated every 25 epochs on 50 paired initial conditions, with success and standard error reported across five seeds.Completion steps average successful rollouts; success-by-budget requires first success by that step, with full budgets of 500 steps for Can and Square and 700 for Transport.Source: §D.6 · p. 22
  • Bridge continuations compare recorded actions with a selected bridge from the same restored source state, then resume the same DP policy from each landing.The comparison uses the fastest successful resume, so fitted completion distributions describe jointly successful resumes.Source: §D.6 · p. 22
  • Failure-state recovery averages success across resumes, sampled reachable-workspace failure states, and five policy seeds.Each task uses 50 shared failure states, 16 resumes per state, and a full task horizon; dropped objects outside the robot's reachable workspace are excluded.Source: §D.6 · p. 22

Related papers

Loading 2610.02339v1…