Efficient Hypergradient Descent for Inverse Reinforcement Learning

Resumen Técnico: Descenso de Hipergradiente Eficiente para la Inferencia de Recompensa por Refuerzo

Planteamiento del Problema
La Inferencia de Recompensa por Refuerzo (IRL, por sus siglas en inglés) busca recuperar una función de recompensa a partir de demostraciones de un experto, de tal manera que la política óptima bajo dicha recompensa reproduzca el comportamiento observado. Los autores se centran en la Máxima Verosimilitud de IRL (ML-IRL), la cual puede formularse como un problema de optimización bilevel. El nivel interno optimiza una política regularizada por entropía bajo una recompensa aprendida $r_\phi$, mientras que el nivel externo minimiza la discrepancia (específicamente, la divergencia KL) entre la política inducida y los datos del experto.

El principal cuello de botella computacional al resolver este problema bilevel mediante diferenciación implícita es el requisito de calcular el hipergradiente. La construcción explícita del Hessiano requiere una memoria cuadrática ($O(d_\theta^2)$), mientras que las aproximaciones iterativas (por ejemplo, gradientes conjugados) pueden ser computacionalmente costosas y sensibles al condicionamiento. Aunque se ha propuesto el Descenso de Hipergradiente Natural (NHGD) para reemplazar el Hessiano con un sustituto de la matriz de información de Fisher, el NHGD estándar asume una distribución de datos fija. En ML-IRL, el objetivo interno implica una divergencia KL inversa sobre una distribución de trayectoria inducida por la política, lo que hace que la aplicación directa de NHGD no sea trivial.

Metodología
El artículo propone un método para calcular eficientemente el hipergradiente para ML-IRL aprovechando la estructura geométrica del problema interno y empleando el esbozado de matrices (matrix sketching).

1. Equivalencia de Fisher-Hessiano: Los autores demuestran que, bajo la optimalidad interna exacta y la realizabilidad (donde la clase de políticas es lo suficientemente rica como para modelar perfectamente la distribución del experto), el Hessiano del objetivo interno es proporcional a la matriz de información de Fisher de la trayectoria con descuento ($F_\theta$). Específicamente, $\nabla^2_\theta L_{inner} = \alpha F_\theta$. Esto permite que el producto inverso-Hessiano-vector requerido para el hipergradiente sea reemplazado por un producto inverso-Fisher-vector.
2. Estimación del Hipergradiente: El método descompone el cálculo del hipergradiente en tres componentes:
   • El gradiente externo con respecto a los parámetros de la política.
   • La solución de un sistema lineal amortiguado que involucra la matriz de Fisher y el gradiente externo.
   • Un término de derivada mixta que involucra el gradiente de la recompensa y la solución del sistema lineal.
3. Esbozado Espectral de Flujo (SCFD): Para evitar el costo de memoria $O(d_\theta^2)$ de almacenar la matriz de Fisher densa, los autores aproximan el producto inverso-Fisher-vector requerido utilizando un esbozado de Direcciones Frecuentes de Compensación Espectral de Flujo (SCFD). En lugar de materializar la matriz de Fisher, el algoritmo transmite vectores de puntuación de la política desde las trayectorias del agente hacia un esbozo compacto de tamaño $m$. Esto reduce la complejidad de almacenamiento a $O(md_\theta)$. El esbozo se utiliza entonces para resolver el sistema lineal amortiguado de forma aproximada.

Contribuciones Clave

• Derivación Teórica: El artículo deriva estimadores basados en muestras para el hipergradiente implícito de ML-IRL y demuestra que el Hessiano interno es proporcional a la matriz de información de Fisher de la trayectoria en un óptimo interno realizable. Esto se presenta como la primera aplicación de esta identidad para el cálculo de hipergradientes implícitos en ML-IRL.
• Solucionador Escalable: Los autores proponen un solucionador basado en SCFD que transmite vectores de puntuación de la política. Este enfoque evita la construcción de la matriz de puntuación o de la matriz de Fisher densa, reduciendo significativamente los requisitos de memoria.
• Validación Empírica: El método se evalúa frente a una línea base de ML-IRL de un solo bucle (single-loop) de primer orden estocástico y un solucionador de Fisher explícito en entornos de control discretos (CartPole) y continuos (LQR).

Resultados
Los experimentos se realizaron bajo presupuestos computacionales equivalentes (límite de entrenamiento de 24 horas) con arquitecturas e hiperparámetros idénticos.

• Rendimiento de la Política: Los métodos basados en Fisher lograron una calidad de política competitiva. En CartPole, todos los métodos alcanzaron retornos de nivel de experto. En LQR, el ML-IRL de un solo bucle logró el mejor Log-Verosimilitud Negativa (NLL) de la política y el mejor Retorno del Entorno, aunque el margen sobre el método de "Fisher con Esbozado" fue pequeño.
• Clasificación de Recompensas: Los métodos basados en Fisher demostraron una fuerte calidad de clasificación de recompensas. En LQR, el método de "Fisher con Esbozado" alcanzó la mayor Correlación de Rango (0.972) entre los retornos de trayectoria de la recompensa aprendida y la real, superando tanto a la línea base de un solo bucle como al solucionador de Fisher explícito.
• Eficiencia:
  • Memoria: El esbozado proporcionó el beneficio más significativo en el entorno de alta dimensión LQR, reduciendo el uso de memoria pico hasta en 1.31$\times$ en comparación con el solucionador de Fisher explícito.
  • Velocidad: En CartPole, el esbozado produjo una aceleración de 1.29$\times$. En LQR, tamaños de esbozo moderados (por ejemplo, $m=32$) ofrecieron una ligera aceleración (1.04$\times$), mientras que los esbozos más grandes ($m=256$) resultaron más lentos que el solucionador explícito debido a la sobrecarga del proceso de esbozado, a pesar de obtener valores de pérdida finales menores.
  • Estabilidad: Los autores encontraron que el solucionador con esbozo se mantuvo estable bajo parámetros de amortiguación ($\lambda$) sustancialmente más débiles que la formulación explícita, lo que sugiere que el proceso de esbozado puede actuar como un regularizador implícito.

Significancia y Reivindicaciones
El artículo afirma abordar el cuello de botella de escalabilidad de la la IRL bilevel estableciendo un vínculo directo entre el Hessiano interno y la matriz de información de Fisher en el contexto de la optimización de la divergencia KL inversa. Al combinar este conocimiento teórico con el esbozado de matrices de flujo, los autores demuestran que es posible calcular hipergradientes implícitos sin el costo de memoria cuadrática de la construcción explícita de las matrices de Hessiano o de Fisher.

Los autores concluyen modestamente que, si bien los métodos basados en Fisher son competitivos con el ML-IRL de un solo bucle en términos de calidad final de la política, ofrecen una ventaja distintiva en el rendimiento de la clasificación de recompensas y en la eficiencia computacional (específicamente en la reducción de memoria) para problemas de mayor dimensión. El trabajo sugiere que los tamaños de esbozo moderados ofrecen el equilibrio óptimo entre la calidad de la optimización, el uso de memoria y el tiempo de ejecución.