Parameter-Efficient Distributional RL via Normalizing Flows and a Geometry-Aware Cramér Surrogate

Riepilogo Tecnico: NFDRL

Enunciato del Problema
L'Apprendimento per Rinforzo Distribuzionale (DistRL) mira a modellare l'intera distribuzione del ritorno anziché solo il ritorno atteso, fornendo segnali più ricchi per incertezza, rischio e variabilità. Tuttavia, i metodi esistenti affrontano un compromesso tra espressività ed efficienza parametrica.

• Metodi categorici (es. C51): Si basano su supporti discreti fissi. Sebbene efficaci, mancano di parsimonia; il numero di parametri scala linearmente con la risoluzione desiderata (numero di atomi). Impongono inoltre limiti rigidi sui ritorni, limitando l'applicabilità in domini con ricompense dinamiche o a coda pesante.
• Metodi quantili (es. QR-DQN, IQN): Imparano le posizioni dei quantili per gestire ritorni illimitati. Tuttavia, approssimano le distribuzioni come miscele di delta di Dirac (funzioni a gradino). Questa rappresentazione può essere inefficiente per distribuzioni lisce, multimodali o a coda pesante, richiedendo spesso molti quantili per approssimare densità semplici e fallendo nel catturare forme distribuzionali dettagliate.
• Metodi a Densità Continua: Sebbene la modellazione di funzioni di densità di probabilità (PDF) continue offra supporto adattivo e liscezza, gli approcci precedenti hanno faticato a definire un obiettivo di formazione che sia simultaneamente teoricamente solido (garantendo la contrazione dell'operatore di Bellman) e praticamente ottimizzabile (fornendo gradienti di campione non distorti).

Metodologia: NFDRL
Gli autori propongono NFDRL, un'architettura DistRL efficiente in termini di parametri che modella le distribuzioni di ritorno utilizzando flussi normalizzanti continui.

1. Architettura:

   • Il modello utilizza un flusso normalizzante condizionale per mappare campioni di rumore di base $z \sim U$ in esiti di ritorno $y = f_\theta(z)$.
   • Il flusso è costruito tramite una Funzione di Distribuzione Cumulativa (CDF) a Miscela Gaussiana. Una rete neurale $h_\theta$ mappa uno stato $x$ ai parametri di una CDF a miscela per ogni azione.
   • Per affrontare la natura limitata delle CDF ($[0, 1]$), una trasformazione affine aggiuntiva ridimensiona l'output a un intervallo dinamico $[-G_{max}, G_{max}]$, consentendo al modello di adattarsi al supporto dei ritorni reali senza limiti fissi.
   • La distribuzione target è costruita componendo il flusso previsto con la funzione di bootstrap ($r + \gamma y$), trattata come un livello di flusso invertibile per preservare la normalizzazione tramite la formula del cambiamento di variabili.

2. Obiettivo di Formazione (Surrogato Cramér Consapevole della Geometria):

   • La divergenza KL standard non è adatta per supporti non sovrapposti, e la minimizzazione esatta di Wasserstein soffre di gradienti distorti.
   • Gli autori introducono una distanza ispirata a Cramér e consapevole della geometria definita sulle masse di probabilità piuttosto che sulle densità grezze.
   • La funzione di perdita è definita come:
     $$D(\eta, \hat{T}\eta) = \left( \frac{1}{N^2} \sum_{i,j} (w_i - v_i)^2 \cdot |y_i - y_j| \right)^{1/2}$$
     dove $w_i$ e $v_i$ sono masse di probabilità nei punti di supporto $y_i$.
   • Per gestire il disallineamento tra i supporti previsto e target (che differiscono a causa del passo di bootstrap), il metodo impiega la Stima della Densità tramite Kernel (KDE) per valutare le densità su supporti condivisi prima di calcolare la perdita.

Contributi Chiave

1. Efficienza Parametrica: A differenza delle basi categoriche dove la dimensione del modello cresce con la risoluzione, NFDRL mantiene un ingombro parametrico costante indipendentemente dalla risoluzione effettiva della distribuzione, offrendo risoluzione effettiva infinita e supporto dinamico.
2. Solidità Teorica: Gli autori dimostrano che la loro distanza basata sulle masse è una vera metrica e che l'operatore distribuzionale di Bellman associato è una contrazione $\sqrt{\gamma}$. Ciò garantisce la convergenza a un punto fisso unico.
3. Gradienti Non Distorti: L'obiettivo proposto ammette gradienti di campione non distorti, una proprietà tipicamente assente nei precedenti metodi DistRL basati su PDF (a differenza di Wasserstein, che spesso richiede approssimazioni distorte).
4. Espressività: La natura continua del flusso permette al modello di catturare paesaggi di ritorno complessi e multimodali che i metodi basati su quantili spesso sfocano o non riescono a risolvere.

Risultati

• MDP Fittizi: In ambienti controllati, NFDRL recupera con successo ricchi paesaggi di ritorno multimodali (es. distribuzioni bimodali) con PDF lisce, mentre le basi quantili (IQN) producono distribuzioni rumorose e sfocate che oscurano la modalità.
• Benchmark Atari-5: NFDRL raggiunge prestazioni competitive con le basi categoriche (C51) e si avvicina alle prestazioni dei migliori metodi quantili (IQN).
  • Sul benchmark Atari-5, NFDRL (variante surrogata) ha ottenuto un punteggio medio normalizzato sull'uomo di 418, superando C51 (330) e DQN (189), sebbene sia leggermente indietro rispetto a IQN (525).
  • In modo notevole, NFDRL-S ha superato IQN in giochi specifici come Double Dunk e QBert*.
• Efficienza: NFDRL eguaglia le prestazioni di C51 con significativamente meno parametri. Mentre il numero di parametri di C51 scala linearmente con il numero di atomi, NFDRL mantiene una dimensione fissa, eguagliando la capacità di C51 con soli 11 atomi.

Significato e Affermazioni
Il documento posiziona NFDRL non come una sostituzione universale per agenti basati su quantili potenti come IQN su Atari, ma come un'alternativa continua e parsimoniosa con compromessi distinti. Il suo significato principale risiede nell'offrire un quadro teoricamente rigoroso (contrazione e gradienti non distorti) per distribuzioni di ritorno continue con un ingombro parametrico costante. Ciò è particolarmente prezioso in scenari in cui le decisioni a valle dipendono dalla forma dettagliata della distribuzione di ritorno (es. multimodalità, rischio di coda) piuttosto che solo dal valore atteso. Gli autori riconoscono le limitazioni, inclusa una maggiore varianza di formazione, una convergenza più lenta rispetto ad alcune basi e il costo computazionale del passo di allineamento KDE, considerandoli aree per un affinamento futuro.