Jetson-PI: Towards Onboard Real-Time Robot Control via Foresight-Aligned Asynchronous Inference

Riepilogo Tecnico: Jetson-PI

Problema

I modelli Vision-Language-Action (VLA) hanno dimostrato prestazioni straordinarie nei compiti incarnati (embodied) sfruttando i modelli Vision-Language (VLM) pre-addestrati e il fine-tuning su dataset di dimostrazione robotica. Tuttavia, l'implementazione di questi modelli su dispositivi onboard a bassa potenza (ad esempio, NVIDIA Jetson Orin) affronta due barriere critiche:

1. Elevata Latenza di Inferenza: La complessità computazionale dei modelli VLA risulta in una latenza di inferenza sostanziale (ad esempio, 1,4s su Jetson Orin), portando a frequenze di controllo basse (0,7 Hz) e tempi di reazione lenti.
2. Limitazioni dell'Inferenza Asincrona Esistente: Sebbene l'inferenza asincrona (parallelizzando l'esecuzione dell'azione con la successiva inferenza) possa mascherare la latenza, essa introduce due problemi specifici sui dispositivi edge:
   • Disallineamento Percezione-Esecuzione: Nel momento in cui un chunk di azione viene predetto ed eseguito, l'ambiente si è evoluto rispetto all'immagine di input utilizzata per la predizione. Questo disallineamento peggiora con tempi di inferenza più lunghi, causando un calo della precisione della predizione.
   • Tempo di Reazione Lungo: La velocità di reazione del robot rimane limitata dal tempo di inferenza del VLA. Inoltre, i metodi asincroni esistenti richiedono spesso ampi chunk di azione per mantenere la stabilità, il che aumenta ulteriormente il tempo medio di reazione ai cambiamenti ambientali.

L'implementazione di modelli VLA su GPU di fascia alta (ad esempio, RTX 4090) non è una soluzione praticabile per la robotica mobile a causa dei severi problemi di consumo energetico (riducendo l'autonomia della batteria di 6,0×) e della necessità di connettività di rete se eseguiti come server.

Metodologia: Jetson-PI

Gli autori propongono Jetson-PI, un framework per l'efficiente implementazione di VLA su dispositivi onboard tramite Correzione Asincrona Allineata alla Previsione (Foresight-Aligned Asynchronous Correction). Il metodo affronta le sfide identificate attraverso tre componenti principali:

1. Correzione Asincrona Allineata alla Previsione

Per risolvere il disallineamento percezione-esecuzione, gli autori introducono un modulo di Correzione Futura (Future Correction Module) leggero.

• Meccanismo: Invece di predire solo immagini future o stati del robot, questo modulo predice lo stato latente futuro (rappresentazione compressa del VLM) al tempo $t + \Delta$ (dove $\Delta$ è la latenza di inferenza), condizionato dalle azioni impegnate per l'esecuzione.
• Addestramento: Viene utilizzato un pipeline di addestramento in due fasi. Primo, l'esperto d'azione apprende a predire le azioni partendo dagli stati latenti futuri. Secondo, il modulo di correzione viene addestrato a predire questi stati futuri a partire dallo stato attuale e dalle azioni impegnate, fornendo una stima della confidenza.
• Beneficio: Ciò consente all'esperto d'azione di predire le azioni basandosi sullo stato futuro dell'ambiente piuttosto che sull'osservazione obsoleta, allineando la predizione con il tempo di esecuzione. Il modulo è estremamente leggero (40 milioni di parametri, circa l'1% della dimensione totale del VLA).

2. Ottimizzazione dello Scheduling Basata sulla Confidenza

Per ridurre il tempo di reazione, gli autori propongono una strategia di scheduling dinamico che bilancia le invocazioni del VLM e dell'Esperto d'Azione.

• Logica: Il Modulo di Correzione Futura funge da scheduler. Se la sua confidenza predetta è superiore a una soglia $\theta$, il sistema salta l'inferenza pesante del VLM e si affida esclusivamente all'Esperto d'Azione per generare i successivi chunk di azione. Se la confidenza scende (indicando alta incertezza o un cambiamento significativo nell'ambiente), il VLM viene invocato per aggiornare la cache KV e gli stati latenti.
• Razionale: Questo evita l'approccio "naïve" di eseguire VLM ed Esperto d'Azione in parallelo, che gli autori dimostrano essere inefficiente sui dispositivi onboard a causa della contesa di banda (entrambi i componenti diventano limitati dalla banda sull'hardware edge, a differenza delle GPU di fascia alta dove il VLM è limitato dal calcolo).

3. Accelerazioni a Livello di Sistema

Gli autori hanno costruito un framework di inferenza basato su llama.cpp adattato alle caratteristiche dei VLA sulle piattaforme edge:

• Riutilizzo dei CUDA Graph: Poiché la lunghezza dei token del VLA è deterministica (risoluzione dell'immagine fissa e istruzioni brevi), il CUDA graph viene costruito una sola volta e riutilizzato, eliminando l'overhead di costruzione del grafo.
• Buffering Intermedio Residente nella GPU: I risultati intermedi tra i componenti (ViT $\to$ LLM $\to$ Esperto d'Azione) sono mantenuti nella memoria della GPU per evitare la latenza di comunicazione CPU-GPU.
• Unrolling del Flow Matching: I passaggi iterativi di denoising dell'Esperto d'Azione vengono srotolati (unrolled) in un unico grafo computazionale unificato, riducendo il numero di chiamate al lancio dei kernel.

Contributi Chiave

1. Identificazione delle Sfide Specifiche dell'Edge: Il documento evidenzia come l'inferenza asincrona standard fallisca sui dispositivi edge a causa della contesa di banda e della natura specifica del disallineamento percezione-esecuzione, che non può essere risolto predendo solo gli stati del robot.
2. Correzione Allineata alla Previsione: Un modulo innovativo che predice i latenti futuri del VLM condizionati dalle azioni impegnate, consentendo all'esperto d'azione di operare su rappresentazioni ambientali future.
3. Scheduling Adattivo: Un meccanismo guidato dalla confidenza che regola dinamicamente la frequenza delle invocazioni del VLM, riducendo significativamente il tempo di reazione senza sacrificare l'accuratezza.
4. Motore di Inferenza Ottimizzato per l'Edge: Un design di sistema che sfrutta il riutilizzo dei grafi, la residenza in memoria e l'unrolling per massimizzare il throughput su dispositivi limitati dalla banda.

Risultati Sperimentali

Gli autori hanno valutato Jetson-PI sul benchmark LIBERO (usando il modello $\pi0.5$) e su robot reali (PrimeBot X2-W).

• Frequenza di Controllo: Su NVIDIA Jetson Orin, Jetson-PI ottiene un miglioramento della frequenza di controllo di 8,66× rispetto all'inferenza PyTorch naive e di 5,41× rispetto a vla.cpp.
• Tasso di Successo: Sul benchmark LIBERO, Jetson-PI supera lo stato dell'arte dei metodi asincroni, VLASH, del 14,8% nel tasso di successo medio. Mantiene un'alta accuratezza anche all'aumentare della latenza di inferenza ($\Delta$), mentre VLASH e RTC (Reinforced Trajectory Correction) degradano significativamente.
• Prestazioni nel Mondo Reale: Negli esperimenti nel mondo reale che coinvolgono compiti complessi come piegare i vestiti, Jetson-PI su un Jetson Orin ha ottenuto prestazioni paragonabili a un setup con RTX 4090, con una perdita di accuratezza trascurabile rispetto alla baseline di fascia alta.
• Suddivisione della Latenza: Le ottimizzazioni di sistema (riutilizzo dei grafi, buffering, unrolling) hanno ridotto il tempo di reazione totale da ~1,4s a ~0,16s su Jetson Orin.

Significato e Rivendicazioni

Il documento afferma che Jetson-PI fornisce una soluzione pratica per l'implementazione di modelli VLA su larga scala su robot mobili alimentati a batteria e a bassa potenza. Affrontando i colli di bottiglia specifici del calcolo edge — ovvero la limitazione della banda e il divario temporale tra percezione ed esecuzione — il metodo consente un controllo robotico ad alta frequenza e reattivo senza le penalità di potenza dei cluster di GPU di fascia alta. Gli autori sottolineano che, sebbene le piattaforme onboard siano fondamentalmente limitate rispetto ai cluster di GPU, Jetson-PI offre una via percorribile per espandere le applicazioni VLA dalle demo di laboratorio a scenari mobili del mondo reale.