Distributional Inverse Reinforcement Learning

Resumen Técnico: Aprendizaje por Refuerzo Inverso Distribucional (DistIRL)

1. El Problema

El Aprendizaje por Refuerzo Inverso (IRL) busca inferir la función de recompensa subyacente y la política de un experto a partir de trayectorias observadas. Sin embargo, la mayoría de los métodos actuales (como MaxEntIRL o IRL Bayesiano estándar) asumen que la función de recompensa es determinista o que solo se necesita recuperar un estimado puntual (la media).

Esta suposición es limitante en escenarios del mundo real donde las señales de recompensa son inherentemente estocásticas:

• Robótica: La incertidumbre en el contacto con objetos deformables o frágiles introduce variabilidad en la recompensa para los mismos pares estado-acción.
• Neurociencia: Las señales de dopamina, que actúan como recompensas biológicas, muestran variaciones significativas entre ensayos, sugiriendo una distribución subyacente estocástica.
• Limitaciones de los métodos existentes: Los enfoques tradicionales solo coinciden con la recompensa esperada (media). Esto ignora momentos de orden superior (varianza, asimetría) y la estructura completa de la distribución de retornos, lo que resulta en políticas que no son conscientes del riesgo y no pueden capturar la verdadera incertidumbre del entorno. Además, en configuraciones offline (sin interacción con el entorno), inferir distribuciones completas es particularmente desafiante.

2. Metodología: El Marco DistIRL

Los autores proponen DistIRL, un marco que modela explícitamente tanto la distribución de la recompensa como la distribución de los retornos (cumulativos). La metodología se basa en tres pilares fundamentales:

A. Aprendizaje de la Distribución de Recompensa mediante Dominancia Estocástica de Primer Orden (FSD)
En lugar de minimizar la diferencia en las recompensas esperadas, DistIRL busca alinear las distribuciones completas de los retornos del agente y del experto.

• Objetivo: Minimizar las violaciones de la Dominancia Estocástica de Primer Orden (FSD). Si la distribución de retornos del experto $Z_E$ domina estocásticamente a la del agente $Z_\pi$, entonces $F_{Z_E}(z) \leq F_{Z_\pi}(z)$ para todo $z$.
• Función de Pérdida: Se define una función de energía basada en la violación de FSD:
  $$L_{FSD}(\pi, r) = \int_{-\infty}^{\infty} [F_{Z_E}(z) - F_{Z_\pi}(z)]_+ dz$$
  Esta pérdida se interpreta como una función de energía en un Modelo Basado en Energía (EBM), donde las recompensas que generan menores violaciones de FSD son exponencialmente más probables bajo los datos del experto.
• Inferencia Variacional: Dado que la inferencia directa es intratable, se utiliza un enfoque variacional para aproximar la distribución posterior de la recompensa $q_\phi(r|s,a)$, optimizando un límite inferior de la evidencia (ELBO) que combina la pérdida de FSD esperada y una regularización KL respecto a un prior.

B. Aprendizaje de Políticas Conscientes del Riesgo (Risk-Aware)
Una vez estimada la distribución de recompensas, el objetivo es aprender una política que maximice un Medida de Riesgo de Distorsión (DRM).

• Relación con FSD: Se demuestra teóricamente que si una distribución domina estocásticamente a otra, cualquier medida de riesgo basada en distorsión (como CVaR) también será mayor.
• Objetivo de la Política: En lugar de maximizar la esperanza, la política maximiza una medida de riesgo $M_\xi(Z_\pi)$ (definida mediante una función de distorsión $\xi$) más la entropía de la política:
  $$\max_\phi \int_0^1 F^{-1}_{Z_\pi}(v) d\tilde{\xi}(v) + H(\pi_\phi)$$
  Esto permite que el agente aprenda comportamientos aversos al riesgo (o propensos al riesgo) alineados con la variabilidad de los datos del experto, no solo con su media.

C. Algoritmo Práctico
El algoritmo alterna entre actualizar la distribución de recompensas y la política:

1. Muestreo: Se generan muestras de retorno utilizando la distribución de recompensa actual y las políticas del experto y del agente.
2. Actualización del Crítico: Se utiliza regresión de cuantiles (Quantile Regression) para estimar la distribución de retornos.
3. Actualización de la Recompensa: Se minimiza la pérdida FSD para ajustar la distribución paramétrica de la recompensa (ej. distribución normal sesgada de Azzalini).
4. Actualización de la Política: Se optimiza la política utilizando el gradiente de la medida de riesgo DRM.

3. Contribuciones Clave

1. Aprendizaje de Distribución de Recompensas: Es el primer trabajo que recupera la distribución completa de la función de recompensa (no solo parámetros de una distribución determinista) en un marco principista utilizando FSD.
2. Políticas Conscientes de la Distribución: Extiende el IRL para aprender políticas que son conscientes de la distribución de retornos, permitiendo el aprendizaje de comportamientos aversos al riesgo en escenarios de IRL offline.
3. Análisis Teórico: Se proporciona un análisis de convergencia que demuestra que el algoritmo converge a un punto estacionario de primer orden con una complejidad de iteración de $O(\epsilon^{-2})$.
4. Validación Empirica: Demostración de recuperación de distribuciones significativas en datos sintéticos, datos neuroconductuales reales y tareas de control robótico de alta dimensión.

4. Resultados Experimentales

Los experimentos se realizaron en tres dominios:

• Gridworld (Entorno Discreto):

  • En un entorno con recompensas estocásticas, DistIRL recuperó con precisión tanto la media como la varianza de las recompensas.
  • En comparación con IRL Bayesiano (BIRL), que solo recuperó la media (y falló en capturar la varianza, generando estimaciones espurias), DistIRL identificó correctamente la incertidumbre y aprendió una política aversa al riesgo que evitaba estados con alta variabilidad.

• Comportamiento Espontáneo de Ratones (Neurociencia):

  • Se aplicó a un conjunto de datos de trazas de dopamina en ratones.
  • Los resultados mostraron que los modelos deterministas fallaban en capturar la forma de la distribución de dopamina.
  • DistIRL con distribución normal sesgada (S-DistIRL) logró la mayor correlación con los niveles reales de dopamina y la menor distancia de Wasserstein-1, recuperando la asimetría y la multimodalidad de las señales biológicas. Esto sugiere que la estructura de recompensa neuronal puede inferirse directamente del comportamiento.

• Benchmarks MuJoCo (Control Continuo Offline):

  • Se evaluó en tareas de D4RL con recompensas estocásticas (penalizaciones por fallos catastróficos).
  • DistIRL superó consistentemente a los métodos state-of-the-art (como ValueDICE, Offline ML-IRL y BC) en tareas de HalfCheetah, Hopper y Walker2D.
  • Los métodos basados en esperanza (ValueDICE) fallaron porque optimizaban para un retorno esperado que no reflejaba el comportamiento averso al riesgo del experto.
  • DistIRL logró un rendimiento cercano al experto, validando su capacidad para aprender políticas robustas en entornos estocásticos.

5. Significado e Impacto

• Superación de Limitaciones del IRL Clásico: DistIRL rompe la barrera de la suposición de recompensas deterministas, permitiendo modelar la incertidumbre inherente en sistemas complejos (robótica, biología).
• Seguridad y Robustez: Al aprender distribuciones completas y optimizar medidas de riesgo, el método permite generar políticas que son intrínsecamente más seguras y robustas ante variaciones en el entorno, algo crítico para la implementación de IA en el mundo real.
• Herramienta para Neurociencia: Proporciona una nueva herramienta computacional para decodificar señales neuromoduladoras (como la dopamina) a partir de datos conductuales, ofreciendo insights sobre los estados motivacionales internos y sus sustratos neuronales.
• Generalización: El marco es lo suficientemente general para funcionar tanto en entornos estocásticos como deterministas, demostrando versatilidad en diversos dominios de aplicación.

En conclusión, DistIRL representa un avance fundamental al unificar el aprendizaje de distribuciones de recompensa, la estimación de distribuciones de retorno y el aprendizaje de políticas conscientes del riesgo, ofreciendo una solución más fiel y expresiva a los problemas de IRL en el mundo real.