Parallel Branch Model Predictive Control on GPUs

Resumo Técnico: Controle Preditivo de Modelo de Ramificação Paralela em GPUs

Definição do Problema
O Controle Preditivo de Modelo de Ramificação (BMPC - Branch Model Predictive Control) é uma estrutura de planejamento poderosa para lidar com a incerteza em ambientes dinâmicos, como a condução automatizada, ao gerar árvores de trajetórias onde as ramificações correspondem a diferentes realizações de incerteza. No entanto, a implementação generalizada do BMPC é dificultada pelo pesado fardo computacional exigido para resolver esses problemas, particularmente ao lidar com horizontes de planejamento longos e numerosos cenários previstos. Os resolvedores existentes frequentemente têm dificuldade em explorar eficientemente a estrutura de árvore inerente ou falham em alcançar o paralelismo temporal, limitando sua adequação para aplicações em tempo real. Além disso, lidar com restrições por estágio gerais dentro de uma estrutura de controle ótimo em árvore em hardware paralelo permanece um desafio.

Metodologia
Os autores propõem um resolvedor baseado em GPU para BMPC que integra uma formulação de multiple-shooting com um método de Lagrangiana Aumentada (AL) para o tratamento de restrições. O núcleo da abordagem baseia-se em dois resolvedores internos de Regulador Quadrático Linear (LQR) adaptados para explorar a estrutura esparsa de árvore:

1. Resolvedores LQR de Árvore Paralelos:

   • SLQR (Paralelismo ao Nível de Cenário): Este resolvedor realiza uma recursão de Riccati modificada dos nós folha para a raiz. Ele agrega funções de valor dos nós filhos em cada estágio, permitindo que problemas de minimização independentes em cada nó sejam resolvidos em paralelo. Esta abordagem requer menos recursos de GPU e é adequada para cenários onde os recursos são limitados.
   • STLQR (Paralelismo de Cenário e Temporal): Este resolvedor utiliza o algoritmo de parallel scan para alcançar tanto o paralelismo ao nível de cenário quanto o temporal, tanto nas passagens de retrocesso (Riccati) quanto de avanço (rollout). Ele utiliza Funções de Valor Condicionais (CVFs) e uma regra de combinação estruturada em árvore para computar funções de valor e leis de controle afins em complexidade de tempo $O(\log N)$. Este método oferece maior paralelismo, mas demanda mais recursos de GPU.

2. Tratamento de Restrições via Lagrangiana Aumentada:
   Para abordar restrições por estágio gerais, os autores empregam um método de Lagrangiana Aumentada (AL). O loop interno utiliza uma abordagem de LQR iterativo (iLQR), onde o problema com restrições é aproximado como um problema LQR de árvore sem restrições usando a função de penalidade de Powell-Hestenes-Rockafellar (PHR). Um rollout linear é usado para computar perturbações ótimas, permitindo a paralelização eficiente em GPUs. O loop externo atualiza adaptativamente os multiplicadores de Lagrange e os pesos de penalidade com base nas violações de restrição, seguindo a regra BCL.

3. Implementação:
   O resolvedor é implementado em JAX, utilizando sua diferenciação automática e o compilador XLA para aceleração em GPU. O framework suporta aritmética de precisão simples (FP32) e dupla (FP64).

Principais Contribuições
O artigo descreve três contribuições primárias:

1. Dualidade de Resolvedores Paralelos: O desenvolvimento de dois resolvedores LQR de árvore paralelos (SLQR e STLQR) que oferecem diferentes níveis de paralelismo, permitindo que os usuários selecionem o método apropriado com base no tamanho do problema e nos recursos computacionais disponíveis.
2. Resolvedor BMPC Não Linear com Restrições: A integração destes resolvedores LQR de árvore em um resolvedor de multiple-shooting iterativo para problemas BMPC não lineares, incorporando um método de Lagrangiana aumentada para um tratamento robusto de restrições e capacidades de warm-start.
3. Benchmarking e Código Aberto: Um benchmarking abrangente do resolvedor proposto contra resolvedores iLQR existentes (TRAJAX, MPX) e um resolvedor baseado em CPU de alto desempenho (HPIPM), juntamente com o lançamento de uma implementação de código aberto.

Resultados Numéricos
Os autores avaliaram o resolvedor em duas tarefas distintas: problemas LQR de árvore sem restrições e planejamento de trajetória restrito para um uniciclo e um quadripêndulo.

• Desempenho em Árvore LQR: O desempenho dos resolvedores baseados em GPU é altamente dependente do tamanho do problema e do hardware. Em tamanhos de problema pequenos (ex: $Z = 2$ caminhos de árvore), os resolvedores são significativamente mais lentos que o resolvedor baseado em CPU HPIPM, com o STLQR sendo mais de 5$\times$ mais lento e o SLQR mais de 20$\times$ mais lento em uma NVIDIA RTX 5060 Ti devido à latência de acesso à memória da GPU e overhead. No entanto, em instâncias de grande escala, o desempenho se inverte: o SLQR pode superar o HPIPM em até 2$\times$ em instâncias de grande escala ($Z \ge 128$) na RTX 5060 Ti. Da mesma forma, em GPUs de alto desempenho como a RTX 4090, o STLQR alcança um aumento de velocidade de até 1.9$\times$ sobre o HPIPM para tamanhos de árvore de moderados a grandes ($Z \ge 16$).
• Tratamento de Restrições: Nas tarefas de planejamento de trajetória, o resolvedor proposto (ILQRJAX) demonstrou comportamento de convergência comparável ao resolvedor de estado da arte em CPU, IPOPT, mas com tempo de computação por iteração significativamente reduzido (ex: reduzindo o tempo médio de iteração de 3.80 ms para 1.87 ms para o uniciclo). O resolvedor lidou com sucesso com todas as instâncias de teste, enquanto outros resolvedores baseados em GPU (TRAJAX, MPX) tiveram dificuldades com instâncias mais desafiadoras, frequentemente falhando em convergir devido a limitações de formulação ou falta de esquemas de atualização adaptativos.

Significância e Alegações
O artigo afirma que a abordagem proposta oferece um caminho viável para o BMPC de grande escala em tempo real, explorando totalmente a estrutura de árvore através de algoritmos paralelos em GPUs. Os autores enfatizam que seu método alcança desempenho superior em comparação com resolvedores de CPU de alto desempenho especificamente em instâncias de grande escala, onde a estrutura de árvore pode ser efetivamente paralelizada. Eles reconhecem, contudo, que o resolvedor baseado em parallel scan possui altas demandas de recursos de GPU, o que pode limitar a escalabilidade se os recursos forem saturados, e que para tamanhos de problema pequenos, os resolvedores baseados em CPU ainda podem superar o proposto. O trabalho posiciona-se como um passo para tornar o planejamento consciente da incerteza viável para aplicações complexas do mundo real, equilibrando eficiência computacional com o tratamento rigoroso de restrições e incertezas. O trabalho futuro identifica a implementação do método em CUDA C++ para otimizar ainda mais a utilização de recursos e, potencialmente, explorar aritmética de precisão mista para melhorar a estabilidade numérica em hardware otimizado para FP32.