💻 computer science

SI-Diff: A Framework for Learning Search and High-Precision Insertion with a Force-Domain Diffusion Policy

This paper presents SI-Diff, a unified force-domain diffusion policy framework that integrates a novel mode-conditioning mechanism and a search teacher policy to simultaneously learn robust search and high-precision insertion, significantly improving misalignment tolerance and zero-shot transferability compared to existing baselines.

Original authors: Yibo Liu, Stanko Oparnica, Simon Shewchun-Jakaitis, Guoyi Fu, Jie Wang, Jun Yang, Anand Jagannathan, Tony Hong-Yau Lo

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yibo Liu, Stanko Oparnica, Simon Shewchun-Jakaitis, Guoyi Fu, Jie Wang, Jun Yang, Anand Jagannathan, Tony Hong-Yau Lo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to put a key into a lock, but you are wearing thick gloves and you can't see the keyhole very well. You have to feel your way around. If the key is slightly off-center, you might just push it straight down and get stuck. If you know exactly where the hole is, you can slide it in perfectly. But what if you need a robot to do this, and the robot doesn't know exactly where the hole is?

This paper introduces SI-Diff, a new "brain" for robots that solves this exact problem. It teaches a robot how to do two very different things using the same set of instructions: searching for the hole when it's lost, and sliding the object in once it's found.

Here is how it works, broken down with simple analogies:

1. The Problem: Two Different Moves

Usually, robots are taught to do these tasks separately.

The Search: If the robot is lost, it needs to wiggle and sweep the area (like a person feeling for a light switch in the dark).
The Insertion: Once found, it needs to be very delicate, wiggling just enough to slide the object in without getting stuck (like threading a needle).

Most robots have to switch "modes" or load different software to do these two things. The authors wanted to build a single robot brain that can do both without switching gears, just like a human worker who can feel around for a hole and then slide a peg in without thinking about changing tools.

2. The Solution: A "Diffusion" Policy

The team used a type of AI called a Diffusion Policy.

The Analogy: Think of diffusion like a sculptor starting with a block of noisy, random clay and slowly chipping away the noise until a perfect statue emerges.
In the Robot: The robot starts with a random guess of how to move its arm. The AI "denoises" this guess, refining it step-by-step based on what the robot's sensors (touch and force) are telling it, until it finds the perfect movement to either search or insert.

3. The Secret Sauce: The "Mode" Switch

The biggest challenge was teaching the AI to know which move to make (search vs. insert) without changing the code.

The Analogy: Imagine a music player that can play both a "Search Song" and an "Insertion Song." Usually, you'd need two different players. SI-Diff uses a Mode Prompt, which is like a "track switch" button.
How it works: The robot is told, "Right now, you are in Search Mode" or "Right now, you are in Insertion Mode." This tiny instruction tells the AI to look at the same sensor data but interpret it differently. In Search Mode, it looks for patterns that mean "keep looking." In Insertion Mode, it looks for patterns that mean "wiggle gently to fit."

4. The Teacher: Learning from a Smart Guide

Robots learn best by watching experts. The authors created a "Teacher Policy" (a set of rules) to generate training data.

The Search Teacher: Instead of just drawing a perfect spiral (which might miss the hole), this teacher is a bit chaotic. It tries eight different search patterns with random speeds and directions. It's like a teacher who doesn't just show one way to find a lost item, but shows you many different ways to sweep the floor so you learn to adapt.
The Insertion Teacher: This teacher knows how to wiggle a stuck peg out of a tight spot, a trick humans use instinctively.

The robot watches thousands of these "successful" attempts by the teacher and learns to copy the feeling of success.

5. The Results: From "Maybe" to "Definitely"

The team tested their robot against the current best methods (like a system called TacDiffusion).

The Old Way: If the peg was off by more than 2 millimeters (about the thickness of a credit card), the robot would usually fail. It was too rigid.
SI-Diff: The new system could handle misalignments up to 5 millimeters. It was much more forgiving of mistakes.
Zero-Shot Magic: Even though the robot was only trained on a square block, it could successfully insert unseen shapes like cylinders, triangles, and even a USB connector. It learned the concept of searching and inserting, not just the specific shape of the block.

Summary

SI-Diff is a robot control system that uses a single AI model to both hunt for a hole and slide an object into it. By using a "mode switch" and learning from a diverse set of "teacher" demonstrations, it makes robots much more robust, allowing them to handle bigger mistakes and work with shapes they've never seen before, all without needing to switch software or reprogram themselves.

Technical Summary: SI-Diff: A Framework for Learning Search and High-Precision Insertion with a Force-Domain Diffusion Policy

Problem Statement
Contact-rich assembly, particularly peg-in-hole tasks, faces significant challenges due to uncertainties in relative poses, such as misalignments and small clearances. Existing approaches typically treat the "search" phase (resolving large x–y misalignments) and the "high-precision insertion" phase (resolving jamming in small clearances) as separate problems requiring distinct strategies or model switches. While visual servoing can achieve coarse alignment, its accuracy is often insufficient for sub-millimeter clearance tasks, necessitating force-based methods. Current state-of-the-art methods like TacDiffusion [1] excel at high-precision insertion but assume the peg and hole are already generally aligned, failing when initial misalignments exceed 2 mm. Conversely, search strategies often lack the adaptability to transition seamlessly into precision insertion. There is a need for a unified framework capable of learning and executing both distinct action patterns without switching models or weights.

Methodology
The authors propose SI-Diff, a framework that unifies search and high-precision insertion within a single force-domain diffusion policy.

System Architecture: The system employs a feedforward-force-based impedance controller. The model learns to generate feedforward forces ( $F_{ff}$ ) based on tactile and kinematic observations. The input state includes external wrench ( $\tau_{ext}$ ), internal wrench ( $\tau_{in}$ ), and end-effector velocities ( $v_e$ ) from both current and previous frames.
Force-Domain Diffusion Policy: The core is a Denoising Diffusion Probabilistic Model (DDPM) that maps Gaussian noise to the target distribution of $F_{ff}$ . The model is trained to predict the noise added at each diffusion step, conditioned on the robot's state.
Mode Conditioning Mechanism: To enable a single model to handle two distinct behaviors (search and insertion), the authors introduce a mode-conditioning mechanism. Inspired by natural language processing tokenization, discrete mode indices (0 for search, 1 for insertion) are embedded into high-dimensional learnable vectors. These mode tokens are prepended to the sequence of robot-state tokens, allowing the diffusion model to learn distinct action patterns based on the specified mode prompt.
Teacher Policies:
- Search Teacher: A novel rule-based policy is developed to generate diverse search trajectories. Unlike fixed spiral searches, this policy randomizes parameters (e.g., search direction, amplitude, frequency) and adapts based on external forces to ensure efficient exploration. It retains only successful and efficient demonstrations (completed within 2.5 seconds) for training.
- Insertion Teacher: The framework utilizes the self-adaptive tactile insertion method from [2] (and TacDiffusion [1]) as the teacher for the insertion phase. This policy detects "stuck" states via residual force monitoring and applies sinusoidal wiggle motions to escape jamming.
Training Strategy: The model is trained on a dataset comprising 300 successful search demonstrations and 1,500 insertion demonstrations. To address the class imbalance, Balanced Batch Sampling (BBS) is employed.

Key Contributions

Unified Framework: Development of a novel framework (SI-Diff) that learns both search and high-precision insertion using a single force-domain diffusion policy, eliminating the need for model switching.
Mode Conditioning: Design of a new conditioning mechanism that allows the model to perform distinct action patterns (exploratory search vs. precision wiggle) based on a mode prompt, enabling simultaneous learning of both tasks.
Diverse Search Teacher: Creation of an effective search teacher policy that generates diverse, randomized trajectories, facilitating efficient data collection and robust learning of wrench-action mappings.

Experimental Results
Experiments were conducted using a 7-DoF Franka Emika robot with a 0.1 mm clearance cuboid peg for training. The method was evaluated against TacDiffusion [1], rule-based variants (Teacher-D, Teacher-R), and FORGE [3] across nine shapes (including unseen geometries like hexagonal/prism, cylinder, and USB-A connectors) and misalignment levels ranging from 2 mm to 5 mm.

Misalignment Tolerance: SI-Diff extended the tolerance for x–y misalignment from 2 mm (the limit of TacDiffusion) to 5 mm.
Success Rates: On the training cuboid shape, SI-Diff achieved a 95.8% success rate across all misalignment levels, significantly outperforming TacDiffusion (29.1%) and other baselines. It maintained high success rates (90%–97%) on unseen shapes, including complex geometries like triangular and pentagonal prisms.
Efficiency: SI-Diff demonstrated superior efficiency, with the lowest average execution times (both failure-included and success-only) compared to competitors.
Zero-Shot Transfer: The model, trained on a single cuboid shape, demonstrated robust zero-shot transferability to unseen shapes without retraining.
Qualitative Behavior: The model learned to adapt its search strategy dynamically, exhibiting varying amplitudes and frequencies that decreased as alignment improved, unlike the fixed patterns of the teacher policy.

Significance and Claims
The paper claims that SI-Diff bridges the technological gap between search and insertion strategies, offering a more integrated solution for intelligent assembly systems. By learning from successful and efficient demonstrations provided by the new teacher policy, the model surpasses both the teacher policy and the state-of-the-art baseline (TacDiffusion) in terms of success rate and efficiency. The primary significance lies in the ability to handle larger initial misalignments (up to 5 mm) while maintaining high-precision insertion capabilities within a single framework, thereby reducing engineering complexity and facilitating integration in real-world industrial applications. The authors note that while the method is robust, mode switching is currently triggered by z-displacement rather than learned tactile cues, and the scope is limited to x–y misalignment and small-clearance insertion.