Parallel Branch Model Predictive Control on GPUs

Resumen Técnico: Control Predictivo de Modelo de Rama Paralela en GPUs

Planteamiento del Problema
El Control Predictivo de Modelo de Rama (BMPC, por sus siglas en inglés) es un marco de planificación potente para manejar la incertidumbre en entornos dinámicos, como la conducción automatizada, mediante la generación de árboles de trayectoria donde las ramas corresponden a diferentes realizaciones de la incertidumbre. Sin embargo, el despliegue generalizado del BMPC se ve obstaculizado por la significativa carga computacional requerida para resolver estos problemas, particularmente cuando se trata de horizontes de planificación largos y numerosos escenarios predichos. Los solvers existentes a menudo tienen dificultades para explotar la estructura de árbol inherente de manera eficiente o no logran alcanzar el paralelismo temporal, lo que limita su idoneidad para aplicaciones en tiempo real. Además, el manejo de restricciones generales por etapa dentro de un marco de control óptimo estructurado en árbol sobre hardware paralelo sigue siendo un desafío.

Metodología
Los autores proponen un solver basado en GPU para BMPC que integra una formulación de tiro múltiple (multiple-shooting) con un método de Lagrangiano Aumentado (AL) para el manejo de restricciones. El núcleo del enfoque se basa en dos solvers de Regulador Cuadrático Lineal (LQR) internos diseñados específicamente para explotar la estructura dispersa de árbol:

1. Solvers LQR de Árbol Paralelos:

   • SLQR (Paralelismo a Nivel de Escenario): Este solver realiza una recursión de Riccati modificada desde los nodos hoja hacia la raíz. Agrega funciones de valor de los nodos hijos en cada etapa, lo que permite resolver problemas de minimización independientes en cada nodo de forma paralela. Este enfoque requiere menos recursos de GPU y es adecuado para escenarios donde los recursos son limitados.
   • STLQR (Paralelismo Temporal y de Escenario): Este solver aprovecha el algoritmo de escaneo paralelo (parallel scan) para lograr paralelismo tanto a nivel de escenario como temporal en las pasadas hacia atrás (Riccati) y hacia adelante (rollout). Utiliza Funciones de Valor Condicionales (CVFs) y una regla de combinación estructurada en árbol para computar funciones de valor y leyes de control afines en una complejidad de tiempo de $O(\log N)$. Este método ofrece un mayor paralelismo pero demanda más recursos de GPU.

2. Manejo de Restricciones mediante Lagrangiano Aumentado:
   Para abordar restricciones generales por etapa, los autores emplean un método de Lagrangiano Aumentado (AL). El bucle interno utiliza un enfoque de LQR iterativo (iLQR) donde el problema restringido se aproxima como un problema LQR de árbol no restringido utilizando la función de penalización de Powell-Hestenes-Rockafellar (PHR). Se utiliza un rollout lineal para calcular las perturbaciones óptimas, permitiendo una paralelización eficiente en GPUs. El bucle externo actualiza adaptativamente los multiplicadores de Lagrange y los pesos de penalización basándose en las violaciones de las restricciones, siguiendo la regla BCL.

3. Implementación:
   El solver está implementado en JAX, utilizando su diferenciación automática y su compilador XLA para la aceleración en GPU. El marco soporta aritmética de precisión simple (FP32) y de doble precisión (FP64).

Contribuciones Clave
El artículo describe tres contribuciones primarias:

1. Solvers Duales Paralelos: El desarrollo de dos solvers LQR de árbol paralelos (SLQR y STLQR) que ofrecen diferentes niveles de paralelismo, permitiendo a los usuarios seleccionar el método apropiado según el tamaño del problema y los recursos computacionales disponibles.
2. Solver BMPC No Lineal con Restricciones: La integración de estos solvers LQR de árbol en un solver iterativo de tiro múltiple para problemas BMPC no lineales, incorporando un método de Lagrangiano Aumentado para un manejo robusto de restricciones y capacidades de warm-start.
3. Benchmarking y Código Abierto: Una evaluación exhaustiva del rendimiento del solver propuesto frente a solvers iLQR existentes (TRAJAX, MPX) y un solver de alto rendimiento basado en CPU (HPIPM), junto con la liberación de una implementación de código abierto.

Resultados Numéricos
Los autores evaluaron el solver en dos tareas distintas: problemas LQR de árbol sin restricciones y planificación de trayectoria restringida para un monociclo y un péndulo-cuádruple.

• Rendimiento en Tree LQR: El rendimiento de los solvers basados en GPU depende altamente del tamaño del problema y del hardware. En tamaños de problema pequeños (por ejemplo, $Z = 2$ rutas de árbol), los solvers son significativamente más lentos que el solver basado en CPU HPIPM, siendo STLQR más de 5 veces más lento y SLQR más de 20 veces más lento en una NVIDIA RTX 5060 Ti debido a la latencia de acceso a la memoria de la GPU y la sobrecarga (overhead). Sin embargo, en instancias de gran escala, el rendimiento se invierte: SLQR puede superar a HPIPM hasta en un 2$\times$ en instancias de gran escala ($Z \ge 128$) en la RTX 5060 Ti. Del mismo modo, en GPUs de gama alta como la RTX 4090, STLQR logra una aceleración de hasta 1.9$\times$ sobre HPIPM para tamaños de árbol de moderados a grandes ($Z \ge 16$).
• Manejo de Restricciones: En las tareas de planificación de trayectoria, el solver propuesto (ILQRJAX) demostró un comportamiento de convergencia comparable al solver de vanguardia basado en CPU, IPOPT, pero con un tiempo de computación por iteración significativamente reducido (por ejemplo, reduciendo el tiempo promedio de iteración de 3.80 ms a 1.87 ms para el monociclo). El solver manejó con éxito todas las instancias de prueba, mientras que otros solvers basados en GPU (TRAJAX, MPX) tuvieron dificultades con las instancias más desafiantes, fallando a menudo en la convergencia debido a limitaciones de formulación o falta de esquemas de actualización adaptativa.

Significancia y Reivindicaciones
El artículo afirma que el enfoque propuesto ofrece un camino viable hacia el BMPC en tiempo real para problemas de gran escala mediante el aprovechamiento total de la estructura de árbol a través de algoritmos paralelos en GPUs. Los autores enfatizan que su método logra un rendimiento superior comparado con los solvers de alto rendimiento basados en CPU específicamente en instancias de gran escala donde la estructura de árbol puede ser paralelizada efectivamente. Reconocen, sin embargo, que el solver basado en parallel scan tiene altas demandas de recursos de GPU, lo que puede limitar la escalabilidad si los recursos se saturan, y que para tamaños de problema pequeños, los solvers basados en CPU pueden seguir siendo superiores. El trabajo se posiciona como un paso hacia la realización de la planificación consciente de la incertidumbre para aplicaciones complejas del mundo real, equilibrando la eficiencia computacional con el manejo riguroso de restricciones e incertidumbre. Se identifica como trabajo futuro la implementación del método en CUDA C++ para optimizar aún más la utilización de recursos y explorar potencialmente la aritmética de precisión mixta para mejorar la estabilidad numérica en hardware optimizado para FP32.