# KOI: Accelerating Online Imitation Learning via Hybrid Key-state Guidance

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

## Key facts

- Cited by 6 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 31 of its 41 references that are in the map.
- First posted on arXiv 6 Aug 2024.

## Placed in

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

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

- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [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)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning](https://mlsanity.com/paper/2306.03310/) (2023)
- [Code as Policies: Language Model Programs for Embodied Control](https://mlsanity.com/paper/2209.07753/) (2022)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Open X-Embodiment: Robotic Learning Datasets and RT-X Models](https://mlsanity.com/paper/2310.08864/) (2023)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [Rethinking Spatiotemporal Feature Learning: Speed-Accuracy Trade-offs in Video Classification](https://mlsanity.com/paper/1712.04851/) (2017)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Language to Rewards for Robotic Skill Synthesis](https://mlsanity.com/paper/2306.08647/) (2023)
- [A Survey of Inverse Reinforcement Learning: Challenges, Methods and Progress](https://mlsanity.com/paper/1806.06877/) (2018)
- [Extrapolating Beyond Suboptimal Demonstrations via Inverse Reinforcement Learning from Observations](https://mlsanity.com/paper/1904.06387/) (2019)
- [Bridge Data: Boosting Generalization of Robotic Skills with Cross-Domain Datasets](https://mlsanity.com/paper/2109.13396/) (2021)
- [RT-H: Action Hierarchies Using Language](https://mlsanity.com/paper/2403.01823/) (2024)
- [A Survey of Imitation Learning: Algorithms, Recent Developments, and Challenges](https://mlsanity.com/paper/2309.02473/) (2023)
- [RoboCLIP: One Demonstration is Enough to Learn Robot Policies](https://mlsanity.com/paper/2310.07899/) (2023)
- [Primal Wasserstein Imitation Learning](https://mlsanity.com/paper/2006.04678/) (2020)
- [Deeply AggreVaTeD: Differentiable Imitation Learning for Sequential Prediction](https://mlsanity.com/paper/1703.01030/) (2017)
- [Teach a Robot to FISH: Versatile Imitation from One Minute of Demonstrations](https://mlsanity.com/paper/2303.01497/) (2023)
- [Watch and Match: Supercharging Imitation with Regularized Optimal Transport](https://mlsanity.com/paper/2206.15469/) (2022)
- [Wasserstein Adversarial Imitation Learning](https://mlsanity.com/paper/1906.08113/) (2019)
- [Language Reward Modulation for Pretraining Reinforcement Learning](https://mlsanity.com/paper/2308.12270/) (2023)
- [Imitation Learning with Sinkhorn Distances](https://mlsanity.com/paper/2008.09167/) (2020)
- [Optimal Transport: Fast Probabilistic Approximation with Exact Solvers](https://mlsanity.com/paper/1802.05570/) (2018)
- [Universal Visual Decomposer: Long-Horizon Manipulation Made Easy](https://mlsanity.com/paper/2310.08581/) (2023)

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

- [Phoenix: A Motion-based Self-Reflection Framework for Fine-grained Robotic Action Correction](https://mlsanity.com/paper/2504.14588/) (2025)
- [Human-assisted Robotic Policy Refinement via Action Preference Optimization](https://mlsanity.com/paper/2506.07127/) (2025)
- [FLARE: A Failure-Aware Framework for Autonomous Correction and Recovery in Visual-Language Robotic Manipulation](https://mlsanity.com/paper/2608.26645/) (2026)
- [FoldNet: Learning Generalizable Closed-Loop Policy for Garment Folding via Keypoint-Driven Asset and Demonstration Synthesis](https://mlsanity.com/paper/2505.09109/) (2025)
- [Continuous Vision-Language-Action Co-Learning with Semantic-Physical Alignment for Behavioral Cloning](https://mlsanity.com/paper/2511.14396/) (2025)
- [Grasp-Then-Plan with Failure Attribution: A Closed Two-Stage Framework for Precise and Generalizable Robotic Manipulation](https://mlsanity.com/paper/2606.03385/) (2026)

## Abstract (arXiv)

Online Imitation Learning struggles with the gap between extensive online exploration space and limited expert trajectories, hindering efficient exploration due to inaccurate reward estimation. Inspired by the findings from cognitive neuroscience, we hypothesize that an agent could estimate precise task-aware reward for efficient online exploration, through decomposing the target task into the objectives of "what to do" and the mechanisms of "how to do". In this work, we introduce the hybrid Key-state guided Online Imitation (KOI) learning method, which leverages the integration of semantic and motion key states as guidance for reward estimation. Initially, we utilize visual-language models to extract semantic key states from expert trajectory, indicating the objectives of "what to do". Within the intervals between semantic key states, optical flow is employed to capture motion key states to understand the mechanisms of "how to do". By integrating a thorough grasp of hybrid key states, we refine the trajectory-matching reward computation, accelerating online imitation learning with task-aware exploration. We evaluate not only the success rate of the tasks in the Meta-World and LIBERO environments, but also the trend of variance during online imitation learning, proving that our method is more sample efficient. We also conduct real-world robotic manipulation experiments to validate the efficacy of our method, demonstrating the practical applicability of our KOI method. Videos and code are available at https://gewu-lab.github.io/Keystate_Online_Imitation/.