Sintesi Tecnica: Token d'Azione Ordinati per l'Apprendimento di Policy Visuomotorie

Definizione del Problema

La tokenizzazione delle azioni funge da interfaccia critica tra il controllo robotico continuo e la modellazione sequenziale nelle moderne policy visuomotorie. Tuttavia, gli approcci esistenti non riescono a soddisfare simultaneamente tre proprietà essenziali richieste per un efficace apprendimento della policy:

1. Alta Compressione: Le sequenze di token devono essere compatte per minimizzare la difficoltà di addestramento e la latenza di inferenza.
2. Decodificabilità Totale: Qualsiasi sequenza di token generata da una policy deve decodificarsi in modo affidabile in azioni continue valide ed eseguibili.
3. Struttura dei Token Ordinata: La sequenza di token deve organizzare le informazioni rilevanti per il controllo in modo tale che i primi token catturino movimenti grossolani ad alto impatto, mentre i token successivi rifiniscano i dettagli residui. Questa struttura "dal macro al micro" (coarse-to-fine) si allinea con il bias induttivo della predizione del next-token nei modelli autoregressivi e fornisce una supervisione significativa per i modelli vision-language (VLM).

Gli attuali metodi soddisfano sottoinsiemi di questi desiderata ma non tutti simultaneamente. Il binning per dimensione produce una decodificabilità totale ma genera sequenze proibitivamente lunghe. I tokenizer nel dominio della frequenza (es. FAST) offrono un ordine ma mancano di decodificabilità totale a causa di problemi di espansione a lunghezza variabile. I tokenizer latenti appresi (es. QueST, ACodec) forniscono compressione ma spesso mancano di una struttura ordinata esplicita che posizioni le informazioni di controllo globale all'inizio, limitando la loro efficacia come target di policy.

Metodologia: Tokenizzazione dell'Azione Ordinata (OAT)

Gli autori propongono la Tokenizzazione dell'Azione Ordinata (OAT), un tokenizer appreso progettato per soddisfare tutti e tre i desiderata. OAT discretizza chunk di azioni continue in sequenze di token compatte, totalmente decodificabili e ordinate.

Architettura

OAT impiega un'architettura encoder-decoder con un collo di bottiglia discreto:

• Collo di Bottiglia a Registri (Register Bottleneck): L'encoder utilizza token "registro" apprendibili (ispirati ai registri ViT) per effettuare la cross-attention su tutte le posizioni nel chunk di azione continua. Questi registri aggregano l'informazione temporale in una sequenza di stati latenti a lunghezza fissa.
• Quantizzazione Scalare Finita (FSQ): Gli stati dei registri latenti sono discretizzati in un vocabolario fisso tramite FSQ, garantendo che ogni indice del vocabolario mappi un valido chunk di azione continua (Decodificabilità Totale).
• Decoder: Un decoder Transformer a cross-attention ricostruisce il chunk di azione continua dai registri quantizzati.

Addestramento: Ricostruzione del Prefisso Nidificato

Per indurre il desiderato ordine dal macro al micro, OAT utilizza la Ricostruzione del Prefisso Nidificato (Nested Prefix Reconstruction) durante l'addestramento:

• Dropout Nidificato: L'obiettivo di addestramento campiona un budget di ricostruzione $K$ da un insieme di budget (es. $\{1, 2, 4, \dots, H_l\}$). Il decoder è addestrato a ricostruire il chunk di azione completo utilizzando solo i primi $K$ token, con la parte finale (suffix) mascherata.
• Allocazione dell'Informazione: Questo forza i primi token a trasportare informazioni di controllo grossolane e globalmente salienti necessarie per un'azione plausibile, mentre i token successivi raffinano i dettagli residui. Ciò crea una rappresentazione "progressiva" dove qualsiasi prefisso decodifica in un'azione eseguibile.
• Vincoli di Flusso dei Registri: Gli autori introducono specifiche maschere di auto-attenzione per i registri per supportare diversi schemi di generazione. Definiscono due varianti:
  • $OAT_{sing}$ (per token): Utilizza la standard causal attention, imponendo un ordine di dipendenza stretto su ogni token.
  • $OAT_{pow2}$ (Potenza di due): Utilizza la block-causal attention dove i registri all'interno di un gruppo potenza di due non possono attendersi tra loro, ma solo ai gruppi precedenti. Ciò consente una decodifica autoregressiva a blocchi scalabile.

Interfacce di Policy

OAT è validato in due paradigmi di policy prevalenti:

1. Policy Autoregressive (AR): La policy genera i token OAT sequenzialmente (o a blocchi), che vengono poi detonizzati in azioni. La struttura ordinata supporta budget di inferenza variabili, permettendo al sistema di scambiare il costo di inferenza con la fedeltà dell'azione.
2. Co-Addestramento dei Token (TC): Il VLM è addestrato a predire i token OAT (supervisione), mentre un esperto di flow-matching separato genera azioni continue partendo dal contesto distaccato del VLM. In questo caso, il primo token OAT agisce come un riassunto "simile a un piano" dell'intero chunk di azione, fornendo un forte obiettivo globale per la rappresentazione di prefill del VLM.

Generazione Scalabile

Il paper introduce un framework di Generazione Autoregressiva a Blocchi (BAR) che unifica la generazione per token, la decodifica parallela one-shot e gli schemi a blocchi intermedi. Nello specifico, la variante $OAT_{pow2}$ combinata con uno schema di generazione a potenza di due riduce la profondità delle chiamate sequenziali della policy da $O(H_l)$ a $O(\log H_l)$, migliorando significativamente l'efficienza dell'inferenza senza sacrificare le prestazioni.

Contributi Chiave

1. Formalizzazione dei Desiderata: Il paper identifica e formalizza l'alta compressione, la decodificabilità totale e la struttura dei token ordinata come i tre requisiti fondamentali per la tokenizzazione dell'azione come interfaccia di policy.
2. Algoritmo OAT: La proposta di un tokenizer appreso che soddisfa questi requisiti attraverso la compressione basata su registri, FSQ e l'addestramento a prefisso nidificato.
3. Framework di Decodifica Scalabile: L'introduzione di un framework BAR che unifica vari schemi di generazione, dimostrando come il raggruppamento in potenze di due possa ridurre drasticamente la latenza di inferenza.
4. Validazione Completa: Esperimenti estesi su tre backbone di policy (Transformer leggero, PaliGemma2, Qwen3VL) e oltre 60 task (simulati e reali) che coprono sia i paradigmi AR che TC.

Risultati Sperimentali

• Rate-Distortion: OAT traccia una curva fluida di errore di ricostruzione rispetto al budget di token, dimostrando che i prefissi di lunghezze variabili decodificano in azioni progressivamente più accurate.
• Policy Leggere: Nel controllo a ciclo chiuso con un backbone Transformer fisso, OAT supera costantemente i baseline (Bin, FAST, QueST, ACodec). Le prestazioni scalano con il budget di token, con $OAT_{sing}^{8}$ che raggiunge i tassi di successo più elevati in LIBERO, RoboMimic, MetaWorld e compiti reali.
• VLM Autoregressivi: OAT migliora i tassi di successo per i grandi VLM (PaliGemma2, Qwen3VL). La variante potenza di due ($OAT_{pow2}$) eguaglia le prestazioni della generazione per token richiedendo significativamente meno chiamate alla policy (es. 5 chiamate contro 16 per la generazione completa).
• Co-Addestramento dei Token (TC): Anche quando i token non vengono decodificati durante l'inferenza, la supervisione OAT porta a rappresentazioni VLM superiori. OAT ottiene il successo medio più alto nei setting TC, superando QueST e ACodec in diversi benchmark, validando l'ipotesi della supervisione "simile a un piano".
• Ablazioni: Rimuovere l'addestramento del prefisso ordinato (nested dropout) degrada significativamente le prestazioni, confermando che la struttura dal macro al micro è critica. Inoltre, un disallineamento tra la struttura di attenzione dei registri e lo schema di generazione (es. usare $OAT_{sing}$ con blocchi potenza di due) riduce le prestazioni, evidenziando la necessità di allineamento tra il design del tokenizer e lo schema di inferenza.

Significato e Rivendicazioni

Il paper sostiene che OAT fornisca un'interfaccia robusta e flessibile per l'apprendimento di policy visuomotorie. Soddisfacendo la triade di compressione, decodificabilità e ordine, OAT abilita:

• Inferenza Flessibile: Le policy possono adattarsi scegliendo budget di token variabili per bilanciare latenza e fedeltà.
• Migliore Apprendimento: La struttura ordinata si allinea con i bias induttivi dei modelli autoregressivi e fornisce una supervisione globale significativa per i VLM nei setup di co-addestramento.
• Scalabilità: L'integrazione con BAR permette un'inferenza efficiente su modelli grandi senza la penalità di profondità lineare della generazione standard per token.

Gli autori concludono che, sebbene lo studio sia limitato a specifici benchmark e backbone, i principi di design di OAT offrono una soluzione generalizzabile al problema della tokenizzazione dell'azione, aprendo la strada a sistemi di IA incarnata (embodied AI) più efficienti e capaci. Notano che il lavoro futuro potrebbe esplorare il calcolo adattivo, dove le policy decidono dinamicamente il budget di token necessario in base alla complessità o incertezza del compito.