Riepilogo Tecnico: DataMIL – Selezione dei Dati per l'Imitation Learning Robotico con i Datamodels

1. Definizione del Problema

La comunità della robotica si è recentemente orientata verso l'addestramento di policy fondazionali generaliste utilizzando dataset massivi e diversificati (ad es., Open-X Embodiment). Sebbene queste policy ottengano una forte performance media su compiti vari, esse spesso sottoperformano su singoli compiti specializzati, rendendo necessario il fine-tuning su nuovi dati specifici per il compito.

La sfida centrale affrontata da questo lavoro è la selezione dei dati: identificare quali sottoinsiemi di grandi dataset pregressi, se combinati con limitate dimostrazioni specifiche per il compito, producano le migliori policy specializzate.

• Limitazioni dei Metodi Attuali: La selezione ingenua basata su nozioni umane di qualità (ad es., somiglianza semantica, somiglianza visiva o prossimità stato-azione) spesso fallisce. Questi euristiche assumono che i dati più "simili" siano i più "utili", ignorando l'effettivo impatto di un punto dati sulla performance della policy a valle.
• La Barriera di Ottimizzazione: Idealmente, si selezionerebbero i dati testando esaustivamente i sottoinsiemi tramite il riaddestramento e valutando la policy risultante. Tuttavia, nella robotica, la valutazione della policy richiede rollouts nel mondo reale, che sono dispendiosi in termini di tempo, pericolosi e computazionalmente impraticabili per una ricerca su larga scala.
• Il Gap: Gli esistenti framework di attribuzione dei dati come i datamodels (che predicono la performance del modello basandosi sui dati di addestramento) hanno avuto successo in NLP e Computer Vision, ma non sono stati direttamente applicabili alla robotica a causa della natura non differenziabile e costosa della valutazione delle policy nel mondo reale.

2. Metodologia: DataMIL

Gli autori introducono DataMIL (Datamodels per l'Imitation Learning), un framework che estende il paradigma dei datamodels alla robotica. DataMIL ragiona sulla selezione dei dati in modo end-to-end, utilizzando la policy stessa per identificare i punti dati che migliorano la performance, invece di fare affidamento su euristiche statiche.

2.1 Framework Core

DataMIL tratta l'algoritmo di apprendimento della policy come una "black box" e addestra un estimatore (un datamodel) per predire la performance di una policy addestrata su un particolare sottoinsieme di dati.

• Obiettivo: Trovare un sottoinsieme $D' \subset D$ che massimizzi la metrica target $M(A(D'))$, dove $A$ è l'algoritmo di apprendimento e $M$ è la metrica di performance.
• Formulazione del Datamodel: Invece di addestrare la policy, DataMIL apprende una funzione $\hat{f}(D') \approx M(A(D'))$. Per i datamodels lineari, ciò si riduce ad assegnare un punteggio scalare $\tau(z_i)$ a ciascun punto dati $z_i$, dove la performance predetta è la somma dei punteggi del sottoinsieme selezionato.

2.2 Adattamenti Chiave per la Robotica

Per rendere i datamodels trattabili per la robotica, gli autori introducono tre modifiche critiche:

1. Funzione di Perdita Sostitutiva (Evitare i Rollout):
   I rollout nel mondo reale sono non differenziabili e costosi. DataMIL sostituisce la vera metrica del tasso di successo con una metrica proxy $\tilde{M}$, definita come la perdita negativa di behavior cloning (BC) su un set di dimostrazioni target tenuto da parte:
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   Questa metrica è completamente differenziabile e non richiede interazione con l'ambiente, consentendo l'uso di metodi di stima basati sul gradiente.

2. Estimatori di Datamodel:
   Il paper impiega due metodi per stimare i punteggi di influenza $\tau(z_i)$:

• Estimatore di Regressione: Campiona sottoinsiemi casuali di dati, addestra le policy, valuta la metrica proxy e adatta un modello lineare per predire la performance basandosi sulla presenza dei dati. Questo metodo è accurato ma computazionalmente pesante.
• Estimatore Basato su Metagradiente: Utilizza funzioni di influenza e metagradienti per calcolare la derivata della metrica target rispetto ai pesi dei dati. Questo metodo è più efficiente per modelli grandi (come Octo) ed evita la necessità di riaddestrare le policy per ogni sottoinsieme.

3. Clustering e Mitigazione dello Shift di Distribuzione:

• Clustering: Per ridurre il rumore nella stima dell'influenza, le singole coppie stato-azione vengono raggruppate in cluster temporali (sub-traiettorie o traiettorie complete). La granularità ottimale del cluster dipende dalla dimensione del dataset; dataset più grandi (come OXE) beneficiano dell'aggregazione a livello di traiettoria.
• Shift di Distribuzione: Per prevenire stime errate causate dagli shift tra i domini prior e target, una piccola frazione di dati del compito target viene inclusa durante la fase di stima del datamodel per allineare il processo di apprendimento al dominio target.

2.3 Pipeline di Addestramento

Una volta stimati i punteggi di influenza:

1. Il top $x\%$ dei dati pregressi (con la più alta influenza positiva) viene selezionato per formare $D_{sel}$.
2. La policy finale viene addestrata tramite co-training sui dati target $D_{target}$ e sui dati selezionati $D_{sel}$. Ad ogni step, un batch viene campionato da $D_{target}$ con probabilità $\alpha$ e da $D_{sel}$ con probabilità $1-\alpha$.

3. Contributi Chiave

• Estensione dei Datamodels alla Robotica: DataMIL è il primo framework che applica con successo i datamodels all'imitation learning robotico, superando la barriera degli costosi rollout nel mondo reale tramite una perdita surrogata differenziabile.
• Consapevolezza della Performance End-to-End: A differenza dei metodi euristici (somiglianza visiva, di movimento o di stato), DataMIL ottimizza direttamente l'impatto sulla performance della policy, permettendogli di selezionare i dati che migliorano la performance e scartare quelli che la degradano, anche se i dati appaiono visivamente simili.
• Estimatori Scalabili: L'introduzione di un estimatore basato su metagradienti permette al framework di scalare a policy grandi e complesse (es. Octo) e a dataset massivi (es. OXE) dove gli approcci basati sulla regressione sarebbero computazionalmente proibitivi.
• Selezione Cross-Embodiment e Multi-Task: Il metodo dimostra la capacità di selezionare dati utili per compiti e robotica non presenti nel dataset precedente, e di curare dataset che supportino simultaneamente molteplici obiettivi a valle.

4. Risultati Sperimentali

Gli autori validano DataMIL attraverso oltre 60 compiti in simulazione e in contesti reali:

• MetaWorld (50 Compiti): DataMIL ha ottenuto un incremento di performance del 10% rispetto ai baseline allo stato dell'arte (Behavior Retrieval, Flow Retrieval, STRAP). Ha filtrato con successo dimostrazioni rumorose e sub-ottimali dai dati di esplorazione autonoma, laddove i baseline basati sulla somiglianza spesso recuperavano dati irrilevanti o di bassa qualità.
• LIBERO (10 Compiti): Utilizzando la policy transformer Octo, DataMIL ha superato costantemente i baseline in compiti diversificati a lungo termine. Mentre i baseline di somiglianza visiva performavano bene in alcuni compiti specifici, DataMIL ha fornito la performance media più robusta.
• Open-X Embodiment (OXE) (Realtà):
  • Franka-Ball & Franka-Pouch: DataMIL ha selezionato dati che hanno migliorato significativamente i tassi di successo rispetto alla selezione casuale e ai baseline euristici.
  • Tiago-Sink: In un setting challenging di cross-embodiment (selezionare dati per un robot Tiago da un dataset che non contiene dati di Tiago), DataMIL ha identificato con successo dimostrazioni rilevanti da altri robot (es. RT-1, BC-Z) che condividevano l'"essenza" del compito (manipolazione su tavolo), mentre i baseline non sono riusciti a generalizzare.
  • Droid-Multitask: DataMIL ha curato un singolo dataset che ha migliorato le performance su tre compiti distinti simultaneamente, superando i baseline in termini di tasso di successo medio.

Risultati Qualitativi:

• DataMIL seleziona un mix diversificato di dataset piuttosto che fare eccessivo affidamento su una singola fonte (a differenza dei baseline che spesso recuperano l'80%+ da un unico dataset).
• Distingue efficacemente tra dati utili e dannosi anche quando appaiono visivamente identici (es. stesso stato ma diverse distribuzioni di azione), un modo di fallimento comune nel retrieval basato sulla somiglianza.

5. Significato e Claim

Il paper sostiene che la selezione dei dati end-to-end e consapevole della performance sia critica per sbloccare il potenziale dei grandi dataset pregressi nella robotica. Gli autori argomentano che, sebbene i grandi modelli fondazionali forniscano una base solida, la loro specializzazione dipende fortemente dalla qualità dei dati di fine-tuning.

• Claim Modesti sulla Scalabilità: Gli autori riconoscono che stimare i datamodels rimane computazionalmente costoso (diverse volte il costo di addestrare una policy su tutti i dati), sebbene l'approccio metagradiente mitighi il problema. Notano inoltre che la sintonizzazione degli iperparametri (es. dimensione del cluster, rapporto di selezione) attualmente manca di una forte intuizione teorica.
• Ambito: Il lavoro si concentra principalmente sulla specializzazione di singolo compito, con il setting Droid-Multitask che funge da passo preliminare verso la curatela multi-task.
• Conclusione: DataMIL dimostra che passare dalle euristiche definite dall'uomo a una selezione guidata dai dati e consapevole del modello produce guadagni consistenti nei tassi di successo, particolarmente in scenari complessi, eterogenei e reali dove i metodi di retrieval tradizionali falliscono.