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

Resumen Técnico: NFDRL

Enunciado del Problema
El Aprendizaje por Refuerzo Distribucional (DistRL) busca modelar la distribución completa de la recompensa en lugar de solo la recompensa esperada, proporcionando señales más ricas para la incertidumbre, el riesgo y la variabilidad. Sin embargo, los métodos existentes enfrentan una compensación entre expresividad y eficiencia de parámetros.

• Métodos categóricos (p. ej., C51): Se basan en soportes discretos fijos. Aunque son efectivos, carecen de parsimonia; la cantidad de parámetros escala linealmente con la resolución deseada (número de átomos). También imponen límites rígidos a las recompensas, limitando su aplicabilidad en dominios con recompensas dinámicas o de cola pesada.
• Métodos de cuantiles (p. ej., QR-DQN, IQN): Aprenden las ubicaciones de los cuantiles para manejar recompensas ilimitadas. Sin embargo, aproximan las distribuciones como mezclas de deltas de Dirac (funciones escalón). Esta representación puede ser ineficiente para distribuciones suaves, multimodales o de cola pesada, a menudo requiriendo muchos cuantiles para aproximar densidades simples y fallando en capturar formas distribucionales detalladas.
• Métodos de Densidad Continua: Aunque modelar funciones de densidad de probabilidad (PDF) continuas ofrece soporte adaptable y suavidad, los enfoques anteriores han luchado por definir un objetivo de entrenamiento que sea simultáneamente teóricamente sólido (garantizando la contracción del operador de Bellman) y prácticamente optimizable (proporcionando gradientes de muestra no sesgados).

Metodología: NFDRL
Los autores proponen NFDRL, una arquitectura DistRL eficiente en parámetros que modela distribuciones de recompensa utilizando flujos normalizadores continuos.

1. Arquitectura:

   • El modelo utiliza un flujo normalizador condicional para mapear muestras de ruido base $z \sim U$ a resultados de recompensa $y = f_\theta(z)$.
   • El flujo se construye mediante una Función de Distribución Acumulada (CDF) de Mezcla Gaussiana. Una red neuronal $h_\theta$ mapea un estado $x$ a los parámetros de una CDF de mezcla para cada acción.
   • Para abordar la naturaleza acotada de las CDFs ($[0, 1]$), una transformación afín adicional reescala la salida a un rango dinámico $[-G_{max}, G_{max}]$, permitiendo que el modelo se adapte al soporte de las recompensas del mundo real sin límites fijos.
   • La distribución objetivo se construye componiendo el flujo predicho con la función de bootstrap ($r + \gamma y$), tratada como una capa de flujo invertible para preservar la normalización mediante la fórmula del cambio de variable.

2. Objetivo de Entrenamiento (Sustituto de Cramér Consciente de la Geometría):

   • La divergencia KL estándar no es adecuada para soportes que no se superponen, y la minimización exacta de Wasserstein sufre de gradientes sesgados.
   • Los autores introducen una distancia inspirada en Cramér y consciente de la geometría definida sobre masas de probabilidad en lugar de densidades crudas.
   • La función de pérdida se define como:
     $$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}$$
     donde $w_i$ y $v_i$ son masas de probabilidad en los puntos de soporte $y_i$.
   • Para manejar la desalineación entre los soportes predicho y objetivo (que difieren debido al paso de bootstrap), el método emplea Estimación de Densidad Kernel (KDE) para evaluar densidades en soportes compartidos antes de calcular la pérdida.

Contribuciones Clave

1. Eficiencia de Parámetros: A diferencia de las líneas base categóricas donde el tamaño del modelo crece con la resolución, NFDRL mantiene una huella de parámetros constante independientemente de la resolución efectiva de la distribución, ofreciendo resolución efectiva infinita y soporte dinámico.
2. Solidez Teórica: Los autores demuestran que su distancia basada en masas es una métrica verdadera y que el operador de Bellman distribucional asociado es una contracción de $\sqrt{\gamma}$. Esto garantiza la convergencia a un punto fijo único.
3. Gradientes No Sesgados: El objetivo propuesto admite gradientes de muestra no sesgados, una propiedad típicamente ausente en métodos DistRL basados en PDF anteriores (a diferencia de Wasserstein, que a menudo requiere aproximaciones sesgadas).
4. Expresividad: La naturaleza continua del flujo permite que el modelo capture paisajes de recompensa complejos y multimodales que los métodos basados en cuantiles a menudo difuminan o no logran resolver.

Resultados

• MDPs de Juguete: En entornos controlados, NFDRL recupera con éxito paisajes de recompensa ricos y multimodales (p. ej., distribuciones bimodales) con PDFs suaves, mientras que las líneas base de cuantiles (IQN) produjeron distribuciones ruidosas y difuminadas que ocultaban la modalidad.
• Benchmark Atari-5: NFDRL logra un rendimiento competitivo con las líneas base categóricas (C51) y se acerca al rendimiento de los mejores métodos de cuantiles (IQN).
  • En el benchmark Atari-5, NFDRL (variante sustituto) logró una puntuación media normalizada por humanos de 418, superando a C51 (330) y DQN (189), aunque quedando ligeramente por detrás de IQN (525).
  • Cabe destacar que NFDRL-S superó a IQN en juegos específicos como Double Dunk y QBert*.
• Eficiencia: NFDRL iguala el rendimiento de C51 con significativamente menos parámetros. Mientras que la cantidad de parámetros de C51 escala linealmente con el número de átomos, NFDRL mantiene un tamaño fijo, igualando la capacidad de C51 con solo 11 átomos.

Significado y Afirmaciones
El artículo posiciona a NFDRL no como un reemplazo universal para agentes basados en cuantiles fuertes como IQN en Atari, sino como una alternativa continua y parsimoniosa con compensaciones distintas. Su principal significado radica en ofrecer un marco teóricamente riguroso (contracción y gradientes no sesgados) para distribuciones de recompensa continuas con una huella de parámetros constante. Esto es particularmente valioso en escenarios donde las decisiones aguas abajo dependen de la forma detallada de la distribución de recompensas (p. ej., multimodalidad, riesgo de cola) en lugar de solo del valor esperado. Los autores reconocen limitaciones, incluida una mayor varianza de entrenamiento, una convergencia más lenta en comparación con algunas líneas base y el costo computacional del paso de alineación KDE, considerándolos áreas para un refinamiento futuro.