Parallel Branch Model Predictive Control on GPUs

Résumé Technique : Contrôle Prédictif par Modèle à Branches Parallèle sur GPU

Énoncé du Problème
Le Contrôle Prédictif par Modèle à Branches (BMPC - Branch Model Predictive Control) est un cadre de planification puissant pour gérer l'incertitude dans les environaux dynamiques, tels que la conduite automatisée, en générant des arbres de trajectoires où les branches correspondent à différentes réalisations de l'incertitude. Cependant, le déploiement généralisé du BMPC est entravé par la charge de calcul considérable requise pour résoudre ces problèmes, particulièrement lorsqu'on traite de longs horizons de planification et de nombreux scénarios prédits. Les solveurs existants peinent souvent à exploiter efficacement la structure arborescente inhérente ou ne parviennent pas à obtenir un parallélisme temporel, ce qui limite leur adéquation aux applications en temps réel. De plus, la gestion de contraintes générales par étape au sein d'un cadre de contrôle optimal structuré en arbre reste un défi sur le matériel parallèle.

Méthodologie
Les auteurs proposent un solveur basé sur GPU pour le BMPC qui intègre une formulation de tir multiple (multiple-shooting) avec une méthode de Lagrangien augmenté (AL) pour la gestion des contraintes. Le cœur de l'approche repose sur deux solveurs LQR (régulateur linéaire quadratique) internes adaptés, conçus pour exploiter la structure creuse de l'arbre :

1. Solveurs LQR d'Arbre Parallèles :

   • SLQR (Parallélisation au Niveau des Scénarios) : Ce solveur effectue une récursion de Riccati modifiée de la racine vers les feuilles. Il agrège les fonctions de valeur des nœuds enfants à chaque étape, permettant de résoudre des problèmes de minimisation indépendants à chaque nœud en parallèle. Cette approche nécessite moins de ressources GPU et convient aux scénaux où les ressources sont limitées.
   • STLQR (Parallélisation Temporelle et de Scénarios) : Ce solveur exploite l'algorithme de balayage parallèle (parallel scan) pour obtenir un parallélisme tant au niveau des scénarios que temporel, lors des passes arrière (Riccati) et avant (rollout). Il utilise des Fonctions de Valeur Conditionnelles (CVF) et une règle de combinaison structurée en arbre pour calculer les fonctions de valeur et les lois de commande affines en une complexité temporelle de $O(\log N)$. Cette méthode offre un parallélisme plus élevé mais exige plus de ressources GPU.

2. Gestion des Contraintes via le Lagrangien Augmenté :
   Pour traiter les contraintes générales par étape, les auteurs utilisent une méthode de Lagrangien augmenté (AL). La boucle interne utilise une approche LQR itérative (iLQR) où le problème contraint est approximé par un problème LQR d'arbre non contraint en utilisant la fonction de pénalité de Powell-Hestenes-Rockafellar (PHR). Un déroulement linéaire (linear rollout) est utilisé pour calculer les perturbations optimales, permettant une parallélisation efficace sur les GPU. La boucle externe met à jour de manière adaptative les multiplicateurs de Lagrange et les poids de pénalité en fonction des violations de contraintes, suivant la règle BCL.

3. Implémentation :
   Le solveur est implémenté dans JAX, utilisant sa différenciation automatique et son compilateur XLA pour l'accélération GPU. Le cadre prend en charge l'arithmétique en simple précision (FP32) et en double précision (FP64).

Contributions Clés
Le document souligne trois contributions principales :

1. Solveurs Parallèles Duaux : Le développement de deux solveurs LQR d'arbre parallèles (SLQR et STLQR) qui offrent différents niveaux de parallélisme, permettant aux utilisateurs de choisir la méthode appropriée en fonction de la taille du problème et des ressources de calcul disponibles.
2. Solveur BMPC Non Linéaire Contraint : L'intégration de ces solveurs LQR d'arbre dans un solveur itératif de tir multiple pour les problèmes BMPC non linéaires, incorporant une méthode de Lagrangien augmenté pour une gestion robuste des contraintes et des capacités de pré-initialisation (warm-starting).
3. Benchmarking et Open Source : Un benchmarking complet du solveur proposé par rapport aux solveurs iLQR existants (TRAJAX, MPX) et à un solveur CPU haute performance (HPIPM), ainsi que la publication d'une implémentation en open-source.

Résultats Numériques
Les auteurs ont évalué le solveur sur deux tâches distinctes : des problèmes LQR d'arbre non contraints et la planification de trajectoire contrainte pour un unicycle et un pendule à quadruple pivot (quad-pendulum).

• Performance sur le LQR d'Arbre : La performance des solveurs basés sur GPU dépend fortement de la taille du problème et du matériel. Sur de petites tailles de problèmes (ex: $Z = 2$ chemins d'arbre), les solveurs sont nettement plus lents que le solveur CPU HPIPM, le STLQR étant plus de 5$\times$ plus lent et le SLQR plus de 20$\times$ plus lent sur un NVIDIA RTX 5060 Ti en raison de la latence d'accès à la mémoire GPU et des surcoûts (overhead). Cependant, sur des instances à grande échelle, la performance s'inverse : le SLQR peut surpasser HPIPM jusqu'à 2$\times$ sur des instances de grande taille ($Z \ge 128$) sur l'RTX 5060 Ti. De même, sur des GPU haut de gamme comme l'RTX 4090, le STLQR atteint une accélération allant jusqu'à 1,9$\times$ par rapport à HPIPM pour des tailles d'arbres modérées à grandes ($Z \ge 16$).
• Gestion des Contraintes : Dans les tâches de planification de trajectoire, le solveur proposé (ILGRJAX) a démontré un comportement de convergence comparable au solveur CPU de pointe IPOPT, mais avec un temps de calcul par itération nettement réduit (réduisant par exemple le temps moyen d'itération de 3,80 ms à 1,87 ms pour l'unicycle). Le solveur a géré avec succès toutes les instances de test, là où d'autres solveurs basés sur GPU (TRAJAX, MPX) ont éprouvé des difficultés, échouant souvent à converger en raison de limitations de formulation ou de l'absence de schémas de mise à jour adaptatifs.

Signification et Revendications
Le papier affirme que l'approche proposée offre une voie viable vers le BMPC à grande échelle en exploitant pleinement la structure d'arbre via des algorithmes parallèles sur GPU. Les auteurs soulignent que leur méthode atteint des performances supérieures aux solveurs CPU haute performance spécifiquement sur les instances à grande échelle où la structure de l'arbre peut être parallélisée efficacement. Ils reconnaissent toutefois que le solveur basé sur le balayage parallèle présente des exigences élevées en ressources GPU, ce qui peut limiter la scalabilité si les ressources sont saturées, et que pour les petits problèmes, les solveurs CPU peuvent encore être plus performants. Ce travail se positionne comme une étape vers la rendre la planification consciente de l'incertitude réalisable pour des applications complexes du monde réel en équilibrant l'efficacité de calcul avec la gestion rigoureuse des contraintes et de l'incertitude. Les travaux futurs identifient l'implémentation de la méthode en CUDA C++ pour optimiser davantage l'utilisation des ressources et potentiellement explorer l'arithmétique en précision mixte pour améliorer la stabilité numérique sur le matériel optimisé pour le FP32.