Resumen Técnico: Un Marco de Regularización KL para Aprender a Planificar con Priors Adaptativos

1. Enunciado del Problema

En el Aprendizaje por Refuerzo Basado en Modelos (MBRL), la exploración efectiva sigue siendo un desafío crítico, particularmente en tareas de control continuo de alta dimensión donde la eficiencia de muestras es primordial. Enfoques recientes, como TD-MPC y BMPC, utilizan políticas aprendidas como distribuciones de propuesta para la planificación de Integral de Trayectoria Predictiva de Modelo (MPPI). Sin embargo, surge una limitación fundamental derivada de la desviación de la política: la política de muestreo ($\pi_{\theta_s}$) a menudo se actualiza independientemente de la distribución de acciones del planificador ($\pi_P$).

Esta desacoplamiento crea una desviación de distribución donde la función de valor de acción bootstrap se entrena sobre trayectorias generadas por el planificador MPPI (que explora una distribución específica), mientras que la política de muestreo se actualiza para maximizar los valores de acción locales, a menudo divergiendo de la distribución del planificador. Esta desalineación conduce a estimaciones de valor poco fiables, particularmente para horizontes cortos, y degrada el rendimiento a largo plazo. Los intentos previos de alinear estas distribuciones, como la minimización de la divergencia KL inversa (BMPC), sufren de alta varianza porque dependen de muestras del planificador "obsoletas" almacenadas en el búfer de repetición, las cuales representan una mezcla de políticas de planificación desactualizadas en lugar de una distribución coherente actual.

2. Metodología: Marco PO-MPC

Los autores proponen Optimización de Políticas–Control Predictivo de Modelo (PO-MPC), un marco unificado que trata la distribución de acciones del planificador como un prior en la optimización de políticas. PO-MPC plantea el paso de aprendizaje de la política de muestreo como una instancia de Aprendizaje por Refuerzo Regularizado con KL.

Formulación Central

El objetivo es aprender una política de muestreo $\pi_{\theta_s}$ que maximice la recompensa esperada mientras permanece cerca de un prior adaptativo $\pi_p$ (derivado del planificador MPPI). La función objetivo es:

$$J(\pi_{\theta_s}) = \mathbb{E} \left[ \sum_{t=0}^{T-1} \gamma^t \left( r(z_t, a_t) - \lambda \text{KL}[\pi_{\theta_s}(\cdot | z_t) \| \pi_p(\cdot | z_t)] \right) \right]$$

Donde:

• $z_t$ es la representación del estado latente.
• $\lambda$ es un hiperparámetro que controla la fuerza de la regularización KL.
• $\pi_p$ es el prior de la política.

Componentes Clave

1. RL Regularizado con KL: El marco unifica los métodos existentes como casos especiales de esta formulación:

   • $\lambda = 0$: Recupera TD-MPC2, donde la política se actualiza únicamente mediante la maximización del valor de acción y la regularización de entropía.
   • $\lambda \to \infty$: Recupera BMPC (con KL inverso), donde la política se restringe a imitar la distribución del planificador.
   • $\lambda$ Intermedio: El artículo explora valores intermedios (por ejemplo, $\lambda \in [0.1, 9]$) para equilibrar la maximización de la recompensa y la alineación de distribuciones.

2. Prior Intermedio Adaptativo ($\pi_{\theta_p}$):
   Para abordar la alta varianza causada por el uso directo de muestras obsoletas del planificador desde el búfer de repetición, PO-MPC introduce un prior intermedio aprendido. En lugar de clonar la distribución cruda de las muestras del planificador almacenadas (que pueden ser una mezcla gaussiana variable en el tiempo), el marco entrena una red neuronal $\pi_{\theta_p}$ para aproximar la distribución actual inducida por MPPI $\pi_P$.

   • Este prior se destila de la salida del planificador, protegiendo las actualizaciones de la política de muestreo de la varianza de las muestras desactualizadas.
   • El prior puede entrenarse utilizando KL Inverso (minimizando $\text{KL}[\pi_{\theta_p} \| \pi_P]$) o KL Directo (minimizando $\text{KL}[\pi_P \| \pi_{\theta_p}]$).
     • KL Inverso: Sesga la política hacia la coincidencia de modos específicos del planificador (convergencia más rápida, potencialmente menos exploración).
     • KL Directo: Sesga la política hacia cubrir el soporte de todas las distribuciones muestreadas (exploración mejorada, potencialmente convergencia más lenta).

3. Bucle del Algoritmo:
   El entrenamiento sigue un ciclo Planificar $\to$ Inferir $\to$ Regularizar:

   • Planificar: MPPI genera una distribución de trayectorias $\pi_P$ utilizando la política de muestreo actual y la función de valor.
   • Inferir: El prior adaptativo $\pi_{\theta_p}$ se actualiza para aproximar $\pi_P$.
   • Regularizar: La política de muestreo $\pi_{\theta_s}$ se actualiza utilizando el objetivo regularizado con KL, equilibrando la maximización del valor Q y la penalización KL hacia $\pi_{\theta_p}$.

3. Contribuciones Clave

• Unificación: El artículo unifica métodos dispares de RL basados en MPPI (TD-MPC2, BMPC) bajo un único marco teórico (PO-MPC), aclarando sus decisiones de diseño y relaciones.
• Priors Adaptativos: Introduce un prior intermedio aprendido para reemplazar la clonación directa de muestras del búfer de repetición, demostrando teórica y empíricamente que esto reduce la varianza del gradiente y estabiliza el entrenamiento.
• Objetivos Flexibles: El marco permite la exploración de variantes previamente no estudiadas, específicamente ajustando la fuerza KL $\lambda$ y eligiendo entre pérdidas KL Directa/Inversa para el prior.
• Rendimiento de Vanguardia: El método logra una eficiencia de muestras superior y un rendimiento final en tareas de control continuo de alta dimensión en comparación con las líneas base existentes.

4. Resultados Experimentales

Los autores evaluaron PO-MPC en 21 tareas de control continuo desafiantes de la Suite de Control de DeepMind (Humanoide y Perro) y HumanoidBench.

• Mejoras de Rendimiento: PO-MPC con valores intermedios de $\lambda$ superó consistentemente tanto a TD-MPC2 ($\lambda=0$) como a BMPC ($\lambda \to \infty$).
  • En tareas de alta dimensión (por ejemplo, HumanoidBench), los valores intermedios de $\lambda$ (particularmente $\lambda=1$ o $\lambda=9$) produjeron mejoras significativas en la recompensa final y la eficiencia de muestras.
  • En algunos entornos donde las líneas base fallaron en aprender (por ejemplo, HumanoidBench "Balance Simple", "Slide", "Stair"), PO-MPC convergió con éxito.
• Eficacia del Prior Intermedio: El uso del prior intermedio aprendido ($\pi_{\theta_p}$) igualó o superó el rendimiento del uso de muestras crudas del planificador desde el búfer de repetición, confirmando la reducción en la varianza de actualización.
• Sensibilidad de la Función de Pérdida: La elección de la pérdida KL para entrenar el prior importó. El KL Directo fue beneficioso en tareas con mucha exploración (por ejemplo, "Stair"), mientras que el KL Inverso aceleró la convergencia en tareas críticas de precisión (por ejemplo, "Balance Simple").
• Ablación: Los resultados mostraron que el rendimiento se degrada si $\lambda$ es demasiado bajo (regresa a la desviación de la política) o demasiado alto (colapso prematuro de modos), destacando la importancia de ajustar este compromiso.

5. Significado y Afirmaciones

El artículo afirma que PO-MPC representa una unificación principista de los enfoques basados en MPPI. Al alinear explícitamente la política de muestreo con el planificador mediante un objetivo regularizado con KL y un prior adaptativo, el marco:

1. Cierra el bucle entre la planificación y el aprendizaje de políticas, asegurando que la función de valor se evalúe sobre distribuciones relevantes para el planificador actual.
2. Mejora la estabilidad al mitigar la varianza introducida por las muestras de planificación desactualizadas en el búfer de repetición.
3. Avanza el estado del arte en RL basado en modelos para control continuo, ofreciendo un marco flexible que puede ajustarse a requisitos específicos de la tarea (exploración vs. explotación).

Los autores señalan que, aunque el marco es poderoso, asume distribuciones gaussianas para las políticas y no explota completamente todas las trayectorias simuladas para el aprendizaje de valores. Se sugiere trabajo futuro para extender las distribuciones de políticas más allá de las gaussianas y ajustar automáticamente el hiperparámetro $\lambda$.