Abstract
Large scale robot learning has recently shown promise in enabling robots to per- form complex tasks by integrating perception, control, and optionally, language understanding into a unified framework. However, they continue to struggle with long-horizon, contact-rich manipulation tasks, such as the handling of deformable objects, where supervision from demonstrations is often inconsistent in quality. In such settings, reward modeling offers a natural solution: by providing grounded progress signals, it can transform noisy demonstrations into stable supervision that generalizes across diverse trajectories. In this work, we introduce a stage-aware, video-based reward modeling framework that jointly predicts the high-level task stage and fine-grained progress within each stage. Reward labels are automatically derived from natural language subtask annotations, enabling consistent progress estimation across variable-length and heterogeneous demonstrations. This design overcomes the limitations of frame-index-based labeling, which collapses in long, variable-duration tasks such as folding a T-shirt. Our reward model demonstrates robustness to demonstration variability, generalization to out-of-distribution sce- narios, and strong utility for downstream policy training. Building upon this re- ward model, we propose the Reward-Aligned Behavior Cloning (RA-BC) frame- work, which selectively filters high-quality data and reweights training samples according to reward estimates. Extensive experiments demonstrate that the reward model outperforms baselines on out-of-distribution real robot policy rollouts and human demonstration validation. Our approach achieves 83% success on folding T-shirts from the flattened state and 67% from the crumpled state—dramatically surpassing vanilla behavior cloning, which attains only 8% and 0% success un- der the same training dataset, respectively. Overall, our results highlight reward modeling as a key enabler for scalable, annotation-efficient, and robust imitation learning in long-horizon robotic manipulation.
Overview of our method
Overview of our method's framework for (a) data processing, (b) reward model training, and (c) policy training with reward signals.
SARM's structure
Overview of SARM, stage-aware reward modeling. Left: SARM overview, which in- cludes both a stage estimator and subtask estimator. First the task stage is predicted from the ob- servations. This prediction is additionally passed into the subtask estimator which predicts a scale value of the progress within the stage. Right: An overview of the estimator architecture which is replicated for both the stage estimator and the subtask estimato
Results of SARM
Demo Data Estimation
Examples of SARM's prediction on demonstration data
Citations
If you find our work useful in your research, please consider citing:
@misc{chen2025sarm,
title={SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation},
author={Qianzhong Chen, Justin Yu, Mac Schwager, Pieter Abbeel, Fred Shentu, Philipp Wu},
year={2025},
eprint={2509.25358},
archivePrefix={arXiv},
primaryClass={cs.RO},
url={https://arxiv.org/abs/2509.25358},
}