Fisher Decorator: Refining Flow Policy via A Local Transport Map

1. Il Problema: Il Disallineamento Geometrico nell'RL Offline

L'apprendimento per rinforzo offline (Offline RL) mira a imparare politiche ottimali da dataset statici senza interazione con l'ambiente. La sfida centrale è bilanciare la massimizzazione della ricompensa con il vincolo di rimanere vicini alla distribuzione della politica comportamentale ($\pi_\beta$) per evitare lo shift distribuzionale.

Il problema fondamentale identificato dagli autori risiede nel divario geometrico tra l'obiettivo teorico e le implementazioni pratiche nelle politiche basate su Flow Matching:

• Obiettivo Teorico: L'ottimizzazione è vincolata dalla divergenza di Kullback-Leibler (KL). La KL è intrinsecamente anisotropa e sensibile alla densità: penalizza fortemente le deviazioni nelle regioni ad alta densità della politica comportamentale e impedisce di spostarsi in regioni non supportate (dove la densità è zero).
• Implementazione Attuale: Metodi recenti come Flow Q-Learning (FQL) e DeFlow approssimano il vincolo KL utilizzando un upper bound della distanza di Wasserstein-2 ($W_2$), che viene spesso rilassato a una penalità isotropa ($L_2$).
• Conseguenze: La regolarizzazione $L_2$ tratta tutte le direzioni e le regioni dello spazio degli stati-azioni allo stesso modo, ignorando la struttura della densità. Questo porta a:
  1. Shift Distribuzionale: La politica appresa può assegnare massa a regioni non supportate dai dati.
  2. Collasso delle Modalità (Mode Collapse): In distribuzioni multimodali, la penalità isotropa tende a "mediare" le modalità, spostando la politica verso regioni a bassa valore (baricentro di Wasserstein) invece di raffinare le modalità esistenti verso valori più alti.

2. Metodologia: Fisher Decorator (FiDec)

Gli autori propongono un nuovo quadro geometrico per la raffinazione della politica, chiamato Fisher Decorator (FiDec), che risolve il disallineamento trattando l'aggiornamento come una mappa di trasporto locale.

A. Riformulazione Geometrica

Invece di apprendere una nuova politica da zero, FiDec parametrizza la politica raffinata $\pi_\theta$ come una mappa di trasporto locale $T_s$ applicata alla politica comportamentale $\pi_\beta$:
$$\pi_\theta(\cdot|s) = (T_s)_\# \pi_\beta(\cdot|s)$$
dove $T_s(a) = a + \delta(s, a)$. Qui, $\delta(s, a)$ è un termine residuo appreso che induce uno spostamento locale.

B. Approssimazione Quadratica Locale

Analizzando la trasformazione della densità di probabilità indotta da questa mappa, gli autori dimostrano che, per piccoli spostamenti $\delta$, la divergenza KL ammette un'approssimazione quadratica di secondo ordine governata dalla Matrice di Informazione di Fisher locale $I(s, a)$:
$$D_{KL}(\pi_\theta \| \pi_\beta) \approx \frac{1}{2} \mathbb{E}_{a \sim \pi_\beta} \left[ \delta(s, a)^\top I(s, a) \delta(s, a) \right]$$
dove $I(s, a) = \nabla_a \log \pi_\beta(a|s) \nabla_a \log \pi_\beta(a|s)^\top$.

Questa formulazione è cruciale perché:

1. Anisotropia: La penalità è direzionale. Le direzioni allineate con il gradiente del log-densità (score function) sono penalizzate diversamente rispetto a quelle ortogonali, preservando la struttura della varietà statistica.
2. Preservazione del Supporto: Poiché l'aggiornamento è un trasporto locale, la massa probabilistica viene "spinta" all'interno del supporto esistente, evitando di creare massa in regioni vuote.

C. Stima Pratica e Algoritmo

Un vantaggio chiave è che la funzione di score $\nabla_a \log \pi_\beta$ è intrinsecamente legata al campo di velocità $v_\beta$ appreso nel modello di Flow Matching. Gli autori derivano una stima stabile della metrica di Fisher utilizzando una perturbazione temporale $\epsilon$ vicino a $t=1$:
$$I(s, a) \approx \frac{(t_\epsilon v_\beta(t_\epsilon, s, a) - a)(t_\epsilon v_\beta(t_\epsilon, s, a) - a)^\top}{(1-t_\epsilon)^2}$$
L'algoritmo ottimizza un problema di trust-region vincolato: massimizzare il valore atteso $Q$ soggetto al vincolo quadratico definito dalla metrica di Fisher. Questo viene risolto efficientemente tramite un metodo primale-duale (aggiornamento di Lagrange).

3. Contributi Chiave

1. Caratterizzazione Geometrica del Divario KL-Wasserstein: Identificano e formalizzano matematicamente perché l'approssimazione $W_2/L_2$ fallisce nell'Offline RL multimodale, dimostrando che la KL induce una geometria anisotropa che le metriche isotrope non possono catturare.
2. Ottimizzazione Anisotropa Tractabile: Derivano una formulazione quadratica gestibile per l'ottimizzazione della politica vincolata da KL, utilizzando la metrica di Fisher locale invece della semplice norma $L_2$.
3. Framework FiDec: Propongono un algoritmo semplice ed efficiente che non richiede la risoluzione di ODE (equazioni differenziali ordinarie) durante l'inferenza né la distillazione iterativa, mantenendo però la ricchezza espressiva dei modelli di flusso.

4. Risultati Sperimentali

Gli autori hanno valutato FiDec su una vasta gamma di benchmark, inclusi OGBench (task multimodali complessi come puzzle, calcio robotico, labirinti) e D4RL.

• Prestazioni Offline: FiDec supera lo stato dell'arte (SOTA) su 73 task diversi, superando metodi basati su Gaussiani, Diffusion e Flow (come FQL e DeFlow). I miglioramenti sono particolarmente evidenti nei task multimodali dove i metodi precedenti soffrono di collasso delle modalità o shift distribuzionale.
• Analisi Visuale: Le visualizzazioni mostrano che mentre FQL tende a collassare verso regioni intermedie a basso valore e DeFlow rimane troppo conservativo, FiDec segue la geometria del supporto dei dati, spostando la massa verso regioni ad alto valore senza uscire dal supporto.
• Fine-tuning Offline-to-Online: FiDec fornisce un'inizializzazione di alta qualità che permette un adattamento online più rapido ed efficace rispetto ai metodi basati su distillazione.
• Efficienza: Nonostante l'aggiunta della metrica di Fisher, FiDec mantiene tempi di addestramento comparabili ai metodi esistenti (circa 2.7 ms/step), evitando il costo computazionale di metodi iterativi complessi.

5. Significato e Impatto

Il lavoro di Fisher Decorator rappresenta un passo significativo verso l'integrazione rigorosa della geometria dell'informazione nell'apprendimento per rinforzo.

• Cambiamento di Paradigma: Sposta l'attenzione dalla semplice minimizzazione della distanza euclidea ($L_2$) alla comprensione della struttura intrinseca della distribuzione dei dati.
• Robustezza: Dimostra che comprendere e rispettare la geometria anisotropa della politica comportamentale è essenziale per evitare lo shift distribuzionale e il collasso delle modalità in scenari offline complessi.
• Scalabilità: Fornisce un metodo che è sia teoricamente fondato (basato su approssimazioni di secondo ordine della KL) che praticamente efficiente, rendendo l'ottimizzazione geometrica accessibile per applicazioni reali di RL.

In sintesi, FiDec risolve il problema fondamentale della "mappatura sbagliata" tra vincoli teorici e regolarizzazioni pratiche, offrendo un framework che preserva la diversità delle modalità e guida l'ottimizzazione verso soluzioni più robuste e performanti.