Gated Memory Policy

1. Il Problema

Le attività di manipolazione robotica presentano requisiti di memoria estremamente variabili:

• Task Markoviani: Richiedono poca o nessuna memoria storica (es. presa e posizionamento semplice).
• Task Non-Markoviani: Richiedono memoria "in-trial" (all'interno di un singolo episodio) o "cross-trial" (tra più tentativi) per inferire proprietà fisiche non osservabili (es. attrito, massa) o adattarsi al contesto.

La sfida principale: Estendere semplicemente la finestra di osservazione storica (input history) di una policy visuomotoria per gestire task complessi porta spesso a un crollo delle prestazioni. Questo accade perché:

1. Overfitting: Aumentare la dimensionalità dell'input con dati storici fissi rende difficile la generalizzazione.
2. Shift di Distribuzione: La presenza di osservazioni rumorose o non familiari in alcuni frame storici può spostare l'intera sequenza di input fuori dalla distribuzione di addestramento.
3. Costo Computazionale: L'attenzione su lunghe storie storiche (specialmente nei Transformer) scala quadraticamente ($O(H^2)$), rendendo l'inferenza lenta e costosa.
4. Ridondanza: In molti momenti di un task, la memoria non è necessaria; forzarne l'uso distrae la policy dall'osservazione corrente e riduce la reattività.

2. Metodologia: Gated Memory Policy (GMP)

Gli autori propongono Gated Memory Policy (GMP), un framework che impara a rispondere a due domande fondamentali: Quando richiamare la memoria e Cosa richiamare.

A. Architettura di Base

GMP si basa su una Diffusion Policy (DiT - Diffusion Transformer). Invece di concatenare direttamente tutti i token storici in una lunga sequenza (come fanno le policy tradizionali), GMP utilizza un modulo di Cross-Attention dedicato.

• Token Storici: Le azioni storiche e le osservazioni visive (estratte da un encoder ViT e aggregati tramite Multi-head Attention Pooling per ridurre i token) vengono memorizzati in una cache a finestra scorrevole.
• Efficienza: Questo approccio riduce la complessità computazionale da quadratica a lineare rispetto alla lunghezza della storia.

B. Meccanismi Chiave

1. Porta di Memoria Appresa (Learned Memory Gate):

   • GMP utilizza un MLP con attivazione sigmoide per produrre un valore binario $\mu_t \in \{0, 1\}$.
   • Se $\mu_t = 1$, l'output dell'attenzione incrociata sulla storia viene attivato; se $\mu_t = 0$, viene ignorato.
   • Calibrazione della Porta: Poiché addestrare la porta in modo end-to-end è difficile (tende a usare sempre la storia per minimizzare l'errore di training), gli autori propongono un processo di calibrazione self-supervisionato:
     1. Si addestrano due policy separate su un set di training: una senza memoria ($\pi$) e una con memoria sempre attiva ($\pi_{mem}$).
     2. Si valuta su un set di validazione: se l'errore di predizione dell'azione di $\pi$ è significativamente maggiore di quello di $\pi_{mem}$ (soglia $\theta$), si etichetta quel timestep come "richiede memoria" ($\mu_t=1$).
     3. Si congela la porta calibrata e si ri-addestra la policy finale su tutto il dataset.

2. Augmentation con Rumore Diffusionale (Diffusion Noise Augmentation):

   • Per migliorare la robustezza e prevenire l'overfitting su storie "pulite", GMP inietta rumore schedulato nelle azioni storiche durante l'addestramento.
   • A ogni passo di diffusione $k$, la policy è condizionata su azioni storiche al passo $k-1$ (leggermente meno rumorose). Questo simula una distribuzione di input più ampia e diversificata, rendendo la policy robusta a errori accumulati o rumore durante l'inferenza reale.

3. Contributi Principali

1. Gated Memory Policy (GMP): Un framework pratico ed efficiente che seleziona dinamicamente la memoria solo quando necessario, mantenendo robustezza e reattività. Funziona sia per task Markoviani che Non-Markoviani senza sacrificare le prestazioni.
2. MemMimic: Un nuovo benchmark (simulazione e mondo reale) progettato specificamente per valutare l'uso della memoria, con un focus sui task "cross-trial" (adattamento contestuale) spesso trascurati dai benchmark esistenti.
3. Studio Empirico Completo: Analisi dettagliata delle scelte progettuali (calibrazione della porta, iniezione di rumore, attenzione incrociata) e del loro impatto su prestazioni, robustezza ed efficienza computazionale.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su MemMimic (task non Markoviani) e RoboMimic (task Markoviani).

• Prestazioni su Task Non-Markoviani: Su MemMimic, GMP ha ottenuto un miglioramento medio del 30.1% nel tasso di successo rispetto alle baseline a lunga storia (Long-hist baselines). GMP eccelle in task come:
  • Iterative Pushing/Flinging/Casting: Adattamento al attrito o massa sconosciuta attraverso tentativi ed errori.
  • Match Color with Delay: Mantenimento della memoria visiva per lunghi periodi.
• Prestazioni su Task Markoviani: Su RoboMimic (Square, Transport, Tool Hang), GMP mantiene prestazioni competitive, evitando il degrado tipico delle policy che forzano l'uso di lunghe storie storiche. La porta di memoria rimane spenta ($\mu_t=0$) nella maggior parte dei casi (es. 73% in Match Color).
• Efficienza Computazionale:
  • L'uso dell'attenzione incrociata riduce la scalatura dei costi da $O(H^2)$ a $O(H)$.
  • Quando la porta è spenta, il costo di inferenza diventa costante e minimo, permettendo una reattività superiore rispetto alle policy che elaborano sempre l'intera storia.
• Robustezza: L'iniezione di rumore nelle azioni storiche ha dimostrato di migliorare significativamente la robustezza in ambienti reali e simulati rumorosi, superando approcci come Diffusion Forcing.

5. Significato e Impatto

Il lavoro di GMP risolve il dilemma fondamentale tra la necessità di memoria per task complessi e la necessità di efficienza e robustezza per task semplici.

• Generalizzazione: Dimostra che non è necessario un compromesso tra task Markoviani e Non-Markoviani; una policy può adattarsi dinamicamente.
• Efficienza: Introduce un meccanismo per "saltare" il calcolo della memoria quando non necessario, rendendo i sistemi robotici più reattivi ed economici da eseguire.
• Nuovo Standard: Con il benchmark MemMimic, fornisce un modo rigoroso per valutare le capacità di memoria a lungo termine (intra- e inter-trial), spingendo la ricerca verso policy robotiche più adattive e capaci di apprendere dall'esperienza passata in modo intelligente.

In sintesi, GMP rappresenta un passo avanti verso policy robotiche che non solo "vedono" il presente, ma sanno anche "ricordare" il passato solo quando è strettamente necessario per il successo del compito.