Parallel Branch Model Predictive Control on GPUs

技術要約：GPU上での並列分岐モデル予測制御

問題提起
分岐モデル予測制御（BMPC）は、自動運転などの動的な環境において、不確実性の実現（realizations）に対応する軌跡ツリーを生成することで、不確実性を扱うための強力なプランニングフレームワークである。しかし、特に長いプランニングホライゾンや多数の予測シナリオを扱う場合、BMPCの広範な展開は、膨大な計算負荷によって阻害されている。既存のソルバーは、多くの場合、ツリー構造を効率的に活用できていなかったり、時間的並列性を達成できなかったりするため、リアルタイムアプリケーションへの適合性に限界がある。さらに、ツリー構造を持つ最適制御フレームワーク内で、一般的なステージごとの制約を並列ハードウェア上で扱うことは依然として課題である。

手法
著者らは、制約処理のために拡張ラグランジュ（AL）法を用いたマルチシューティング定式化を統合した、BMPC用のGPUベースのソルバーを提案している。このアプローチの中核は、ツリーのスパース構造を活用するように設計された2つの特化した内部線形二次レギュレータ（LQR）ソルバーにある。

1. 並列ツリーLQRソルバー:

   • SLQR (シナリオレベルの並列化): このソルバーは、リーフノードからルートノードに向かって修正リッカチ再帰を実行する。各ステージにおいて子ノードからの価値関数を集約することで、各ノードにおける独立した最小化問題を並列に解くことを可能にする。この手法は、より少ないGPUリソースを必要とし、リソースが限られているシナリオに適している。
   • STLQR (シナリオおよび時間的並列化): このソルバーは、パラレルスキャンアルゴリズムを利用して、後方パス（リッカチ）および前方パス（ロールアウト）の両方において、シナリオレベルおよび時間的並列性を達成する。条件付き価値関数（CVF）とツリー構造の結合規則を用いて、$O(\log N)$ の時間計算量で価値関数とアフィン制御則を計算する。この手法は、より高い並列性を提供するが、より多くのGPUリソースを要求する。

2. 拡張ラグランジュによる制約処理:
   一般的なステージごとの制約に対処するため、著者らは拡張ラグランジュ（AL）法を採用している。内部ループでは、制約付き問題を制約なしのツリーLQR問題として近似する反復LQR（iLQR）アプローチを用い、Powell-Hestenes-Rockafellar（PHR）ペナルティ関数を使用する。最適な摂動を計算するために線形ロールアウトが使用され、これによりGPU上での効率的な並列化が可能となる。外部ループでは、BCLルールに従い、制約違反に基づいてラグランジュ乗数とペナルティ重みを適応的に更新する。

3. 実装:
   ソルバーはJAXで実装されており、自動微分とXLAコンパイラを利用してGPU加速を実現している。本フレームワークは、単精度（FP32）および倍精度（FP64）演算の両方をサポートしている。

主な貢献
本論文は、主に3つの貢献を述べている。

1. デュアル並列ソルバーの開発: 異なるレベルの並列性を提供する2つの並列ツリーLQRソルバー（SLQRおよびSTLQR）を開発し、ユーザーが問題の規模と利用可能な計算リソースに基づいて適切な手法を選択できるようにした。
2. 制約付き非線形BMPCソルバー: これらのツリーLQRソルバーを、非線形BMPC問題のためのマルチシューティング反復ソルバーへと統合し、堅牢な制約処理とウォームスタート機能のための拡張ラグランジュ法を組み込んだ。
3. ベンチマークとオープンソース化: 提案されたソルバーを、既存のiLQRソルバー（TRAJAX, MPX）および高性能なCPUベースのソルバー（HPIPM）と比較した包括的なベンチマークを実施し、オープンソースの実装を公開した。

数値結果
著者らは、無制約のツリーLQR問題と、ユニサイクルおよびクアッド・ペンデュラム（四重振り子）の制約付き軌跡計画という2つの異なるタスクでソルバーを評価した。

• Tree LQRにおける性能: GPUベースのソルバーの性能は、問題のサイズとハードウェアに強く依存する。小規模な問題（例：$Z = 2$ のツリーパス）では、GPUのメモリ・アクセス・レイテンシとオーバーヘッドにより、STLQRはHPIPMより5倍以上、SLQRは20倍以上遅く、CPUベースのHPIPMソルバーよりも大幅に遅い。しかし、大規模なインスタンスでは、この関係が逆転する。RTX 5060 Tiにおいて、大規模なインスタンス（$Z \ge 128$）では、SLQRはHPIPMを最大2倍上回る性能を示す。同様に、RTX 4090のようなハイエンドGPUでは、中規模から大規模なツリーサイズ（$Z \ge 16$）において、STLQRはHPIPMに対して最大1.9倍のスピードアップを達成する。
• 制約処理: 軌跡計画タスクにおいて、提案されたソルバー（ILCLRJAX）は、最先端のCPUソルバーであるIPOPTと同等の収束挙動を示したが、1イテレーションあたりの計算時間は大幅に短縮された（例：ユニサイクルの場合、平均イテレーション時間を3.80 msから1.87 msに短縮）。本ソルバーはすべてのテストインスタンスを正常に処理したが、他のGPUベースのソルバー（TRAJAX, MPX）は、定式化の制限や適応的な更新スキームの欠如により、より困難なインスタンスで収束に失敗することが多かった。

意義と主張
本論文は、提案されたアプローチが、GPU上の並列アルゴリズムを通じてツリー構造を完全に活用することにより、大規模な問題に対するリアルタイムBMPCへの実行可能な道筋を提供すると主張している。著者らは、大規模なインスタンスにおいて、提案手法が（ツリー構造を効果的に並列化できるため）高性能なCPUベースのソルバーに対して優れた性能を発揮することを強調している。ただし、並列スキャンベースのソルバーはGPUリソースの要求が高く、リソースが飽和した場合にスケーラビリティを制限する可能性があること、また小規模な問題サイズでは依然としてCPUベースのソルバーが優れている可能性があることを認めている。本研究は、計算効率と制約および不確実性の厳密な取り扱いのバランスを取ることで、複雑で現実世界のアプリケーションにおける不確実性を考慮したプランニングを可能にするためのステップとして位置づけられている。今後の課題としては、リソース利用率をさらに最適化するためのCUDA C++での実装や、FP32に最適化されたハードウェア上での数値的安定性を向上させるための混合精度演算の探索が挙げられている。