Policy Contrastive Decoding for Robotic Foundation Models

1. Enunciato del Problema

I modelli fondazionali robotici (policy robotiche generaliste) hanno dimostrato un grande potenziale per la manipolazione flessibile e a scopo generale. Tuttavia, gli autori identificano un difetto critico: questi modelli apprendono spesso correlazioni spurie dalle traiettorie di pre-addestramento. Invece di concentrarsi sulle caratteristiche degli oggetti rilevanti per il compito, le policy si affidano a indizi di sfondo irrilevanti (ad esempio, illuminazione, texture o oggetti specifici di sfondo) per prevedere le azioni.

• Conseguenza: Quando l'ambiente di deployment cambia (ad esempio, variazioni dello sfondo, regolazioni dell'illuminazione o riposizionamento degli oggetti), le prestazioni della policy si degradano significativamente perché le "scorciatoie" apprese non sono più valide.
• Limitazioni Attuali: Le soluzioni esistenti richiedono spesso il riaddestramento o il fine-tuning, che sono computazionalmente costosi e richiedono l'accesso ai pesi del modello. Manca un metodo senza addestramento, plug-and-play per correggere questi fallimenti al momento dell'inferenza.

2. Metodologia: Policy Contrastive Decoding (PCD)

Gli autori propongono Policy Contrastive Decoding (PCD), un plugin di inferenza senza addestramento progettato per reindirizzare l'attenzione di una policy robotica dalle caratteristiche spurie agli indizi visivi rilevanti per gli oggetti.

Meccanismo Principale

PCD opera contrastando le distribuzioni di probabilità delle azioni generate da due diversi input visivi:

1. Osservazione Originale ($o_i$): L'input visivo standard.
2. Osservatione con Maschera sull'Oggetto ($\hat{o}_i$): Lo stesso input visivo in cui gli oggetti target rilevanti per il compito sono mascherati (rimossi/inpainted).

Il metodo calcola una nuova distribuzione di probabilità delle azioni contrastiva ($\pi^*$) utilizzando la seguente formula:
$$\pi^*_\theta(a_i | \ell, o_i) = \frac{1}{C} \cdot \pi_\theta(a_i | \ell, o_i) \cdot \left( \frac{\pi_\theta(a_i | \ell, o_i)}{\pi_\theta(a_i | \ell, \hat{o}_i)} \right)^\alpha$$

• Logica: Se la policy si affida a caratteristiche spurie (sfondo), la distribuzione di probabilità sarà simile sia per $o_i$ che per $\hat{o}_i$, risultando in un rapporto vicino a 1 (nessun cambiamento). Se la policy si affida all'oggetto, la distribuzione per $\hat{o}_i$ differirà significativamente, amplificando la probabilità dell'azione corretta nell'output finale.
• Iperparametro ($\alpha$): Controlla la forza del contrasto. $\alpha=0$ ripristina la policy di base.

Componenti Tecnici Chiave

Per rendere PCD applicabile a diverse policy robotiche, gli autori introducono due strategie specifiche:

1. Tracking-to-Mask (Track2Mask):

   • Obiettivo: Generare automaticamente osservazioni con maschera sull'oggetto per ogni fotogramma di una traiettoria senza intervento manuale.
   • Processo:
     • Annotazione: Gli oggetti target nel fotogramma iniziale sono identificati utilizzando prompt umani (Punto/Box) o rilevatori open-vocabulary pronti all'uso (ad esempio, Grounding DINO).
     • Tracking: SAM2 (Segment Anything Model 2) traccia e segmenta questi oggetti nei fotogrammi successivi.
     • Inpainting: Gli oggetti segmentati vengono inpainted (rimossi) per creare l'osservazione mascherata $\hat{o}_i$.

2. Modellazione Probabilistica basata su KDE (KDE-PM):

   • Sfida: Le policy autoregressive (ad esempio, OpenVLA) escono direttamente le probabilità dei token, ma le policy basate su diffusione (ad esempio, Octo, $\pi_0$) generano azioni tramite un processo di denoising e non escono nativamente una distribuzione di probabilità $\pi_\theta(a|o)$.
   • Soluzione: Gli autori utilizzano la Stima della Densità di Kernel (KDE) per approssimare la distribuzione di probabilità delle azioni. Campionano $N$ azioni candidate dalla policy di diffusione e adattano una densità di kernel gaussiana a questi campioni per stimare $\pi_\theta(a|o)$ e $\pi_\theta(a|\hat{o})$, permettendo alla formula di decoding contrastivo di funzionare con i modelli di diffusione.

3. Contributi Chiave

1. Prima Soluzione Senza Addestramento: PCD è il primo lavoro a proporre un approccio plug-and-play, senza addestramento, per mitigare le correlazioni spurie nelle policy robotiche durante l'inferenza.
2. Agnostico rispetto al Modello: Funziona come un plugin black-box compatibile sia con policy autoregressive (OpenVLA) che basate su diffusione (Octo, $\pi_0$) senza richiedere l'accesso ai pesi pre-addestrati o al fine-tuning.
3. Robustezza allo Shift di Distribuzione: Contrastando esplicitamente input rilevanti per l'oggetto rispetto a input irrilevanti per l'oggetto, PCD costringe la policy a concentrarsi su caratteristiche causali, migliorando significativamente la generalizzazione a scenari non visti.

4. Risultati Sperimentali

Gli autori hanno valutato PCD su 15 compiti diversi attraverso simulazione (SIMPLER) e ambienti reali utilizzando tre policy all'avanguardia: OpenVLA, Octo e $\pi_0$.

• Prestazioni in Simulazione:
  • OpenVLA: Miglioramento del tasso di successo del 50,6% in media.
  • Octo: Miglioramento del tasso di successo del 29,7% in media.
  • $\pi_0$: Miglioramento del tasso di successo dell'8,9% in media (nonostante $\pi_0$ fosse già una baseline forte).
• Prestazioni nel Mondo Reale:
  • Su compiti di manipolazione nel mondo reale utilizzando la policy $\pi_0$, PCD ha ottenuto un enorme miglioramento del 108% nel tasso di successo rispetto alla baseline.
  • Il metodo si è dimostrato robusto contro vari distrattori (cambiamenti di illuminazione, disordine dello sfondo, spostamenti della posizione degli oggetti).
• Studi di Ablazione:
  • Il metodo è robusto rispetto a diverse strategie di rilevamento degli oggetti (Prompt umani vs Grounding DINO vs YOLO World).
  • È robusto rispetto a diversi algoritmi di inpainting (LaMa ha ottenuto i risultati migliori).
  • L'iperparametro $\alpha$ mostra un impatto positivo chiaro quando $>0$.

5. Significato e Limitazioni

• Significato: PCD affronta una debolezza fondamentale nei modelli fondazionali robotici attuali (dipendenza da correlazioni spurie) senza il pesante costo del riaddestramento. Dimostra che semplici interventi al momento dell'inferenza possono migliorare drasticamente la generalizzazione e la sicurezza nel deployment nel mondo reale.
• Limitazioni:
  • Sovraccarico Computazionale: PCD richiede un secondo passaggio in avanti (per l'immagine mascherata) e il tracking/inpainting degli oggetti, aumentando la latenza di inferenza di circa 24-100% (a seconda della policy e del parallelismo). Tuttavia, gli autori sostengono che ciò è accettabile data la massiccia crescita nei tassi di successo.
  • Ambito: Attualmente si concentra sulla correzione al momento dell'inferenza; i lavori futuri mirano ad affrontare le correlazioni spurie durante la fase di addestramento.

Conclusione: Il paper presenta un framework altamente efficace e versatile (PCD) che migliora significativamente l'affidabilità e la generalizzazione dei modelli fondazionali robotici costringendoli a prestare attenzione alle caratteristiche visive corrette, colmando il divario tra le prestazioni in simulazione e la robustezza nel mondo reale.