Technische Samenvatting: DataMIL – Selectie van Data voor Robot Imitatie-leren met Datamodellen

1. Probleemstelling

De robotica-gemeenschap is onlangs verschoven naar het trainen van algemene fundamentele beleidssystemen (foundation policies) met behulp van enorme, diverse datasets (bijv. Open-X Embodiment). Hoewel deze beleidssystemen een sterke gemiddelde prestatie bereiken over uiteenlopende taken, presteren ze vaak ondermaats op individuele, gespecialiseerde taken, wat het finetunen op nieuw verworven, taakspecifieke data noodzakelijk maakt.

De kernuitdaging die door dit werk wordt aangepakt, is dataselectie: het identificeren van welke subsets van grote voorafgaande datasets, wanneer gecombineerd met beperkte taakspecifieke demonstraties, de beste gespecialiseerde beleidssystemen opleveren.

• Beperkingen van Huidige Methoden: Naïeve selectie op basis van menselijke concepten van kwaliteit (bijv. semantische gelijkenis, visuele gelijkenis of staat-actie nabijheid) faalt vaak. Deze heuristieken gaan ervan uit dat de meest "gelijke" data ook de meest "nuttige" data is, waarbij het werkelijke effect van een datapunt op de downstream prestaties van het beleid wordt genegeerd.
• De Optimalisatiebarrière: Ideaal gezien zou men data selecteren door subsets exhaustief te testen via hertraining en de resulterende prestaties van het beleid te evalueren. Echter, in de robotica vereist de evaluatie van beleid real-world rollouts, wat tijdrovend, gevaarlijk en computationeel onhaalbaar is voor grootschalige zoekopdrachten.
• De Kloof: Bestaande kaders voor data-attributie zoals datamodels (die de prestaties van een model voorspellen op basis van trainingsdata) zijn succesvol geweest in NLP en Computer Vision, maar zijn niet direct toepasbaar op robotica vanwege de niet-differentieerbare en kostbare aard van real-world evaluatie van beleid.

2. Methodologie: DataMIL

De auteurs introduceren DataMIL (Datamodels voor Imitatie-leren), een framework dat het datamodel-paradigma uitbreidt naar robotica. DataMIL redeneert over dataselectie op een end-to-end wijze, waarbij het beleid zelf wordt gebruikt om datapunten te identificeren die de prestaties verbeteren, in plaats van te vertrouwen op statische heuristieken.

2.1 Kernframework

DataMIL behandelt het leeralgoritme van het beleid als een "black box" en traint een estimator (een datamodel) om de prestaties te voorspellen van een beleid dat getraind is op een specifieke subset van data.

• Doel: Vind een subset $D' \subset D$ die de doelmetriek $M(A(D'))$ maximaliseert, waarbij $A$ het leeralgoritme is en $M$ de prestatie-metriek.
• Datamodel Formulering: In plaats van het beleid te trainen, leert DataMIL een functie $\hat{f}(D') \approx M(A(D'))$. Voor lineaire datamodellen komt dit neer op het toekennen van een scalaire score $\tau(z_i)$ aan elk datapunt $z_i$, waarbij de voorspelde prestatie de som is van de scores voor de geselecteerde subset.

2.2 Belangrijke Aanpassingen voor Robotica

Om datamodels hanteerbaar te maken voor robotica, introduceren de auteurs drie cruciale modificaties:

1. Surrogaat Verliesfunctie (Vermijden van Rollouts):
   Real-world rollouts zijn niet-differentieerbaar en kostbaar. DataMIL vervangt de ware succesrate-metriek door een proxy-metriek $\tilde{M}$, gedefinieerd als de negatieve Behavior Cloning (BC) loss op een uitgesloten set van doel-demonstraties:
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   Deze metriek is volledig differentieerbaar en vereist geen interactie met de omgeving, wat het gebruik van gradiënt-gebaseerde schattingsmethoden mogelijk maakt.

2. Datamodel Estimators:
   Het paper gebruikt twee methoden om de invloedscores $\tau(z_i)$ te schatten:

• Regressie Estimator: Sample willekeurige subsets van data, traint beleidssystemen, evalueert de proxy-metriek en past een lineair model toe om prestaties te voorspellen op basis van de aanwezigheid van data. Dit is accuraat maar computationeel zwaar.
• Metagradient-gebaseerde Estimator: Gebruikt invloedsfuncties en metagradiënten om de afgeleide van de doelmetriek te berekenen ten opzichte van de datagewichten. Deze methode is efficiënter voor grote modellen (zoals Octo) en voorkomt de noodzaak om voor elke subset beleidssystemen te hertrainen.

3. Clustering en Mitigatie van Distributieverschuiving:

• Clustering: Om ruis in de invloedschatting te verminderen, worden individuele staat-actie paren gegroepeerd in temporele clusters (sub-trajecten of volledige trajecten). De optimale cluster-granulariteit hangt af van de datasetgrootte; grotere datasets (zoals OXE) profiteren van aggregatie op trajectniveau.
• Distributieverschuiving: Om slechte schattingen te voorkomen die worden veroorzaakt door verschuivingen tussen de voorafgaande en de doel-domeinen, wordt een klein deel van de doel-taakdata opgenomen tijdens de datamodel-schattingsfase om het leerproces af te stemmen op het doel-domein.

2.3 Trainingspipeline

Zodra de invloedscores zijn geschat:

1. De top $x\%$ van de voorafgaande data (hoogste positieve invloed) wordt geselecteerd om $D_{sel}$ te vormen.
2. Het uiteindelijke beleid wordt getraind via co-training op de doel-data $D_{target}$ en de geselecteerde data $D_{sel}$. Bij elke stap wordt een batch gesampled uit $D_{target}$ met waarschijnlijkheid $\alpha$ en uit $D_{sel}$ met waarschijnlijkheid $1-\alpha$.

3. Belangrijkste Bijdragen

• Uitbreiding van Datamodels naar Robotica: DataMIL is het eerste framework dat er succesvol in is geslaagd om datamodels toe te passen op robotica-imitatie-leren, waarbij de barrière van kostbare real-world rollouts wordt overwonnen via een differentieerbare surrogaatverliesfunctie.
• End-to-End Prestatiebewustzijn: In tegen tegenoverheid van heuristische methoden (visuele, bewegings- of staat-gelijkenis), optimaliseert DataMIL direct voor de impact op de succesrate van het beleid, waardoor het data kan selecteren die prestaties verbeteren en data kan verwerpen die prestaties verslechteren, zelfs als de data visueel zeer gelijkaand lijkt.
• Schaalbare Estimators: De introductie van de metagradient-gebaseerde estimator stelt het framework in staat om te schalen naar grote, complexe beleidssystemen (zoals Octo) en massale datasets (zoals OXE), waar regressie-gebaseerde benaderingen computationeel onhaalbaar zouden zijn.
• Cross-Embodiment en Multi-Task Selectie: De methode demonstreert het vermogen om nuttige data te selecteren voor taken en robot-embodiments die niet aanwezig zijn in de voorafgaande dataset, en om datasets te cureren die meerdere downstream doelstellingen tegelijkertijd ondersteunen.

4. Experimentele Resultaten

De auteurs valideren DataMIL over meer dan 60 taken in simulatie en real-world settings:

• MetaWorld (50 Taken): DataMIL behaalde een prestatieverbetering van 10% ten opzichte van state-of-the-art baselines (Behavior Retrieval, Flow Retrieval, STRAP). Het filterde succesvol ruisende, suboptimale demonstraties uit autonome exploratie-data, terwijl gelijkenis-gebaseerde baselines vaak irrelevante of kwalitatief lage data ophaalden.
• LIBERO (10 Taken): Gebruikmakend van het Octo transformer beleid, presteerde DataMIL consistent beter dan baselines over diverse langdurige taken. Hoewel visuele gelijkenis-baselines goed presteerden in enkele specifieke taken, bood DataMIL de meest robuuste gemiddelde prestatie.
• Open-X Embodiment (OXE) (Real-World):
  • Franka-Ball & Franka-Pouch: DataMIL selecteerde data die de succesrates aanzienlijk verbeterde vergeleken met willekeurige selectie en heuristische baselines.
  • Tiago-Sink: In een uitdagende cross-embodiment setting (het selecteren van data voor een Tiago-robot uit een dataset die geen Tiago-data bevat), identificeerde DataMIL succesvol relevante demonstraties van andere robots (bijv. RT-1, BC-Z) die de "essentie" van de taak deelden (tabletop manipulatie), terwijl baselines faalden om te generaliseren.
  • Droid-Multitask: DataMIL cureerde een enkele dataset die de prestaties op drie verschillende taken tegelijkertijd verbeterde, waarbij het de baselines overtrof in de gemiddelde succesrate.

Kwalitatieve Bevindingen:

• DataMIL selecteert een diverse mix van datasets in plaats van te leunen op één bron (in tegenstelling tot baselines die vaak 80%+ uit één dataset halen).
• Het onderscheidt effectief tussen behulpzame en schadelijke data, zelfs wanneer ze visueel identiek zijn (bijv. dezelfde staat maar verschillende actie-distributies), een foutmodus die veel voorkomt bij gelijkenis-gebaseerde retrieval.

5. Betekenis en Claims

Het paper claimt dat end-to-end, prestatiebewuste dataselectie cruciaal is voor het ontsluiten van het potentieel van grote voorafgaande datasets in robotica. De auteurs stellen dat hoewel grote fundamentele modellen een sterke basis bieden, hun specialisatie sterk afhangt van de kwaliteit van de fine-tuning data.

• Bescheiden Claims over Schaalbaarheid: De auteurs erkennen dat het schatten van datamodellen computationeel duur blijft (meerdere malen de kosten van het trainen van een beleid op alle data), hoewel de metagradient-benadering dit verzacht. Ze merken op dat hyperparameter tuning (bijv. clustergrootte, selectieratio) momenteel nog geen sterke theoretische intuïtie heeft.
• Scope: Het werk richt zich primair op single-task specialisatie, waarbij de Droid-Multitask setting dient als een voorlopige stap richting multi-task curatie.
• Conclusie: DataMIL demonstreert dat het bewegen van menselijke gedefinieerde heuristieken naar data-gedreven, model-bewuste selectie consistente winst oplevert in succesrates, met name in complexe, heterogene en real-world scenario's waar traditionele retrieval-methoden tekortschieten.