Amortizing Trajectory Diffusion with Keyed Drift Fields

1. Il Problema

I pianificatori basati su modelli di diffusione (Diffusion Models) hanno dimostrato grande successo nella generazione di sequenze di azioni multimodali per l'apprendimento per rinforzo offline (Offline RL) e il controllo robotico. Tuttavia, presentano un limite critico: richiedono un processo iterativo di "denoising" (rimozione del rumore) composto da molti passaggi sequenziali ($T$) per generare una singola traiettoria.

• Costo computazionale: Questo approccio comporta un'alta latenza di inferenza e un elevato numero di valutazioni della rete neurale per ogni passo di controllo.
• Impatto sul controllo in tempo reale: In scenari con vincoli computazionali stretti o requisiti di controllo ad alta frequenza, l'iterazione sequenziale rende il pianificatore troppo lento per essere utilizzato in loop di controllo chiusi (closed-loop).
• Fallimento delle soluzioni naive: Tentare di addestrare un generatore "one-step" (in un solo passaggio) semplicemente allineando la distribuzione delle traiettorie generate con i dati nel pieno spazio delle finestre temporali porta a un collasso della diversità delle azioni. Le traiettorie tendono a convergere verso una media statica perché le dimensioni future non vincolate dominano le metriche di distanza, ignorando le condizioni attuali.

2. Metodologia: Keyed Drifting Policies (KDP)

Gli autori propongono KDP, un generatore di traiettorie condizionato che opera in un singolo passaggio di inferenza, ma che mantiene la capacità di campionare piani diversificati e adattarsi allo stato corrente.

Concetto Chiave: Vicinato Consapevole della Condizione

L'osservazione fondamentale è che la generazione condizionata fallisce se la similarità tra le traiettorie è misurata sull'intera finestra temporale. In spazi ad alta dimensionalità, le componenti future non vincolate dominano la distanza, trascinando il campione verso la media.
KDP risolve questo introducendo uno spazio delle chiavi (key space):

• Spazio delle Chiavi ($\kappa$): Una mappa compatta che estrae la parte della traiettoria rilevante per la condizione (es. lo stato iniziale $s_0$). La distanza tra le traiettorie per la selezione dei vicini è calcolata solo in questo spazio.
• Spazio della Traiettoria: Gli aggiornamenti effettivi (spostamenti) avvengono nell'intero spazio della traiettoria (stato + azione per tutto l'orizzonte temporale).

Meccanismo di Addestramento: Campo di Deriva (Drift Field)

KDP non imita direttamente i dati, ma viene addestrato per evolvere la sua distribuzione verso un obiettivo dinamico, "amortizzando" (spalmare) il processo di raffinamento iterativo tipico della diffusione direttamente nei parametri del modello durante l'addestramento.
Il processo di addestramento utilizza un campo di deriva con due forze opposte:

1. Attrazione: Spinge le traiettorie generate verso le finestre del dataset che corrispondono alla condizione (misurata nello spazio delle chiavi).
2. Repulsione: Allontana le traiettorie generate dalle altre traiettorie generate nello stesso batch (per mantenere la diversità e prevenire il collasso modale).

Componenti Tecniche Critiche

• Obiettivo con Stop-Gradient: Invece di applicare iterativamente la deriva durante l'inferenza, il modello viene addestrato per predire un "target derivato" ($\tilde{x}_i$) calcolato come $x_i + V(x_i)$, dove $V$ è il campo di deriva. L'uso dello stop-gradient sul target impedisce che il gradiente fluisca attraverso il calcolo della deriva stessa, stabilizzando l'addestramento.
• Vincoli Rigidi (Hard Constraints): Le dimensioni vincolate (es. lo stato iniziale deve essere esattamente lo stato osservato) vengono "clamate" (bloccate) e non modificate dalla deriva.
• Mascheramento dei Negativi Self: Per evitare che un campione si respinga da se stesso (collassando la repulsione), il termine diagonale viene mascherato.
• Media Multi-Temperatura: La deriva viene calcolata su diverse temperature ($\tau$) e mediata per rendere il modello robusto alla scelta di un singolo parametro di vicinato.

3. Contributi Chiave

1. Formulazione Pratica del Campo di Deriva Chiavato: Un metodo per la generazione condizionata che evita il collasso delle azioni sotto vincoli rigidi, utilizzando una metrica di vicinato consapevole della condizione.
2. Ricetta di Pianificazione One-Step: Un approccio che elimina il ciclo sequenziale di denoising ($T$ passaggi), permettendo di generare finestre di traiettoria complete in un singolo forward pass, pur mantenendo l'interfaccia dei pianificatori a diffusione.
3. Valutazione End-to-End: Dimostrazione delle prestazioni sia su benchmark di simulazione standard (D4RL) che su implementazioni hardware in tempo reale, sotto vincoli di latenza rigorosi.

4. Risultati Sperimentali

Gli esperimenti coprono locomozione, navigazione a lungo raggio e manipolazione destrezza (Adroit).

• Efficienza Computazionale: KDP riduce drasticamente la latenza di pianificazione rispetto ai metodi a diffusione (es. Diffuser). Mentre Diffuser richiede 20 passaggi sequenziali per campione, KDP ne richiede solo 1.
  • Su hardware (NVIDIA A100), KDP riduce la latenza di pianificazione da ~150-280 ms (Diffuser) a ~2-5 ms, permettendo frequenze di ripianificazione molto più elevate (es. 38 Hz contro 4 Hz su un drone Crazyflie).
• Prestazioni:
  • Locomozione (D4RL): KDP è competitivo con i migliori metodi a basso NFE (Number of Function Evaluations) e supera significativamente Diffuser in termini di efficienza, mantenendo prestazioni simili o migliori.
  • Pianificazione a Lungo Orizzonte (Maze2D, AntMaze): KDP supera Diffuser, dimostrando che l'amortizzazione del raffinamento nell'addestramento è un bias induttivo migliore per problemi condizionati rispetto all'iterazione online.
  • Manipolazione (Adroit): KDP mostra guadagni significativi sia in velocità che in punteggio di successo rispetto a Diffuser, specialmente in compiti che richiedono precisione sotto contatto.
• Hardware Reale: Su un drone quadrotore e un braccio robotico (SO-100), KDP mantiene un successo del compito simile o leggermente superiore ai metodi a diffusione, ma con una latenza di controllo molto inferiore, permettendo al sistema di reagire più velocemente agli imprevisti.

5. Significato e Implicazioni

Il lavoro di KDP è significativo perché risolve il compromesso fondamentale tra la ricchezza multimodale offerta dai modelli generativi (come la diffusione) e la velocità di inferenza richiesta dal controllo robotico in tempo reale.

• Cambio di Paradigma: Sposta il costo computazionale dall'inferenza (dove è critico per il tempo reale) all'addestramento (dove può essere parallelizzato e non vincola il loop di controllo).
• Robustezza: Dimostra che è possibile ottenere comportamenti di pianificazione complessi e adattivi senza il pesante onere computazionale del denoising iterativo.
• Applicabilità: Apre la strada all'uso di pianificatori generativi avanzati in scenari reali con vincoli di sicurezza e tempo stringenti, dove i ritardi di pianificazione possono portare a fallimenti catastrofici.

In sintesi, KDP offre un modo per "comprimere" l'intelligenza di un pianificatore a diffusione iterativa in un modello one-step, preservando la diversità e la capacità di condizionamento attraverso un meccanismo innovativo di campi di deriva basati su chiavi.