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

1. Formulação do Problema

O artigo aborda o desafio crítico da manipulação robótica de longo horizonte e rica em contato, especificamente envolvendo objetos deformáveis (por exemplo, dobrar uma camiseta). Embora os Modelos de Comportamento Robótico (RBMs) recentes tenham mostrado promessa, eles lutam com:

• Qualidade Inconsistente de Demonstrações: Grandes conjuntos de dados frequentemente contêm trajetórias ruidosas, subótimas ou falhas de operadores não especialistas.
• Rotulagem Frágil: Métodos existentes frequentemente dependem de índices de quadro ou rótulos de progresso baseados em tempo simples. Em tarefas com durações variáveis e progresso não linear (por exemplo, achatar uma camiseta leva quantidades diferentes de tempo dependendo do tecido), a rotulagem baseada em quadros introduz ruído severo, tornando impossível distinguir entre um passo "lento mas correto" e um passo "falho".
• Generalização: Modelos treinados em dados curados falham em generalizar para cenários fora da distribuição (OOD), como recuperação de erros ou manipulação de oclusões em objetos deformáveis.

O problema central é como derivar sinais de recompensa consistentes e de alta granularidade a partir de demonstrações diversas e de comprimento variável para guiar o aprendizado da política de forma eficaz, sem exigir funções de recompensa perfeitas e criadas manualmente.

2. Metodologia

Os autores propõem um framework de duas partes: SARM (Stage-Aware Reward Modeling) para gerar recompensas, e RA-BC (Reward-Aligned Behavior Cloning) para utilizar essas recompensas para treinar políticas.

A. SARM: Stage-Aware Reward Modeling

O SARM é um framework de modelagem de recompensa baseado em vídeo projetado para prever o progresso da tarefa de forma robusta.

• Anotações em Linguagem Natural: Em vez de índices de quadro, o método usa anotações de sub-tarefas em linguagem natural (por exemplo, "agarrar", "achatar", "dobrar") para decompor tarefas em estágios semânticos.
• Arquitetura de Dupla Cabeça:
  1. Estimador de Estágio: Uma cabeça de classificação que prevê o estágio atual de alto nível da tarefa (por exemplo, "Fase de Achatar") com base em observações visuais (quadros RGB) e estados das juntas.
  2. Estimador de Sub-tarefa: Uma cabeça de regressão que prevê o progresso de alta granularidade ($\in [0, 1]$) condicionado ao estágio previsto.
• Estratégia de Rotulagem:
  • Os autores computam uma proporção a priori no nível do conjunto de dados ($\bar{\alpha}_k$) para cada sub-tarefa, baseada na duração temporal média dessa sub-tarefa em todo o conjunto de dados.
  • Rótulos de progresso são gerados por interpolação linear dentro das sub-tarefas, garantindo que estados semânticos idênticos (por exemplo, uma camiseta totalmente achatada) recebam valores de progresso consistentes, independentemente de quantos quadros levaram para chegar lá.
• Processamento de Entrada: O modelo usa um encoder CLIP congelado para características visuais e um encoder Transformer para capturar dependências temporais. Emprega aumentação de retrocesso (anexando quadros revertidos de timestamps anteriores) para melhorar a robustez contra casos de falha e cenários OOD.

B. RA-BC: Reward-Aligned Behavior Cloning

O RA-BC é um framework de treinamento que aproveita o modelo de recompensa SARM aprendido para filtrar e reponderar dados de demonstração durante o treinamento da política.

• Cálculo de Delta de Progresso: Para cada amostra de treinamento, o modelo calcula a melhoria de progresso ($\Delta \phi$) entre o estado atual e o próximo estado usando o modelo de recompensa SARM.
• Ponderação Dinâmica:
  • O sistema mantém estatísticas de execução online (média $\mu$ e desvio padrão $\sigma$) dos deltas de progresso.
  • Amostras com progresso positivo (acima de um limiar) são reponderadas para cima, enquanto aquelas com progresso negativo (regressões/falhas) são reponderadas para baixo ou filtradas.
  • Um mecanismo de "anulação de priori" atribui um peso de 1 a amostras claramente boas e 0 a amostras claramente ruins para garantir filtragem decisiva.
• Objetivo: A perda padrão de Clonagem de Comportamento (BC) é substituída por uma perda ponderada, ensinando efetivamente à política a imitar apenas os segmentos de alta qualidade e progressivos do conjunto de dados diverso.

3. Principais Contribuições

1. Framework SARM: Um modelo de recompensa baseado em vídeo, novo e consciente de estágio, que prevê conjuntamente estágios de tarefa e progresso de alta granularidade usando anotações em linguagem natural. Supera a fragilidade da rotulagem por índice de quadro e generaliza bem para cenários OOD.
2. Framework RA-BC: Um método de treinamento de política que usa o modelo de recompensa aprendido para identificar e reponderar automaticamente demonstrações de alta qualidade, permitindo aprendizado eficaz a partir de conjuntos de dados ruidosos e diversos sem filtragem manual.
3. Validação Empírica em Objetos Deformáveis: Aplicação bem-sucedida à tarefa altamente desafiadora de dobrar camisetas, demonstrando melhorias significativas sobre a Clonagem de Comportamento padrão e outras bases de modelagem de recompensa.

4. Resultados Experimentais

Os autores avaliaram sua abordagem em duas tarefas: dobrar camisetas (longo horizonte, deformável) e descarregar pratos (multi-estágio, rígido).

• Desempenho do Modelo de Recompensa:

  • O SARM superou significativamente as bases (LIV, VLC, GVL, REDS, ReWiND) na previsão de progresso em demonstrações humanas e execuções reais de robôs.
  • Na dobra de camisetas, o SARM alcançou um MSE de 0,008 em dados de validação (vs. 0,019 para a melhor base seguinte) e uma pontuação de classificação de 0,94 em execuções reais de robôs (vs. 0,78 para o ReWiND).
  • O SARM foi robusto a variações de iluminação e modos de falha onde outros modelos produziram sinais ruidosos ou otimistas.

• Desempenho da Política (Dobra de Camisetas):

  • Bases: BC Vanilla treinado no conjunto de dados diverso completo ($D_{all}$) alcançou 0% de sucesso na tarefa "Difícil" (dobrar a partir de um estado amassado).
  • Filtragem Ingênua: Filtrar por duração ($D_{2min}$) melhorou o desempenho, mas ainda falhou na tarefa Difícil (0% de sucesso).
  • RA-BC com SARM:
    • Tarefa Média (Achatado para Dobrado): 83% de taxa de sucesso.
    • Tarefa Difícil (Amassado para Dobrado): 67% de taxa de sucesso.
  • Comparação: O RA-BC com SARM superou significativamente o RA-BC com o modelo de recompensa ReWiND (67% vs. 25% na tarefa Difícil), provando que a qualidade do modelo de recompensa é o gargalo para o RA-BC.

• Estudos de Ablação:

  • Consciência de Estágio: Remover o estimador de estágio degradou o desempenho, confirmando a necessidade de modelagem hierárquica.
  • Aumentação de Retrocesso: Essencial para generalizar para casos de falha do mundo real; sem ela, o modelo tornou-se excessivamente otimista.
  • Escala do Modelo: Um modelo de 60M de parâmetros foi encontrado como o equilíbrio ótimo entre capacidade e sobreajuste.

5. Significado

Este artigo fornece uma solução escalável e eficiente em anotações para manipulação robótica de longo horizonte.

• Eficiência de Dados: Demonstra que o aprendizado de política de alta qualidade é possível mesmo com conjuntos de dados ruidosos e diversos, se o sinal de recompensa puder distinguir com precisão entre "progresso lento" e "falha".
• Escalabilidade: Ao aproveitar anotações em linguagem natural (que são mais fáceis de obter do que funções de recompensa densas) e automatizar a derivação de recompensa, o método escala para tarefas complexas como manipulação de objetos deformáveis.
• Impacto Prático: Os resultados na dobra de camisetas (67% de sucesso a partir de um estado amassado) representam um salto significativo para a robótica doméstica, um domínio onde os modelos atuais geralmente falham devido à complexidade das interações ricas em contato e deformáveis.

O trabalho estabelece que a modelagem de recompensa é um gargalo crítico no aprendizado por imitação e que a estimação de progresso semântico e consciente de estágio é a chave para desbloquear comportamentos robóticos de longo horizonte robustos.