Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics

1. Planteamiento del Problema

El trabajo aborda el problema de la Optimización de Trayectorias y Políticas (TPO) en sistemas dinámicos discretos con horizonte finito. El objetivo es encontrar una secuencia de controles o una política de retroalimentación que minimice el costo esperado de una trayectoria.

• Contexto: Se asume que las dinámicas del sistema $f(x_t, u_t)$ son diferenciables respecto al estado y al control. Esto permite el uso de métodos basados en gradientes.
• Desafío: Aunque los métodos basados en gradientes (como la programación no lineal o los métodos de primer orden en RL) son eficientes, tienden a quedar atrapados en mínimos locales no deseados debido a la no convexidad del problema. Por otro lado, los métodos de muestreo (como MPPI) exploran mejor el espacio pero a menudo requieren un número prohibitivo de muestras y no aprovechan la información del gradiente.
• Objetivo Unificado: Diseñar un algoritmo unificado que combine las ventajas de la exploración basada en muestreo con la refinación basada en gradientes, formulando el diseño del controlador como un problema de inferencia.

2. Metodología Propuesta

El autor propone un marco basado en la minimización de un costo esperado regularizado por la divergencia KL, lo que lleva a una distribución óptima de tipo "Boltzmann-tilted" (sesgada por Boltzmann) sobre los parámetros del controlador $\theta$:

$$p^*(\theta) \propto p_0(\theta) \exp\left(-\frac{1}{\lambda} E(\theta)\right)$$

Donde $E(\theta)$ es el costo esperado de la trayectoria, $p_0$ es una distribución previa y $\lambda$ es una temperatura. A medida que $\lambda \to 0$, la distribución se concentra en los mínimos globales.

Para muestrear eficientemente de esta distribución objetivo (que puede ser aguda y multimodal), se introduce Tempered Sequential Monte Carlo (TSMC).

Componentes Clave de TSMC:

1. Ruta de Templado (Tempering Path):

   • En lugar de muestrear directamente de la distribución objetivo difícil, se construye una secuencia de distribuciones intermedias $p_k$ que deforman suavemente la previa $p_0$ hacia la objetivo $p^*$.
   • Se utiliza una secuencia de escalares $\beta_k$ (temperaturas) donde $p_k(\theta) \propto p_0(\theta) \exp(-\beta_k E(\theta)/\lambda)$.

2. Muestreo Secuencial (Importance Sampling y Resampling):

   • Se mantiene un conjunto de partículas ponderadas.
   • En cada paso de templado, las partículas se reponderan según la razón de verosimilitud incremental.
   • Se utiliza el Tamaño de Muestra Efectivo (ESS) para adaptar dinámicamente el siguiente nivel de temperatura, evitando el colapso de los pesos.
   • Se realiza un resampling para eliminar partículas de bajo peso y duplicar las de alto peso.

3. Rejuvenecimiento con Hamiltonian Monte Carlo (HMC):

   • Para evitar que el conjunto de partículas colapse en un número pequeño de modos y para explorar el espacio de parámetros localmente, se aplica un paso de HMC después del resampling.
   • Ventaja clave: Dado que las dinámicas son diferenciables, TSMC puede calcular gradientes exactos ($\nabla_\theta E(\theta)$) mediante diferenciación automática a través de la simulación de la trayectoria (rollout). Esto permite a HMC proponer movimientos grandes y eficientes en el espacio de parámetros, preservando la diversidad de las partículas.

Adaptaciones para Optimización de Políticas (PO):

Para la optimización de políticas (donde el costo es una expectativa sobre estados iniciales), el autor propone dos variantes prácticas:

• Aproximación Determinista: Aproximar la distribución de estados iniciales con un conjunto fijo de muestras empíricas, permitiendo el cálculo de gradientes deterministas.
• Construcción de Espacio Extendido: Tratar la aleatoriedad de la simulación (estados iniciales) como variables auxiliares dentro del estado de la partícula. Esto permite usar TSMC en un espacio aumentado $(\theta, x_0)$, donde se actualizan tanto los parámetros de la política como los estados iniciales mediante cadenas de Markov.

3. Contribuciones Principales

1. Marco Unificado TPO: Se presenta un algoritmo que trata tanto la Optimización de Trayectorias (TO) como la Optimización de Políticas (PO) bajo el mismo marco de inferencia, aprovechando la diferenciabilidad de las dinámicas.
2. TSMC con Gradientes Exactos: Se integra HMC con diferenciación automática a través de simulaciones de control, permitiendo una exploración eficiente en paisajes de costos no convexos y multimodales.
3. Mecanismos para PO: Se desarrollan estrategias específicas (aproximación determinista y espacio extendido) para manejar la intractabilidad de la expectativa en la optimización de políticas sin sacrificar la capacidad de usar gradientes.
4. Rendimiento Superior: Validación experimental que demuestra que TSMC supera a los métodos de última generación (basados en gradientes puros, MCMC directo y muestreo sin gradientes) en tareas complejas.

4. Resultados Experimentales

El método se evaluó en varios benchmarks de control clásico y robótica:

• Optimización de Trayectorias (TO):

  • Péndulo Invertido: TSMC alcanzó costos cercanos al límite inferior certificado (SDP), superando a MPPI y mostrando mayor robustez que MALA.
  • Acrobot: TSMC fue el único método capaz de lograr un "swing-up" exitoso (levantar el péndulo a la posición vertical) con un costo significativamente menor que los baselines (NUTS, MALA, MPPI, IPOPT).
  • Pusher Slider (Dinámicas no suaves): En un problema con contactos no suaves, TSMC fue el único método que alcanzó la pose objetivo, superando ampliamente a los métodos basados en gradientes puros que fallaron debido a la no suavidad.

• Optimización de Políticas (PO):

  • Se comparó contra PPO y SAC (algoritmos estándar de RL).
  • En tareas difíciles como el Acrobot y el Doble Péndulo sobre un carrito, PPO y SAC fallaron consistentemente en encontrar políticas exitosas, incluso con grandes presupuestos de entrenamiento.
  • TSMC logró encontrar políticas con costos de trayectoria significativamente menores, demostrando una capacidad superior para escapar de mínimos locales en la optimización de parámetros de redes neuronales.

• Eficiencia: Aunque TSMC es computacionalmente más costoso por iteración que métodos simples como MPPI, su capacidad para encontrar soluciones de menor costo con menos "ensayos" globales lo hace competitivo en términos de calidad de solución.

5. Significado e Impacto

• Puente entre Muestreo y Optimización: El trabajo demuestra que es posible combinar la exploración global de los métodos de muestreo con la eficiencia local de los métodos basados en gradientes, superando las limitaciones de ambos enfoques por separado.
• Aprovechamiento de la Diferenciabilidad: Refuerza la tendencia de utilizar simuladores diferenciables (física diferenciable) no solo para el cálculo de gradientes directos, sino para guiar procesos de inferencia bayesiana complejos como el muestreo de distribuciones multimodales.
• Robustez en Control No Convexo: Ofrece una solución robusta para problemas de control donde la no convexidad y los mínimos locales son barreras críticas, especialmente en tareas de contacto y sistemas subactuados complejos.
• Escalabilidad: Aunque actualmente limitado por la memoria GPU al mantener muchas partículas, el marco es inherentemente paralelizable y se beneficia de la aceleración hardware moderna.

En conclusión, el artículo presenta TSMC como una herramienta poderosa y general para el control óptimo, capaz de resolver problemas de optimización de trayectorias y políticas que son intratables para los métodos convencionales de optimización local o muestreo ciego.