Search arXiv⌕ Search
Home/Archive/arXiv

arXiv · 2610.40427

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

Tanmay Dokania·Yashwanth Kumar Nakka
View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

eess.SYcs.ROcs.SY

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Tanmay Dokania, Yashwanth Kumar Nakka. 2026-10-29. socx: GPU-Accelerated Chance-Constrained Stochastic Optimal Control using JAX. https://arxiv.org/abs/2610.40427

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Robust Taylor-Lagrange Control for Safety-Critical Systems with Guaranteed Feasibility

The Control Barrier Function (CBF) method has been widely adopted in solving safety-critical control problems. However, the existence of a CBF is only a sufficient condition for system safety. The recently proposed Taylor-Lagrange Control (TLC) method addresses this limitation by providing a necessary and sufficient condition, but it is still vulnerable to the feasibility preservation problem (e.g., due to the inter-sampling effect when time is discretized). In this paper, we propose a robust TLC (rTLC) method to address the feasibility preservation problem and show how to use it to solve optimal control problems. Specifically, the rTLC method expands the safety function to an order higher than its relative degree using Taylor's expansion with Lagrange remainder, which allows the control to explicitly show up at the current time instead of a future time in the original TLC method. This implies that the safety constraint remains satisfied over the entire discretization time interval that follows. Moreover, the rTLC method naturally addresses the feasibility preservation problem with only a single hyper-parameter (the discretization time interval size during implementation), compared to significantly more parameters for its counterparts. Finally, we illustrate the effectiveness of the proposed rTLC method through an adaptive cruise control problem, and compare it with existing safety-critical control methods.

eess.SY↗
arXiv2026

Reduced-Order Data Assimilation for Thermospheric Density Using Physics-informed SINDyc Models

Orbit prediction and space situational awareness require accurate thermospheric mass density, which responds nonlinearly to solar and geomagnetic forcing across several orders of magnitude. Physics-based general circulation models resolve that response but are computationally expensive, while empirical models run cheaply and carry no time-evolving atmospheric state. An autoregressive Sparse Identification of Nonlinear Dynamics with control (SINDyc-AR) model, derived from the Thermosphere-Ionosphere-Electrodynamics General Circulation Model (TIE-GCM), captures the dominant modes of variability and their dependence on the drivers at a fraction of the parent model's cost. An extended Kalman filter assimilates in situ observations from CHAMP, GRACE, GRACE-FO, GOCE, and Swarm into that model across several orbital configurations and geomagnetic conditions, with a linear DMDc model as reference. Assimilation reduces density error relative to open-loop prediction, most visibly during geomagnetic storms and under single-satellite coverage. The two reduced-order models perform comparably throughout. NRLMSIS 2.1 and HASDM can outperform the assimilated model far from the assimilated track, so results are framed as improvements over the open-loop forecast. An observation-calibrated global density record spanning August 2000 to December 2025 is released alongside this work.

eess.SY↗
arXiv2026

Executor-Side Progressive Risk-Gated Actuation for Agentic AI in Wireless Supervisory Control

Agentic artificial intelligence built on large language models shows promise for automating Open Radio Access Network supervisory control. However, action-level safeguards alone may be insufficient: evidence acquisition consumes time and shared resources, and network changes can invalidate approvals before release. We propose Progressive Risk-Gated Actuation (PRGA), an executor-side model and contract governing intents from assessment to release or termination. PRGA combines local checks with selective retrieval of coordination evidence to reduce acquisition costs while preserving mandatory evidence requirements. Version and lease rechecks prevent release under invalidated approvals, while reevaluation requires recorded progress within each transaction's retry budget. Building on this contract, we formulate evidence allocation to maximize the expected whole-period rate of timely valid releases under risk, quality-of-service, deadline, and shared byte-budget constraints. For energy-saving and slice-protection workloads with stable daily load cycles, thresholds are selected offline and interpolated using an online load-phase estimate. On average over modeled start phases and quota levels, the allocation policy improves timely valid release rates over arrival-order fetching and an offline-selected global threshold under identical quotas and release checks. In OpenAirInterface radio emulation, an immediate scripted resend after a version-mismatch refusal raises on-time receiver delivery from 77.3% to 89.9% at modest additional command cost.

eess.SY↗