PolicyFlow: Policy Optimization with Continuous Normalizing Flow in Reinforcement Learning

Resumen Técnico: PolicyFlow

Declaración del Problema

La Optimización de Política Próxima (PPO) es un algoritmo de aprendizaje por refuerzo (RL) de política on-policy estándar conocido por su estabilidad y simplicidad. Sin embargo, típicamente depende de distribuciones de política gaussianas, las cuales tienen dificultades para representar distribuciones de acciones complejas, multimodales o altamente sesgadas. Aunque los modelos generativos expresivos como los Flujos Normalizantes Continuos (CNF) y los modelos de difusión pueden capturar tales distribuciones complejas, la integración de estos en PPO presenta desafíos significativos.

El cuello de botella principal es el requisito de la evaluación de la verosimilitud (likelihood evaluation) para calcular las razones de importancia ($\pi(a|s)/\hat{\pi}(a|s)$) para el objetivo de la función de sustitución de PPO. Para los CNF, la evaluación exacta de la verosimilitud requiere simular la trayectoria completa de la Ecuación Diferencial Ordinaria (ODE) y realizar la retropropagación por camino (path-wise backpropagation). Este proceso es computacionalmente costoso, intensivo en memoria y, a menudo, numéricamente inestable debido al estallido o desvanecimiento de los gradientes. Los intentos existentes de aplicar actualizaciones de tipo PPO a políticas basadas en flujos (por ejemplo, FPO, DPPO) sufren de un sesgo de estimación asimétrico, requieren tamaños de lote grandes o degradan su rendimiento al entrenar desde cero debido a la dificultad de la exploración fuera de la variedad (off-manifold). Además, las técnicas estándar de regularización de entropía son difíciles de aplicar a las políticas basadas en flujos porque el cálculo explícito de la entropía de la política requiere una costosa integración de la divergencia.

Metodología

Los autores proponen PolicyFlow, un algoritmo de RL on-policy que integra políticas CNF expresivas con objetivos de tipo PPO sin requerir la evaluación de la verosimilitud de la trayectoria completa. El método consta de dos componentes principales:

1. Aproximación de la Razón de Importancia mediante Variaciones del Campo de Velocidad

En lugar de simular la trayectoria completa de la ODE para calcular el desplazamiento terminal entre la política actual $\phi_1$ y la política de referencia $\hat{\phi}_1$, PolicyFlow aproxima la razón de importancia utilizando las variaciones del campo de velocidad a lo largo de una trayectoria de interpolación lineal simple.

• Representación de la Política: La política genera acciones integrando un flujo $\phi_t(z; s)$ gobernado por un campo de velocidad neuronal $v_t$, añadiendo ruido gaussiano en el tiempo terminal: $a = \phi_1(z; s) + n$.
• Aproximación: El desplazamiento terminal $\delta\phi_1 = \phi_1 - \hat{\phi}_1$ se aproxima integrando la diferencia en los campos de velocidad $\delta v_t = v_t - \hat{v}_t$ a lo largo de una trayectoria de interpolación lineal $x_t = (1-t)z + t\hat{\phi}_1(z; s)$.
• Objetivo de Sustitución: La razón de importancia $\rho$ se estima utilizando la razón de verosimilitud gaussiana de las variaciones del campo de velocidad en lugar de las trayectorias completas del flujo. Esto permite que el objetivo de entrenamiento se estime eficientemente a lo largo de la trayectoria de interpolación sin necesidad de retropropagar a través de la simulación completa de la ODE durante el paso de actualización. El análisis teórico indica que esto introduce solo un error de primer orden en el logaritmo de la verosimilitud, el cual está naturalmente acotado por el mecanismo de recorte (clipping) de PPO.

2. Regularizador Browniano para la Entropía

Para prevenir el colapso de modo y fomentar una exploración diversa sin calcular explícitamente la entropía intratable del flujo, los autores introducen el Regularizador Browniano.

• Concepto: Inspirado en el movimiento browniano y la ecuación del calor, este regularizador fomenta que el campo de velocidad aprendido siga la dinámica de un proceso que aumenta la entropía.
• Mecanismo: Aprovecha la relación entre el campo de velocidad y la función de puntuación (score function $\nabla \log p_t$) en flujos rectificados. El regularizador penaliza las desviaciones del campo de velocidad aprendido respecto a un campo de vectores objetivo derivado de la puntuación del flujo de referencia.
• Formulación: El término del regularizador $J_{Reg}$ incluye un término que minimiza la diferencia entre el campo de velocidad aprendido y un término de alineación objetivo (derivado del campo de velocidad de referencia y el estado), más un término que fomenta la varianza del ruido gaussiano inyectado. Este enfoque evita los costosos cálculos de log-verosimilitud y la inyección de ruido heurística, proporcionando una forma ligera y fundamentada para mantener la diversidad de las trayectorias.

Contribuciones Clave

1. Integración Eficiente de CNF-PPO: PolicyFlow permite el uso de políticas CNF expresivas en RL on-policy mediante la aproximación de las razones de importancia a través de las variaciones del campo de velocidad, eliminando la necesidad de la costosa retropropagación por camino a través de trayectorias completas de ODE.
2. Regularizador Browniano: Un regularizador de entropía novedoso y computacionalmente ligero, inspirado en el movimiento browniano, que mitiga el colapso de modo y fomenta comportamientos diversos sin requerir el cálculo explícito de la entropía.
3. Validación Empírica: El método se valida en diversos entornos, demostrando un rendimiento competitivo o superior comparado con PPO Gaussiano y con las bases de estado del arte para flujos (FPO, DPPO).

Resultados Experimentales

Los autores evaluaron PolicyFlow en los entornos de referencia MultiGoal, MuJoCo Playground e IsaacLab:

• MultiGoal (Prueba de Colapso de Modo): En un entorno con seis metas simétricas, el PPO estándar (Gaussiano) falló en alcanzar todas las metas. Las bases FPO y DPPO colapsaron en unos pocos modos debido a la falta de una regularización de entropía efectiva. PolicyFlow, con el regularizador browniano, capturó con éxito una distribución multimodal equilibrada, alcanzando todas las metas con una probabilidad aproximadamente igual.
• MuJoCo Playground: PolicyFlow logró un rendimiento comparable o superior al de FPO y DPPO en diversas tareas (por ejemplo, WalkerWalk, CheetahRun). A menudo convergió más rápido, lo que indica una mayor eficiencia de muestreo.
• IsaacLab: PolicyFlow fue comparado contra PPO en tareas de robótica (locomoción, manipulación). Consistentemente igualó o superó el rendimiento asintótico de PPO.
• Eficiencia Computacional: Aunque PolicyFlow implica la simulación de ODE para el muestreo, el paso de entrenamiento (que evita la retropropagación completa de la ODE) aumentó el tiempo de entrenamiento por iteración en menos del 50% en comparación con PPO en tareas estándar, y se mantuvo por debajo del doble de costo incluso con el aumento de las dimensiones de incrustación (embedding).

Significancia y Reivindicaciones

El artículo afirma que PolicyFlow ofrece un marco práctico y expresivo para el RL on-policy. Al desacoplar la necesidad de la evaluación de la verosimilitud de la trayectoria completa del paso de actualización de la política, cierra la brecha entre el poder de representación de los flujos normalizantes continuos y la estabilidad de PPO.

Los autores enfatizan que:

• La aproximación de las razones de importancia está teóricamente acotada y es suficiente para un entrenamiento estable dentro del régimen de recorte de PPO.
• El regularizador browniano proporciona una alternativa fundamentada a la inyección de ruido heurística o a la costosa integración de la divergencia para mantener la exploración.
• PolicyFlow aborda las limitaciones específicas de los métodos de RL basados en flujos existentes (FPO, DPPO) con respecto a la estabilidad de la optimización, el costo computacional y el sesgo, haciendo que las políticas generativas expresivas sean más accesibles para el aprendizaje por refuerzo en línea.

El trabajo sugiere que PolicyFlow sirve como base para unir el modelado generativo y el aprendizaje por refuerzo, con posibles direcciones futuras que incluyen el ajuste fino (fine-tuning) de políticas de correspondencia de flujo (flow-matching), la extensión al RL offline y la exploración de conexiones con los modelos de difusión.