OSTRA: Learning Multi-Stage Robot Manipulation through Object-State Transition Supervision

Sparse supervision is a key challenge in multi-stage manipulation, with the lack of supervision over the past-to-future evolution of objects required by the policy being particularly critical. To address this issue, we introduce the Object-State Transition Map (**OST-Map**), a unified object-centric representation that organizes task-relevant objects as identity-consistent graphs and represents historical and future object states in a shared structured space. Building on OST-Map, we propose **OSTRA**, a diffusion-based Object-State Transition Reasoning Architecture that hierarchically predicts future object states, target TCP poses, and continuous action trajectories. OSTRA combines Transition-Guided Attention, which prioritizes historical object-state transitions during future-state prediction, with Staged Denoising, which progressively translates predicted object changes into robot subgoals and executable actions. Experiments on two simulation benchmarks and five diverse real-world tasks demonstrate consistent improvements over representative imitation learning baselines. Further ablations validate the contributions of OST-Map and the proposed architectural designs.