Multivariate Distributional Reinforcement Learning Using Sliced Divergences

Resumen Técnico: Aprendizaje por Refuerzo Distribucional Multivariante mediante Divergencias Segmentadas (Sliced Divergences)

Planteamiento del Problema

El Aprendizaje por Refuerzo Distribucional (DRL) modela la distribución completa de los retornos en lugar de solo su esperanza, ofreciendo ventajas teóricas y empíricas sobre el RL clásico basado en la esperanza. Sin embargo, extender el DRL a entornos multivariantes (donde los retornos son vectores de $d$ dimensiones, $d > 1$) sigue siendo un desafío significativo.

Los enfoques existentes enfrentan tres limitaciones principales:

1. Intratabilidad Computacional: Métricas comunes como la distancia de Wasserstein sufren de altos costos computacionales en altas dimensiones (típicamente $O(n^3 \log n)$ para resolvedores de transporte óptimo) y tasas estadísticas deficientes que dependen de la dimensión ambiental.
2. Falta de Garantías de Contracción: Aunque el DRL escalar posee resultados de contracción bien establecidos bajo el operador de Bellman, los entornos multivariantes con un descuento matricial general (donde el factor de descuento es una matriz dependiente del estado-acción $\Gamma(s, a) \in \mathbb{R}^{d \times d}$) carecen de garantías de contracción rigurosas. Los métodos existentes suelen fallar al no contraer bajo descuentos anisotrópicos (no uniformes).
3. Sesgo de Gradiente en el Entrenamiento Estocástico: Muchas divergencias, particularmente la de Wasserstein, fallan en la propiedad de gradiente de muestra no sesgado (U) cuando el objetivo de Bellman se instancia a partir de un único sucesor muestreado (aprendizaje TD estándar). Esto conduce a gradientes sesgados y un ajuste distribucional subóptimo.

Metodología: Aprendizaje por Refuerzo Distribucional Segmentado (SDRL)

Los autores introducen el Aprendizaje por Refuerzo Distribucional Segmentado (SDRL), un marco que eleva las divergencias unidimensionales tratables a distribuciones de retorno multivariantes mediante proyecciones aleatorias (segmentación o slicing).

Mecanismo Central

En lugar de calcular directamente una divergencia de alta dimensión, el SDRL proyecta las distribuciones multivariantes sobre direcciones aleatorias $\theta \in S^{d-1}$, calcula una divergencia base $\Delta$ en una dimensión y agrega los resultados.

• Segmentación Uniforme (Uniform Slicing): Promedia la divergencia base sobre $L$ direcciones aleatorias:
  $$S\Delta_p^p(\mu, \nu) = \int_{S^{d-1}} \Delta^p((P_\theta)_\#\mu, (P_\theta)_\#\nu) d\sigma(\theta)$$
• Segmentación Máxima (Max-Slicing/MSDRL): Optimiza sobre la dirección más discriminativa para obtener garantías de contracción más fuertes para descuentos matriciales generales:
  $$MS\Delta(\mu, \nu) = \sup_{\theta \in S^{d-1}} \Delta((P_\theta)_\#\mu, (P_\theta)_\#\nu)$$

El marco admite una amplia clase de divergencias base, analizando específicamente Wasserstein ($W_p$), distancia de Cramér ($C_2$) y Máximo Desacuerdo Medio (MMD).

Contribuciones Teóricas

El artículo establece fundamentos teóricos rigurosos para el SDRL:

1. Propiedades Métricas: Demuestra que tanto las elevaciones segmentadas uniformes como las máximas preservan los axiomas métricos si la divergencia base es una métrica.
2. Garantías de Contracción:
   • Segmentación Uniforme: Prueba la contracción de Bellman bajo un descuento escalar compartido ($\Gamma = \gamma I$) para cualquier divergencia base que satisfaga la no-expansión de traslación, la contracción de escala y la $p$-convexidad de mezcla.
   • Segmentación Máxima: Introduce un nuevo resultado de contracción para descuentos matriciales densos generales (actualizaciones anisotrópicas). Demuestra que la segmentación máxima produce contracción bajo la norma del operador de la matriz de descuento, un escenario donde la segmentación uniforme y el MMD estándar suelen fallar en la contracción.
3. Complejidad de Muestreo: Demuestra que la segmentación uniforme hereda tasas de convergencia independientes de la dimensión de la divergencia base unidimensional, evitando la maldición de la dimensionalidad. Las tasas de la segmentación máxima escalan polinómicamente con la dimensión ($O(d \log n / n)$), lo cual es significativamente mejor que el transporte óptimo exacto de alta dimensión.
4. Análisis del Sesgo de Gradiente:
   • Segmentación Uniforme: Preserva la propiedad de gradiente de muestra no sesgado (U) si la divergencia base la satisface. Esto hace que la Cramér Segmentada y el MMD Segmentado sean compatibles con el bootstrapping TD de una sola muestra (estándar).
   • Segmentación Máxima: Falla en la propiedad (U) debido al sesgo de selección introducido por el paso de maximización. La dirección seleccionada depende del lote de muestra específico, lo que genera gradientes sesgados en entornos TD estándar.

Resultados Experimentales

Los autores evalúan el SDRL en tres entornos de referencia: una cadena MDP tabular, un entorno de laberinto basado en píxeles y un subconjunto de juegos de Atari.

1. Entorno de Cadena (Evaluación de Política):

   • Compara TD de una sola muestra (bootstrapping estándar) frente a TD casi exacto (construcción explícita de mezclas).
   • Hallazgo: Los objetivos que satisfacen la propiedad (U) (ej. Cramér Segmentada, MMD Segmentado) logran un ajuste distribucional preciso en el régimen de una sola muestra. Los objetivos que violan (U) (ej. Wasserstein Segmentada, variantes de Segmentación Máxima) se degradan significativamente a menos que se utilice el objetivo de mezcla completo.
   • Segmentación Máxima: Funciona bien en el caso degenerado de una sola dimensión (donde el sesgo de selección desaparece), pero falla en entornos multivariantes generales bajo TD de una sola muestra.

2. Laberinto Basado en Píxeles (Evaluación de Política):

   • Evalúa la precisión distribucional utilizando la distancia de Wasserstein-2 empírica entre los retornos de Monte Carlo y la predicción del crítico.
   • Hallazgo: La distancia de Cramér Segmentada y el MMD Segmentado rinden a la par con sus contrapartes no segmentadas y superan significativamente a la de Wasserstein Segmentada en el régimen de TD de una sola muestra, confirmando la importancia de la propiedad (U).

3. Juegos de Atari (Control):

   • Evalúa el rendimiento de control sobre retornos multidimensionales descompuestos.
   • Hallazgo: Aunque los objetivos que satisfacen (U) proporcionan estimaciones distribucionales precisas, la Wasserstein-2 Segmentada logra sorprendentemente un fuerte rendimiento de control a pesar de tener un ajuste de la distribución y gradientes sesgados. Esto sugiere que para el control, la estimación precisa de la esperanza del retorno puede ser más crítica que la fidelidad de la distribución completa, incluso si el objetivo distribucional está sesgado.

Significancia y Reivindicaciones

El artículo afirma proporcionar un marco principiado para extender el DRL a entornos multivariantes manteniendo la tratabilidad computacional y las garantías teóricas.

• Recomendación Práctica: Los autores identifican la Cramér Segmentada como una opción sólida por defecto para el aprendizaje distribucional multivariante. Es computacionalmente eficiente ($O(n \log n)$), satisface la propiedad de gradiente no sesgado (permitiendo el entrenamiento TD estándar) y ofrece garantías de contracción en el entorno isotrópico estándar.
• Perspectiva Teórica: El trabajo clarifica los compromisos entre diferentes elecciones de divergencia. Destaca que, si bien la Segmentación Máxima resuelve el problema de la contracción para descuentos matriciales generales, introduce un sesgo de gradiente que la hace inadecuada para el TD de una sola muestra estándar sin modificaciones.
• Identificación de Brechas: El artículo concluye que una divergencia que sea simultáneamente demostrable como contractiva bajo descuentos anisotrópicos generales (en un sentido de solo norma), satisfaga la propiedad de gradiente de muestra no sesgado y ofrezca una estimación tratable, sigue siendo una cuestión abierta.

En resumen, el SDRL ofrece un enfoque escalable y teóricamente fundamentado para el DRL multivariante, emergiendo la Cramér Segmentada como la opción práctica más robusta para los flujos de trabajo estándar de aprendizaje por refuerzo.