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

1. Probleemstelling

Het artikel adresseert de kritieke uitdaging van langdurige, contactrijke robotmanipulatie, specifiek met betrekking tot vervormbare objecten (bijvoorbeeld het vouwen van een T-shirt). Hoewel recente Robot Behavior Models (RBMs) veelbelovend zijn, kampen ze met de volgende problemen:

• Inconsistentie in Demonstratiekwaliteit: Grote datasets bevatten vaak ruis, suboptimale of gefaalde trajecten van niet-expert operators.
• Kwetsbare Labeling: Bestaande methoden vertrouwen vaak op frame-indices of eenvoudige op tijd gebaseerde voortgangslabels. Bij taken met variabele duur en niet-lineaire voortgang (bijvoorbeeld het platleggen van een shirt kost verschillende hoeveelheden tijd, afhankelijk van de stof), introduceert frame-gebaseerde labeling ernstige ruis, waardoor het onmogelijk wordt om onderscheid te maken tussen een "traag maar correct" stap en een "gefaald" stap.
• Generalisatie: Modellen getraind op gecureerde data falen in generalisatie naar out-of-distribution (OOD) scenario's, zoals herstel van fouten of het hanteren van occlusies bij vervormbare objecten.

Het kernprobleem is hoe men consistent, fijnmazige beloningssignalen kan afleiden uit diverse, variabele-lengte demonstraties om beleidsleer effectief te sturen zonder perfecte, handgemaakte beloningsfuncties te vereisen.

2. Methodologie

De auteurs stellen een tweeledig kader voor: SARM (Stage-Aware Reward Modeling) voor het genereren van beloningen, en RA-BC (Reward-Aligned Behavior Cloning) voor het benutten van die beloningen om beleidsmodellen te trainen.

A. SARM: Stage-Aware Reward Modeling

SARM is een op video gebaseerd beloningsmodelleringskader dat is ontworpen om taakvoortgang robuust te voorspellen.

• Natuurlijke Taalannotaties: In plaats van frame-indices gebruikt de methode natuurlijke taal subtaakannotaties (bijvoorbeeld "grijpen", "platleggen", "vouwen") om taken op te splitsen in semantische fasen.
• Dual-Head Architectuur:
  1. Faseschatter: Een classificatiekop die de huidige hoog-niveau taakfase voorspelt (bijvoorbeeld "Fase van platleggen") op basis van visuele observaties (RGB-frames) en gewointoestanden.
  2. Subtaakschatter: Een regressiekop die fijnmazige voortgang ($\in [0, 1]$) voorspelt, geconditioneerd op de voorspelde fase.
• Labelstrategie:
  • De auteurs berekenen een dataset-niveau priorverhouding ($\bar{\alpha}_k$) voor elke subtaak op basis van de gemiddelde temporale duur van die subtaak over de dataset.
  • Voortgangslabels worden gegenereerd door lineaire interpolatie binnen subtaken, zodat identieke semantische toestanden (bijvoorbeeld een volledig platgelegd shirt) consistente voortgangswaarden ontvangen, ongeacht hoeveel frames het kostte om daar te komen.
• Invoerverwerking: Het model gebruikt een bevroren CLIP-encoder voor visuele kenmerken en een Transformer-encoder om temporale afhankelijkheden te vangen. Het hanteert rewind-augmentatie (het toevoegen van omgekeerde frames van eerdere tijdstippen) om robuustheid te verbeteren tegen falende gevallen en OOD-scenario's.

B. RA-BC: Reward-Aligned Behavior Cloning

RA-BC is een trainingskader dat het geleerde SARM-beloningsmodel benut om demonstratiedata te filteren en opnieuw te wegen tijdens beleids training.

• Berekening van Voortgangsdelta: Voor elke trainingsstalen berekent het model de verbetering in voortgang ($\Delta \phi$) tussen de huidige staat en de volgende staat met behulp van het SARM-beloningsmodel.
• Dynamische Weging:
  • Het systeem onderhoudt online lopende statistieken (gemiddelde $\mu$ en standaardafwijking $\sigma$) van voortgangsdelta's.
  • Stalen met positieve voortgang (boven een drempelwaarde) worden opgewogen, terwijl die met negatieve voortgang (regressies/falen) worden afgewogen of gefilterd.
  • Een "prior override"-mechanisme wijst een gewicht van 1 toe aan duidelijk goede stalen en 0 aan duidelijk slechte stalen om beslissende filtering te garanderen.
• Doel: De standaard Behavior Cloning (BC) loss wordt vervangen door een gewogen loss, wat het beleid effectief leert om alleen de hoogwaardige, progressieve segmenten van de diverse dataset te imiteren.

3. Belangrijkste Bijdragen

1. SARM Kader: Een nieuw, fase-bewust, op video gebaseerd beloningsmodel dat taakfasen en fijnmazige voortgang gezamenlijk voorspelt met behulp van natuurlijke taalannotaties. Het overwint de kwetsbaarheid van frame-index labeling en generaliseert goed naar OOD-scenario's.
2. RA-BC Kader: Een beleids trainingsmethode die het geleerde beloningsmodel gebruikt om automatisch hoogwaardige demonstraties te identificeren en opnieuw te wegen, waardoor effectief leren mogelijk is vanuit ruisige, diverse datasets zonder handmatige filtering.
3. Empirische Validatie op Vervormbare Objecten: Succesvolle toepassing op de uiterst uitdagende T-shirt vouwtaken, met aanzienlijke verbeteringen ten opzichte van standaard Behavior Cloning en andere beloningsmodelleringsbaselines.

4. Experimentele Resultaten

De auteurs evalueerden hun aanpak op twee taken: T-shirt vouwen (langdurig, vervormbaar) en Afwas legen (meerfasig, stijf).

• Prestaties van het Beloningsmodel:

  • SARM presteerde aanzienlijk beter dan baselines (LIV, VLC, GVL, REDS, ReWiND) bij het voorspellen van voortgang op menselijke demonstraties en echte robotrollouts.
  • Bij het vouwen van T-shirts behaalde SARM een 0,008 MSE op validatiedata (tegenover 0,019 voor de volgende beste baseline) en een 0,94 classificatiescore op echte robotrollouts (tegenover 0,78 voor ReWiND).
  • SARM was robuust tegen lichtvariaties en faalmodi waarbij andere modellen ruisige of optimistische signalen produceerden.

• Prestaties van het Beleid (T-shirt Vouwen):

  • Baselines: Vanille BC getraind op de volledige diverse dataset ($D_{all}$) behaalde 0% succes op de "Moeilijke" taak (vouwen vanuit een gekreukelde staat).
  • Naïeve Filtering: Filtering op duur ($D_{2min}$) verbeterde de prestaties, maar faalde nog steeds op de Moeilijke taak (0% succes).
  • RA-BC met SARM:
    • Gemiddelde Taak (Plat naar Gevouwen): 83% successcore.
    • Moeilijke Taak (Gekreukeld naar Gevouwen): 67% successcore.
  • Vergelijking: RA-BC met SARM presteerde aanzienlijk beter dan RA-BC met het ReWiND-beloningsmodel (67% vs. 25% op de Moeilijke taak), wat bewijst dat de kwaliteit van het beloningsmodel de bottleneck is voor RA-BC.

• Ablatiestudies:

  • Fasebewustzijn: Het verwijderen van de faseschatter verslechterde de prestaties, wat de noodzaak van hiërarchische modellering bevestigt.
  • Rewind Augmentatie: Essentieel voor generalisatie naar real-world faalgevallen; zonder dit werd het model te optimistisch.
  • Modelgrootte: Een model van 60M parameters bleek de optimale balans tussen capaciteit en overfitting.

5. Betekenis

Dit artikel biedt een schaalbare en annotatie-efficiënte oplossing voor langdurige robotmanipulatie.

• Data-efficiëntie: Het toont aan dat hoogwaardig beleidsleren mogelijk is, zelfs met ruisige, diverse datasets, als het beloningssignaal nauwkeurig onderscheid kan maken tussen "trage voortgang" en "falen".
• Schaalbaarheid: Door gebruik te maken van natuurlijke taalannotaties (die gemakkelijker te verkrijgen zijn dan dichte beloningsfuncties) en het automatiseren van de beloningsafleiding, schaalt de methode naar complexe taken zoals manipulatie van vervormbare objecten.
• Praktische Impact: De resultaten bij het vouwen van T-shirts (67% succes vanuit een gekreukelde staat) vertegenwoordigen een aanzienlijke sprong voorwaarts voor huishoudelijke robotica, een domein waar huidige modellen typisch falen vanwege de complexiteit van contactrijke, vervormbare interacties.

Het werk stelt vast dat beloningsmodellering een kritieke bottleneck is in imitatielearning en dat fase-bewuste, semantische voortgangsschatting de sleutel is tot het vrijmaken van robuuste langdurige robotgedragingen.