Spacetime Optimal-Transport Attention for Visuo-Haptic Imitation Learning of Contact-Rich Manipulation

Abstract: Contact-rich manipulation tasks such as tight-clearance insertion, connector mating, polishing, and surface-conforming wiping remain difficult for data-driven controllers because they couple discontinuous contact dynamics, partial observability, and strict safety constraints. No single sensing modality suffices: vision supplies global context before contact, force/torque (F/T) feedback governs interaction after contact, and proprioceptive pose provides a consistent kinematic backbone. Most prior imitation-learning policies for contact-rich tasks operate on uni- or bi-modal signals, and the few that fuse three modalities typically adopt off-the-shelf attention modules with no explicit prior on how attention mass should be distributed across task-relevant regions. We present Spacetime Optimal-Transport Attention (SO-TA), a tri-modal fusion backbone that replaces softmax-normalized patch attention by an entropy-regularized Optimal Transport (OT) alignment between force-pose-derived sub-queries and visual patches. Explicit marginal constraints act as a structured inductive bias for contact-rich tasks, encouraging conditioning-aware spatial selection that is stable across illumination, distractors, and partial occlusion. SO-TA is paired with a diffusion-based sequence policy mapping observation windows to pose-action chunks. We evaluate SO-TA on three real-robot tasks: tight peg-in-hole assembly, BCM wiring-connector insertion, and curved-surface mark erasing. With ~200 rollouts per condition, SO-TA reaches 100% success on tight peg-in-hole versus 93% for cross-attention at matched capacity, and retains 82.5% success under illumination, distractor, and partial-occlusion perturbations where a concatenation baseline drops to 43.5%. OT-derived patch heatmaps and leave-one-out modality-influence ratios provide interpretable, phase-dependent diagnostics.

• The paper introduces SO-TA, formulating visuo-haptic fusion as an entropy-regularized Optimal Transport problem to enhance imitation learning.
• The method integrates synchronized RGB, force/torque, and pose signals via a diffusion-based policy, achieving high precision in tasks like peg insertion and connector mating.
• The approach demonstrates superior robustness against visual perturbations, outperforming standard fusion techniques in real-world contact-rich manipulation.

Spacetime Optimal-Transport Attention for Tri-Modal Visuo-Haptic Imitation Learning

Introduction

This paper proposes Spacetime Optimal-Transport Attention (SO-TA), a tri-modal fusion architecture for imitation learning in contact-rich robot manipulation scenarios. The approach addresses the limitations of existing visuo-haptic policy learning frameworks, in which conventional uni- or bi-modal networks, and even off-the-shelf attention-based fusion methods, fail to optimally exploit task phase-dependence and are highly vulnerable to spurious saliency, partial observability, and domain shift. This is especially relevant in settings such as tight-clearance insertion, connector mating, and surface-conforming operations where phase- and modality-aware reasoning is crucial for robustness and performance. The paper substantiates SO-TA using systematic real-robot evaluations targeting three industry-relevant manipulation domains.

Problem Formulation and Architecture

The central technical innovation is the cast of visuo-haptic fusion as an entropy-regularized Optimal Transport (OT) problem between force/pose-derived sub-queries and visual image patches, with the row-marginal supply explicitly parameterized by the haptic context. This design replaces the generic softmax-normalized attention frequently adopted in cross-modal Transformer systems. The architectural pipeline comprises synchronized vision (RGB), force/torque (F/T), and pose signals processed via learned encoders, fused using SO-TA, and finally consumed by a diffusion-based sequence policy for pose-action chunking and closed-loop control.

Upon ingesting observation windows, vision features are spatially aggregated according to an OT transport plan, which enforces explicit mass constraints representing haptic state—injecting a strong structured prior crucial for robustness and interpretability.

Figure 1: The overall tri-modal imitation learning pipeline, integrating preprocessing, fusion (SO-TA), and diffusion-based policy inference.

Spacetime Optimal-Transport Attention (SO-TA)

Figure 3: The SO-TA module: force/pose features generate sub-queries and supply; retrieval uses an entropy-regularized OT plan over image patches. Supply is conditionally determined by the haptic context.

The optimal-transport-based attention module operates as follows:

• Sub-query and marginal prediction: FC networks project force/pose context into a set of sub-queries and corresponding non-uniform supply distributions (“row marginals”).
• Key/Value mapping: Visual patches are projected into the OT-attention space.
• Cost computation: Cosine similarities between normalized sub-queries and patch keys parametrize the transport cost.
• Sinkhorn OT alignment: The transport plan matches the predicted supply (based on haptics) to the (uniform) patch capacity, producing a doubly-stochastic, entropy-regularized matching.
• Aggregation and merging: Messages are aggregated per sub-query and then globally merged with weights tied to the predicted supply vector, ensuring preservation of contextual dependence.

This attention design embeds strong inductive bias: it enforces phase-appropriate spatial selection aligned with physical task structure and adapts dynamically to changes in visual context, distractors, or occlusion.

Experimental Evaluation

Three real-robot tasks were used to benchmark SO-TA: tight peg-in-hole insertion (sub-millimeter clearance), automotive wiring-connector insertion, and curved-surface mark erasing. Each was collected with several hundred human demonstrations, using tri-modal synchronized sensing.

Peg-in-Hole Insertion

Figure 2: Tight peg-in-hole experimental setup (third-person and robot view).

SO-TA achieved $100\%$ success over $\sim200$ trials in standard conditions, outperforming cross-attention (93%) and matching/ slightly exceeding concatenation (99.5%). Importantly, SO-TA outperformed other fusion schemes under severe visual perturbations (illumination shifts, distractors, occlusion), retaining 82.5% success where simple concatenation collapsed to 43.5%.

Figure 4: First-step $\Delta z^L$ predictions from SO-TA diffusion policy closely match ground-truth with $\geq$2 Sinkhorn iterations, justifying the network’s low-latency deployment.

Ablation reveals that the F/T stream dominates post-contact, with vision and pose providing essential context-dependent support. The SO-TA fusion backbone delivers markedly more consistent success and latency profiles, especially in visually perturbed deployment regimes.

Figure 5: Success rate and completion time KDE for modality ablations (concatenation baseline). Tri-modal input is essential for optimal performance.

Figure 6: Success rate and completion time KDE across fusion methods. SO-TA gives the highest aggregate reliability and tightest temporal density tails.

Figure 7: SO-TA’s robustness is clear under visual distribution shift (illumination noise, occlusion). It retains 82.5% success—concatenation drops to 43.5%.

Interpretability analyses highlight that SO-TA’s marginal supply and patch weights align with expected task phases: visual share peaks pre-contact, while F/T influence dominates during search and insertion. Patch attention heatmaps reflect plausible region relevance, although some diffuseness remains due to the absence of explicit spatial priors in the cost function.

Figure 8: SO-TA interpretability: patch heatmaps and leave-one-out modality ratios across phases; visual attention dominates before contact, yielding to F/T after.

Automotive Wiring-Connector Insertion

Figure 9: Automotive BCM connector insertion (setup and robot view), with clutter and minimal visual contrast.

The BCM connector task, characterized by visual ambiguity and background clutter, further validates the value of structured fusion: SO-TA closely matches human performance in mean completion time, with only modest long-tail failures—tractable using a lightweight cutoff-and-reset runtime wrapper.

Figure 10: Success rate and completion time KDE for SO-TA vs. humans in wiring-connector insertion; policy-induced delay tail is mitigated by wrapper.

(Figure 12), (Figure 13)

Figure 11: Cutoff-and-reset wrapper selection: optimal $12\,\mathrm{s}$ abort threshold minimizes tail latency, achieving $>99.5\%$ success.

Curved-Surface Mark Erasing

Erasing (wiping) requires visual localization, force-modulation, and non-trivial trajectory synthesis on a curved surface with partial occlusion. SO-TA achieves mean completion time near human levels, with minor degradation on under-represented mark placements.

Figure 12: Curved-surface mark erasing; the eraser occludes the target region, challenging visual+force fusion.

Figure 13: SO-TA achieves high success rates and close human parity in wipe task completion time.

Modality and attention analyses again confirm phase- and task-relevant allocation: vision dominates before contact, force contributes substantially during wiping, and pose governs retreat to canonical locations.

Figure 14: SO-TA interpretability in mark erasing; phase-consistent gating and spatial focus even with moving targets and occlusion.

Implications and Future Directions

This work demonstrates that explicit marginal-constrained fusion outperforms standard MLP or vanilla Transformer-based alternatives for tri-modal visuo-haptic imitation learning, especially in shifting deployment domains and under partial observability. The inductive bias introduced by SO-TA’s OT attention is key for generalization, sample efficiency, and interpretability in contact-rich contexts, and aligns with emerging theoretical and practical needs identified in recent literature (2605.20433).

The approach opens several avenues: incorporating spatial priors in the OT cost for sharper and more localized attention allocation, addressing compounding error accumulation with more advanced decision-time recovery/head-tail truncation strategies, and scaling tri-modal architectures to broader, more varied manipulation skillsets using online adaptation and internet-scale dataset co-training.

Conclusion

Spacetime Optimal-Transport Attention (SO-TA) introduces a principled, scalable, and robust tri-modal fusion strategy for complex contact-rich manipulation. By structuring attention via haptics-conditioned OT, SO-TA realizes predictable, interpretable, and efficient multi-modal policy learning, outperforming extant generic fusion methods under matched capacity and harsh deployment shifts. This work positions structured fusion as a central axis for future progress in robot perception and control.