Analytic Planning under Uncertainty with Moment Closure

Resumen Técnico: Planificación Analítica bajo Incertidumbre con Cierre de Momentos

Planteamiento del Problema

El aprendizaje por refuerzo basado en modelos (MBRL) efectivo en entornos estocásticos requiere una planificación que considere la incertidumbre predictiva. Los enfoques tradicionales enfrentan un compromiso:

1. Muestreo Estocástico: El MBRL profundo moderno suele depender del muestreo de trayectorias imaginadas (Monte Carlo) para propagar la incertidumbre. Aunque es escalable, esto introduce una varianza significativa en el objetivo, lo que hace que la estimación del valor sea sensible al presupuesto de muestreo y potencialmente inestable.
2. Aproximaciones Deterministas: Los métodos que colapsan los modelos estocásticos a su media o aprenden dinámicas deterministas descartan la covarianza predictiva. Esto conduce a un sesgo sistemático cuando las funciones de valor son no lineales.
3. Propagación Analítica: Aunque la propagación analítica de distribuciones de estado completas es un principio sólido, históricamente ha requerido estructuras de política o recompensa restrictivas (por ejemplo, entornos lineales-cuadráticos) para seguir siendo tratable.

La pregunta central abordada es si la planificación consciente de la distribución es posible sin estas restricciones, específicamente mediante el cálculo analítico de la esperanza de Bellman bajo dinámicas estocásticas aprendidas.

Metodología: Planificación Analítica Compatible con Momentos (MoCA)

Los autores proponen un marco que desacopla la tratabilidad de la selección de acciones codiciosas de la tratabilidad de la esperanza predictiva. El método se basa en dos componentes estructurales clave:

1. Parametrización de Valor-Acción Cuadrática (Desacoplamiento de la Maximización)

Para evitar la intratabilidad de integrar un operador de maximización ($\max_{a'} Q(s', a')$) sobre una distribución estocástica, los autores utilizan una parametrización de valor-acción cuadrática (inspirada en NAF y BNAF).

• Estructura: La función Q se descompone como:
  $$Q(s, a) = V(s) - \frac{1}{2}(a - \tilde{\mu}(s))^\top P(s) (a - \tilde{\mu}(s))$$
  donde $V(s)$ es el valor del estado, $\tilde{\mu}(s)$ es la acción codiciosa sin restricciones, y $P(s)$ es una matriz de curvatura definida positiva.
• Efecto: Esta estructura permite que la selección de la acción codiciosa se resuelva en forma cerrada ($\arg\max_a Q(s,a) = \tilde{\mu}(s)$) y que el valor máximo sea exactamente $V(s)$. En consecuencia, el respaldo de Bellman se reduce de una esperanza sobre una función Q maximizada a una simple esperanza sobre la función de valor del estado únicamente:
  $$y(s, a) = r(s, a) + \gamma \mathbb{E}_{s' \sim p(\cdot|s,a)}[V(s')]$$
  Esto elimina la maximización interna de la esperanza, dejando solo el requisito de calcular $\mathbb{E}[V(s')]$ analíticamente.

2. Principio de Compatibilidad de Momentos

La tratabilidad de la esperanza restante $\mathbb{E}[V(s')]$ depende de la compatibilidad entre la distribución de transición predictiva y la clase de la función de valor.

• Modelo de Dinámica: Las dinámicas del entorno se aprenden como una mezcla de Gaussianas heterocedásticas (un ensamble de $K$ modelos). Para evitar el muestreo durante la planificación, el ensamble se aproxima mediante una única distribución Gaussiana cuya media y covarianza coinciden con los dos primeros momentos de la mezcla (capturando tanto la incertidumbre aleatoria como la epistémica).
• Representación de Valor: La función de valor del estado $V(s)$ se parametriza como una mezcla de funciones de base radial (RBF) Gaussianas.
• Cierre Analítico: Debido a que la distribución predictiva es Gaussiana y la función de valor es una mezcla de Gaussianas, la esperanza $\mathbb{E}[V(s')]$ admite una solución en forma cerrada en términos de los momentos de la distribución (media y covarianza). Esto permite que el objetivo de Bellman se compute exactamente sin muestreo de Monte Carlo.

Procedimiento de Aprendizaje

El algoritmo (MoCA) integra estos componentes en un bucle de aprendizaje fuera de la política (off-policy):

1. Aprendizaje de la Dinámica: Se entrena un ensamble de modelos de dinámica Gaussianos mediante máxima verosimilitud sobre remuestreos de Bootstrap del búfer de repetición (replay buffer).
2. Construcción del Objetivo: Para cada transición, las predicciones del ensamble se agregan mediante el ajuste de momentos para formar una Gaussiana predictiva $\mathcal{N}(\bar{\mu}, \bar{\Sigma})$. El objetivo de Bellman se calcula analíticamente utilizando la esperanza de forma cerrada de la función de valor RBF bajo esta Gaussiana.
3. Actualización del Crítico: Se actualizan dos críticos (para mitigar el sesgo de sobreestimación) utilizando la estimación doble recortada sobre los objetivos analíticos.

Contribuciones Clave

1. Respaldo de Bellman Analítico: El artículo demuestra que es posible lograr esperanzas de Bellman en forma cerrada bajo dinámicas estocásticas aprendidas al desacoplar la maximización (mediante la parametrización Q cuadrática) de la esperanza.
2. Marco de Compatibilidad de Momentos: Introduce un principio de diseño donde la representación del valor se elige específicamente para ser compatible con los momentos de la distribución predictiva, permitiendo la propagación analítica de la incertidumbre sin restringir la clase de la política.
3. Planificación Consciente de la Incertidumbre: Propaga tanto la media predictiva como la covarianza de manera determinista, eliminando la varianza de Monte Carlo de la construcción del objetivo mientras preserva los beneficios del modelado estocástico.

Resultados Empíricos

Los experimentos se realizaron en tareas de control continuo (Cartpole Balance y Pendulum Swingup) bajo ruido de observación Gaussiano controlado (tanto homocedástico como heterocedástico).

• Robustez al Ruido: MoCA superó consistentemente a los modelos de referencia (SAC, BNAF y una versión de MoCA de equivalencia de certeza que ignora la covarianza) a medida que aumentaba el ruido de observación. Mantuvo retornos asintóticos más altos y una dinámica de aprendizaje más estable.
• Analítico vs. Monte Carlo: Al compararse con variantes de Monte Carlo de la misma arquitectura (MoCA-MC) con diferentes conteos de muestras ($K=4, 16$), el método analítico (MoCA) logró una convergencia más rápida y un rendimiento asintótico superior. Los métodos de Monte Carlo sufrieron de alta varianza, especialmente con muestras limitadas, lo que llevó a una selección de acciones subóptima.
• Calibración: El enfoque analítico produjo una incertidumbre predictiva bien calibrada (la cobertura de $1\sigma$ empírica se mantuvo cerca del nivel nominal del 68%) y una desviación estándar epistémica estable.

Significado y Reivindicaciones

El artículo afirma proporcionar un marco fundamentado para la planificación con modelos de distribución aprendidos que evita la varianza del muestreo y el sesgo de la sustitución por la media.

• Estabilidad: Al eliminar el ruido de Monte Carlo de la construcción del objetivo, el método proporciona una planificación estable incluso cuando aumenta el ruido de observación.
• Generalidad: El principio de compatibilidad de momentos no está ligado a una única forma funcional; la combinación Gaussiana-RBF se presenta como una instancia concreta que equilibra la tratabilidad y la expresividad.
• Practicidad: El enfoque ofrece una alternativa computacionalmente eficiente a la planificación basada en muestreo en control continuo, demostrando que la propagación analítica de momentos puede servir como un sustituto práctico de las trayectorias estocásticas.

Los autores reconocen limitaciones, señalando que la instancia RBF enfrenta la "maldición de la dimensionalidad" respecto al número de bases requeridas en espacios de estado de alta dimensión, y que la implementación actual depende de supuestos de covarianza diagonal para la eficiencia computacional. Se sugiere como trabajo futuro explorar representaciones de espacio latente para escalar el método.