Flex-$\pi$: A Multi-Stream World-Action Model with Compute Flexibility

Problematica

I modelli world-action (WAM) sono emersi come una potente alternativa ai modelli vision-language-action (VLA), poiché prevedono congiuntamente le osservazioni e le azioni future, apprendendo così rappresentazioni spazio-temporali robuste ed ereditando i priori dai modelli di generazione video su larga scala. Tuttavia, gli attuali WAM generalisti prevedono quasi esclusivamente i latenti RGB delle immagini future. Questi latenti, addestrati puramente per la ricostruzione dei pixel, catturano i dettagli dell'aspetto ma spesso non riescono ad allinearsi con la geometria 3D esplicita e la semantica degli oggetti necessaria per una manipolazione robotica precisa.

Supervisionare i WAM direttamente sulla geometria e sulla semantica comporta tipicamente tre costi significativi:

1. Sensori Aggiuntivi: Richiede nuovi hardware per modalità (es. sensori di profondità).
2. Nuovo Pre-addestramento: Necessità di addestrare nuovi priori visivi da zero.
3. Latenza di Inferenza: Rallenta l'implementazione a causa della complessa elaborazione multi-modale.

Il paper pone la seguente domanda: Come possiamo amplificare i punti di forza dei WAM (predizione congiunta e forti priori visivi) per includere la geometria 3D e la semantica senza questi sacrifici?

Metodologia: FLEX-π

Gli autori introducono FLEX-π, un WAM da 6 miliardi di parametri che predice simultaneamente osservazioni RGB future, pointmap 3D e semantica DINO centriche sugli oggetti. L'innovazione principale è ottenere questo "pasto gratis" di grounding multi-modale senza nuovi sensori o pre-addestramento.

1. Spazio Latente Condiviso tramite VAE Congelato

FLEX-π utilizza un singolo VAE (Variational Autoencoder) congelato proveniente da un modello di generazione video pre-addestrato (Wan-2.2).

• Codifica RGB: Il VAE codifica le immagini RGB come di consueto.
• Codifica Pointmap 3D: Sorprendentemente, gli autori hanno scoperto che questo stesso VAE congelato, addestrato solo su pixel RGB, può codificare e decodificare pointmap 3D quasi senza perdite. Ciò consente alla geometria 3D di essere proiettata esattamente nello stesso spazio latente dell'RGB senza addestrare un nuovo encoder.
• Codifica Semantica: La semantica degli oggetti è derivata utilizzando un encoder DINOv3 congelato.

2. Backbone Multi-Stream Mixture-of-Transformers (MoT)

Il modello incorpora tre stream visivi (RGB, Pointmap, DINO) in uno spazio latente condiviso di stream di token.

• Architettura: Un backbone Mixture-of-Transformers (5 miliardi di parametri) elabora questi stream. Utilizza un backbone visual transformer condiviso ma mantiene reti di proiezione feedforward indipendenti per ogni modalità.
• Action Expert: Un separato e più stretto Action Expert Transformer (1 miliardo di parametri) effettua la cross-attenzione agli stream visivi per generare chunk di azioni.
• Generazione Congiunta Causale: I token di azione effettuano l'attenzione verso le osservazioni correnti e la cross-attenzione verso i token visivi futuri attivi. Ciò consente alla generazione delle azioni di beneficiare della rappresentazione evolutiva del futuro fornita dal modello.

3. Addestramento Flessibile tramite Stream Dropout e Cross-Modality Forcing

Per consentire a un singolo checkpoint di supportare qualsiasi combinazione di input e output durante l'inferenza, FLEX-π impiega due strategie chiave di addestramento:

• Visual Stream Dropout: Durante l'addestramento, gli stream di input vengono rimossi casualmente (con probabilità 0.5), garantendo che il modello possa operare con osservazioni parziali (es. solo RGB).
• Cross-Modality Forcing: Fondamentalmente, anche se uno stream viene rimosso dall'input, il modello è comunque costretto a denoisare e predire quello stream all'output. Ad esempio, se la pointmap 3D manca dall'input, il modello deve sintetizzare la pointmap futura basandosi sugli input RGB e DINO. Questo forza il modello a interiorizzare una rappresentazione unificata dove aspetto, geometria e semantica sono mutuamente predittivi.

4. Flessibilità di Inferenza

Al momento del deployment, gli utenti possono scambiare latenza e prestazioni selezionando quali stream generare:

• Modalità Action-Only: Il modello predice solo le azioni (inferenza più veloce).
• Generazione Congiunta Completa: Il modello predice azioni insieme a futuri RGB, pointmap e feature DINO (prestazioni più elevate).
• Generazione Parziale: Qualsiasi sottoinsieme di stream può essere generato.

Contributi Chiave

1. Grounding Geometrico e Semantico Gratuito: Dimostrare che un VAE video pre-addestrato su RGB può codificare senza perdite le pointmap 3D, permettendo ai WAM di essere supervisionati sulla geometria 3D e sulla semantica DINO senza nuovi sensori o pre-addestramento.
2. Flessibilità Computazionale: Un singolo checkpoint addestrato supporta qualsiasi sottoinsieme di stream di input/output, consentendo ai professionisti di scegliere il proprio punto di equilibrio tra velocità e prestazioni al momento del deployment.
3. Cross-Modality Forcing: Un regime di addestramento che forza il modello a predire le modalità mancanti da quelle osservate, migliorando significativamente la rappresentazione interna e la performance della policy.

Risultati Sperimentali

Gli autori valutano FLEX-π su benchmark di simulazione (RoboTwin, LIBERO, LIBERO-Plus) e su compiti di manipolazione bimanuale nel mondo reale su un robot YAM.

• Performance nel Mondo Reale: Su cinque compiti destri e precisi (inclusi "Self-Repair Gripper" e "Soft-Bag Zipping"), FLEX-π supera i forti baseline (π0.5, ManiFlow, Fast-WAM) di 2–7× nei tassi di successo.
  • Notevolmente, la variante action-only di FLEX-π (che non genera futuri visivi durante l'inferenza) supera già tutti i baseline ed è più veloce di π0.5.
  • La generazione congiunta completa aumenta ulteriormente le prestazioni, ottenendo fino a un miglioramento assoluto del 42,7% nel completamento dei compiti rispetto al più forte baseline.
• Efficienza dei Dati: FLEX-π è significativamente più efficiente nei dati. In RoboTwin, supera i baseline di 1.9–4.5× in regimi a bassa disponibilità di dati (50 dimostrazioni). Anche quando addestrato su solo il 50% dei dati reali, FLEX-π supera i baseline addestrati sull'intero dataset.
• Generalizzazione (OOD): FLEX-π dimostra robustezza agli spostamenti out-of-distribution (oggetti non visti, clutter, distrattori). Mentre i baseline come ManiFlow e π0.5 subiscono cali significativi di performance (25–26%) sotto spostamento di distribuzione, FLEX-π cala di meno del 5%.
• Ablazioni:
  • Rimuovere gli input DINO o Pointmap durante l'addestramento riduce significativamente i tassi di successo (rispettivamente del 6,8% e del 20%).
  • Rimuovere il Cross-Modality Forcing penalizza i tassi di successo del 21%, confermando che forzare il modello a sintetizzare le modalità mancanti è critico per apprendere una forte rappresentazione condivisa.

Significato e Rivendicazioni

Il paper sostiene che FLEX-π dimostra che un WAM può ottenere un grounding semantico e geometrico essenzialmente gratuitamente. Sfruttando lo spazio latente di un modello video pre-addestrato e impiegando il cross-modality forcing, il modello ottiene:

• Migliori Performance della Policy: Tassi di successo significativamente più alti in compiti complessi, ricchi di contatto e a lungo termine.
• Nessun Overhead Hardware: Nessuna necessità di sensori aggiuntivi oltre l'RGB.
• Nessun Overhead di Pre-addestramento: Nessuna necessità di addestrare nuovi priori visivi.
• Flessibilità di Deployment: La capacità di eseguire una modalità rapida (solo azioni) o una modalità a piena prestazione (generazione congiunta completa) utilizzando lo stesso checkpoint.

Gli autori concludono che, sebbene FLEX-π sia altamente efficace, richiede comunque più epoche di fine-tuning per convergere rispetto a modelli più semplici, e la sua modalità a generazione congiunta completa è più lenta rispetto ai VLAs con numero di parametri comparabile. Suggeriscono che i miglioramenti futuri potrebbero derivare dall'unire FLEX-π con backbone VLM più forti e l'accesso a dataset robotici ancora più ampi.