Riepilogo Tecnico: Trust Region Q-Adjoint Matching (TRQAM)

Enunciato del Problema

Il documento affronta la sfida del fine-tuning off-policy per politiche di flusso preaddestrate. Sebbene le politiche di flow matching offrano distribuzioni di azioni ricche e multi-modali superiori alle politiche Gaussiane unimodali convenzionali, la loro definizione implicita tramite processi di denoising multi-step rende instabile il miglioramento della politica basato sul gradiente. La retropropagazione diretta attraverso la catena di campionamento è computazionalmente costosa e soggetta a instabilità.

Le soluzioni esistenti tentano di aggirare questo problema apprendendo residui additivi (ignorando le dinamiche multi-step) o operando nello spazio del rumore (limitato dall'espressività della politica congelata). Un approccio recente, Q-learning con Adjoint Matching (QAM), riformula il fine-tuning come un problema di Controllo Ottimo Stocastico (SOC) senza memoria, utilizzando un critico appreso per guidare il processo di campionamento tramite adjoint matching. Tuttavia, gli autori identificano una fragilità fondamentale in QAM: in contesti off-policy, il critico appreso è inevitabilmente imperfetto. Piccoli errori di approssimazione nel critico vengono amplificati esponenzialmente quando l'aggiornamento della politica si basa su una "temperatura" fissa (o temperatura inversa $\beta$), portando a una deriva distruttiva rispetto al prior preaddestrato e al collasso finale del modello. Questa instabilità persiste anche con il clipping del gradiente.

Metodologia: Trust Region Q-Adjoint Matching (TRQAM)

TRQAM propone un algoritmo stabile di fine-tuning off-policy che integra un meccanismo di trust region direttamente nelle dinamiche di campionamento SOC. L'innovazione centrale è il controllo adattivo della divergenza di Kullback-Leibler (KL) nello spazio dei percorsi tra la politica fine-tunata e il prior preaddestrato.

Principali Intuizioni Teoriche

1. Amplificazione Esponenziale degli Errori: Gli autori formalizzano (Lemma 1) che l'adjoint matching a temperatura fissa soffre di un'amplificazione esponenziale degli errori del critico. Un $\beta$ fisso non può sfruttare simultaneamente un critico affidabile e proteggersi da uno inaffidabile.
2. $\lambda$ come Parametro di Trust Region: TRQAM introduce un parametro scalare $\lambda$ che scala il coefficiente di diffusione nell'SDE SOC di $\sqrt{\lambda}$.
3. Identità KL Esatta (Teorema 1): Utilizzando il teorema di Girsanov, gli autori dimostrano che scalare la diffusione di $\sqrt{\lambda}$ rende la divergenza KL nello spazio dei percorsi una funzione esatta e in forma chiusa di $\lambda$:
   $$D_{KL}(P_u \parallel P_{base}) = \mathbb{E}_{X \sim P_u} \left[ \frac{1}{2\lambda} \int_0^1 \|u(X_\tau, \tau)\|^2 d\tau \right]$$
   Questo stabilisce un legame strutturale diretto tra il parametro $\lambda$ e la deviazione dalla politica di base.
4. Limite Terminal KL (Proposizione 1): La KL nello spazio dei percorsi limita superiormente la KL della distribuzione delle azioni terminali, garantendo che il controllo della KL nello spazio dei percorsi vincoli efficacemente la deviazione finale della politica.

Quadro Algoritmico

Invece di aggiungere un termine KL come penalità di perdita esterna (che può essere sovrascritta da gradienti forti del critico), TRQAM internalizza $\lambda$ all'interno delle dinamiche di campionamento:

• Campionamento: L'SDE controllato utilizza un coefficiente di diffusione $\sqrt{\lambda}\sigma(\tau)$. Aggiustare $\lambda$ rimodella lo stesso SDE, avvicinando o allontanando le dinamiche controllate dalle dinamiche di base.
• Adattamento: L'algoritmo impiega una discesa duale proiettata per adattare $\lambda$ dinamicamente. Stima la KL realizzata nello spazio dei percorsi ($\hat{D}_n$) tramite un surrogato Monte Carlo derivato dal campionatore senza memoria OT discretizzato.
• Regola di Aggiornamento: $\lambda$ viene aggiornato per imporre un budget KL target $\epsilon_{KL}$:
  $$\lambda_{n+1} \leftarrow \max\{0, \lambda_n + \eta_\lambda (D_n - \epsilon_{KL})\}$$
  Se la KL realizzata supera il budget, $\lambda$ aumenta, rendendo le dinamiche più conservative; se scende sotto, $\lambda$ diminuisce, consentendo un miglioramento più aggressivo.

Principali Contributi

1. Identificazione della Fragilità: Il documento identifica e formalizza l'amplificazione esponenziale degli errori del critico nell'adjoint matching a temperatura fissa, dimostrando la sua tendenza a causare il collasso del modello su benchmark come Robomimic.
2. Collegamento Teorico: Dimostra che scalare il coefficiente di diffusione di $\sqrt{\lambda}$ crea un'identità esatta tra il costo di controllo e la KL nello spazio dei percorsi, trasformando $\lambda$ in un parametro di trust region principiato.
3. Progettazione dell'Algoritmo: TRQAM è proposto come un metodo che internalizza il parametro di trust region nelle dinamiche SOC e lo adatta tramite discesa duale, imponendo il limite KL a livello di campionamento piuttosto che come penalità di perdita soft.
4. Superiorità Empirica: Estesi esperimenti dimostrano che TRQAM raggiunge una convergenza stabile dove i metodi precedenti falliscono e supera le baseline state-of-the-art.

Risultati Sperimentali

Gli autori valutano TRQAM su 50 task dal benchmark OGBench e dal benchmark di manipolazione Robomimic.

• Prestazioni RL Offline: Su OGBench, TRQAM raggiunge un tasso di successo complessivo del 68%, superando significativamente la baseline più forte (DSRL al 46%) e i precedenti metodi di adjoint matching (QAM al 35%, QAM-E al 45%). I guadagni sono particolarmente pronunciati nei task a lungo orizzonte e combinatori.
• Stabilità: Su Robomimic, i metodi a temperatura fissa (QAM, QAM-E) mostrano perdite adjoint esplosive (fino a $10^{20}$) e tassi di successo che collassano vicino allo zero attraverso varie impostazioni di iperparametri. Al contrario, TRQAM rimane stabile su tutti i budget testati.
• Validazione del Meccanismo: Studi di ablazione confermano che:
  • L'adattamento è necessario: I budget KL adattivi superano quelli fissi.
  • L'internalizzazione è necessaria: Internalizzare $\lambda$ nell'SDE (TRQAM) traccia strettamente il budget KL, mentre la regolarizzazione KL esterna permette alla KL realizzata di deviare ben oltre l'obiettivo, portando a un degrado delle prestazioni.
• Sensibilità: Il budget KL target $\epsilon_{KL}$ agisce come una manopola di controllo prevedibile; budget più stretti generalmente producono prestazioni migliori, sebbene il valore ottimale vari sistematicamente con la struttura del task (ad esempio, budget più ampi sono benefici per task con spazi di stato più grandi come puzzle-4x4).

Significato e Affermazioni

Il documento afferma che TRQAM fornisce un'alternativa principiale e stabile ai metodi esistenti di fine-tuning off-policy per le politiche di flusso. Internalizzando il vincolo di trust region nelle dinamiche di campionamento tramite il teorema di Girsanov, TRQAM evita la "deriva distruttiva" causata da critici imperfetti in metodi come QAM.

Gli autori posizionano TRQAM come l'analogo off-policy dell'ottimizzazione della politica Trust Region (TRPO) e dell'ottimizzazione della politica Prossimale (PPO) per RL on-policy. Sostengono che per le politiche di flusso preaddestrate, che fungono da priors comportamentali preziosi, mantenere una deviazione controllata tramite un'identità KL esatta è cruciale per prevenire l'oblio catastrofico e il collasso del modello. Il metodo è presentato come una soluzione robusta sia per l'RL offline che per la transizione al fine-tuning online, offrendo un meccanismo affidabile per sfruttare le abilità preaddestrate senza sacrificare la stabilità.