Source-linked AI summary
RodForesight: A World Model Enhanced Diffusion Policy for Slender Rod Insertion
Chuanbo Yu, Mingyu Yue, Yan Lyu, Chuhan Song, Peng Wang
TL;DR
Slender rods can bend during manipulation, making precise insertion difficult when distal motion depends on configuration, grasp, material properties, and contact. RodForesight combines coarse visual-servo approach with diffusion-proposed insertion actions evaluated by a world model. Predictive selection raises success from 88.9% to 96.7% compared with the same diffusion policy.
Problem
Slender-rod insertion requires precise distal position and orientation control despite bending, reduced volumes, and material- and contact-dependent motion.
Source: §I · p. 1Method
RodForesight uses visual servoing for coarse approach, then evaluates diffusion-policy action chunks with an action-conditioned world model predicting tilt and radial alignment errors.
Source: §I · p. 2Results
Predictive selection raises success from 88.9% to 96.7% compared with the same diffusion policy, with gains across all three training materials.
Source: 2) Predictive Action Selection · p. 6Takeaways & Limitations
Experiments support transfer to lower-stiffness materials without retraining, including dry PA66 represented by an isotropic surrogate.
Source: F. Ablation and Generalisation · p. 7
Abstract
from arXiv · showhide
Slender rod insertion arises in precision manufacturing, where millimetre scale diameter and tight clearances demand accurate perception and control. Conventional peg-in-hole methods assume a rigid object whose tip pose is fixed relative to the gripper. This assumption breaks down for a high aspect ratio rod, which can bend during manipulation, making its tip motion dependent on the rod configuration, grasp, material properties, and contact. We present RodForesight, a learning framework that factorises the task into two stages: 1) coarse approaching, which uses visual servoing to map diverse initial configurations into a compact near hole hand-off region; and 2) predictive insertion, which performs fine alignment and completes the insertion. It is worth noting that the two stages can be wrapped into an end-to-end design. During insertion, a diffusion policy generates candidate action chunks, while an action conditioned world model predicts their effects on rod-hole alignment. This pre-execution evaluation enables RodForesight to select the best action chunk based on predicted tilt and radial errors before execution. Experiments investigate the performance of different stages and the end-to-end setting, where RodForesight improves the success rate from 88.9% to 96.7%, compared to baseline methods such as diffusion policy.
I. INTRODUCTION
Slender-rod insertion challenges rigid-peg assumptions because bending, reduced volume, and material-dependent responses complicate perception and control. RodForesight separates coarse approach from predictive fine insertion, evaluating candidate actions before execution.
I. INTRODUCTION
- High aspect ratio rods can bend, making distal motion depend on configuration, grasp, material response, and contact while reducing volumes make detection and contact-based methods challenging.The task involves millimetre-scale diameters and precise distal position and orientation control.Source: §I · p. 1
- Rigid-peg learning methods do not transfer readily to slender rods, and limited demonstrations cannot enumerate diverse initial poses, motivating a shared near-hole hand-off region.Most insertion trajectories share the near-hole area, supporting a two-stage approach.Source: §I · p. 1
- RodForesight uses visual servoing to reach a compact near-hole hand-off, then uses a diffusion policy and action-conditioned world model to predict tilt and radial alignment effects before execution.The stages can also be wrapped into an end-to-end design.Source: §I · p. 2
- The framework explicitly accounts for deformation and material variation, and experiments evaluate both individual stages and the end-to-end setting across materials and unseen stiffnesses.The test suite includes a 100 mm by 4 mm rod entering a 5 mm diameter, 20 mm deep hole with 0.5 mm radial clearance.Source: Figure 1 · p. 1, §I · p. 2
II. RELATED WORK
Prior work addresses insertion through contact feedback, flexible-object geometry or flexibility estimates, and diffusion-based action generation. RodForesight combines candidate motion proposals with a world model that predicts distal alignment effects.
A. Insertion and Assembly with Contact
- Contact sensing and compliant motion support rigid peg-in-hole localisation and insertion,, while flexible-object methods use approximate Jacobians or visual flexibility estimates,.RodForesight focuses on slender rods across different material responses [4].Source: A. Insertion and Assembly with Contact · p. 2
B. Diffusion Policies for Robotic Manipulation
- Diffusion Policy generates observation-conditioned action sequences and replans after partial execution; RodForesight uses diffusion to propose near-hole corrections for comparison before execution.Related work also combines reinforcement learning and diffusion for deformable-object routing.Source: B. Diffusion Policies for Robotic… · p. 2
C. World Models and Predictive Action Selection
- Predictive methods model action effects using images, latent representations, graphs, or physical constraints,,,,,; RodForesight predicts distal geometry relative to the hole.Its training branches share the same committed Cosserat state, so predictions compare candidate motions from identical initial conditions.Source: C. World Models and Predictive Action… · p. 2
III. PROBLEM FORMULATION
The problem is vision-guided insertion of a flexible rod into a cylindrical hole despite pose variation, deformation, and contact. The formulation defines hole-relative geometry and errors, while a Cosserat-based simulator supplies physical transitions.
III. PROBLEM FORMULATION
- The task represents distal rod geometry in a hole-fixed frame using transverse displacement, insertion depth, and two tilt components; radial and tilt errors quantify alignment.This geometry defines task progress and supplies supervision and evaluation targets.Source: §III · p. 2
- Synchronized RGB-D observations from side and wrist cameras provide semantic masks, and each action is a relative Cartesian gripper motion with three translations and three rotations.The objective is to advance the distal rod while reducing radial and tilt errors and accommodating contact-induced deformation.Source: §III · p. 3
- Approach and near-hole insertion have different control regimes: explicit geometry supports large-error correction before contact, while small motions near the hole interact with deformation and contact.Overall success includes Stage 1 failures.Source: §III · p. 3
A. Method Overview
- Stage 1 maps varied initial poses to a compact hand-off region; Stage 2 proposes local insertion motions, predicts alignment effects, executes a prefix, and replans after reobservation.The Cosserat model supplies physical transitions for data generation and evaluation, while the learned world model predicts candidates during control.Source: A. Method Overview · p. 3
B. Cosserat Rod Transition Model
- The simulator represents the rod by a centreline and cross-section orientations, with element strain relative to its undeformed state describing shear, extension, bending, and twist.The material law maps strain to internal forces and tangent stiffness; axial and bending responses are bilinear elastoplastic, while shear and torsion remain elastic.Source: B. Cosserat Rod Transition Model · p. 3
- Hole contact accounts for rod radius, resists penetration with a logarithmic barrier, and uses Coulomb friction; the grasp constrains a rod band while the rest deforms.Transitions use backward Euler and Newton iterations with line search and adaptive step reduction; converged deformation and contact histories are committed for the next step.Source: B. Cosserat Rod Transition Model · p. 3
C. Stage 1: Coarse Approaching
Stage 1 estimates rod-hole alignment and uses bounded corrections to approach a stationary hand-off. Stage 2 proposes insertion chunks, predicts their alignment effects, selects a prefix, and replans after observation.
C. Stage 1: Coarse Approaching
- Stage 1 estimates alignment from segmented rod and hole masks, tracked across observations, and uses calibrated depth to reconstruct the hole annulus and visible rod centreline.Endpoint and axis estimates are propagated using commanded tilt increments, separately smoothed for position and tilt, and angular updates are bounded.Source: C. Stage 1: Coarse Approaching · p. 3
- Bounded proportional corrections reduce lateral and angular errors while approaching a target depth; poor alignment attenuates axial speed, and the controller waits at zero motion for valid observations once tolerances are met.Stage 2 begins from the resulting physical state with a fresh semantic history.Source: C. Stage 1: Coarse Approaching · p. 3
D. Stage 2: Predictive Insertion
- At each decision, Stage 2 compares predicted geometry after the executed prefix horizon, choosing the smallest tilt error and using radial error to break ties.It executes the selected prefix, discards the rest, then reobserves and replans; invalid observations trigger holds and reacquisition, while exhausted budgets or solver nonconvergence end in failure.Source: D. Stage 2: Predictive Insertion · pp. 3–4
- The diffusion policy conditions on recent side and wrist semantic-observation histories and generates action sequences learned from expert insertion trajectories.A shared ResNet-18 encodes views and times; a temporal U-Net denoises action sequences, with axial motion restricted to the insertion direction.Source: D. Stage 2: Predictive Insertion · p. 3
- The world model predicts geometry changes along candidate action prefixes, adding each predicted change to the shared source geometry; training branches start from the same expert state and receive per-step geometry-change supervision.The model uses a separate mask encoder and initializes candidate rollouts from the same hidden state.Source: D. Stage 2: Predictive Insertion · pp. 3–4
V. EXPERIMENTS
Experiments instantiate Stage 1 with bounded visual servoing and train Stage 2 using expert insertions and matched Cosserat rollouts.
A. Experimental Instantiation
The experimental setup specifies the task geometry, evaluation range, controller horizons, and physical-model parameters.
A. Experimental Instantiation
- Training and principal evaluations span three materials, three grasp locations, and a 4 × 2 × 4 grid of hole positions.Source: A. Experimental Instantiation · p. 4
- The Cosserat rod is discretised into 20 elements, with contact stiffness kc of 104 N/m and static and dynamic friction coefficients of 0.20 and 0.15.Source: A. Experimental Instantiation · p. 4
- Stage 1 allows 100 motions and uses position and tilt update gains of 0.65 and 0.35, respectively.Source: A. Experimental Instantiation · p. 4
- Stage 1 caps tilt innovation at 1° and uses command ranges of [−1, 1] mm and [−1, 1]° for translation and rotation.Source: A. Experimental Instantiation · p. 5
- Stage 2 permits 160 controls at 50 ms intervals, with transverse, axial, and angular ranges of [−0.5, 0.5] mm, [0, 0.25] mm, and [−1, 1]°.Source: A. Experimental Instantiation · p. 5
- The policy uses a 256-dimensional image condition and 256-channel temporal U-Net; the world model uses 128-dimensional action embeddings and a 256-unit GRU.Source: A. Experimental Instantiation · p. 5
B. Expert Data and Diffusion Policy Training
The diffusion policy is trained on geometry-expert insertion trajectories, with action chunks corrupted by scheduled Gaussian noise and noise recovery learned from visual conditioning.
B. Expert Data and Diffusion Policy Training
- A geometry expert provides 300 successful Stage 2 trajectories, processed into 6,654 samples containing two semantic frames and a target sequence of 16 actions.The split holds out 623 states by disjoint hole-position cells.Source: B. Expert Data and Diffusion Policy… · p. 5
- The encoded expert action chunk A has dimensions H × da, with da denoting action dimension, and is scaled to [−1, 1].Source: B. Expert Data and Diffusion Policy… · p. 5
- The noisy action chunk is formed as Aτ = cτA + sτϵ using Gaussian noise and signal and noise coefficients from a cumulative squared cosine schedule.Source: B. Expert Data and Diffusion Policy… · p. 5
- The policy learns to recover sampled noise from the noisy chunk and visual condition, which encodes a history of length Ho.Source: B. Expert Data and Diffusion Policy… · p. 5
- The policy is trained for 80 epochs with AdamW, batch size 64, learning rate 10−4, and mixed precision, then frozen for branch generation.Source: B. Expert Data and Diffusion Policy… · p. 5
C. Matched Branches and World Model Training
The world model is trained to compare policy-generated action branches from matched states, and validation loss selects the replica used at inference.
C. Matched Branches and World Model Training
- At each expert state, eight proposals of 16 actions are rolled out from the same observation and committed Cosserat state, yielding 53,232 matched branches.Holding the source state fixed makes differences in outcomes reflect candidate actions; controller-reached states are also evaluated for prediction accuracy.Source: C. Matched Branches and World Model… · p. 5
- Each replica receives two mask frames and all eight action chunks, with source geometry and changes supervised at all 16 prediction steps.Source: C. Matched Branches and World Model… · p. 5
- Table II reports success rates for controller ablations and generalization tests, with Stage 2 rates conditional on successful Stage 1 handoffs and overall rates covering all initial cases.Source: Table II · p. 5
- Figure 4 presents successful executions across five materials, including unseen hole positions and initial rod poses, 6061-T6 aluminium, and dry PA66.Its columns show the initial state, Stage 1 approach and handoff, Stage 2 insertion, and success.Source: Figure 4 · p. 6
- Three replicas share architecture and hyperparameters but use independent initialization and data order; the lowest-validation-loss replica is retained for inference.Source: C. Matched Branches and World Model… · p. 6
D. Evaluation Protocol
The evaluation compares direct and staged controllers on fixed initial configurations, then tests generalisation across hole positions, rod poses, diameters, and materials.
D. Evaluation Protocol
- The principal paired comparison uses 90 fixed configurations: 29 AISI 304, 31 C11000 copper, and 30 carbon epoxy cases.Direct policies start from initial poses, while staged variants restore corresponding Stage 1 handoffs; failures remain in the overall denominator.Source: D. Evaluation Protocol · p. 6
- Generalisation uses 256 cases with unseen hole positions and broader initial rod positions and orientations, including 88 AISI 304, 87 C11000 copper, and 81 carbon epoxy cases.Diameter tests reuse the 90 principal configurations; models train at 2, 4, and 6 mm and are evaluated at 3 and 5 mm with 0.5 mm radial clearance.Source: D. Evaluation Protocol · p. 6
- Material tests use 100 mm long, 4 mm diameter 6061-T6 aluminium or dry PA66 rods, initialized straight and unstrained without motion, contact history, or external load.Both learned models remain frozen, and overall success includes every initial case.Source: D. Evaluation Protocol · p. 6
- Code, model weights, and evaluation data are available in an anonymised repository.Source: D. Evaluation Protocol · p. 6
E. Results
Stage separation raises success before predictive selection, and evaluating candidate actions with a world model further improves insertion success. Prediction also reduces selection regret, although prediction error increases on controller-reached states.
1) Stage Separation:
- 88.9% success with the staged temporal U-Net diffusion policy versus 1.1% for direct control, while the deterministic Transformer achieves 0.0%.Within this demonstration budget, the results support assigning broad pose correction to Stage 1 so the policy can focus on local insertion motions.Source: 1) Stage Separation · p. 6
2) Predictive Action Selection:
- 96.7% success with predictive selection versus 88.9% for the same diffusion policy and initial cases; the paired McNemar p value is 0.0156.Gains occur across all three training materials; the Table III comparison retains every initial case in the denominator.Source: 2) Predictive Action Selection · p. 6, Table III · p. 7
- Replacing only the GRU transition with a Transformer yields 93.3% success, but the paired p value of 0.375 does not establish architectural superiority.Source: 2) Predictive Action Selection · p. 6
- Prediction reduces mean regret from 0.031 to 0.025 across test panels and from 0.534 to 0.248 on 28 panels selected by an action-diversity threshold.Regret measures the selected-versus-oracle outcome gap normalized by success tolerances; exact index agreement can count near-identical outcomes as different.Source: 2) Predictive Action Selection · p. 7
- Expert-branch radial MAE is 0.026–0.029 mm and tilt MAE is 0.101–0.130° across training materials, but radial MAE rises to 0.576 mm in-distribution and 0.761 mm in generalisation on controller-reached states.Table IV reports prediction errors after four actions; generalisation covers 256 cases with unseen hole positions and initial rod poses.Source: 2) Predictive Action Selection · p. 7, Table IV · p. 7
F. Ablation and Generalisation
The staged controller transfers to softer materials and unseen hole positions and rod poses, while diameter transfer is less uniform; the conclusion summarizes broader generalisation across these settings.
F. Ablation and Generalisation
- Stage 1 succeeds in every 6061-T6 aluminium and dry PA66 case, and overall success is 93.3% for each material.All failures occur during Stage 2; the rods have Young’s moduli of 68.9 and 3.5 GPa, respectively, versus a training range of 117–193 GPa.Source: F. Ablation and Generalisation · p. 7
- For unseen hole positions and initial rod poses, median radial error falls from 19.39 to 0.38 mm and median tilt from 20.06° to 1.20° among successful handoffs.Stage 1 succeeds in 95.7% of cases, while the complete controller succeeds in 87.5%, with all initial cases retained.Source: F. Ablation and Generalisation · p. 7
- After training at 2, 4, and 6 mm, success is 68.9% at 3 mm and 86.7% at 5 mm with radial clearance fixed; perception recovery accounts for 89.3% of failures at 3 mm.Source: F. Ablation and Generalisation · p. 7
VI. CONCLUSION
- The conclusion reports generalisation across unseen hole positions, initial rod poses, rod diameters, and materials, including a highly flexible material beyond the training range.It characterizes Stage 1 as contracting initial variation into a compact hand-off region so Stage 2 can focus on contact-sensitive corrections.Source: §VI · p. 7