SARM2: Multi-Task Stage Aware Reward Modeling for Self Improving Robotic Manipulation

2026-06-08

AI Summary

Existing reward models struggle to balance generality and fine-grained phase awareness in long-horizon robotic manipulation, limiting the generalization capability of reinforcement learning. This work proposes a multitask phase-aware reward model (RM) that integrates an action-primitive phase estimator with a multi-gated mixture-of-experts (MMoE) value head, enabling for the first time cross-task dense reward modeling without task-specific annotations. Embedded within the SPIRAL framework, the approach leverages inexpensive autonomous trajectories to drive self-improving policy learning. Evaluated across ten tasks, the method reduces mean squared error in value estimation by 80% and achieves success rates of 100% and 90% on Folding Shorts and Cleaning Whiteboard tasks, respectively, establishing a critical pathway toward high-quality reward learning and a stable data flywheel.

Abstract

Fine-tuning vision-language-action (VLA) policies for long-horizon manipulation still relies heavily on behavior cloning, which requires costly high-quality demonstrations and keeps policies near the demonstration distribution. Reward models can reduce this dependence by reweighting demonstrations and providing dense supervision for on-robot reinforcement learning (RL), but they must be dense, accurate, and general. Existing methods fall short: task-specific stage-aware models are accurate but require per-task annotations, while general vision-language-model (VLM) reward models are broadly applicable but too coarse for fine-grained long-horizon progress. We introduce RM, a multi-task stage-aware reward model that combines an action-primitive-based stage estimator with a multi-gate Mixture-of-Experts (MMoE) value head to produce dense per-step rewards across manipulation tasks. Building on RM, we further propose SPIRAL (Self-Policy Improvement via Reward-Aligned Learning), an on-policy reward-guided framework that improves VLA policies from cheap autonomous rollouts. On a 10-task benchmark, RM reduces value-estimation MSE by 80% over the strongest baselines; when used in SPIRAL, it improves task success from around 50% to near-perfect performance on Folding Shorts (58% to 100%) and Cleaning Whiteboard (50% to 90%), showing that high-quality dense rewards are key to a stable robot data flywheel. Project website: https://qianzhong-chen.github.io/sarm2.github.io/.

Problem

• reward modeling
• robotic manipulation
• long-horizon tasks
• stage-awareness
• vision-language-action policies

Innovation

• stage-aware reward modeling
• multi-task reinforcement learning
• Mixture-of-Experts
• robotic manipulation
• self-improving policies