# Robust Offline Imitation Learning Through State-level Trajectory Stitching

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Imitation learning methods, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 29 of its 40 references that are in the map.
- First posted on arXiv 28 Mar 2025.

## Placed in

- Offline RL and imitation learning from data / [Imitation learning methods](https://mlsanity.com/offlinerl/imitation-learning-methods/)

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

- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Vision-Language Foundation Models as Effective Robot Imitators](https://mlsanity.com/paper/2311.01378/) (2023)
- [RoboAgent: Generalization and Efficiency in Robot Manipulation via Semantic Augmentations and Action Chunking](https://mlsanity.com/paper/2309.01918/) (2023)
- [Toward the Fundamental Limits of Imitation Learning](https://mlsanity.com/paper/2009.05990/) (2020)
- [Discriminator-Weighted Offline Imitation Learning from Suboptimal Demonstrations](https://mlsanity.com/paper/2207.10050/) (2022)
- [STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning](https://mlsanity.com/paper/2310.09615/) (2023)
- [Q-attention: Enabling Efficient Learning for Vision-based Robotic Manipulation](https://mlsanity.com/paper/2105.14829/) (2021)
- [BAKU: An Efficient Transformer for Multi-Task Policy Learning](https://mlsanity.com/paper/2406.07539/) (2024)
- [Data Quality in Imitation Learning](https://mlsanity.com/paper/2306.02437/) (2023)
- [How to Leverage Unlabeled Data in Offline Reinforcement Learning](https://mlsanity.com/paper/2202.01741/) (2022)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [Versatile Offline Imitation from Observations and Examples via Regularized State-Occupancy Matching](https://mlsanity.com/paper/2202.02433/) (2022)
- [Optimal Transport for Offline Imitation Learning](https://mlsanity.com/paper/2303.13971/) (2023)
- [Universal Visual Decomposer: Long-Horizon Manipulation Made Easy](https://mlsanity.com/paper/2310.08581/) (2023)
- [Discriminator-Guided Model-Based Offline Imitation Learning](https://mlsanity.com/paper/2207.00244/) (2022)
- [SEABO: A Simple Search-Based Method for Offline Imitation Learning](https://mlsanity.com/paper/2402.03807/) (2024)
- [Offline Imitation Learning Through Graph Search and Retrieval](https://mlsanity.com/paper/2407.15403/) (2024)
- [How to Leverage Diverse Demonstrations in Offline Imitation Learning](https://mlsanity.com/paper/2405.17476/) (2024)
- [Offline Imitation Learning with Model-based Reverse Augmentation](https://mlsanity.com/paper/2406.12550/) (2024)

## Built on by (most cited first, 1 of 1)

- [NEEDLEWORK: Offline Rewriting of Robot Data with Verified Local Stitches](https://mlsanity.com/paper/2610.02339/) (2026)

## Abstract (arXiv)

Imitation learning (IL) has proven effective for enabling robots to acquire visuomotor skills through expert demonstrations. However, traditional IL methods are limited by their reliance on high-quality, often scarce, expert data, and suffer from covariate shift. To address these challenges, recent advances in offline IL have incorporated suboptimal, unlabeled datasets into the training. In this paper, we propose a novel approach to enhance policy learning from mixed-quality offline datasets by leveraging task-relevant trajectory fragments and rich environmental dynamics. Specifically, we introduce a state-based search framework that stitches state-action pairs from imperfect demonstrations, generating more diverse and informative training trajectories. Experimental results on standard IL benchmarks and real-world robotic tasks showcase that our proposed method significantly improves both generalization and performance.