Parallel Branch Model Predictive Control on GPUs

Technische Samenvatting: Parallelle Branch Model Predictive Control op GPU's

Probleemstelling
Branch Model Predictive Control (BMPC) is een krachtig planningsframework voor het omgaan met onzekerheid in dynamische omgevingen, zoals geautomatiseerd rijden, door trajectbomen te genereren waarbij takken (branches) overeenkomen met verschillende realisaties van onzekerheid. Echter, de brede inzetbaarheid van BMPC wordt gehinderd door de aanzienlijke computationele last die nodig is om deze problemen op te lossen, met name bij lange planningshorizonten en talrijke voorspelde scenario's. Bestaande solvers worstelen vaak met het efficiënt exploiteren van de inherente boomstructuur of falen in het bereiken van temporele parallellisme, wat hun geschiktheid voor real-time toepassingen beperkt. Bovendien blijft het afhandelen van algemene fase-gewijze restricties binnen een boomgestructureerd optimal control framework op parallelle hardware een uitdaging.

Methodologie
De auteurs stellen een GPU-gebaseerde solver voor BMPC voor die een multiple-shooting formulering integreert met een augmented Lagrangian (AL) methode voor het afhandelen van restricties. De kern van de aanpak rust op twee op maat gemaakte innerlijke Linear Quadratic Regulator (LQR) solvers die ontworpen zijn om de boom-sparse structuur te exploiteren:

1. Parallelle Tree LQR Solvers:

   • SLQR (Scenario-Level Parallelization): Deze solver voert een gemodificeerde Riccati-recursie uit van de bladknopen naar de wortel. Het aggregeert waarde-functies van kindknopen bij elke fase, waardoor onafhankelijke minimalisatieproblemen op elk knooppunt parallel kunnen worden opgelost. Deze aanpak vereist minder GPU-bronnen en is geschikt voor scenario's waarin de middelen beperkt zijn.
   • STLQR (Scenario & Temporal Parallelization): Deze solver maakt gebruik van het parallel scan-algoritme om zowel scenario-niveau als temporele parallellisme te bereiken in zowel de achterwaartse (Riccati) als voorwaartse (rollout) passes. Het maakt gebruik van Conditional Value Functions (CVF's) en een boomgestructureerde combinatieregel om waarde-functies en affiene controlewetten te berekenen in $O(\log N)$ tijdcomplexiteit. Deze methode biedt hogere parallellisme maar vereist meer GPU-bronnen.

2. Restrictieafhandeling via Augmented Lagrangian:
   Om algemene fase-gewijze restricties aan te pakken, gebruiken de auteurs een Augmented Lagrangian (AL) methode. De innerlijke lus gebruikt een iteratieve LQR (iLQR) benadering waarbij het beperkte probleem wordt benaderd als een onbeperkt tree LQR-probleem met behulp van de Powell-Hestenes-Rockafellar (PHR) penalty functie. Een lineaire rollout wordt gebruikt om optimale perturbaties te berekenen, wat efficiënte parallellisme op GPU's mogelijk maakt. De buitenste lus werkt Lagrange-multiplicatoren en penalty-gewichten adaptief bij op basis van restrictie-schendingen, volgens de BCL-regel.

3. Implementatie:
   De solver is geïmplementeerd in JAX, waarbij gebruik wordt gemaakt van automatische differentiatie en de XLA-compiler voor GPU-acceleratie. Het framework ondersteunt zowel single-precision (FP32) als double-precision (FP64) rekenkunde.

Belangrijkste Bijdragen
Het artikel schetst drie primaire bijdragen:

1. Duale Parallelle Solvers: De ontwikkeling van twee parallelle tree LQR solvers (SLQR en STLQR) die verschillende niveaus van parallellisme bieden, waardoor gebruikers de geschikte methode kunnen selecteren op basis van de probleemgrootte en beschikbare computationele bronnen.
2. Beperkte Nietlineaire BMPC Solver: De integratie van deze tree LQR solvers in een multiple-shooting iteratieve solver voor nietlineaire BMPC-problemen, inclusief een augmented Lagrangian methode voor robuuste restrictieafhandeling en warm-start mogelijkheden.
3. Benchmarking en Open Source: Een uitgebreide benchmarking van de voorgestelde solver tegen bestaande iLQR solvers (TRAJAX, MPX) en een high-performance CPU-gebaseerde solver (HPIPM), samen met de release van een open-source implementatie.

Numerieke Resultaten
De auteurs hebben de solver geëvalueerd op twee verschillende taken: onbeperkte tree LQR-problemen en beperkte trajectplanning voor een unicycle en een quad-pendulum.

• Prestaties op Tree LQR: De prestaties van de GPU-gebaseerde solvers zijn sterk afhankelijk van de probleemgrootte en hardware. Bij kleine probleemomvang (bijv. $Z = 2$ tree-paden) zijn de solvers aanzienlijk trager dan de CPU-gebaseerde HPIPM solver, waarbij STLQR meer dan 5$\times$ trager is en SLQR meer dan 20$\times$ trager op een NVIDIA RTX 5060 Ti door GPU-geheugen-access latentie en overhead. Echter, bij grootschalige instanties keert de prestatie om: SLQR kan HPIPM met wel 2$\times$ verslaan op grootschalige instanties ($Z \ge 128$) op de RTX 5060 Ti. Op vergelijkbare wijze bereikt STLQR op high-end GPU's zoals de RTX 4090 een versnelling tot 1.9$\times$ ten opzichte van HPIPM voor matige tot grote boomgroottes ($Z \ge 16$).
• Restrictieafhandeling: In traject-plannings-taken vertoonde de voorgestelde solver (ILQRJAX) een convergentiegedrag vergelijkbaar met de state-of-the-art CPU-solver IPOPT, maar met een aanzienlijk lagere berekeningstijd per iteratie (bijv. het reduceren van de gemiddelde iteratietijd van 3.80 ms naar 1.87 ms voor de unicycle). De solver handelde alle testinstanties succesvol af, terwijl andere GPU-gebaseerde solvers (TRAJAX, MPX) moeite hadden met meer uitdagende instanties en vaak niet convergeren door formuleringbeperkingen of een gebrek aan adaptieve update-schema's.

Betekenis en Claims
Het artikel claimt dat de voorgestelde aanpak een levensvatbaar pad biedt naar real-time BMPC voor grootschalige problemen door de boomstructuur volledig te exploiteren via parallelle algoritmen op GPU's. De auteurs benadrukken dat hun methode superieure prestaties levert vergeleken met high-performance CPU-gebaseerde solvers, specif