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

1. Problem Statement

Vision-Language-Action (VLA) models have established end-to-end policy learning as a dominant paradigm in robotic manipulation. However, these models typically struggle with long-horizon tasks without extensive fine-tuning on in-domain data. Supervised Fine-Tuning (SFT) relies on costly, high-quality human demonstrations and remains prone to out-of-distribution (OOD) failures. While Reinforcement Learning (RL) offers a path to self-improvement, it faces two critical bottlenecks:

1. Sparse Rewards: Existing online RL methods often rely on sparse terminal rewards, which fail to provide the dense, step-level supervision necessary for long-horizon, contact-rich tasks.
2. Reward Model Limitations: Current reward modeling approaches fall into two suboptimal categories:
   • Task-Specific Models: Accurate but require per-task annotations and retraining for every new task (e.g., SARM), limiting scalability.
   • General-Purpose VLM Models: Broadly applicable but produce coarse-grained, noisy predictions at the step level, making them unsuitable as dense supervisory signals for fine-grained manipulation progress.

The paper posits that scaling VLA-RL requires a reward model that is simultaneously dense (per-step), accurate (faithful to task progress), and general (applicable across diverse tasks without retraining).

2. Methodology

The authors propose SARM2 (Multi-Task Stage Aware Reward Modeling) and SPIRAL (Self-Policy Improvement via Reward-Aligned Learning).

3.1 SARM2: Multi-Task Stage-Aware Reward Model

SARM2 addresses the trade-off between generalization and accuracy by combining a task-agnostic stage estimator with a specialized value head.

• Action-Primitive Stage Estimator:

  • Instead of task-specific stage annotations, the model uses a shared vocabulary of action primitives (e.g., "Pick," "Place," "Rotate," "Erase") as an intermediate representation.
  • Based on an analysis of 10,000+ hours of manipulation data, the authors identified K=21 action primitives that cover >90% of task duration, plus a "null" class for uncertainty.
  • The estimator is a 4-layer causal Transformer that takes multi-view images (wrist-left, wrist-right, top-down) and proprioceptive states as input. It classifies the current segment into one of K+1=22 primitive classes.
  • This estimator is trained on a balanced dataset of 66 hours of segments spanning 100 tasks, allowing it to generalize across domains.

• Multi-Gate Mixture-of-Experts (MMoE) Value Head:

  • To produce dense value estimates (specifically, normalized remaining steps to completion), SARM2 employs an MMoE decoder.
  • Architecture: The value model uses a separate 6-layer causal Transformer backbone. It receives the same visual/proprioceptive inputs as the stage estimator, plus the task embedding and the predicted primitive embedding from the stage estimator.
  • Routing Mechanism: The predicted primitive selects a specific gate in the MMoE. Each gate routes tokens to a shared pool of experts (3-layer MLPs). This design allows the model to specialize in the visual dynamics and progress signatures of specific primitive groups while sharing capacity across tasks.
  • Objective: The model predicts a dense reward $r_t \in [-1, 0]$ representing the remaining progress. The loss function combines Mean Squared Error (MSE) on the goal distance with auxiliary losses for gate balance and entropy to prevent router collapse.

3.2 SPIRAL: Self-Policy Improvement Framework

SPIRAL is an on-policy, reward-guided framework designed to create a "robot data flywheel" using SARM2's dense rewards.

• Residual RL Substrate: SPIRAL builds upon the DICE-RL framework, using a residual learning approach where a residual policy corrects the outputs of a base VLA policy.
• Hybrid Objective:
  • Dense Component: Uses TD3 with dense per-step rewards provided by SARM2.
  • Sparse Component: Combines this with a Monte Carlo (MC) objective driven by sparse episode-level rewards to ensure alignment with the ultimate task goal.
• Autonomous Loop: The framework operates in a closed loop:
  1. Initialization: Fine-tune a base VLA policy via Behavior Cloning (BC).
  2. One-Time Adaptation: A small set of human-annotated rollouts adapts the reward model to the policy's rollout distribution.
  3. Self-Improvement: The system autonomously collects rollouts, relabels them with the adapted SARM2, and updates the policy via residual RL. This cycle repeats without further human intervention.

3. Key Contributions

1. SARM2: A multi-task stage-aware reward model that couples a generalizable action-primitive stage estimator with a multi-gate MMoE value head. This architecture produces accurate, dense rewards across diverse long-horizon tasks within a single model, eliminating the need for per-task retraining.
2. SPIRAL: A reward-aligned on-policy residual RL framework that leverages SARM2's dense rewards to enable continual policy refinement from cheap autonomous rollouts, closing an autonomous robot data flywheel.
3. Empirical Validation: Demonstration that high-quality dense rewards are critical for stable self-improvement, showing substantial gains over sparse-reward and VLM-based baselines on real-world tasks.

4. Experimental Results

The authors evaluated SARM2 and SPIRAL on a 10-task benchmark and two specific real-world long-horizon tasks: Folding Shorts and Cleaning Whiteboard.

Reward Model Performance

• Accuracy: SARM2 reduced the Mean Squared Error (MSE) on held-out demonstrations by 80% compared to the strongest baselines (ReWiND, TOPReward, Robometer).
• Generalization: On the "Cleaning Whiteboard" task (T2), SARM2 achieved a rollout classification correlation ($\rho$) of 0.667, significantly outperforming general-purpose VLM-based models (which often yielded negative correlations due to over-optimism).
• Ablations: Removing the stage estimator increased demo loss by ~70%; removing the multi-gate design increased it by ~30%. This confirms that both stage awareness and MMoE routing are synergistic.
• Router Health: The MMoE routing remained balanced (Density score ~0.10), whereas single-gate variants suffered from router collapse (Density > 0.80).

Policy Learning Performance (SPIRAL)

• Folding Shorts:
  • Flat: Success rate improved from 58% (initial BC/RL-Dense) to 100% after 3 rounds of SPIRAL with SARM2.
  • Crumpled (Harder): Success rate improved from 33% to 67%. Notably, sparse-reward baselines degraded performance on this task (dropping to 17%), highlighting the necessity of dense rewards.
• Cleaning Whiteboard:
  • Success rate improved from 50% to 90%.
  • Average progress score increased from 0.813 to 0.975.
• Comparison: SPIRAL with SARM2 consistently outperformed all baselines, including LoRA-fine-tuned VLM reward models (Robometer-FT) and sparse-reward RL. The results indicate that while density is necessary, the quality and stage-awareness of the reward are the bottlenecks for effective self-improvement.

5. Significance and Claims

The paper claims that SARM2 and SPIRAL demonstrate that high-quality dense rewards are the "load-bearing factor" for a self-sustaining robot data flywheel. The work challenges the notion that general-purpose VLMs are sufficient for fine-grained robotic control, showing instead that a structured, primitive-based approach is required for accurate progress estimation.

The authors emphasize that their method enables continual policy refinement on long-horizon tasks with minimal human intervention after an initial adaptation phase. By decoupling the stage estimation (generalizable primitives) from the value estimation (specialized experts), they achieve a model that is both scalable across tasks and precise enough to guide on-policy RL.

Limitations Acknowledged by the Authors:

• Embodiment Scope: The action-primitive vocabulary is derived from bimanual table-top data; extending to mobile manipulation requires rebuilding the primitive set.
• Residual RL Ceiling: The framework refines policies via residual learning rather than updating VLA weights, meaning it cannot correct errors at the intention or high-level task-decomposition level.
• Sample Efficiency: While autonomous, the on-policy rollout collection still requires human operators for routine resets between episodes.