socx：基于 JAX 的 GPU 加速机会约束随机最优控制求解器
socx: GPU-Accelerated Chance-Constrained Stochastic Optimal Control using JAX
Tanmay Dokania, Yashwanth Kumar Nakka
AI 速览：作者提出并实现 JAX 版 GPU 加速求解器 socx，用于线性机会约束随机最优控制的提升凸重构，并在四旋翼悬停线性化模型上报告了相对 MOSEK 与 SCS 的显著加速。
In this article, we present socx, a GPU-accelerated solver for lifted convex reformulations of linear stochastic optimal control under chance constraints, implemented in JAX. We demonstrate that by exploiting the sparse and independent structure of the problem, distributed optimization can reduce it to simpler problems: independent linear-quadratic tracking (LQT) problems and Euclidean projections onto sets. These simpler problems can be solved in parallel and then coupled using the constraints. To validate our solver, we run it on a hover-linearized 3D quadrotor and show a speed-up of up to $98\times$ over MOSEK and up to $2224\times$ over SCS. socx also solves horizons up to $N=1120$, whereas MOSEK runs out of memory at $N=320$. We present ablations to isolate the key components that enable the speed-up.