# How to Utilize Failure Demo Data?: Effective Data Selection for Imitation Learning Using Distribution Differences in Attention Mechanism

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Language-conditioned policies and VLAs, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 17 of its 23 references that are in the map.
- First posted on arXiv 8 May 2026.

## Placed in

- Offline RL and imitation learning from data / [Language-conditioned policies and VLAs](https://mlsanity.com/offlinerl/language-conditioned-policies-and-vlas/)

## Builds on (most cited first, 17 of 17)

- [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)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [AgiBot World Colosseo: A Large-scale Manipulation Platform for Scalable and Intelligent Embodied Systems](https://mlsanity.com/paper/2503.06669/) (2025)
- [Learning Hand-Eye Coordination for Robotic Grasping with Deep Learning and Large-Scale Data Collection](https://mlsanity.com/paper/1603.02199/) (2016)
- [π*0.6: a VLA That Learns From Experience](https://mlsanity.com/paper/2511.14759/) (2025)
- [A Survey of Imitation Learning: Algorithms, Recent Developments, and Challenges](https://mlsanity.com/paper/2309.02473/) (2023)
- [AHA: A Vision-Language-Model for Detecting and Reasoning Over Failures in Robotic Manipulation](https://mlsanity.com/paper/2410.00371/) (2024)
- [Learning from Imperfect Demonstrations with Self-Supervision for Robotic Manipulation](https://mlsanity.com/paper/2401.08957/) (2024)
- [Real-Time Out-of-Distribution Failure Prevention via Multi-Modal Reasoning](https://mlsanity.com/paper/2505.10547/) (2025)
- [Fail2Progress: Learning from Real-World Robot Failures with Stein Variational Inference](https://mlsanity.com/paper/2509.01746/) (2025)
- [Learning from Successful and Failed Demonstrations via Optimization](https://mlsanity.com/paper/2107.11918/) (2021)
- [Sensorimotor Attention and Language-based Regressions in Shared Latent Variables for Integrating Robot Motion Learning and LLM](https://mlsanity.com/paper/2407.09044/) (2024)
- [Aligning Human Intent from Imperfect Demonstrations with Confidence-based Inverse soft-Q Learning](https://mlsanity.com/paper/2312.11194/) (2023)
- [Motion ReTouch: Motion Modification Using Four-Channel Bilateral Control](https://mlsanity.com/paper/2502.20982/) (2025)

## Abstract (arXiv)

Imitation learning for robotic tasks has relied primarily on policies trained only on successful demonstrations, although failures are unavoidable during human data collection. Many existing approaches for exploiting failure data require additional data processing or iterative policy updates through autonomous rollouts, making it difficult to directly and stably utilize failure data accumulated during data collection. In this work, we propose a method that learns latent representations of success-failure discrepancies and incorporates them into the attention mechanism. During inference, an appropriate latent mode is selected from the initial observation to improve action stability. Furthermore, we introduce a post-training metric that quantifies the attention discrepancy between each failure sample and successful demonstrations to select failure data. Simulation results show that the proposed method improves task success rates when trained with failure data and that the proposed metric identifies failure samples that are beneficial for learning when combined with successful demonstrations. These results suggest that the proposed method can support more efficient use of collected demonstrations in robotic data collection pipelines.