PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations

1. Enunciato del Problema

I modelli attuali Vision-Language-Action (VLA) hanno avanzato il controllo robotico sfruttando i modelli Vision-Language (VLM) come backbones percettivi, tipicamente utilizzando un'architettura a doppio sistema (Sistema 2 per il ragionamento, Sistema 1 per il controllo). Tuttavia, gli approcci esistenti soffrono di una limitazione fondamentale:

• Statico vs Temporale: Inquadrano prevalentemente il pre-addestramento come clonazione comportamentale supervisionata (BC), focalizzandosi su correlazioni semantiche visivo-linguistiche statiche. Manca loro una comprensione intrinseca della raggiungibilità dell'obiettivo—la probabilità quantitativa che una coppia stato-azione corrente porti a un obiettivo specificato in linguaggio nel tempo.
• Colli di bottiglia nella stima del valore: I tentativi di incorporare l'apprendimento del valore spesso si affidano a reti di valore separate, richiedono annotazioni manuali delle ricompense o utilizzano l'apprendimento Temporal-Difference (TD) che soffre di errori di bootstrapping e alta varianza in compiti a lungo orizzonte.
• Divario di generalizzazione: Senza una consapevolezza temporale esplicita, questi modelli faticano nella pianificazione a lungo orizzonte, nelle istruzioni nuove zero-shot e nel recupero da interventi umani perché non riescono a distinguere tra stati "facili ma errati" e "difficili ma raggiungibili".

2. Metodologia: PRTS

Gli autori propongono PRTS, un modello fondazionale VLA che riformula il pre-addestramento attraverso l'Apprendimento per Rinforzo Condizionato all'Obiettivo (GCRL) utilizzando l'Apprendimento per Rinforzo Contrastivo (CRL).

A. Concetto Centrale: RL Contrastivo per VLA

Invece di prevedere direttamente le azioni o utilizzare una rete di valore separata, PRTS impara uno spazio di embedding unificato in cui il prodotto interno di un embedding stato-azione ($\phi(s, a)$) e di un embedding obiettivo ($\psi(l)$) approssima l'occupazione log-scontata dell'obiettivo.

• Base Teorica: Basata sull'equivalenza tra la funzione Q per le ricompense di raggiungimento dell'obiettivo e la probabilità di occupazione condizionata (Lemma 1).
• Obiettivo: Il modello impara a discriminare tra coppie stato-azione che possono raggiungere un obiettivo linguistico specifico (campioni positivi) e quelle che non possono (campioni negativi).

B. Innovazioni Tecniche Chiave

1. RL Contrastivo Condizionato al Linguaggio:

   • A differenza del CRL standard che campiona stati futuri geometricamente, PRTS tratta l'istruzione linguistica come l'obiettivo fisso per un'intera traiettoria.
   • Problema Multi-Positivo: Poiché un singolo obiettivo linguistico si applica a tutti i timestep in una traiettoria, un batch contiene molte coppie positive per un obiettivo.
   • Ponderazione Temporale: Per recuperare la distribuzione di campionamento geometrica del CRL standard, gli autori introducono uno schema di ponderazione temporale. Le coppie stato-azione più vicine alla fine di una traiettoria (più vicine all'obiettivo) sono ponderate più alto ($\gamma^{T-t}$). Ciò permette al modello di apprendere la struttura temporale del raggiungimento dell'obiettivo senza etichette di ricompensa esplicite.
   • Obiettivi Bidirezionali:
     • Stato-Azione $\to$ Linguaggio: Allinea $(s, a)$ con la sua istruzione specifica.
     • Linguaggio $\to$ Stato-Azione: Costringe il modello a classificare le azioni-stato all'interno di una traiettoria in base alla loro vicinanza all'obiettivo, codificando il progresso temporale.

2. Architettura a Singolo Passaggio in Avanti (Sovraccarico Trascurabile):

   • PRTS aggiunge due blocchi di token ausiliari alla sequenza di input: <CRL_action> e <CRL_goal>.
   • Maschera Causale Consapevole del Ruolo: Una maschera di attenzione personalizzata (implementata tramite un kernel CuTe-FlashAttention fuso) garantisce l'isolamento delle informazioni:
     • I token <CRL_action> si focalizzano solo su visione/propriocezione (non sul linguaggio).
     • I token <CRL_goal> si focalizzano solo sull'istruzione.
     • I token di azione autoregressivi standard mantengono l'attenzione causale standard.
   • Risultato: Il modello produce logit di azione discreti, embedding stato-azione e embedding obiettivo in un singolo passaggio in avanti, con un sovraccarico computazionale trascurabile rispetto alla clonazione comportamentale classica.

3. Pipeline di Addestramento in Due Fasi:

   • Pre-addestramento: Ottimizza congiuntamente il backbone VLM sulla clonazione comportamentale (predizione del prossimo token), su compiti ausiliari (Robot-QA) e sulle perdite contrastive bidirezionali. Non è ancora attaccato alcun esperto di azione.
   • Post-addestramento: Viene attaccato e fine-tunato su dati a valle un Esperto di Azione a Corrispondenza di Flusso (basato su Diffusion Transformer). Il backbone pre-addestrato fornisce i prior di raggiungibilità dell'obiettivo.

3. Contributi Chiave

1. Riformulazione del Pre-addestramento VLA: Primo a scalare l'RL Contrastivo senza ricompense e condizionato al linguaggio direttamente nella fase di pre-addestramento VLA, eliminando la necessità di annotazioni manuali delle ricompense o reti di valore separate.
2. Consapevolezza Intrinseca della Raggiungibilità dell'Obiettivo: Le rappresentazioni apprese codificano quantitativamente la probabilità di raggiungere un obiettivo, colmando il divario tra ragionamento semantico e progresso temporale del compito.
3. Architettura Efficiente: La "Maschera Causale Consapevole del Ruolo" e il kernel CuTe personalizzato abilitano una supervisione densa sulla raggiungibilità dell'obiettivo con un sovraccarico computazionale vicino allo zero, rendendolo scalabile a modelli di grandi dimensioni (4B parametri) e dataset massicci (167B token).
4. Framework Unificato: Unifica il ragionamento semantico, la predizione di azioni discrete e la stima del valore all'interno di un'unica architettura VLM.

4. Risultati Sperimentali

PRTS è stato pre-addestrato su 167,8 miliardi di token da dataset diversificati (AgiBotWorld, RoboMind, Open X-Embodiment, ecc.) e valutato su benchmark di simulazione e mondo reale.

• Benchmark di Simulazione (LIBERO, LIBERO-Plus, LIBERO-Pro, SimplerEnv):

  • Prestazioni SOTA: PRTS ottiene risultati all'avanguardia, eguagliando o superando baselines solide (es. $\pi^0.5$, ABot-M0) utilizzando significativamente meno risorse di calcolo per il post-addestramento (es. 1/8 del budget di $\pi^0.5$ su LIBERO).
  • Compiti a Lungo Orizzonte: Guadagni significativi sulla suite "Long" di LIBERO (96,6% vs 94,8% per baselines con budget corrispondente), dimostrando l'efficacia della consapevolezza temporale.
  • Robustezza: Supera le baselines nella generalizzazione zero-shot sotto spostamenti di distribuzione (camera, stato del robot, illuminazione, disposizione degli oggetti) e perturbazioni di Compito/Posizione in LIBERO-Pro.
  • Istruzione Zero-Shot: Miglioramento massiccio su istruzioni nuove (31,5% vs 0,8% per $\pi^0.5$ sull'asse "Task" di LIBERO-Pro).

• Valutazione nel Mondo Reale (RealMan Dual-Arm & Flexiv Single-Arm):

  • Manipolazione Complessa: Ha raggiunto un 95,9% di tasso di successo medio su 11 compiti a due bracci e 90,0% su 3 compiti a braccio singolo a grana fine.
  • Recupero da Intervento Umano: PRTS ha dimostrato una robustezza superiore quando gli umani interrompevano i compiti (es. rimozione di oggetti durante l'esecuzione). A differenza delle baselines che fallivano o riproducevano traiettorie obsolete, PRTS ha riesaminato con successo lo stato e ripianificato per raggiungere l'obiettivo.
  • Generalizzazione: Ha mantenuto alti tassi di successo sotto cambiamenti di illuminazione, identità degli oggetti e layout spaziale.

• Visualizzazione del Valore:

  • È stato dimostrato che il prodotto interno $\phi(s, a)^\top \psi(l)$ correla fortemente con il progresso del compito, aumentando mentre il robot si avvicina all'obiettivo e rimanendo basso per istruzioni non corrispondenti, confermando che il modello apprende una funzione di valore condizionata all'obiettivo valida.

5. Significato

PRTS rappresenta un cambiamento di paradigma nei modelli fondazionali robotici:

• Dall'Imitazione al Ragionamento: Si sposta oltre la semplice clonazione comportamentale (imitazione) per apprendere la consapevolezza della raggiungibilità dell'obiettivo, permettendo al robot di capire quanto è vicino a un obiettivo e quanto è fattibile un percorso.
• Scalabilità: Dimostra che gli obiettivi contrastivi, precedentemente limitati a domini decisionali più piccoli, possono scalare a backbones VLA di grandi dimensioni (4B+ parametri) e dataset massicci senza costi proibitivi.
• Intelligenza a Scopo Generale: Iniettando prior temporali e consapevoli dell'obiettivo durante il pre-addestramento, PRTS ottiene una generalizzazione zero-shot e una robustezza notevoli, suggerendo una strada verso un'intelligenza robotica veramente a scopo generale, orientata agli obiettivi e capace di gestire ambienti reali complessi, a lungo orizzonte e dinamici.