NEEDLEwork: Offline Rewriting of Robot Data with Verified Local Stitches

What could stitching offer?

A within-trajectory shortcut (S1) bypasses an inefficient detour inside one demonstration. A cross-trajectory transplant (S2) combines complementary behaviors from two successful demonstrations that were never recorded together. A failure funnel (S3) connects a state on a failed attempt to a successful continuation. The animations illustrate these opportunities; realizing them requires finding connections that are both useful and feasible.

Which connections would improve the data?

Connecting two nearby observations may simply add a detour. NEEDLE therefore looks for targets with better recorded continuations. From a successful demonstration, the connection and target continuation must together take fewer steps than the original route. From a failed attempt, the target must offer a successful continuation where none was recorded. Among accepted bridges from the same source, NEEDLE keeps the one with the shortest combined bridge and remaining successful continuation. For successful sources, this also maximizes the steps saved. These criteria identify useful destinations; they do not establish that the robot can reach them.

How should we best extract useful information from existing demonstrations for stitching?

Key insight. Rather than learning a globally accurate dynamics model or task-value function and asking it to reason over long trajectories, NEEDLE reduces stitching in high-dimensional observation spaces to a series of short, locally supervised reachability problems. Short recorded segments already teach both feasibility and duration: they show which observations a sequence of actions connects, and in how many steps, even in an episode that ultimately fails.

How does NEEDLE stitch demonstrations together?

The recordings contain the useful segments, but may contain no actions connecting them. A goal-conditioned inverse dynamics model (IDM) proposes those missing actions from a source observation to a target. Plausible actions can still miss the target, so an action-conditioned verifier, trained on these short recorded segments, assesses whether each prefix reaches it and after how many steps. Its training pairs share one recorded action chunk: the chunk is a positive example for the goal it reaches within the horizon, and a negative example for a goal beyond the horizon or on another trajectory. A negative label says only that this chunk does not reach that goal, not that no action sequence could; finding one that does is the objective of stitching.

From proposal to bridge. For each source–target pair, NEEDLE selects the shortest prefix that passes the verifier’s threshold as a bridge. This requires neither a simulator nor reconstructed images along the bridge.

How can the policy learn from a bridge?

A bridge supplies new actions but no intermediate images. NEEDLE pairs these actions with recorded observations at and before the source, using recorded actions to lead into the bridge. Varying the starting observation lets the policy learn the new continuation from several preceding situations.

The bridge may end slightly away from its target, so the target’s next actions may not apply there. Training therefore stops at the bridge endpoint; the target continuation is learned separately from its own observations. Because a stitch creates an alternative local continuation rather than an exact synthetic trajectory, the original actions at the bridge source and along the skipped segment stay in the training set, at a controlled sampling weight, so the policy keeps supervision for situations it may still encounter.

Does stitching improve task success?

NEEDLE succeeds in 44 of 50 dish trials and 48 of 50 sweater trials. Each policy is evaluated on the same 50 starting states per task.

Does stitching reduce completion time?

Mean completion time among successes can be misleading, since a policy that solves fewer difficult starts can look faster. We therefore also count, for each time limit, the fraction of all 50 trials completed within it, failures included. NEEDLE has the lowest mean on sweater folding. On dish racking, AWR has a slightly lower mean, while NEEDLE succeeds on more trials.

Does the policy repeat less motion, and where do failures stop?

The demonstrations contain hesitations, repeated grasps, and re-attempts. These clips show baselines repeating a motion without completing the task, while NEEDLE completes them. The figure below marks, for every failed trial, the stage the policy never completed.

Does the success gain hold across simulated tasks?

NEEDLE has the highest success rate on all three tasks: on average 8.4 percentage points above DP and 6.9 above the strongest baseline on each task.

Do bridges lead to faster completion?

We test whether individual bridges reduce the remaining steps to completion, and whether the trained policy succeeds within a given step budget.

Can failed demonstrations teach recovery?

A failed trajectory can contain useful intermediate behavior even though it never reaches the goal. A failure funnel connects such a state to a successful continuation, turning a state that had no successful action target into one with explicit supervision toward success.

Does checking the bridges matter?

A useful connection must pass two separate tests. It must be useful: following the target’s recorded actions should give the policy a better route to completion. It must also be feasible: the proposed bridge actions must actually reach the target. Candidate selection estimates usefulness from recorded continuations; the verifier judges feasibility from the proposed actions. We test whether its scores depend on the actions, whether accepted bridges reach their targets in simulation, and whether verification improves policy success.

What if the demonstrations are already good?

On three DexMimicGen tasks, where the demonstrations are near-optimal, NEEDLE matches DP on Threading and improves by 1.3 and 4.0 percentage points on Coffee and Three Piece Assembly. AWR performs best on Threading.

Which kinds of connection help?

NEEDLE uses three kinds of bridge: skipping ahead within a demonstration (S1), connecting two demonstrations (S2), and connecting a failed attempt to a successful demonstration (S3). Each row trains NEEDLE with a subset of them; the full method is S1+S2+S3.

What is the best way to train on stitched data?

A bridge gives the policy a second way to act from a recorded observation. Each row below changes one choice about how that is turned into training examples. Under the chart, each change opens into a short animation of what it does to the windows built from one bridge.

Do the bridge actions themselves matter?

This removes every training example whose window contains bridge actions, while keeping the original demonstrations' examples at the weights NEEDLE gives them. What remains is the recorded data, reweighted by where the bridges were found.