Scaling Short-Term Memory of Visuomotor Policies for Long-Horizon Tasks

Sintesi Tecnica: Scalare la Memoria a Breve Termine delle Politiche Visuomotorie per Task a Lungo Orizzonte

1. Definizione del Problema

I compiti robotici richiedono spesso una memoria a breve termine (che copra secondi o minuti) per gestire l'osservabilità parziale, come il recupero di oggetti che non sono più visibili o la temporizzazione di azioni come spegnere un elettrodomestico. Tuttavia, la maggior parte delle attuali politiche visuomotorie addestrate tramite apprendimento per imitazione sono miopi, poiché si affidano esclusivamente agli input sensoriali immediati senza sfruttare le esperienze passate.

Estendere le standard politiche basate su transformer per includere finestre di contesto più ampie al fine di affrontare questo problema presenta due sfide critiche:

1. Correlazioni Spurie: L'attenzione ingenua verso cronologie lunghe introduce informazioni irrilevanti (ad esempio, posizioni passate di distrattori), portando i modelli a imparare false correlazioni che degradano le prestazioni quando il contesto cambia.
2. Scalabilità Computazionale: L'impronta computazionale e di memoria dei meccanismi di attenzione standard scala linearmente o quadraticamente con la lunghezza del contesto. Elaborare input visivi ad alta dimensionalità su sequenze temporali lunghe (ad esempio, minuti) diventa proibitivo in termini di costi.

2. Metodologia: PRISM

Gli autori propongono PRISM (Policy Architecture with Short-Term Memory), un'architettura basata su transformer progettata per utilizzare efficacemente la memoria a breve termine attraverso due innovazioni chiave:

• Gated Attention (Filtraggio delle Informazioni): Per mitigare le correlazioni spurie, PRISM introduce un meccanismo di gating post-attenzione. Questo componente filtra selettivamente le informazioni recuperate modulando l'output di ogni blocco di attenzione tramite un gate apprendibile condizionato dalle caratteristiche dell'input corrente. Ciò consente alla politica di sopprimere adattivamente i dettagli irrilevanti dalla cronologia, riducendo il rumore nella predizione dell'azione.
• Architettura Gerarchica (Efficienza): Per affrontare la scalabilità, PRISM impiega un design gerarchico.
  • Codifica Locale: Un encoder di memoria locale elabora gli input sensoriali e li comprime in token compatti utilizzando un resampler in stile Perceiver.
  • Attenzione Globale: Un encoder di memoria globale calcola l'attenzione su questi token compressi attraverso il tempo.
  • Fusione: L'informazione globale risultante viene trasmessa nuovamente e fusa con i token locali originali tramite connessioni residue.
  • Beneficio: Questo approccio riduce la complessità computazionale da $O((n \cdot k_{tot})^2)$ a circa $O(n \cdot k_{tot} \cdot m + (n \cdot k_{tot}/m)^2)$, dove $m$ è il fattore di compressione. Riduce inoltre significativamente l'impronta di memoria durante il test memorizzando solo i token riassuntivi compressi.

PRISM è addestrato utilizzando il classico behavior cloning (apprendimento per imitazione) senza loss ausiliarie o moduli di memoria esterni, mappando sequenze di input sensoriali e azioni passate verso azioni future in un contesto a ciclo chiuso.

3. Contributi Chiave

• Architettura PRISM: Una nuova architettura di policy che combina l'attenzione a gate e la riassunzione gerarchica per scalare la memoria a breve termine fino a due minuti (a 2–3 Hz) mantenendo l'efficienza computazionale e la robustezza al rumore.
• REMEMBENCH: Un nuovo benchmark composto da otto diversi compiti di manipolazione domestica che spaziano su quattro categorie di memoria a breve termine derivate dalle scienze cognitive:
  • Spaziale: Ricordare le posizioni degli oggetti (ad esempio, recuperare frutta nascosta).
  • Prospettica: Trattenere le intenzioni durante un ritardo (ad esempio, cuocere la carne per una durata specifica).
  • Associazione Oggetto-Oggetto: Ricordare le associazioni tra oggetto e posizione (ad esempio, restituire un oggetto lavato al suo piattino originale specifico).
  • Insieme di Oggetti: Mantenere e aggiornare insiemi di più oggetti (ad esempio, contare le bacchette di pane spostate in un microonde).
  • A differenza dei benchmark precedenti che si concentrano su singoli aspetti o test di memoria indiretti, REMEMBENCH richiede esplicitamente la memoria per avere successo, poiché gli input sensoriali immediati sono insufficienti.

4. Risultati Sperimentali

Il paper valuta PRISM rispetto ad architetture ricorrenti (LSTM, Mamba), transformer ricorrenti a livello di segmento (Transformer-XL) e altre varianti di transformer (Linear Attention, Scene Memory Transformer) su REMEMBENCH e benchmark standard.

• Prestazioni su REMEMBENCH: PRISM raggiunge un tasso di successo medio del 41%, superando i baseline più forti (SMT al 36%, PTP al 28%, Mamba al 25%) con un margine assoluto del 5%–12%.
  • Scalabilità: Le prestazioni scalano con la dimensione della memoria; aumentare la finestra da 1 a 512 timestep migliora il tasso di successo dal 16% al 42% senza saturazione.
  • Ablation: Rimuovere il meccanismo di gating causa un calo di 16 punti nelle prestazioni, confermando la sua necessità per filtrare la cronologia irrilevante.
• Benchmark Standard (ROBOCASA & LIBERO):
  • Su ROBOCASA, PRISM con memoria (n=256) raggiunge il 43% di successo, un miglioramento assoluto dell'11% rispetto alla sua variante senza memoria e un miglioramento dell'11% rispetto al baseline GR00T-N1-2B fine-tuned.
  • Su LIBERO, PRISM raggiunge l'89% di successo medio, un miglioramento del 15% rispetto al baseline OpenVLA fine-tuned e un miglioramento del 12% rispetto alla sua variante senza memoria.
  • Gli autori notano che l'aggiunta della memoria aiuta a risolvere l'ambiguità dell'azione in stati visivamente identici (ad esempio, distinguere tra "appena afferrato" e "per essere posizionato"), riducendo la multimodalità anche in compiti non esplicitamente progettati per testare la memoria.
• Valutazione nel Mondo Reale: In un adattamento reale del compito "Lavare e Restituire al Contenitore", PRISM ha raggiunto il 30% di successo, rispetto al 15% di un baseline transformer senza memoria e allo 0% di una policy senza memoria.
• Efficienza: Con una lunghezza di contesto di 512 step, PRISM riduce l'uso di memoria GPU di picco dell'80% e il tempo di esecuzione del 40% rispetto all'attenzione globale completa e agli approcci ibridi a finestra scorrevole.

5. Significato e Rivendicazioni

Il paper sostiene che PRISM e REMEMBENCH stabiliscono una base per lo sviluppo e la valutazione di politiche visuomotorie aumentate dalla memoria a breve termine capaci di gestire compiti a lungo orizzonte.

• Generalizzazione: Coprendo quattro distinte categorie di memoria, REMEMBENCH incoraggia lo sviluppo di meccanismi di memoria generali piuttosto che soluzioni isolate e specifiche per il compito.
• Scalabilità: PRISM dimostra che la memoria a breve termine può essere scalata a minuti di interazione senza i costi proibitivi associati alle estensioni naive della finestra di contesto.
• Robustezza: Il meccanismo di gated attention si rivela essenziale per prevenire il modello dall'apprendere correlazioni spurie in dati a lungo orizzonte rumorosi.

Gli autori concludono che, sebbene i tassi di successo attuali su REMEMBENCH (intorno al 30–40%) indichino che i compiti a lungo orizzonte rimangono difficili, i costanti guadagni dall'aumento della memoria attraverso diverse architetture e benchmark validano l'approccio. Identificano direzioni future come l'integrazione della memoria persistente a lungo termine e l'utilizzo del pre-training su scala internet per i moduli di memoria, ma sottolineano che il loro lavoro attuale si concentra strettamente sulla meccanica della memoria a breve termine all'interno del framework di apprendimento per imitazione.