SARM: Modelagem de Recompensa Consciente de Etapas

SARM (Stage-Aware Reward Modeling) é um framework de modelagem de recompensas baseado em vídeo para tarefas de manipulação robótica de longo horizonte. Este guia aborda como treinar modelos de recompensa SARM e, opcionalmente, como usá-los com Reward-Aligned Behavior Cloning (RA-BC).

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

Por que Modelos de Recompensa?

O clonagem de comportamento padrão trata todos os quadros de demonstração de forma igual, mas os conjuntos de dados de robôs do mundo real são desorganizados. Eles contêm hesitações, correções e trajetórias de qualidade variável. Os modelos de recompensa resolvem isso aprendendo uma noção generalizável de progresso da tarefa a partir das demonstrações: dados os quadros de vídeo e uma descrição da tarefa, eles preveem o quão perto o robô está de concluir a tarefa (0→1). Esse “sinal de progresso” aprendido pode ser usado de várias maneiras, sendo duas aplicações promissoras: (1) aprendizado por imitação ponderado (RA-BC), em que os quadros com alto progresso recebem mais peso durante o treinamento da política, e (2) aprendizado por reforço, em que o modelo de recompensa fornece recompensas densas para o aprimoramento online ou offline da política.

Visão geral

O SARM possui os seguintes recursos:

1. Arquitetura consciente de etapas: Prediz simultaneamente a etapa de alto nível da tarefa e o progresso detalhado dentro de cada etapa
2. Anotações de subtarefas: Usa anotações de subtarefas em linguagem natural para derivar rótulos de progresso consistentes
3. Proporções temporais: Calcula os priors no nível do conjunto de dados (α̅_k) para cada subtarefa, a fim de normalizar o progresso entre demonstrações de comprimento variável

O SARM treina em um alvo compacto stage+tau para cada quadro:

• stage: índice do estágio inteiro k ∈ {0, ..., K-1}
• τ (tau): progresso dentro do estágio τ ∈ [0, 1]
• codificação de alvo: y = k + τ (é isso que o processador de conjunto de dados produz)

No momento da inferência (e no RA-BC de nível inferior), o SARM converte o valor bruto k + τ em um progresso normalizado em [0, 1] usando as proporções temporais de nível de conjunto de dados α̅_k (armazenadas em meta/temporal_proportions_*.json).

Isso corresponde à Fórmula (2) do artigo:

progress_t = P_{k-1} + α̅_k × τ_t

Onde:

• τ_t = (t - s_k) / (e_k - s_k) é o tempo normalizado dentro da subtarefa
• P_{k-1} é o anterior cumulativo (soma das proporções das subtarefas anteriores)
• α̅_k é a proporção temporal para a subtarefa k

Isso garante que estados de tarefa idênticos sejam mapeados para valores de progresso consistentes, mesmo em demonstrações de diferentes comprimentos.

Entradas e Alvos (O que o novo código espera)

O SARM é treinado por meio de seu processador (src/lerobot/rewards/sarm/processor_sarm.py), que:

• Codifica imagens e texto da tarefa com CLIP (ViT-B/32) em video_features e text_features
• Preenche/trunca o estado do robô em state_features (até max_state_dim)
• Constrói os alvos como sparse_targets (e dense_targets em dense_only/dual) usando a codificação stage+tau y = k + τ
• Máscaras reprocessam quadros usando um tensor lengths por amostra (reprocessar é uma aumentação no momento do treinamento)

No mínimo, cada amostra de treinamento precisa de:

• task (string): descrição da tarefa
• imagens de reward_model.image_key e estados de reward_model.state_key do conjunto de dados

Modos de Anotação

Você pode escolher entre 3 modos de anotação que determinam como os rótulos de progresso são calculados:

Detalhes do Modo

pip install -e ".[sarm]"

1. Train SARM → 2. Visualize predictions → 3. (Optional) Train policy with RA-BC

1. Annotate (dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

1. Annotate (sparse+dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

Etapa 1: Anotação de subtarefas

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --dense-only \
  --dense-subtasks "Bring robot arms up from starting position,Grab near side and do 1st fold,Grab side and do 2nd fold,Grab side and do 3rd fold to finish folding" \
  --video-key observation.images.base \
  --num-workers 4 \
  --push-to-hub

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --sparse-subtasks "Bring arms up from starting position,Fold the towel (3 folds in total)" \
  --dense-subtasks "Bring robot arms up from starting position,Grab near side and do 1st fold,Grab side and do 2nd fold,Grab side and do 3rd fold to finish folding" \
  --video-key observation.images.base \
  --num-workers 4 \
  --push-to-hub

Argumentos de Anotação

Observação: Após a conclusão da anotação, 5 episódios são visualizados automaticamente por padrão. Use --num-visualizations 0 para pular esta etapa.

Etapa 2: Verificar Anotações

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --visualize-only \
  --visualize-type dense \
  --num-visualizations 5 \
  --video-key observation.images.base \
  --output-dir ./subtask_viz

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --visualize-only \
  --visualize-type both \
  --num-visualizations 5 \
  --video-key observation.images.base \
  --output-dir ./subtask_viz

Isso gera visualizações mostrando os quadros do vídeo com as fronteiras das subtarefas sobrepostas e a linha do tempo das subtarefas.

Argumentos de Visualização

Dica: Se as anotações forem imprecisas, ajuste as descrições dos subtarefas para serem mais específicas e execute novamente.

Etapa 3: Treinar SARM

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=single_stage \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_single \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=dense_only \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_dense \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=dual \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_dual \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

Treinamento Multi-GPU

Adicione accelerate launch --multi_gpu --num_processes=4 para usar múltiplas GPUs para treinamento.

Argumentos de Treinamento

Etapa 4: Visualizar Previsões

Use compute_rabc_weights.py com --visualize-only para visualizar as previsões do modelo (e, se disponíveis, os alvos derivados das anotações) sem gravar um arquivo parquet.

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode sparse \
  --output-dir ./sarm_viz

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode dense \
  --output-dir ./sarm_viz

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode both \
  --output-dir ./sarm_viz

A visualização mostra:

• Gráfico de progresso: Progresso previsto (e anotação opcional derivada de “GT” quando disponível e --stride 1)
• Probabilidades de estágio: Gráfico de área empilhado das probabilidades de estágio previstas
• Quadros de amostra: Quadros-chave do episódio com rótulos de progresso/etapa

Argumentos de Visualização

Etapa 5 (Opcional): Treinar a Política com RA-BC

O Behavior Cloning Alinhado a Recompensas (RA-BC) usa o modelo SARM treinado para ponderar as amostras de treinamento com base na melhoria de progresso prevista. Isso requer duas etapas:

1. Pré-compute progress values para todos os frames usando o modelo SARM treinado
2. Treinar a política com ponderação RA-BC usando os valores pré-computados

Como o RA-BC funciona

Para cada amostra de treinamento, o RA-BC calcula o delta de progresso:

r_i = φ(o_{t+Δ}) - φ(o_t)

Onde φ é a previsão de progresso do SARM e Δ é o chunk_size da política. Amostras com progresso positivo (boas demonstrações) recebem pesos maiores, enquanto amostras com progresso negativo ou zero têm seus pesos reduzidos.

A ponderação segue as Equações 8-9 do artigo:

• Peso suave: w̃_i = clip((r_i − (μ − 2σ)) / (4σ + ε), 0, 1)
• Peso final: w_i = 𝟙{r_i > κ} + 𝟙{0 ≤ r_i ≤ κ} × w̃_i

Etapa 5a: Calcule os valores de progresso do SARM

Primeiro, execute o modelo SARM em todos os frames do seu conjunto de dados para calcular os valores de progresso:

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --head-mode sparse \
  --num-visualizations 5 \
  --push-to-hub

Este script:

• Processa todos os quadros e calcula os valores de progresso
• Salva os valores de progresso em um arquivo parquet ao lado do conjunto de dados no disco (padrão: <dataset_root>/sarm_progress.parquet)
• Gera visualizações dos primeiros N episódios (padrão: 5)

Argumentos:

Formato de saída (sarm_progress.parquet):

Etapa 5b: Treinar a Política com RA-BC

Depois de obter o arquivo de progresso, treine sua política com ponderação RA-BC. O arquivo de progresso é detectado automaticamente a partir do caminho do dataset (sarm_progress.parquet) se não for fornecido explicitamente. Atualmente, PI0, PI0.5, SmolVLA e FastWAM são suportados com RA-BC:

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=pi0 \
  --sample_weighting.type=rabc \
  --sample_weighting.head_mode=sparse \
  --sample_weighting.kappa=0.01 \
  --output_dir=outputs/train/policy_rabc \
  --batch_size=32 \
  --steps=40000

O script de treinamento faz automaticamente:

• Carrega os valores de progresso pré-computados do arquivo parquet
• Usa o chunk_size da política para calcular os deltas de progresso (Δ)
• Calcula pesos de amostras com base na melhoria do progresso
• Aplica perda ponderada durante o treinamento

Argumentos do RA-BC:

Ajuste do Kappa RA-BC

O parâmetro kappa é o limiar que determina quais amostras recebem peso total (w=1). Compreender como ajustá-lo é fundamental para que o RA-BC funcione de forma eficaz.

Como a ponderação funciona:

Diagnóstico de problemas de kappa:

Monitore estas métricas do Trackio durante o treinamento:

Se sample_weight_mean_weight ≈ 1.0: seu kappa está muito baixo. A maioria das amostras tem delta > kappa e ignora completamente o ponderamento suave. O RA-BC se torna equivalente ao BC vanilla.

Definindo kappa com base nos seus dados:

O valor padrão kappa=0.01 foi ajustado para a tarefa de dobrar camisetas do artigo (episódios de ~90s a 30fps). Para o seu conjunto de dados, verifique os valores registrados em sample_weighting/delta_mean e sample_weighting/delta_std:

# If delta_mean ≈ 0.03 and delta_std ≈ 0.02:
# Most deltas fall in range [0.01, 0.05]

# Option 1: Set kappa = delta_mean (medium selectivity)
--sample_weighting.kappa=0.03

# Option 2: Set kappa = delta_mean + delta_std (high selectivity)
--sample_weighting.kappa=0.05

# Option 3: Set kappa = delta_mean + 2*delta_std (very selective)
--sample_weighting.kappa=0.07

Quando o RA-BC pode não ajudar:

Se o seu conjunto de dados já for de alta qualidade (progresso consistente em todas as demonstrações), o RA-BC não trará muitos benefícios, pois não há nada para filtrar.

Treinamento Multi-GPU com RA-BC

accelerate launch \
  --multi_gpu \
  --num_processes=4 \
  src/lerobot/scripts/lerobot_train.py \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=pi0 \
  --sample_weighting.type=rabc \
  --sample_weighting.kappa=0.01 \
  --output_dir=outputs/train/policy_rabc \
  --batch_size=32 \
  --steps=40000

Dicas e Boas Práticas

Escolhendo um Modo

• Comece com single_stage para experimentos rápidos - sem sobrecarga de anotação
• Use dense_only quando você deseja um acompanhamento detalhado do progresso, mas as tarefas não possuem etapas de alto nível bem definidas
• Use dual para tarefas complexas em que o progresso em nível grosseiro e fino é significativo

Qualidade da Anotação

1. Seja específico com os nomes das subtarefas: em vez de “dobrar”, use “pegue o lado próximo e dobre em direção ao centro”
2. Verifique com visualização: Sempre verifique alguns episódios antes do treinamento
3. Nomenclatura consistente: Use os mesmos nomes de subtarefas em todos os episódios

RA-BC

1. Treine o SARM primeiro: a qualidade do RA-BC depende inteiramente da qualidade do SARM
2. Monitore sample_weight_mean_weight: se for ≈ 1.0, aumente kappa (veja Ajuste do Kappa do RA-BC)

Citação

@article{chen2025sarm,
  title={SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation},
  author={Chen, Qianzhong and Yu, Justin and Schwager, Mac and Abbeel, Pieter and Shentu, Yide and Wu, Philipp},
  journal={arXiv preprint arXiv:2509.25358},
  year={2025}
}