socx: GPU-Accelerated Chance-Constrained Stochastic Optimal Control using JAX

Tanmay Dokania; Yashwanth Kumar Nakka

eess.SYOct 29, 2026

socx: GPU-Accelerated Chance-Constrained Stochastic Optimal Control using JAX

Authors: Tanmay Dokania, Yashwanth Kumar Nakka

Organizations: Authors are with the Daniel Guggenheim School of Aerospace Engineering at the Georgia Institute of Technology, Atlanta, GA, 30332 USA.

Abstract

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.

Figures & tables

Panel of Table 1 — TABLE I : Comparison with related work. CS: covariance steering. Problem class: Gaussian noise (Gau.), exact linear chance constraints (CC), covariance LMI (LMI), convex (Cvx). Solver structure: exploits OCP structure (OCP), time-parallel $O(\log N)$ span ( $\log N$ ), parallel within one problem (Par.), batches across problems (Bat.), GPU. ✓: yes, : partial or only some works in the group, $\times$ : no, –: not applicable.

Panel of Figure 2 — Fig. 1 : Steps of ADMM ( 28 ) in socx . In one iteration, the $\theta$ -update solves mean and stochastic columns’ LQT problems in parallel, followed by the $z$ -update, which projects every chance-constraint cone and the terminal covariance ball in parallel.

	mean channel	stochastic channel $c$
$(\chi_{k},\omega_{k})$	$(\mu_{k},\nu_{k})$	$(\varsigma^{(c)}_{k},\upsilon^{(c)}_{k})$
$\mathcal{Q}$ : $(\tilde{Q}_{k},\tilde{R}_{k})$	$(\tilde{Q}^{\mu}_{k},\tilde{R}^{\mu}_{k})$	$(\tilde{Q}^{\varsigma}_{k},\tilde{R}^{\varsigma}_{k})$
$\mathcal{L}$ : $k_{0}$	$0$	$k_{c}$ ( 26 )
$\mathcal{L}$ : $\chi_{k_{0}}$	$\mu_{0}^{\star}$	$\varsigma^{(c)}_{k_{c}}$ ( 26 )
$\mathcal{L}$ : $(\tilde{q}_{k},\tilde{r}_{k})$	$(\tilde{q}^{\mu}_{k},\tilde{r}^{\mu}_{k})$	$(\tilde{q}^{\varsigma,(c)}_{k},\tilde{r}^{\varsigma,(c)}_{k})$
count	$1$	$n_{e}(N)$

TABLE II : p-ric and p-lqt arguments with shared

(A_{k},B_{k})

Panel of Figure 4 — Fig. 2 : 3D quadrotor trajectories computed with socx , with random wrench perturbations, chance constraints on position & actuator limits and terminal covariance constraint $(n_{x}=12,n_{w}=6,n_{u}=4,N=80,\Delta t=0.1)$ .

Fig. 3 : Solve time vs. (a) tolerance at fixed

N=160

and (b) horizon at fixed

\epsilon=10^{-5}

		socx (Ours) , GPU						MOSEK, CPU			SCS+cuDSS, GPU
$N$	$n_{e}(N)$	time (s)		iterations		obj. err.	$P_{\mathrm{viol}}/\epsilon$	time (s)		speed-up	time (s)		speed-up
40	252	0.38 $\pm$	0.02	377 $\pm$	20	$2{\times}10^{-7}$	1.000	2.9 $\pm$	0.1	7.2–7.9 $\times$	26 $\pm$	11	47–104 $\times$
80	492	0.66 $\pm$	0.06	483 $\pm$	53	$10^{-7}$	1.000	14.9 $\pm$	1.5	20–27 $\times$	179 $\pm$	79	145–369 $\times$
160	972	2.20 $\pm$	0.38	750 $\pm$	135	$2{\times}10^{-5}$	1.002	91.7 $\pm$	5.5	37–53 $\times$	1976 $\pm$	2077	199– 2224 $\times$
240	1,452	4.33 $\pm$	0.30	798 $\pm$	57	$10^{-4}$	1.002	311.7 $\pm$	82.9	57– 98 $\times$	3848 $\pm$	2123	386–1342 $\times$
320	1,932	7.10 $\pm$	0.26	808 $\pm$	31	–	1.000	out of memory (50GB cap)			5789 $\pm$	1633	576–1085 $\times$

TABLE III : 3D quadrotor at tolerance

10^{-5}

: mean

\pm

max. deviation over the three setups of Figure 2 .

	time (ms)			scratch memory (MB)
Variant	$N{=}80$	$160$	$320$	$N{=}80$	$160$	$320$
$\theta$ -update : per iteration
Parallel LQT (Ours)	1.13	2.62	8.17	39	124	390
Sequential LQT	4.48	9.16	20.66	8	40	159
Condensed KKT	2.46	9.54	61.19	10	13	43
Per-column loop	7.78	30.63	120.72	188	747	3768

TABLE IV : Ablations on the helix setup at tolerance

10^{-5}

Stochastic Optimal Control

Stochastic optimal control studies how to choose actions over time when system dynamics, observations, or outcomes are uncertain.

See all 101 papers

Exact and Approximate Convex Reformulation of Linear Stochastic Optimal Control with Chance Constraints
Mar 19, 2026Tanmay Dokania, Yashwanth Kumar NakkaStochastic Optimal Control Chance-Constrained Optimization
Authors are with the Daniel Guggenheim School of Aerospace Engineering at the Georgia Institute of Technology, Atlanta, GA, 30332 USA.
Parallel-in-Time Nonlinear Optimal Control via GPU-native Sequential Convex Programming
Mar 11, 2026Yilin Zou, Zhong Zhang, Maxime Robic +1Optimal Control GPU Acceleration
School of Aerospace Engineering, Tsinghua University, 100084 Beijing, China. · Department of Aerospace Science and Technology, Politecnico di Milano, 20156 Milan, Italy.
WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization
Jul 13, 2026Henrik Hose, Se Hwan Jeon, Charles Khazoom +2GPU Kernel Optimization GPU Acceleration
Institute for Data Science in Mechanical Engineering (DSME), RWTH Aachen University, Germany · Department of Mechanical Engineering, Massachusetts Institute of Technology, Cambridge MA, USA.

Mar 19, 2026J/K move · Enter open · S save

Exact and Approximate Convex Reformulation of Linear Stochastic Optimal Control with Chance Constraints

Tanmay Dokania, Yashwanth Kumar Nakka

Authors are with the Daniel Guggenheim School of Aerospace Engineering at the Georgia Institute of Technology, Atlanta, GA, 30332 USA.

Stochastic Optimal Control Chance-Constrained Optimization Stochastic Linear Dynamical Systems

Open paper AskPDF arXiv

In this paper, we present an equivalent convex optimization formulation for discrete-time stochastic linear systems subject to linear chance constraints, alongside a tight convex relaxation for quadratic chance constraints. By lifting the state vector to encode moment information explicitly, the formulation captures linear chance constraints on states and controls across multiple time steps exactly, without conservatism, yielding strict improvements in both feasibility and optimality. For quadratic chance constraints, we derive convex approximations that are provably less conservative than existing methods. We validate the framework on minimum-snap trajectory generation for a quadrotor, demonstrating that the proposed approach remains feasible at noise levels an order of magnitude beyond the operating range of prior formulations.

Abstract

Stochastic Optimal Control

Explore similar work

Exact and Approximate Convex Reformulation of Linear Stochastic Optimal Control with Chance Constraints

Parallel-in-Time Nonlinear Optimal Control via GPU-native Sequential Convex Programming

WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled LDL⊤LDL^\topLDL⊤ Factorization

WarpMPC: Large-Batch MPC on GPU via ADMM with Unrolled $LDL^\top$ Factorization