Riepilogo Tecnico: Latent Memory Palace (LMP)

Definizione del Problema

Il processo decisionale umano presenta uno spettro che va dalle azioni riflessive al ragionamento deliberato e multi-step, dove il tempo e l'impegno investiti scalano con la complessità del compito. Sebbene i Large Language Models (LLM) abbiano dimostrato con successo capacità di "ragionamento" simili attraverso la generazione iterativa di token (ad esempio, Chain-of-Thought), trasferire questa capacità alle policy di controllo robotico continuo rimane una sfida. Ragionare direttamente nello spazio del linguaggio spesso manca della granularità necessaria per la comprensione spaziale e il controllo preciso del movimento. Inoltre, gli approoli standard di inferenza variazionale per il controllo si basano tipicamente su variabili latenti a dimensionalità fissa, che non possono adattare la propria profondità computazionale alle specifiche esigenze di ogni input. Il problema centrale è come consentire alle policy robotiche di eseguire un calcolo iterativo e adattivo in uno spazio latente appreso end-to-end per la predizione dell'azione, bilanciando così l'efficienza decisionale con un migliore generalismo.

Metodologia

Gli autori propongono il Latent Memory Palace (LMP), un framework che formula il ragionamento per il controllo come inferenza variazionale autoregressiva.

1. Formulazione Core

LMP tratta un'azione $a$ non come una predizione diretta da un'osservazione $o$, ma come il risultato di un calcolo latente intermedio $z$. La policy è definita da un limite inferiore variazionale (ELBO):
$$\log p(a | o) \geq \mathbb{E}_{z \sim q_\theta(z|o,a)} [\log p_\phi(a | o, z)] - D_{KL}(q_\theta(z | o, a) \parallel p_\theta(z | o))$$
A differenza degli approcci standard, la variabile latente $z$ non è un singolo vettore, ma una sequenza autoregressiva a lunghezza variabile $z = (z_1, z_2, \dots, z_{T(z)})$ estratta da un vocabolario che include un token di Fine Sequenza (EOS). La sequenza termina quando viene generato l'EOS, permettendo al modello di allocare adattivamente i passi di calcolo in base alla complessità della decisione.

2. Ragionamento Adattivo tramite Compressione

Per elicitare un comportamento adattivo in cui il modello utilizza più passi solo quando necessario, LMP introduce un meccanismo di compressione. Il decoder dell'azione $p_\phi$ è parametrizzato come una Gaussiana isotropa la cui varianza $\sigma(T(z))$ decade esponenzialmente con la lunghezza della sequenza latente $T(z)$.
$$\sigma(T(z)) = \gamma^{T(z)} \sigma_0, \quad \gamma \in (0, 1)$$
Questo impone una penalità sulla lunghezza: ogni passo latente aggiuntivo restringe la distribuzione Gaussiana. Di conseguenza, il modello può "permettersi" di generare un altro passo solo se la precisione dell'azione predetta giustifica la riduzione della varianza. Ciò costringe la policy a dedicare più calcolo alle decisioni complesse e meno a quelle semplici o intrinsecamente incerte (ad esempio, il timing esatto della chiusura di una pinza).

3. Ottimizzazione Trattabile tramite RL nello Spazio Latente

L'ottimizzazione dell'ELBO attraverso il campionamento autoregressivo non differenziabile è intrattabile tramite la backpropagation standard. Gli autori derivano un obiettivo di Reinforcement Learning (RL) sulle traiettorie latenti:

• Episodio Latente: Una traiettoria viene campionata dalla posterior $q_\theta$ fino all'EOS.
• Ricompensa: Composta da una ricompensa di ricostruzione ($\log p_\phi(a|o,z)$) e una penalità KL rispetto alla prior.
• Riduzione della Varianza: Il termine KL è decomposto in ricompense per singolo step utilizzando la struttura autoregressiva, permettendo l'assegnazione del credito (credit assignment) per step.
• Ottimizzazione: Gli autori utilizzano un obiettivo surrogato clippato di tipo PPO per stabilizzare l'addestramento e prevenire il collasso latente, utilizzando rollout on-policy raccolti in un buffer.

4. Tokenizzazione dell'Azione (LMP-tok)

Rimuovendo il condizionamento dell'osservazione dal framework variazionale, la stessa architettura istanzia LMP-tok, un tokenizer di azioni a lunghezza variabile. Questo comprime chunk di azioni continue in token discreti, che possono poi essere utilizzati per addestrare policy autoregressive a valle. A differenza dei tokenizer che richiedono la ricostruzione esatta di trasformazioni discrete (come FAST), LMP-tok permette a sequenze di lunghezza arbitraria di essere decodificate in azioni valide.

Contributi Chiave

1. Latent Memory Palace (LMP): Una nuova formulazione del ragionamento come inferenza variazionale con una distribuzione latente autoregressiva, che consente un calcolo iterativo e adattivo per il controllo robotico.
2. LMP-$\pi$ (Policy): Una policy di controllo che genera tracce di "ragionamento" a lunghezza variabile durante il test, allocando adattivamente il calcolo in base alla difficoltà e all'incertezza del compito.
3. LMP-tok (Tokenizer): Un tokenizer di azioni a lunghezza variabile che migliora significativamente le performance delle policy autoregressive a valle comprimendo le azioni in spazi latenti discreti.
4. Framework di Ottimizzazione: Un metodo basato su RL trattabile per ottimizzare il limite inferiore variazionale per il campionamento autoregressivo a lunghezza variabile, inclusi i tecniche per stabilizzare l'addestramento contro il collasso latente.

Risultati Sperimentali

Gli autori hanno valutato LMP sia su benchmark del mondo reale (DROID) che simulati (LIBERO, D3IL, RoboMimic).

• Performance della Policy (LMP-$\pi$):

  • Zero-Shot & Fine-tuning: Sulla piattaforma DROID, LMP-$\pi$ ha superato Diffusion Policy (DP) in scenari zero-shot e ha raggiunto tassi di successo più elevati (es. 95% nel cleaning multi-oggetto) in compiti ad alta precisione sottoposti a fine-tuning.
  • Generalizzazione: Su LIBERO-90, LMP-$\pi$ ha mostrato un "pavimento" di performance significativamente più alto nei 10 compiti più difficili rispetto a DP, indicando una riduzione dell'interferenza tra i compiti e una migliore capacità di recupero.
  • Iterativo vs Non-Iterativo: LMP-$\pi$ ha costantemente superato le baseline non iterative (VAE Policy, VQ-BeT), dimostrando che le rappresentazioni latenti a lunghezza fissa mancano della granularità richiesta per la manipolazione complessa.

• Performance del Tokenizer (LMP-tok):

  • Utilizzato per addestrare policy autoregressive a valle, LMP-tok ha superato significativamente le baseline (VQ-VAE, FAST, OAT) nei compiti RoboMimic e LIBERO.
  • In particolare, LMP-tok è stato l'unico metodo a ottenere prestazioni non trascurabili nel compito ad alta precisione "tool-hang" di RoboMimic, dove altri tokenizer sono falliti a causa di vincoli di precisione o ricostruzione.

• Analisi del Calcolo Adattivo:

  • Interpretabilità: Le visualizzazioni mostrano che il numero di passi di ragionamento latente correla con la complessità del compito. Il modello utilizza meno passi durante i movimenti della pinza (dove il timing è intrinsecamente variabile) e più passi durante compiti di allineamento preciso.
  • Compressione: Gli studi di ablazione hanno confermato che una compressione più forte (decadimento più rapido della varianza) porta a una migliore generalizzazione, a patto che sia bilanciata per evitare un eccessivo errore di ricostruzione.
  • Mode-Seeking: Troncare i passi latenti durante il test non ha degradato significativamente le performance; al contrario, il modello ha mostrato un comportamento di "mode-seeking", impegnandosi verso specifici modi d'azione validi invece di mediarli.

Significato e Rivendicazioni

Il paper sostiene di stabilire l'inferenza variazionale autoregressiva come un framework pratico e fondato per il calcolo iterativo e adattivo nel robot learning. Tracciando un'analogia con il metodo mnemonico del "palazzo della memoria", il lavoro dimostra che organizzare le informazioni in uno spazio latente a lunghezza variabile permette alle policy di:

1. Imitare la Variabilità Decisionale Umana: Scalare lo sforzo computazionale con la complessità della decisione.
2. Migliorare la Generalizzazione: Raggiungere performance robuste attraverso diversi compiti e ambienti imparando ad allocare il calcolo in modo efficiente.
3. Potenziare la Tokenizzazione: Fornire un metodo superiore per discretizzare le azioni continue per la modellazione autoregressiva.

Gli autori posizionano LMP come un ponte tra le capacità di ragionamento dei modelli di fondazione e i requisiti di precisione del controllo robotico, offrendo una nuova prospettiva in cui il "ragionamento" è un processo latente implicito ottimizzato attraverso l'inferenza variazionale e il reinforcement learning.