SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation

1. Problem Statement

The paper addresses the critical challenge of long-horizon, contact-rich robot manipulation, specifically involving deformable objects (e.g., T-shirt folding). While recent Robot Behavior Models (RBMs) have shown promise, they struggle with:

• Inconsistent Demonstration Quality: Large datasets often contain noisy, suboptimal, or failed trajectories from non-expert operators.
• Brittle Labeling: Existing methods often rely on frame indices or simple time-based progress labels. In tasks with variable durations and non-linear progress (e.g., flattening a shirt takes different amounts of time depending on the fabric), frame-based labeling introduces severe noise, making it impossible to distinguish between a "slow but correct" step and a "failed" step.
• Generalization: Models trained on curated data fail to generalize to out-of-distribution (OOD) scenarios, such as recovery from errors or handling occlusions in deformable objects.

The core problem is how to derive consistent, fine-grained reward signals from diverse, variable-length demonstrations to guide policy learning effectively without requiring perfect, hand-crafted reward functions.

2. Methodology

The authors propose a two-part framework: SARM (Stage-Aware Reward Modeling) for generating rewards, and RA-BC (Reward-Aligned Behavior Cloning) for utilizing those rewards to train policies.

A. SARM: Stage-Aware Reward Modeling

SARM is a video-based reward modeling framework designed to predict task progress robustly.

• Natural Language Annotations: Instead of frame indices, the method uses natural language subtask annotations (e.g., "grab," "flatten," "fold") to decompose tasks into semantic stages.
• Dual-Head Architecture:
  1. Stage Estimator: A classification head that predicts the current high-level task stage (e.g., "Flattening Phase") based on visual observations (RGB frames) and joint states.
  2. Subtask Estimator: A regression head that predicts fine-grained progress ($\in [0, 1]$) conditioned on the predicted stage.
• Labeling Strategy:
  • The authors compute a dataset-level prior proportion ($\bar{\alpha}_k$) for each subtask based on the average temporal duration of that subtask across the dataset.
  • Progress labels are generated by linearly interpolating within subtasks, ensuring that identical semantic states (e.g., a fully flattened shirt) receive consistent progress values regardless of how many frames it took to get there.
• Input Processing: The model uses a frozen CLIP encoder for visual features and a Transformer encoder to capture temporal dependencies. It employs rewind augmentation (appending reversed frames from earlier timestamps) to improve robustness against failure cases and OOD scenarios.

B. RA-BC: Reward-Aligned Behavior Cloning

RA-BC is a training framework that leverages the learned SARM reward model to filter and reweight demonstration data during policy training.

• Progress Delta Calculation: For each training sample, the model calculates the progress improvement ($\Delta \phi$) between the current state and the next state using the SARM reward model.
• Dynamic Weighting:
  • The system maintains online running statistics (mean $\mu$ and standard deviation $\sigma$) of progress deltas.
  • Samples with positive progress (above a threshold) are up-weighted, while those with negative progress (regressions/failures) are down-weighted or filtered out.
  • A "prior override" mechanism assigns a weight of 1 to clearly good samples and 0 to clearly bad ones to ensure decisive filtering.
• Objective: The standard Behavior Cloning (BC) loss is replaced with a weighted loss, effectively teaching the policy to imitate only the high-quality, progressive segments of the diverse dataset.

3. Key Contributions

1. SARM Framework: A novel stage-aware, video-based reward model that jointly predicts task stages and fine-grained progress using natural language annotations. It overcomes the brittleness of frame-index labeling and generalizes well to OOD scenarios.
2. RA-BC Framework: A policy training method that uses the learned reward model to automatically identify and reweight high-quality demonstrations, enabling effective learning from noisy, diverse datasets without manual filtering.
3. Empirical Validation on Deformable Objects: Successful application to the highly challenging T-shirt folding task, demonstrating significant improvements over standard Behavior Cloning and other reward modeling baselines.

4. Experimental Results

The authors evaluated their approach on two tasks: T-shirt folding (long-horizon, deformable) and Dish unloading (multi-stage, rigid).

• Reward Model Performance:

  • SARM significantly outperformed baselines (LIV, VLC, GVL, REDS, ReWiND) in predicting progress on human demonstrations and real robot rollouts.
  • On T-shirt folding, SARM achieved a 0.008 MSE on validation data (vs. 0.019 for the next best baseline) and a 0.94 classification score on real robot rollouts (vs. 0.78 for ReWiND).
  • SARM was robust to lighting variations and failure modes where other models produced noisy or optimistic signals.

• Policy Performance (T-shirt Folding):

  • Baselines: Vanilla BC trained on the full diverse dataset ($D_{all}$) achieved 0% success on the "Hard" task (folding from a crumpled state).
  • Naive Filtering: Filtering by duration ($D_{2min}$) improved performance but still failed on the Hard task (0% success).
  • RA-BC with SARM:
    • Medium Task (Flat to Folded): 83% success rate.
    • Hard Task (Crumpled to Folded): 67% success rate.
  • Comparison: RA-BC with SARM significantly outperformed RA-BC with the ReWiND reward model (67% vs. 25% on the Hard task), proving that reward model quality is the bottleneck for RA-BC.

• Ablation Studies:

  • Stage Awareness: Removing the stage estimator degraded performance, confirming the need for hierarchical modeling.
  • Rewind Augmentation: Essential for generalizing to real-world failure cases; without it, the model became overly optimistic.
  • Model Scale: A 60M parameter model was found to be the optimal balance between capacity and overfitting.

5. Significance

This paper provides a scalable and annotation-efficient solution for long-horizon robotic manipulation.

• Data Efficiency: It demonstrates that high-quality policy learning is possible even with noisy, diverse datasets if the reward signal can accurately distinguish between "slow progress" and "failure."
• Scalability: By leveraging natural language annotations (which are easier to obtain than dense reward functions) and automating the reward derivation, the method scales to complex tasks like deformable object manipulation.
• Practical Impact: The results on T-shirt folding (67% success from a crumpled state) represent a significant leap forward for household robotics, a domain where current models typically fail due to the complexity of contact-rich, deformable interactions.

The work establishes that reward modeling is a critical bottleneck in imitation learning and that stage-aware, semantic progress estimation is the key to unlocking robust long-horizon robot behaviors.