CusADi: A GPU Parallelization Framework for Symbolic Expressions and Optimal Control (2408.09662v1)

Abstract: The parallelism afforded by GPUs presents significant advantages in training controllers through reinforcement learning (RL). However, integrating model-based optimization into this process remains challenging due to the complexity of formulating and solving optimization problems across thousands of instances. In this work, we present CusADi, an extension of the CasADi symbolic framework to support the parallelization of arbitrary closed-form expressions on GPUs with CUDA. We also formulate a closed-form approximation for solving general optimal control problems, enabling large-scale parallelization and evaluation of MPC controllers. Our results show a ten-fold speedup relative to similar MPC implementation on the CPU, and we demonstrate the use of CusADi for various applications, including parallel simulation, parameter sweeps, and policy training.

• The paper presents a GPU parallelization framework that extends CasADi to accelerate symbolic computation for optimal control problems with dramatic speedups.
• It demonstrates robust integration with MPC and RL in robotics, showcasing applications on platforms like the MIT Humanoid and quadcopter simulations.
• The approach leverages vectorized CUDA kernels to mitigate CPU-GPU data transfer overhead, enabling scalable and real-time optimization in high-dimensional control tasks.

CusADi: A GPU Parallelization Framework for Symbolic Expressions and Optimal Control

This paper introduces "CusADi," an extension of the CasADi symbolic framework designed to enable large-scale parallelization of arbitrary closed-form expressions on Graphics Processing Units (GPUs) using CUDA. The primary motivation is to harness the massive computational power and parallelism of GPUs to accelerate the evaluation and optimization of complex control problems, particularly in the context of Model Predictive Control (MPC) and Reinforcement Learning (RL) applications in robotics. The research demonstrates significant computational speedups and showcases the utility of CusADi in a variety of scenarios, from parallel simulation to RL policy training.

Framework and Methodology

CasADi and Symbolic Framework: CasADi is a widely used software stack for gradient-based numerical optimization, particularly in the formulation of optimal control problems (OCPs). The symbolic expressions in CasADi can be represented as directed graphs of atomic operations, which are exploited by CusADi for vectorized parallel computation on the GPU.

GPU Parallelization: GPUs, with their SIMD (Single Instruction, Multiple Data) architecture, are well-suited for tasks involving repetitive operations on large datasets. CusADi leverages this architecture by transforming CasADi’s symbolic expressions into vectorized CUDA kernels, enabling the parallel evaluation of thousands of function instances. This framework eliminates the overhead associated with CPU-GPU data transfers by maintaining data exclusively on the GPU.

Optimization Approaches: The research utilizes both exact and approximation methods for solving optimization problems. For instance, the Sequential Quadratic Programming (SQP) method is adapted to a penalty-based approach for handling inequality constraints, which can then be solved using LDL$^T$ factorization. This enables closed-form solutions that can be efficiently parallelized on the GPU.

Results and Benchmarks

The paper provides comprehensive benchmarks comparing the performance of CusADi with serial CPU evaluation, parallel CPU evaluation using OpenMP, and PyTorch-based implementations. The results show that CusADi can achieve speedups ranging from 100x to 1000x, particularly for large batch sizes and complex functions. The benchmarks also highlight the substantial effect of data transfer overhead between the CPU and GPU, which CusADi effectively mitigates.

Applications in Robotics

MPC Parallelization for MIT Humanoid: The paper demonstrates the application of CusADi for MPC in a high-dimensional robotic system, the MIT Humanoid. By formulating a closed-form approximation to the OCP, the researchers are able to deploy MPC across thousands of environments in NVIDIA's IsaacGym, achieving a significant reduction in training iteration times (11x faster than CPU-based MPC implementations).

Reinforcement Learning Enhancement: CusADi is used to compute model-based dynamic quantities such as centroidal momentum in parallel, which can be used to augment the observations and rewards in RL environments. For example, incorporating centroidal angular momentum tracking leads to more natural robot behaviors, such as arm swinging during locomotion.

Parallelized Rollouts for Quadcopter: The framework is also applied to a planar quadcopter to perform Monte Carlo simulations and parameter sweeps. These parallelized evaluations provide insights into the region of attraction for a given controller and the effects of varying system and controller parameters on the optimal trajectory.

Implications and Future Directions

The practical implications of CusADi are substantial, particularly in the development and training of advanced control strategies for robotic systems. The ability to leverage GPU parallelism can significantly reduce the computational burden and enable real-time applications of complex controllers like MPC in RL training loops.

Future research could explore extending CusADi to exploit parallelization within individual expressions, potentially utilizing graph-theoretical approaches to identify parallelizable substructures within CasADi expression graphs. Additionally, the promising results with MPC and RL integration suggest that further work could focus on learning residual policies or value functions to enhance the effectiveness of model-based techniques in complex robotic tasks.

In conclusion, CusADi represents a significant advancement in the computational efficiency of symbolic optimization and control frameworks, with broad applicability in robotics and beyond. The integration of GPU acceleration into CasADi opens new avenues for scalable and real-time optimization in high-dimensional control problems.