Improving Diffusion Planners by Self-Supervised Action Gating with Energies

Il paper propone SAGE, un metodo di riordinamento a inferenza che migliora la robustezza dei pianificatori basati su diffusione nell'apprendimento per rinforzo offline penalizzando i piani dinamicamente incoerenti tramite un segnale di consistenza latente, senza richiedere nuove interazioni con l'ambiente o riaddestramento della policy.

L'essenziale

Immagina di dover insegnare a un robot come camminare, afferrare oggetti o navigare in un labirinto, ma hai un solo modo per farlo: dargli un vecchio album di foto di qualcuno che ha già fatto queste cose. Non puoi fargli provare nulla di nuovo, non puoi fargli sbagliare e riprovare, devi solo imparare da quelle foto. Questo è il mondo dell'Apprendimento per Rinforzo Offline.

Il problema? A volte, guardando le foto, il robot immagina scenari fantastici che sembrano perfetti sulla carta, ma che nella realtà sono impossibili. È come se il robot leggesse un libro di ricette, decidesse di fare una torta di cioccolato, ma poi provasse a mettere il cioccolato nel forno insieme a un uovo intero non rotto e a un mattone. Il piano sembra buono sulla carta (il "punteggio" è alto), ma fallisce miseramente appena inizia a cucinare.

Gli scienziati hanno creato dei "pianificatori" basati su una tecnologia chiamata Diffusione (simile a quella usata per generare immagini artistiche) che creano molti possibili futuri e scelgono il migliore. Ma questi pianificatori a volte scelgono piani che sembrano ottimi ma che il robot non può eseguire fisicamente.

Ecco che entra in gioco SAGE (Self-supervised Action Gating with Energies), la soluzione proposta in questo articolo.

L'Analogia del "Controllore di Sicurezza"

Immagina che il tuo robot sia un architetto visionario (il pianificatore a diffusione). Questo architetto è bravissimo a disegnare 100 progetti di case bellissimi e costosi. Tuttavia, l'architetto a volte disegna scale che portano al nulla, o finestre che si aprono su un muro solido.

Prima di SAGE, un ispettore (il sistema di valutazione) guardava questi 100 progetti e sceglieva quello che costava di più o sembrava più lussuoso. Risultato? Spesso si sceglieva un palazzo che crollava appena costruito.

SAGE è come un nuovo ispettore, un "Controllore di Sicurezza" super-intelligente, che lavora in coppia con l'ispettore originale.

Ecco come funziona SAGE, passo dopo passo:

1. L'Apprendimento Silenzioso (Il "Senso Comune"):
   Prima ancora di vedere un piano, SAGE guarda migliaia di foto di persone che camminano o muovono oggetti (i dati offline). Non impara cosa è bello, ma impara cosa è fisicamente possibile.

   • Metafora: È come se SAGE avesse letto milioni di manuali di fisica e di movimento. Sa che se spingi un oggetto pesante, questo non si muove come una piuma. Sa che un'auto non può girare di 90 gradi istantaneamente senza sbandare.

2. Il Test del "Primo Passo" (La Gating):
   Quando l'architetto visionario (il pianificatore) propone 100 piani, SAGE non guarda l'intero piano. Guarda solo i primi pochi passi.

   • L'Energia: SAGE calcola un "livello di energia" (o di sforzo). Se il primo passo del piano richiede che il robot compia un'azione che contraddice la fisica che ha imparato (es. "camminare attraverso un muro"), SAGE dice: "Attenzione! Questo piano ha un'energia troppo alta, è incoerente!".
   • Se il piano è fisicamente possibile, l'energia è bassa.

3. La Selezione Intelligente:
   SAGE non butta via i piani. Fa una selezione a due livelli:

   • Prima, scarta tutti i piani che hanno un'energia troppo alta (quelli che il robot non può eseguire).
   • Poi, tra quelli rimasti (quelli che possono essere eseguiti), lascia che l'ispettore originale scelga quello che promette il miglior risultato (il punteggio più alto).

Perché è una Rivoluzione?

Fino a ora, per evitare questi errori, bisognava ridisegnare tutto il sistema di apprendimento del robot, rendendolo lento e complicato.

SAGE è diverso perché è modulare e non invasivo:

• Non devi ridisegnare l'architetto visionario.
• Non devi fargli fare nuove prove nel mondo reale (che sarebbe pericoloso o costoso).
• SAGE è un "filtro" che si aggiunge alla fine, come un controllore di sicurezza in aeroporto. L'aereo (il piano) può essere bellissimo, ma se il controllore vede che manca un motore, non lo fa decollare.

In Sintesi

Il paper ci dice che i robot spesso falliscono non perché non sanno cosa vogliono fare, ma perché scelgono piani che sembrano ottimi ma sono fisicamente impossibili da iniziare.

SAGE è come un guardiano del tempo che usa la sua conoscenza della fisica (imparata da solo guardando i dati) per dire: "Aspetta, questo piano è bello, ma il primo passo è impossibile. Passiamo al prossimo".

Grazie a questo semplice ma potente filtro, i robot diventano molto più affidabili, meno fragili e capaci di eseguire compiti complessi (come camminare su terreni difficili o cucinare) senza rompersi o cadere, semplicemente perché smettono di scegliere piani che non possono funzionare nella realtà.

Sommario tecnico

1. Il Problema: Fragilità dei Planner Diffusivi nell'RL Offline

L'apprendimento per rinforzo offline (Offline RL) mira ad apprendere politiche da dataset fissi senza interazione online con l'ambiente. Un approccio recente e potente utilizza modelli diffusivi per la pianificazione: invece di apprendere una politica puntuale, questi modelli generano distribuzioni su sequenze di azioni (traiettorie) e selezionano i candidati migliori basandosi su una funzione di valore appresa.

Tuttavia, i planner basati su diffusione soffrono di un fallimento specifico:

• Incoerenza Dinamica Locale: Il processo di selezione si basa spesso su un critico (funzione di valore) che valuta la "bontà" a lungo termine di una traiettoria. Questo può portare a selezionare traiettorie che hanno un alto punteggio di valore ma sono localmente incoerenti con la dinamica dell'ambiente (es. un'azione iniziale che non è fisicamente realizzabile dallo stato corrente).
• Conseguenze: Quando un piano inizia con una transizione irrealizzabile, l'esecuzione diventa fragile. L'agente si impegna in un prefisso irrealistico, e il successivo ripianificazione (replanning) può fallire a cascata, portando a comportamenti instabili o a collisioni.
• Limitazione delle Soluzioni Esistenti: I metodi precedenti che tentano di correggere questo problema spesso richiedono modelli aggiuntivi complessi, guida durante il processo di denoising (che può distorcere la distribuzione appresa) o interazione con l'ambiente, limitando la scalabilità.

2. Metodologia: SAGE (Self-supervised Action Gating with Energies)

Gli autori propongono SAGE, un metodo di riordinamento (re-ranking) a tempo di inferenza che separa esplicitamente il segnale di "fattibilità" (feasibility) dal segnale di "valore". SAGE non modifica il generatore diffusivo né il critico di valore, ma agisce come un filtro intelligente sui candidati generati.

L'architettura si basa su due componenti apprese puramente offline:

A. Rappresentazione Predittiva (Stage I - JEPA)

Viene addestrato un JEPA (Joint-Embedding Predictive Architecture) su sequenze di stati offline.

• Obiettivo: Imparare una rappresentazione latente in cui gli stati futuri sono prevedibili dal contesto locale.
• Meccanismo: Un encoder $e_\theta$ mappa finestre di stati mascherati in uno spazio latente. Un "teacher" (EMA) fornisce i target per stati futuri. Un predittore $g_\phi$ cerca di prevedere l'embedding futuro dal contesto.
• Risultato: Un encoder EMA congelato che cattura la dinamica coerente con il dataset.

B. Predittore Latente Condizionato all'Azione (Stage II)

Sulla rappresentazione latente congelata, viene addestrato un predittore $f_\eta$ che modella le transizioni a breve termine.

• Input: Stato latente corrente $z_t$ e azione $a_t$.
• Output: Predizione dello stato latente successivo $\hat{z}_{t+1}$.
• Loss: Include un termine per la predizione forzata dal teacher, un termine di coerenza per rollout brevi e un hinge loss negativo che penalizza il predittore se rimane accurato anche quando le azioni vengono permutate (assicurando che il predittore sia sensibile alle azioni).

C. Meccanismo di Gating (Inferenza)

Al momento dell'inferenza, SAGE valuta i candidati generati dal planner diffusivo:

1. Calcolo dell'Energia: Per ogni traiettoria candidata $\hat{\tau}$, viene calcolata un'"energia" basata sull'errore di predizione latente sui primi $K$ passi (prefisso):
   $$E(\hat{\tau}) = \frac{1}{K} \sum_{k=0}^{K-1} \| f_\eta(z_{t+k}, a_{t+k}) - z_{t+k+1} \|_1$$
   Un'energia bassa indica che la transizione è coerente con la dinamica appresa dal dataset; un'energia alta segnala incoerenza.
2. Selezione Ibrida:
   • Si filtrano i candidati mantenendo solo la frazione $P$ con l'energia più bassa (fattibilità).
   • Tra i rimanenti, si seleziona quello con il punteggio di valore $J$ più alto, applicando una penalità soft basata sull'energia:
     $$i^* \in \arg \max_{i} (J(\hat{\tau}^{(i)}) - \lambda E(\hat{\tau}^{(i)}))$$

3. Contributi Chiave

1. Separazione Fattibilità/Valore: SAGE introduce un segnale di fattibilità auto-supervisionato distinto dal critico di valore, risolvendo la tensione tra l'esplorazione di valori alti e la necessità di coerenza dinamica locale.
2. Nessuna Interazione Online: L'intero sistema (encoder e predittore) è addestrato esclusivamente su dati offline, senza bisogno di rollouts ambientali o riaddestramento del planner diffusivo.
3. Modularità: SAGE è un modulo "plug-and-play" che può essere integrato in qualsiasi pipeline di pianificazione diffusiva esistente che supporta il campionamento e la selezione basata su punteggio.
4. Scalabilità: Utilizza l'apprendimento auto-supervisionato su larga scala, rendendolo adatto a dataset eterogenei e complessi.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su una vasta gamma di benchmark D4RL, coprendo locomozione (MuJoCo), manipolazione (Franka Kitchen) e navigazione (AntMaze, Maze2D).

• Performance Generale: SAGE migliora costantemente le prestazioni dei planner diffusivi più avanzati (in particolare DV, un planner "generate-and-rank" stato dell'arte).
  • Locomozione: Miglioramento medio da 82.9 (DV) a 84.4.
  • Manipolazione (Kitchen): Aumento significativo, raggiungendo 96.6 su kitchen-partial-v0 (vs 90.0 di DV).
  • Navigazione (AntMaze/Maze2D): SAGE supera i baseline di fattibilità precedenti e migliora la robustezza nella navigazione a lungo termine.
• Validazione del Segnale di Energia:
  • Test di corruzione: Quando le azioni in una traiettoria vengono corrotte (shufflate), l'energia di SAGE mostra un picco localizzato preciso, confermando che il segnale rileva efficacemente le violazioni della dinamica.
  • AUROC: SAGE ottiene un'AUROC molto alta (es. 0.98 su MuJoCo) nel distinguere transizioni reali da quelle incoerenti, superando modelli dinamici diretti (ridge/MLP) e spazi latenti casuali.
• Efficienza Computazionale: L'overhead computazionale aggiunto da SAGE è minimo (~6.8% in più rispetto al solo planner), rendendolo pratico per applicazioni reali.

5. Significato e Impatto

Il lavoro di SAGE è significativo perché affronta una delle principali cause di fallimento nell'RL offline: l'illusione di pianificazione. I planner possono "sognare" traiettorie ottimali che non sono fisicamente eseguibili.

• Affidabilità: SAGE rende i planner diffusivi più robusti, prevenendo l'esecuzione di azioni iniziali che porterebbero a fallimenti immediati.
• Paradigma di Selezione: Dimostra che la fase di selezione (ranking) è cruciale quanto la generazione. Aggiungere un filtro di fattibilità auto-supervisionato permette di sfruttare appieno la capacità generativa dei modelli diffusivi senza comprometterne la stabilità.
• Generalizzazione: Poiché non richiede riaddestramento del generatore, SAGE può essere applicato immediatamente per migliorare qualsiasi sistema di pianificazione basato su diffusione esistente, offrendo un percorso pratico verso un'RL offline più sicura e affidabile per la robotica.