ARM: Advantage Reward Modeling for Long-Horizon Manipulation

1. Problem Statement

Long-horizon robotic manipulation tasks (e.g., folding a towel, assembling parts) pose significant challenges for Reinforcement Learning (RL) and Vision-Language-Action (VLA) models due to the credit assignment problem.

• Sparse Rewards: Standard binary success/failure rewards provide insufficient guidance for learning complex, multi-step sequences.
• Dense Reward Bottleneck: While dense progress rewards are necessary, obtaining them is costly. Existing methods rely on:
  • Absolute Progress Modeling: Assuming a strict monotonic relationship between time and progress. This fails in real-world scenarios involving backtracking, error recovery, or non-linear maneuvers.
  • Zero-shot VLMs: These suffer from high inference costs, spatial-geometric grounding issues, and non-monotonic oscillations in reward signals.
  • Manual Annotation: Requires high cognitive load, leading to inter-annotator inconsistency and scalability issues.
• Suboptimal Data: Existing datasets often contain noisy, suboptimal human demonstrations that standard Behavior Cloning (BC) cannot effectively filter.

2. Methodology: Advantage Reward Modeling (ARM)

The authors propose ARM, a framework that shifts from estimating absolute progress to estimating relative advantage between states. The system consists of three core components:

A. Advantage Reward Model (The Core)

Instead of predicting a continuous scalar progress value (MISO - Multi-Input Single-Output), ARM uses a MIMO (Multi-Input Multi-Output) Temporal Transformer.

• Architecture: It processes a causal window of historical observations (visual features via CLIP, robot proprioceptive states, and task instructions) in parallel.
• Dual-Head Learning:
  1. Multi-frame Advantage Classification: Predicts the relative gain ($\Delta \hat{y}$) between consecutive time steps as a discrete classification: Progressing (+1), Regressing (-1), or Stagnant (0). This avoids the noise of continuous regression.
  2. Task Completion Head: Predicts the probability of the current state being a successful terminal state. This acts as an "anchor" to align relative gains with absolute task completion.
• Training: Supervised by a lightweight Tri-state Labeling Strategy (detailed below) and optimized using Focal Loss to handle class imbalance in successful terminal states.

B. Global Progress Reconstruction

ARM transforms discrete, local advantage predictions into a dense, globally consistent progress trajectory ($P_t$).

• Process: The system integrates the predicted relative gains ($\Delta \hat{y}$) and accumulates them, anchored by the Task Completion Head's prediction (setting $P_T = 1.0$ at success).
• Efficiency: By using a MIMO architecture with video clipping, the model processes non-overlapping segments in parallel, significantly accelerating inference compared to sliding-window approaches.

C. Advantage-Weighted Behavior Cloning (AW-BC)

This is the policy optimization algorithm that utilizes the reconstructed rewards.

• Length-Adaptive Gain: To handle datasets with varying episode lengths, the gain is normalized by the average episode length to prevent bias toward shorter sequences.
• Statistical Reweighting: Samples are weighted based on their advantage.
  • Regressive/Failed samples: Weighted near 0 (filtered out).
  • High-advantage samples: Weighted near 1 (prioritized).
• Objective: Minimizes the weighted negative log-likelihood, effectively performing offline policy improvement by prioritizing high-value recovery trajectories and suppressing suboptimal behaviors.

3. Key Contributions

1. Tri-State Advantage Labeling Strategy:

   • Replaces complex continuous progress labeling with a simple categorical classification: Progressing (+1), Regressing (-1), Stagnant (0).
   • Benefits: Drastically reduces human cognitive load, ensures high cross-annotator consistency, and naturally handles non-monotonic behaviors (backtracking/recovery) that break monotonic assumptions.

2. MIMO Advantage Reward Model:

   • Introduces a MIMO Transformer architecture that predicts sequences of relative advantages rather than single scalar values.
   • Enables automated pseudo-labeling at scale and provides a cost-effective alternative to expensive VLM inference or manual dense annotation.

3. AW-BC Algorithm:

   • Extends Reward-Aligned Behavior Cloning (RA-BC) by incorporating adaptive scaling and dense relative advantage signals.
   • Successfully filters suboptimal samples and prioritizes error-correction behaviors, enabling stable learning from fragmented and noisy datasets (including DAgger-style data).

4. Experimental Results

The framework was evaluated on a challenging long-horizon bimanual towel-folding task (8 distinct stages, 120s limit).

• Policy Performance:
  • Success Rate: Achieved 99.4%, compared to 62.1% for standard Behavior Cloning and 78.5% for the SARM baseline (which uses absolute progress).
  • Task Throughput: 32 episodes/hour (vs. 18 for baseline).
  • Folding Precision: Highest alignment score (3.6) among all methods.
• Reward Model Accuracy:
  • MSE: ARM achieved an MSE of 0.0014 against Ground Truth, significantly outperforming SARM (0.0059).
  • Success/Failure Identification: 100% accuracy in classifying terminal states, compared to ~83-91% for SARM.
• Efficiency:
  • Labeling: The tri-state strategy allowed human annotators to label 250 samples/shift (2.5x faster than subtask segmentation). The automated pipeline processes >2,000 samples/shift (20x faster than humans).
  • Inference: The MIMO architecture achieved 14.1 iterations/second, a 13.7x speedup over zero-shot VLM labeling and a 3.6x speedup over SARM.

5. Significance

• Solving the Reward Engineering Bottleneck: ARM provides a scalable, cost-effective solution to the "Reward Engineering Bottleneck" in robotics, removing the need for expensive, task-specific dense reward engineering.
• Handling Non-Monotonicity: By focusing on relative advantage rather than absolute progress, the framework robustly handles real-world robotic behaviors like backtracking, retries, and recovery, which previously caused reward signal oscillations.
• Data Efficiency: The ability to learn from fragmented, noisy, and suboptimal data (including error-correction episodes) without human intervention during training marks a significant step toward autonomous policy refinement in unstructured environments.
• Scalability: The combination of low-cost labeling and high-speed inference makes this approach viable for large-scale deployment in complex manipulation tasks.