Sintesi Tecnica: Lagrangian Perturbation Diffusion Steering (LP-DS)

1. Definizione del Probleamento

Le policy generative ad alta capacità, come quelle basate su modelli di diffusione e flow-matching, si sono affermate come paradigmi potenti per il controllo continuo e la robotica. Questi modelli eccellono nel rappresentare distribuzioni di azione espressive e multimodali apprese da dati offline. Tuttavia, affrontano due limitazioni primarie:

1. Vincoli di Behavior Cloning (BC): Le policy addestrate tramite BC ereditano i limiti dei loro dataset di dimostrazione, inclusa la copertura incompleta dello spazio degli stati, l'omissione di comportamenti di recupero rari e la fragilità sotto shift della distribuzione.
2. Instabilità del Reinforcement Learning (RL): Sebbene l'RL offra un meccanismo per il miglioramento online, il fine-tuning diretto di grandi decoder generativi è frequentemente instabile e inefficiente in termini di campionamento. L'interazione tra gli obiettivi di ottimizzazione e le dinamiche di denoising o integrazione a lungo orizzonte porta spesso all'instabilità dell'addestramento.

Le alternative esistenti, come la Diffusion Steering via Reinforcement Learning (DSRL), tentano di migliorare le policy congelate apprendendo una policy nello spazio del rumore latente. Tuttavia, DSRL sostituisce efficacemente il prior di pre-addestramento con una policy latente appresa. Questo approccio soffre di due modalità di fallimento:

• Deriva Off-Manifold: Il rumore guidato può deviare dal supporto gaussiano standard utilizzato per addestrare il decoder, portando a un decoding instabile.
• Collasso del Modo (Mode Collapse): La policy latente appresa può far collassare il comportamento multimodale del backbone in un sottoinsieme ristretto di modi, riducendo la diversità comportamentale.

2. Metodologia: Lagrangian Perturbation Diffusion Steering (LP-DS)

Gli autori propongono LP-DS, un framework di adattamento leggero che migliora le policy generative congelate apprendendo un residuo condizionato allo stato nello spazio del rumore latente, piuttosto che modificare i pesi del decoder o sostituire interamente il prior.

Meccanismo Core

LP-DS tratta il backbone generativo congelato $\Phi$ come un decoder black-box che mappa il rumore latente $w$ e lo stato $s$ in azioni. Invece di campionare il rumore direttamente da un prior $p_0$ (ad es., $\mathcal{N}(0, I)$), LP-DS introduce una rete di perturbazione apprendibile $\Delta_\theta(s)$:
$$w = \epsilon + \Delta_\theta(s), \quad \text{dove } \epsilon \sim \mathcal{N}(0, I)$$
Quando $\Delta_\theta(s) = 0$, la policy recupera esattamente la distribuzione originale di behavior cloning.

Ottimizzazione Vincolata tramite Relaxation Lagrangiana

Per evitare che la perturbazione spinga il rumore fuori dalla manifold dove il decoder è stato addestrato, LP-DS formula il problema di apprendimento come un'ottimizzazione vincolata. L'obiettivo è massimizzare il valore Q atteso delle azioni generate soggetto a un vincolo di regione di fiducia (trust-region) sull'ampiezza della perturbazione.

Gli autori approssimano la divergenza di Kullback–Leibler (KL) attesa tra la distribuzione del rumore guidata e il prior utilizzando l'ampiezza al quadrato dello shift:
$$\mathbb{E}_{s \sim D}[\|\Delta_\theta(s)\|^2_2] \leq \delta$$
dove $\delta$ è un iperparametro che controlla la deviazione consentita.

Questo viene risolto utilizzando un approccio di relaxation Lagrangiana. Una variabile duale (moltiplicatore di Lagrange) $\alpha$ viene introdotta per imporre il vincolo dinamicamente:
$$\mathcal{L}(\theta, \alpha) = \mathbb{E}[Q^W(s, w) - \alpha(\|\Delta_\theta(s)\|^2_2 - \delta)]$$
L'algoritmo alterna tra:

1. Aggiornamento dell'Attore: Massimizzazione di $\mathcal{L}$ rispetto a $\theta$ (gradiente ascendente).
2. Aggiornamento del Duale: Massimizzazione di $\mathcal{L}$ rispetto ad $\alpha$ (gradiente ascendente proiettato) per imporre $\alpha \geq 0$. Se la perturbazione viola la regione di fiducia, $\alpha$ aumenta, penalizzando l'attore e costringendolo ad aderire più strettamente al prior.

Componenti dell'Algoritmo

La procedura di addestramento (Algoritmo 1) prevede:

• Raccolta Dati: Campionamento di azioni usando la policy guidata $w = \epsilon + \Delta_\theta(s)$.
• Aggiornamenti del Critic: Aggiornamento di un critic del valore d'azione $Q^A$ tramite apprendimento Temporal Difference (TD) e un critic del valore latente $Q^W$ tramite distillazione da $Q^A$ sulla distribuzione del rumore base.
• Aggiornamenti Primal-Dual: Aggiornamento della rete di perturbazione $\theta$ e del moltiplicatore $\alpha$ per bilanciare la massimizzazione del reward con la preservazione del prior.

3. Contributi Chiave

Il paper delinea quattro contributi principali:

1. Framework LP-DS: Un metodo di adattamento leggero che migliora le policy generative congelate tramite il controllo del rumore residuo condizionato dallo stato, evitando l'instabilità e il costo del fine-tuning completo del decoder.
2. Obiettivo di Trust-Region Lagrangiano: Un obiettivo innovativo che vincola dinamicamente l'ampiezza della perturbazione, fornendo un meccanismo controllabile per bilanciare la massimizzazione del reward rispetto alla preservazione del prior latente pre-addestrato.
3. Mitigazione del Collasso del Modo: Evidenza empirica che mostra come LP-DS preservi una maggiore diversità nello spazio delle azioni e mitighi il collasso del modo rispetto ai metodi di steering nello spazio del rumore che sostituiscono il prior (come DSRL).
4. Generalizzabilità: Validazione del metodo oltre le policy di diffusione compatte, includendo backbone di flow-matching, un grande modello vision-language-action ($\pi_0$) ed esperimenti fisici con il robot Franka.

4. Risultati Sperimentali

Gli autori hanno valutato LP-DS attraverso i benchmark RoboMimic (manipolazione), OpenAI Gym (locomozione) e Adroit (manipolazione destre), confrontandolo con DSRL, DPPO (fine-tuning della policy-gradient) e baseline di offline-to-online diffusion Q-learning.

• Performance: LP-DS ha migliorato costantemente i tassi di successo e i ritorni episodici. In particolare, sull'ambiente WALKER2D, ha ottenuto un miglioramento di circa il 25% nel ritorno rispetto al baseline più forte.
• Efficienza di Campionamento: Aggiornando solo un modulo di perturbazione leggero invece dell'intero decoder, LP-DS ha dimostrato una superiore efficienza di campionamento e dinamiche di apprendimento più stabili.
• Preservazione della Diversità:
  • In un task di navigazione multi-goal toy, LP-DS ha permesso il controllo esplicito del trade-off tra performance e diversità tramite il bound della regione di fiducia $\delta$. Valori piccoli di $\delta$ preservavano il comportamento multimodale, mentre valori grandi permettevano la specializzazione. Al contrario, DSRL è collassato immediatamente su un singolo modo.
  • Nell'ambiente AVOIDING, LP-DS ha preservato un ampio set di rotte fattibili per l'evitamento degli ostacoli (multimodalità a livello di traiettoria) con un piccolo $\delta$, mentre DSRL è collassato su un singolo percorso dominante.
  • Analisi dell'Entropia: LP-DS ha mantenuto costantemente un'entropia dello spazio d'azione più alta rispetto a DSRL durante l'adattamento online, indicando che il vincolo Lagrangiano scoraggia efficacemente il collasso prematuro pur migliorando le performance.
• Studi di Ablazione: Rimuovere la penalità Lagrangiana ha portato a un apprendimento instabile e a una degradazione del successo. Rimuovere sia la penalità Lagrangiana che il bound del rumore ha causato un comportamento altamente instabile e bassi tassi di successo, confermando che il vincolo della regione di fiducia è il meccanismo primario per prevenire la deriva off-distribution.
• Sensibilità: Il metodo si è rivelato robusto rispetto al valore specifico del target della regione di fiducia $\delta$, con valori superiori a 0.1 che generalmente garantivano forti performance.

5. Significato e Rivendicazioni

Il paper sostiene che LP-DS offra un'alternativa stabile ed efficiente al fine-tuning completo per le policy generative. Trattando il backbone congelato come un prior fisso e apprendendo solo uno shift residuo limitato, il metodo:

• Previene la Deriva Off-Manifold: Mantiene le query latenti entro l'intervallo operativo valido del decoder, evitando il comportamento erratico osservato nello steering del rumore non vincolato.
• Preserva la Multimodalità: Mitiga il collasso del modo spesso osservato quando si sostituisce il prior di pre-addestramento, mantenendo così la diversità comportamentale inerente al backbone congelato.
• Si Generalizza tra le Architetture: L'approccio non è limitato ai modelli di diffusione, ma si applica a qualsiasi policy generativa con una mappatura latente-azione, inclusi i modelli di flow-matching e i grandi modelli vision-language-action.
• Permette il Deployment Fisico: Il metodo è stato validato con successo in esperimenti fisici con il robot Franka utilizzando l'adattamento latente basato su simulazione, dimostrando la sua applicabilità oltre i benchmark simulati.

Gli autori posizionano LP-DS come una soluzione pratica per il "miglioramento online" di policy generative ad alta capacità, bilanciando la necessità di ottimizzazione specifica per il task con la necessità di mantenere la robusta struttura multimodale del modello pre-addestrato.