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

arXiv: https://arxiv.org/abs/2605.20433 · HTML: https://mlsanity.com/paper/2605.20433/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 38 of its 45 references that are in the map.
- First posted on arXiv 19 May 2026.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [On the Continuity of Rotation Representations in Neural Networks](https://mlsanity.com/paper/1812.07035/) (2018)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Multimodal Machine Learning: A Survey and Taxonomy](https://mlsanity.com/paper/1705.09406/) (2017)
- [QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/1806.10293/) (2018)
- [Multimodal Transformer for Unaligned Multimodal Language Sequences](https://mlsanity.com/paper/1906.00295/) (2019)
- [An Algorithmic Perspective on Imitation Learning](https://mlsanity.com/paper/1811.06711/) (2018)
- [Multimodal Learning with Transformers: A Survey](https://mlsanity.com/paper/2206.06488/) (2022)
- [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022)
- [RH20T: A Comprehensive Robotic Dataset for Learning Diverse Skills in One-Shot](https://mlsanity.com/paper/2307.00595/) (2023)
- [Making Sense of Vision and Touch: Self-Supervised Learning of Multimodal Representations for Contact-Rich Tasks](https://mlsanity.com/paper/1810.10191/) (2018)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [See, Hear, and Feel: Smart Sensory Fusion for Robotic Manipulation](https://mlsanity.com/paper/2212.03858/) (2022)
- [Visuo-Tactile Transformers for Manipulation](https://mlsanity.com/paper/2210.00121/) (2022)
- [Towards Generalist Robot Learning from Internet Video: A Survey](https://mlsanity.com/paper/2404.19664/) (2024)
- [ForceMimic: Force-Centric Imitation Learning with Force-Motion Capture System for Contact-Rich Manipulation](https://mlsanity.com/paper/2410.07554/) (2024)
- [Deep Generative Models in Robotics: A Survey on Learning from Multimodal Demonstrations](https://mlsanity.com/paper/2408.04380/) (2024)
- [A Survey of Robot Manipulation in Contact](https://mlsanity.com/paper/2112.01942/) (2021)
- [Play to the Score: Stage-Guided Dynamic Multi-Sensory Fusion for Robotic Manipulation](https://mlsanity.com/paper/2408.01366/) (2024)
- [Unlocking Slot Attention by Changing Optimal Transport Costs](https://mlsanity.com/paper/2301.13197/) (2023)
- [A Survey on Imitation Learning for Contact-Rich Tasks in Robotics](https://mlsanity.com/paper/2506.13498/) (2025)
- [Robotic Compliant Object Prying Using Diffusion Policy Guided by Vision and Force Observations](https://mlsanity.com/paper/2503.03998/) (2025)
- [Towards Forceful Robotic Foundation Models: a Literature Survey](https://mlsanity.com/paper/2504.11827/) (2025)
- [MOMA-Force: Visual-Force Imitation for Real-World Mobile Manipulation](https://mlsanity.com/paper/2308.03624/) (2023)

## Abstract (arXiv)

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.