PolicyFlow: Policy Optimization with Continuous Normalizing Flow in Reinforcement Learning

Riepilogo Tecnico: PolicyFlow

Definizione del Problema

L'Ottimizzazione della Politica Prossimale (PPO) è un algoritmo di apprendimento per rinforzo (RL) on-policy standard, noto per la sua stabilità e semplicità. Tuttavia, tipicamente si basa su distribuzioni di policy gaussiane, che faticano a rappresentare distribuzioni di azioni complesse, multimodali o fortemente asimmetriche. Sebbene modelli generativi espressivi come i Continuous Normalizing Flows (CNF) e i modelli di diffusione possano catturare tali distribuzioni complesse, l'integrazione di questi nei PPO presenta sfide significative.

Il collo di bottiglia principale è il requisito della valutazione della verosimiglianza (likelihood evaluation) per calcolare i rapporti di importanza ($\pi(a|s)/\hat{\pi}(a|s)$) per l'obiettivo surrogato PPO. Per i CNF, la valutazione esatta della verosimiglianza richiede la simulazione dell'intera traiettoria dell'equazione differenziale ordinaria (ODE) e la backpropagation path-wise. Questo processo è computazionalmente costoso, intensivo in termini di memoria e spesso numericamente instabile a causa di gradienti esplosivi o evanescenti. I tentativi esistenti di applicare aggiornamenti in stile PPO a policy basate su flow (ad esempio, FPO, DPPO) soffrono di una stima asimmetrica del bias, richiedono grandi batch di dimensioni o degradano nelle prestazioni quando si addestra da zero a causa della difficoltà di esplorazione off-manifold. Inoltre, le tecniche standard di regolarizzazione dell'entropia sono difficili da applicare alle policy basate su flow perché il calcolo esplicito dell'entropia della policy richiede costose integrazioni di divergenza.

Metodologia

Gli autori propongono PolicyFlow, un algoritmo di RL on-policy che integra policy CNF espressive con obiettivi in stile PPO senza richiedere la valutazione completa della verosimiglianza della traiettoria. Il metodo consiste in due componenti principali:

1. Approssimazione del Rapporto di Importanza tramite Variazioni del Campo di Velocità

Invece di simulare l'intera traiettoria ODE per calcolare lo spostamento terminale tra la politica corrente $\phi_1$ e la politica di riferimento $\hat{\phi}_1$, PolicyFlow approssima il rapporto di importanza utilizzando le variazioni del campo di velocità lungo un semplice percorso di interpolazione lineare.

• Rappresentazione della Policy: La policy genera azioni integrando un flow $\phi_t(z; s)$ governato da un campo di velocità neurale $v_t$, aggiungendo rumore gaussiano al tempo terminale: $a = \phi_1(z; s) + n$.
• Approssimazione: Lo spostamento terminale $\delta\phi_1 = \phi_1 - \hat{\phi}_1$ è approssimato integrando la differenza nei campi di velocità $\delta v_t = v_t - \hat{v}_t$ lungo un percorso di interpolazione lineare $x_t = (1-t)z + t\hat{\phi}_1(z; s)$.
• Obiettivo Surrogato: Il rapporto di importanza $\rho$ è stimato utilizzando il rapporto di verosimiglianza gaussiana delle variazioni del campo di velocità piuttosto che le intere traiettorie del flow. Ciò consente di stimare efficientemente l'obiettivo di addestramento lungo il percorso di interpolazione senza eseguire la backpropagation attraverso l'intera simulazione ODE durante il passo di aggiornamento. L'analisi teorica indica che ciò introduce solo un errore del primo ordine nel log-likelihood, che è naturalmente limitato dal meccanismo di clipping di PPO.

2. Regolarizzatore Browniano per l'Entropia

Per prevenire il collasso della moda (mode collapse) e incoraggiare un'esplorazione diversificata senza calcolare esplicitamente l'entropia del flow, gli autori introducono il Regolarizzatore Browniano.

• Concetto: Ispirato al moto browniano e all'equazione del calore, questo regolarizzatore incoraggia la velocità appresa a seguire la dinamica di un processo che aumenta l'entropia.
• Meccanismo: Sfrutta la relazione tra il campo di velocità e la funzione di score ($\nabla \log p_t$) nei rectified flows. Il regolarizzatore penalizza le deviazioni del campo di velocità appreso da un campo vettoriale target derivato dallo score del flow di riferimento.
• Formulazione: Il termine del regolarizzatore $J_{Reg}$ include un termine che minimizza la differenza tra il campo di velocità appreso e un termine di allineamento target (derivato dal campo di velocità di riferimento e dallo stato), oltre a un termine che incoraggia la varianza del rumore gaussiano iniettato. Questo approccio evita costose computazioni di log-likelihood e l'iniezione di rumore euristica, fornendo un modo leggero e principato per mantenere la diversità delle traiettorie.

Contributi Chiave

1. Integrazione Efficiente CNF-PPO: PolicyFlow consente l'uso di policy CNF espressive nell'RL on-policy approssimando i rapporti di importanza tramite variazioni del campo di velocità, eliminando la necessità di costosa backpropagation path-wise attraverso intere traiettorie ODE.
2. Regolarizzatore Browniano: Un regolarizzatore dell'entropia nuovo e computazionalmente leggero, ispirato al moto browniano, che mitiga il collasso della moda e incoraggia comportamenti diversificati senza richiedere il calcolo esplicito dell'entropia.
3. Validazione Empirica: Il metodo è validato attraverso diversi ambienti, dimostrando prestazioni competitive o superiori rispetto a PPO gaussiano e ai migliori baselines basati su flow (FPO, DPPO).

Risultati Sperimentali

Gli autori hanno valutato PolicyFlow sui benchmark MultiGoal, MuJoCo Playground e IsaacLab:

• MultiGoal (Test di Collasso della Moda): In un ambiente con sei obiettivi simmetrici, lo standard PPO (Gaussiano) non è riuscito a raggiungere tutti gli obiettivi. I baselines FPO e DPPO sono collassati su pochi modi a causa della mancanza di una efficace regolarizzazione dell'entropia. PolicyFlow, con il regolarizzatore browniano, ha catturato con successo una distribuzione multimodale bilanciata, raggiungendo tutti gli obiettivi con probabilità circa uguale.
• MuJoCo Playground: PolicyFlow ha ottenuto prestazioni comparabili o superiori a FPO e DPPO in varie attività (ad esempio, WalkerWalk, CheetahRun). Spesso ha converguto più velocemente, indicando una maggiore efficienza campionaria.
• IsaacLab: PolicyFlow è stato confrontato con PPO in compiti di robotica (locomozione, manipolazione). Ha costantemente eguagliato o superato le prestazioni asintotiche di PPO.
• Efficienza Computazionale: Sebbene PolicyFlow comporti la simulazione ODE per il campionamento, il passo di addestramento (che evita la backpropagation completa dell'ODE) ha aumentato il tempo di addestramento per iterazione di meno del 50% rispetto a PPO su compiti standard, e rimaneva al di sotto del doppio del costo anche con l'aumento delle dimensioni di embedding.

Significato e Rivendicazioni

Il paper sostiene che PolicyFlow offra un framework pratico ed espressivo per l'RL on-policy. Decoupling (disaccoppiando) la necessità di valutazione della verosimiglianza della traiettoria completa dall'aggiornamento della policy, esso colma il divario tra il potere rappresentativo dei continuous normalizing flows e la stabilità di PPO.

Gli autori sottolineano che:

• L'approssimazione dei rapporti di importanza è teoricamente limitata e sufficiente per un addestramento stabile all'interno del regime di clipping di PPO.
• Il regolarizzatore browniano fornisce un'alternativa fondata all'iniezione di rumore euristica o all'integrazione costosa della divergenza per mantenere l'esplorazione.
• PolicyFlow affronta i limiti specifici degli esistenti metodi di RL basati su flow (FPO, DPPO) riguardanti la stabilità dell'ottimizzazione, il costo computazionale e il bias, rendendo le policy generative espressive più accessibili per l'apprendimento per rinforzo online.

Il lavoro suggerisce che PolicyFlow serva come fondamento per unire la modellazione generativa e l'apprendimento per rinforzo, con potenziali direzioni future che includono il fine-tuning di policy di flow-matching, l'estensione all'RL offline ed esplorazioni sulle connessioni con i modelli di diffusione.