# SAGE:State-Aware Guided End-to-End Policy for Multi-Stage Sequential Tasks via Hidden Markov Decision Process

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 24 Sep 2025.

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

- [Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning](https://mlsanity.com/paper/2610.10857/) (2026)

## Abstract (arXiv)

Multi-stage sequential (MSS) robotic manipulation tasks are prevalent and crucial in robotics. They often involve state ambiguity, where visually similar observations correspond to different actions. We present SAGE, a state-aware guided imitation learning framework that models tasks as a Hidden Markov Decision Process (HMDP) to explicitly capture latent task stages and resolve ambiguity. We instantiate the HMDP with a state transition network that infers hidden states, and a state-aware action policy that conditions on both observations and hidden states to produce actions, thereby enabling disambiguation across task stages. To reduce manual annotation effort, we propose a semi-automatic labeling pipeline combining active learning and soft label interpolation. In real-world experiments across multiple complex MSS tasks with state ambiguity, SAGE achieved 100% task success under the standard evaluation protocol, markedly surpassing the baselines. Ablation studies further show that such performance can be maintained with manual labeling for only about 13% of the states, indicating its strong effectiveness.