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

1. Enunciato del Problema

Il paper affronta la sfida critica della manipolazione robotica a lungo orizzonte e ricca di contatti, specificamente coinvolgente oggetti deformabili (ad esempio, la piegatura di una maglietta). Sebbene i recenti Modelli di Comportamento Robotico (RBM) abbiano mostrato promesse, faticano con:

• Qualità Incoerente delle Dimostrazioni: Grandi dataset contengono spesso traiettorie rumorose, subottimali o fallite provenienti da operatori non esperti.
• Etichettatura Fragile: I metodi esistenti si basano spesso su indici di frame o semplici etichette di progresso basate sul tempo. In compiti con durate variabili e progresso non lineare (ad esempio, appiattire una maglietta richiede quantità di tempo diverse a seconda del tessuto), l'etichettatura basata sui frame introduce un rumore severo, rendendo impossibile distinguere tra un passaggio "lento ma corretto" e un passaggio "fallito".
• Generalizzazione: I modelli addestrati su dati curati non riescono a generalizzare a scenari fuori distribuzione (OOD), come il recupero da errori o la gestione di occlusioni in oggetti deformabili.

Il problema centrale è come derivare segnali di ricompensa coerenti e a grana fine da dimostrazioni diverse e di lunghezza variabile per guidare efficacemente l'apprendimento della politica senza richiedere funzioni di ricompensa perfette e create manualmente.

2. Metodologia

Gli autori propongono un framework in due parti: SARM (Stage-Aware Reward Modeling) per generare ricompense, e RA-BC (Reward-Aligned Behavior Cloning) per utilizzare tali ricompense per addestrare le politiche.

A. SARM: Stage-Aware Reward Modeling

SARM è un framework di modellazione della ricompensa basato su video progettato per prevedere il progresso del compito in modo robusto.

• Annotazioni in Linguaggio Naturale: Invece degli indici di frame, il metodo utilizza annotazioni di sottocompiti in linguaggio naturale (ad esempio, "afferra", "appiattisci", "piega") per scomporre i compiti in fasi semantiche.
• Architettura a Doppia Testa:
  1. Stima della Fase: Una testa di classificazione che prevede la fase di alto livello del compito corrente (ad esempio, "Fase di Appiattimento") basandosi su osservazioni visive (frame RGB) e stati delle giunture.
  2. Stima del Sottocompito: Una testa di regressione che prevede il progresso a grana fine ($\in [0, 1]$) condizionato alla fase prevista.
• Strategia di Etichettatura:
  • Gli autori calcolano una proporzione a priori a livello di dataset ($\bar{\alpha}_k$) per ogni sottocompito basandosi sulla durata temporale media di quel sottocompito attraverso il dataset.
  • Le etichette di progresso sono generate interpolando linearmente all'interno dei sottocompiti, garantendo che stati semantici identici (ad esempio, una maglietta completamente appiattita) ricevano valori di progresso coerenti indipendentemente dal numero di frame necessari per arrivarci.
• Elaborazione degli Input: Il modello utilizza un encoder CLIP congelato per le caratteristiche visive e un encoder Transformer per catturare le dipendenze temporali. Impiega augmentazione di riavvolgimento (aggiunta di frame invertiti da timestamp precedenti) per migliorare la robustezza contro i casi di fallimento e gli scenari OOD.

B. RA-BC: Reward-Aligned Behavior Cloning

RA-BC è un framework di addestramento che sfrutta il modello di ricompensa SARM appreso per filtrare e ridistribuire i dati di dimostrazione durante l'addestramento della politica.

• Calcolo della Variazione di Progresso: Per ogni campione di addestramento, il modello calcola il miglioramento del progresso ($\Delta \phi$) tra lo stato corrente e il successivo utilizzando il modello di ricompensa SARM.
• Ponderazione Dinamica:
  • Il sistema mantiene statistiche di esecuzione online (media $\mu$ e deviazione standard $\sigma$) delle variazioni di progresso.
  • I campioni con progresso positivo (sopra una soglia) vengono sovrappesati, mentre quelli con progresso negativo (regressioni/fallimenti) vengono sottopesati o filtrati.
  • Un meccanismo di "sovrascrittura a priori" assegna un peso di 1 ai campioni chiaramente buoni e 0 a quelli chiaramente cattivi per garantire un filtraggio deciso.
• Obiettivo: La funzione di perdita standard del Behavior Cloning (BC) viene sostituita con una perdita ponderata, insegnando efficacemente alla politica di imitare solo i segmenti di alta qualità e progressivi del dataset diversificato.

3. Contributi Chiave

1. Framework SARM: Un modello di ricompensa basato su video e consapevole della fase, innovativo, che prevede congiuntamente le fasi del compito e il progresso a grana fine utilizzando annotazioni in linguaggio naturale. Supera la fragilità dell'etichettatura basata sugli indici di frame e generalizza bene agli scenari OOD.
2. Framework RA-BC: Un metodo di addestramento della politica che utilizza il modello di ricompensa appreso per identificare e ridistribuire automaticamente le dimostrazioni di alta qualità, abilitando un apprendimento efficace da dataset rumorosi e diversificati senza filtraggio manuale.
3. Validazione Empirica su Oggetti Deformabili: Applicazione riuscita al compito altamente impegnativo della piegatura della maglietta, dimostrando miglioramenti significativi rispetto al Behavior Cloning standard e ad altre baseline di modellazione della ricompensa.

4. Risultati Sperimentali

Gli autori hanno valutato il loro approccio su due compiti: piegatura della maglietta (lungo orizzonte, deformabile) e scarico dei piatti (multifase, rigido).

• Prestazioni del Modello di Ricompensa:

  • SARM ha superato significativamente le baseline (LIV, VLC, GVL, REDS, ReWiND) nella previsione del progresso su dimostrazioni umane e roll-out di robot reali.
  • Sulla piegatura della maglietta, SARM ha raggiunto un MSE di 0,008 sui dati di validazione (contro 0,019 per la baseline successiva migliore) e un punteggio di classificazione di 0,94 sui roll-out di robot reali (contro 0,78 per ReWiND).
  • SARM è stato robusto alle variazioni di illuminazione e ai modi di fallimento in cui altri modelli producevano segnali rumorosi o ottimistici.

• Prestazioni della Politica (Piegatura della Maglietta):

  • Baseline: BC puro addestrato sull'intero dataset diversificato ($D_{all}$) ha raggiunto il 0% di successo sul compito "Difficile" (piegatura da uno stato accartocciato).
  • Filtraggio Ingenuo: Il filtraggio per durata ($D_{2min}$) ha migliorato le prestazioni ma ha comunque fallito sul compito Difficile (0% di successo).
  • RA-BC con SARM:
    • Compito Medio (Da Piatto a Piegato): 83% di tasso di successo.
    • Compito Difficile (Da Accartocciato a Piegato): 67% di tasso di successo.
  • Confronto: RA-BC con SARM ha superato significativamente RA-BC con il modello di ricompensa ReWiND (67% contro 25% sul compito Difficile), dimostrando che la qualità del modello di ricompensa è il collo di bottiglia per RA-BC.

• Studi di Ablazione:

  • Consapevolezza della Fase: Rimuovere lo stimatore della fase ha degradato le prestazioni, confermando la necessità di una modellazione gerarchica.
  • Augmentazione di Riavvolgimento: Essenziale per generalizzare ai casi di fallimento del mondo reale; senza di essa, il modello diventava eccessivamente ottimistico.
  • Scala del Modello: Un modello da 60 milioni di parametri è stato trovato come il miglior equilibrio tra capacità e sovrapposizione.

5. Significato

Questo paper fornisce una soluzione scalabile ed efficiente in termini di annotazione per la manipolazione robotica a lungo orizzonte.

• Efficienza dei Dati: Dimostra che l'apprendimento di politiche di alta qualità è possibile anche con dataset rumorosi e diversificati se il segnale di ricompensa può distinguere accuratamente tra "progresso lento" e "fallimento".
• Scalabilità: Sfruttando annotazioni in linguaggio naturale (che sono più facili da ottenere rispetto a funzioni di ricompensa dense) e automatizzando l'derivazione della ricompensa, il metodo scala a compiti complessi come la manipolazione di oggetti deformabili.
• Impatto Pratico: I risultati sulla piegatura della maglietta (67% di successo da uno stato accartocciato) rappresentano un significativo passo avanti per la robotica domestica, un dominio in cui i modelli attuali falliscono tipicamente a causa della complessità delle interazioni ricche di contatti e deformabili.

Il lavoro stabilisce che la modellazione della ricompensa è un collo di bottiglia critico nell'apprendimento per imitazione e che la stima del progresso semantica e consapevole della fase è la chiave per sbloccare comportamenti robotici a lungo orizzonte robusti.