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

A PyTorch implementation of SARM (Stage-Aware Reward Modeling for Long Horizon Robot Manipulation) for robotics tasks with variable-length sequences, stage classification, and progress regression.

Key Features:

• Multi-modal input processing (RGB images, joint states, task embeddings)
• Variable sequence length handling with attention masking
• Dual prediction heads for stage classification and progress regression
• CLIP integration for robust visual feature extraction
• HuggingFace Hub integration for easy model sharing
• Model supports multiple tasks, but best results so far from fine-tuning on a specific task

TL;DR

The idea here is that we can use data about stage progress to upsample good windows and downsample bad windows (where robot is fumbling or undoing progress). We use dataset annotations to create a ground truth (GT) of robot progress, based on stages and progress within the stage. The progress that the stage makes is based on the relative proportion of the stage in the demonstration data for the task.

We have:

• SARM weights in the weighted ranges format without the model based on annotations only (outputs/sarm_weights)
• progress -> progress SARM weights from a trained model checkpoint in outputs/progress_weights

We have additional weighted data loader with weighted sampler implementation outside of this repo that supports that weighed ranges format.

Important files:

• train_simple.py - training a SARM progress model for the configuration in config.json
• inference.py - use the SARM model checkpoint to generate progress estimates for a given dataset → 📖 INFERENCE_README
• generate_weights.py - transform the progress from the inference into weighted ranges → 📖 WEIGHTS_README
• dataloader.py - configures training, validation and inference data loaders

Misc:

• profile_training.py - profile where the training is slow → 📖 Profiling Guide
• hf_model_hub.py - download / upload existing models → 📖 HUGGINGFACE_README

Exploration/Visualization/Helpers:

• download_annotations.ipynb - useful if the annotations are missing
• ground_truth_vis.ipynb - visualize the data used to train SARM as well as the data generated from inference
• task_stages.ipynb - preparing the data for training + visualizations of annotation data

Setup

This repo is known to work on both MacOS and Linux. You will need MPS or GPU. Not tested on Windows.

1. Environment Setup

Create a conda environment with Python 3.10 or 3.11 (required for OmniGibson compatibility):

conda create -n sarm python=3.10
conda activate sarm

2. Install BEHAVIOR Repository and Dependencies

Clone and install the BEHAVIOR repository which includes BDDL and OmniGibson:

# Clone BEHAVIOR repository
git clone https://github.com/StanfordVL/BEHAVIOR-1K.git
cd BEHAVIOR-1K

# Install BDDL
cd bddl
pip install -e .
cd ..

# Install OmniGibson  
cd OmniGibson
pip install -e .
cd ..

3. Install Core Dependencies

# PyTorch (adjust CUDA version as needed)
pip install torch torchvision torchaudio

# CLIP for visual features
pip install git+https://github.com/openai/CLIP.git

# Core ML libraries
pip install pandas numpy pillow

# Configuration and utilities
pip install omegaconf dm_tree

# Audio/video processing
pip install av requests

# HuggingFace ecosystem
pip install datasets huggingface_hub google-auth

# LeRobot (specific version)
pip install lerobot==0.3.2

# Optional: Weights & Biases for experiment tracking
pip install wandb

# Optional: Development tools
pip install tqdm

4. Handle PyMeshLab Issues (if needed)

If you encounter issues with PyMeshLab version during the instalation of Omnigibson, just don't install it (remove it from the omnigibson setup.py dependencies).

Download

Pre-trained Models

Download pre-trained SARM models from HuggingFace Hub using hf_model_hub.py:

# Download a specific model
python hf_model_hub.py download \
    --repo-id username/sarm-model \
    --output-dir ./models

# Or use the HuggingFace CLI directly
huggingface-cli download username/sarm-model --local-dir ./models

→ 📖 HuggingFace Guide for complete model sharing documentation

Dataset

The SARM model is trained on the BEHAVIOR dataset (224x224 resolution). Use the download_annotations.ipynb notebook for data exploration:

jupyter notebook download_annotations.ipynb

Dataset: "IliaLarchenko/behavior_224_rgb"

from datasets import load_dataset
dataset = load_dataset("IliaLarchenko/behavior_224_rgb")

Training

1. Data Setup

The model works with the BEHAVIOR dataset (224x224 resolution). The dataloader.py is pre-configured - no manual preparation needed.

2. Configure Training

Edit config.json to match your setup:

{
    "model": {
        "d_model": 768,
        "n_heads": 12, 
        "n_layers": 8,
        "num_stages": 100,
        "d_state": 256,
        "num_tasks": 50
    },
    "training": {
        "max_steps": 10000,
        "learning_rate": 1e-4,
        "batch_size": 16,
        "gradient_accumulation_steps": 4
    },
    "data": {
        "max_sequence_length": 13,
        "tasks": [8],
        "train_episodes": [1, 2, 3, 4, 5],
        "val_episodes": [91, 92, 93, 94, 95]
    }
}

3. Start Training

# Basic training (recommended)
python train_simple.py

# With custom config
python train_simple.py --config my_config.json

# Resume training with different seed (useful to avoid repeating same data)
python train_simple.py --resume checkpoint_step_400.pt --seed 123

# Test dataloader first (recommended)
python dataloader_test.py

Training Script Options:

usage: train_simple.py [-h] [--config CONFIG] [--resume RESUME] [--max-steps MAX_STEPS] [--seed SEED]

4. Monitor Progress

• Loss convergence: Check validation metrics in terminal output
• Checkpoints: Saved at intervals specified by save_steps in config.json
• Best model: Automatically saved as best_model.pt
• Checkpoint directory: Configured in config.json under training settings

5. Performance Profiling

Use profile_training.py to identify bottlenecks:

# Profile current setup
python profile_training.py baseline

# Test multiprocessing optimization
python profile_training.py multiprocessing

# Analyze results
python profile_training.py analyze

→ 📖 Profiling Guide for detailed performance optimization

Use the Model

Inference on Episodes

Use inference.py to generate progress predictions:

# Using configuration from config.json
python inference.py \
    --checkpoint best_model.pt \
    --config config.json \
    --episodes 1 2 3 \
    --sampling-interval 5.0 \
    --jitter-frames 30 \
    --run-name my_inference

# Specify task manually
python inference.py \
    --checkpoint best_model.pt \
    --task 8 \
    --episodes 1 2 3 \
    --run-name my_inference

→ 📖 Inference Guide for detailed usage and options

Generate Weights for Data Sampling

Transform inference results into weighted ranges using generate_weights.py:

python generate_weights.py --help

Output goes to outputs/progress_weights directory.

→ 📖 Weights Guide for weight generation and analysis

Programmatic Usage

import torch
from model import SARM  # See model.py

# Load trained model
model = SARM(d_model=768, num_stages=100, d_state=256, num_tasks=50)
checkpoint = torch.load('best_model.pt')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Run inference
with torch.no_grad():
    stage_logits, progress_pred = model(images, states, tasks, padding_mask)
    
    # Get predictions
    stage_predictions = torch.argmax(stage_logits, dim=-1)  # (B, N)
    progress_values = progress_pred  # (B, N)

Share Models

Upload trained models to HuggingFace Hub using hf_model_hub.py:

python hf_model_hub.py upload \
    --checkpoint best_model.pt \
    --repo-id your-username/sarm-model-name \
    --config config.json

→ 📖 HuggingFace Guide for complete sharing workflow

Configuration

Edit config.json to customize your setup. Here are the key parameters:

🏗️ Model Architecture

Parameter	Default	Description
d_model	768	Hidden dimension for transformers and embeddings
n_heads	12	Number of attention heads in transformer layers
n_layers	8	Number of transformer encoder layers
num_stages	100	Number of task stages for classification (should match your data)
d_state	256	Dimension of robot joint state input
num_tasks	50	Number of different task types in your dataset

🎯 Training Settings

Parameter	Default	Description
max_steps	10000	Maximum training steps
learning_rate	1e-4	Adam optimizer learning rate
batch_size	16	Batch size per GPU
gradient_accumulation_steps	4	Effective batch size = batch_size × this
stage_loss_weight	1.0	Weight for stage classification loss
progress_loss_weight	1.0	Weight for progress regression loss
save_steps	200	Frequency of model checkpointing

📊 Data Configuration

Parameter	Default	Description
tasks	[8]	Task numbers: 8="rearranging_kitchen_furniture", 25="clearing_food_from_table_into_fridge"
max_sequence_length	13	Maximum length of input sequences
train_episodes	[1,2,3...]	Episode indices for training
val_episodes	[91,92,93...]	Episode indices for validation
num_workers	10	Number of dataloader worker processes

Example configurations:

• config.json - Main configuration
• profiling/config_profiling.json - For performance profiling
• config_exp1_multiprocessing.json - Multiprocessing optimization
• config_exp5_large_batches.json - Large batch testing

Concepts

Stage-Aware Reward Modeling

SARM addresses the challenge of long-horizon robot manipulation by modeling task completion as a combination of:

1. Stage Classification: Identifying which stage of the task the robot is currently in
2. Progress Regression: Estimating how much progress has been made within the current stage

This dual approach allows the model to understand both the discrete phases of a task and the continuous progress within each phase.

Multi-Modal Input Processing

The model processes three types of input:

• Visual Information: RGB images processed through a frozen CLIP encoder
• Proprioceptive Information: Robot joint states and end-effector positions
• Task Context: Learned task embeddings that capture task-specific patterns

Variable-Length Sequence Handling

Because we use the REWIND augmentation, we sometimes have longer sequences. SARM handles this through:

• Padding shorter sequences to a fixed length
• Attention masking to ignore padded positions
• Loss computation only on valid (non-padded) timesteps

Implementation Details

🏗️ Model Architecture

See model.py for complete implementation:

Input Processing:
├── CLIP Visual Encoder (frozen) → Linear Projection
├── Joint State → LayerNorm → Linear Projection  
└── Task ID → Embedding Lookup

↓

Transformer Backbone:
├── Positional Bias (learnable)
├── Multi-Head Self-Attention × N layers
└── Feed-Forward Networks

↓

Prediction Heads:
├── Stage Classification → CrossEntropy Loss
└── Progress Regression → MSE Loss

📊 Data Format

The dataloader.py handles all data formatting automatically. For reference, batches have this structure:

We feed data as B, N, individual_data_dim, where N is padded to 13 but actual sequences are variable length (9-13 frames). We don't give the model an individual frame to predict progress, but a set of frames. The first frame is frame 0 of a given episode (the 'anchor'). For timestamp T, frame 8 will have data for timestamp T, and the preceding 7 frames will have data for timestamps T-7, T-6, ..., T-1. Generally we space the preceding timestamps by 1s but it's configurable.

So as an example, we will have:

0, T-7, T-6, T-5, T-4, T-3, T-2, T-1, T

Occasionally (5% of the time - configurable) we add some 'rewinding' (adding timestamps from the past in reverse order to undo progress):

0, T-7, T-6, T-5, T-4, T-3, T-2, T-1, T, T-1, T-2, T-3

Since rewinding can extend our sequence from the base 9 frames to up to 13 frames, we pad all sequences to length 13 and use attention masking to handle the variable lengths.

{
    'images': torch.Tensor,        # (B, N, 3, 224, 224) - RGB sequences
    'states': torch.Tensor,        # (B, N, D_state) - Joint states
    'tasks': torch.Tensor,         # (B,) - Task IDs  
    'stage_labels': torch.Tensor,  # (B, N) - Ground truth stages
    'progress_labels': torch.Tensor, # (B, N) - Ground truth progress [0,1]
    'padding_mask': torch.Tensor   # (B, N) - True = padded position
}

🎯 Loss Function

Implemented in model.py as SARMWithLoss:

total_loss = α × stage_loss + β × progress_loss

where:
stage_loss = CrossEntropyLoss(stage_logits, stage_labels, ignore_padded=True)
progress_loss = MSELoss(progress_pred, progress_labels, ignore_padded=True)

Configure α and β via stage_loss_weight and progress_loss_weight in config.json.

⚡ Performance Features

• Frozen CLIP: Visual encoder frozen to reduce memory/computation (see model.py)
• Attention Masking: Padding positions ignored in attention (see model.py)
• Gradient Accumulation: Larger effective batch sizes (see train_simple.py)
• Mixed Precision: AMP for faster training (see train_simple.py)
• Profiling Tools: Use profile_training.py → 📖 Profiling Guide

📈 Evaluation Metrics

Tracked during training in train_simple.py:

• Stage Accuracy: % correctly classified stages (non-padded positions only)
• Progress MAE: Mean absolute error between predicted/true progress values
• Combined Loss: Weighted sum of classification + regression losses

🔧 Key Files

File	Purpose
model.py	Core SARM model implementation
train_simple.py	Training script with profiling support
dataloader.py	Data loading and preprocessing
inference.py	Model inference and evaluation
generate_weights.py	Convert progress to weighted ranges

Troubleshooting

🔧 Installation Issues

CLIP Import Error: Use the official OpenAI repository:

pip install git+https://github.com/openai/CLIP.git

OmniGibson Issues: Ensure Python 3.10/3.11 and install BDDL first. If PyMeshLab fails, remove it from OmniGibson's setup.py dependencies.

🚀 Performance Issues

GPU Memory: Reduce batch size or increase gradient accumulation in config.json:

{
    "training": {
        "batch_size": 8,
        "gradient_accumulation_steps": 8
    }
}

Slow Training: Use profile_training.py to identify bottlenecks:

python profile_training.py baseline
python profile_training.py analyze

📊 Data Issues

Episode Not Found: Check that episode indices in config.json match available data

Tensor Size Mismatch: Verify task tensor batch size matches other inputs - check dataloader.py

Training Doesn't Start: Test dataloader first:

python dataloader_test.py

🔗 Additional Resources

• 📖 Inference Guide - Detailed inference usage
• 📖 HuggingFace Guide - Model sharing workflow
• 📖 Weights Guide - Weight generation and analysis
• 📖 Profiling Guide - Performance optimization

📁 Key Output Directories

• outputs/sarm_weights/ - Annotation-based SARM weights
• outputs/progress_weights/ - Model-generated progress weights
• outputs/visualizations/ - Generated plots and analysis
• profiling/profiling_results/ - Performance profiling data