OptimalControl.jl: GPU-Enabled Optimal Control DSL

• OptimalControl.jl is a high-level Julia DSL that defines continuous-time optimal control problems in Bolza form and transcribes them directly into structured sparse nonlinear programs.
• It discretizes dynamic systems into banded Jacobians and block-sparse Hessians, enabling efficient GPU-accelerated evaluation and optimization using ExaModels.jl and MadNLP.jl.
• The workflow achieves end-to-end speed-ups of 2×–5× for large-scale problems by leveraging GPU kernels for parallel computation over time stages.

OptimalControl.jl is a high-level Julia DSL for defining continuous-time optimal control problems almost as one would write them on paper and compiling them, via direct transcription, into structured sparse nonlinear programs suitable for GPU execution. In the workflow described for large-scale sparse nonlinear optimal control problems, continuous-time dynamics are modeled in Julia, discretized with OptimalControl.jl, compiled into GPU kernels through ExaModels.jl, and solved entirely on GPU with the interior-point solver MadNLP.jl and the GPU sparse linear solver cuDSS (Montoison et al., 4 Oct 2025).

1. Continuous-time problem class and transcription

OptimalControl.jl is presented for continuous-time optimal control problems in Bolza form. The continuous problem is written as

$$\min_{u(\cdot)}\;g\bigl(x(0),x(t_f)\bigr)\;+\;\int_0^{t_f}f^0\bigl(x(t),u(t)\bigr)\,\mathrm{d}t$$

subject to

$$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$

The discretization introduces a grid

$0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$

with stage variables

$X_i\approx x(t_i),\quad U_i\approx u(t_i).$

A representative transcription enforces, for example, forward-Euler dynamics,

$$X_{i+1}-X_i - h_i\,f(X_i,U_i)=0,\quad i=0,\dots,N-1,$$

and approximates the cost by

$$g(X_0,X_N)\;+\;\sum_{i=0}^{N-1}h_i\,f^0(X_i,U_i)\to\min.$$

Boundary and path constraints are discretized similarly as

$b(X_0,X_N)\le0,\quad c(X_i,U_i)\le0\;\forall i.$

The resulting finite-dimensional problem is a structured sparse NLP in

$\{X_0,\dots,X_N,\;U_0,\dots,U_{N-1}\},$

with block-banded Jacobian and Hessian. This structure is central to the package’s computational model, because the sparsity pattern is inherited directly from the stage-wise dependence of the dynamics and constraints (Montoison et al., 4 Oct 2025).

2. Modeling language and user-facing workflow

The package exposes a Julia DSL intended to let an OCP be specified nearly in continuous form. A canonical declaration is

$$X_{i+1}-X_i - h_i\,f(X_i,U_i)=0,\quad i=0,\dots,N-1,$$3

This syntax expresses the time domain, state and control spaces, endpoint conditions, differential equations, and integral objective within a single declarative block. The package then parses this DSL and compiles it into a sparse nonlinear program through direct transcription.

A minimal GPU-oriented workflow uses the same model declaration and then explicitly discretizes and instantiates the resulting model:

$$X_{i+1}-X_i - h_i\,f(X_i,U_i)=0,\quad i=0,\dots,N-1,$$4

The use of scheme=:trapeze, grid_size=500, and backend=CUDABackend() makes explicit that discretization and execution backend selection are part of the front-end API. The package’s stated design therefore combines an Expr-level DSL with a direct route to GPU-resident NLP evaluation and solution (Montoison et al., 4 Oct 2025).

3. Sparse structure and GPU kernel generation

The sparsity pattern of the direct transcription is stage-local. Because each dynamic constraint at stage $i$ depends only on $(X_i,U_i,X_{i+1})$, the constraint Jacobian has banded form,

$$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$0

where $$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$1 and $$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$2. The Hessian of the Lagrangian is described as block-sparse with the same “two-block” bandwidth.

OptimalControl.jl hands the transcription to ExaModels.jl. In that handoff, ExaModels.jl preserves the stage-wise sparsity pattern, generates SIMD-friendly Julia kernels via KernelAbstractions.jl, and evaluates all $$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$3 together with their first and second derivatives in parallel over $$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$4. The generated GPU code computes the objective, constraints, gradient, Jacobian and Hessian in a single pass over the grid.

This compilation strategy is significant because it ties the mathematical structure of direct transcription to the execution model of the GPU. The package does not merely offload dense algebra; it preserves stage-wise sparsity and exploits SIMD parallelism over time stages for residual and derivative evaluation (Montoison et al., 4 Oct 2025).

4. Solver stack and device-resident optimization

Once compiled to an ExaModels.jl model, the optimization phase is handled by MadNLP.jl, described as a GPU-aware interior-point solver. In this workflow, MadNLP.jl calls the GPU kernels for residuals and derivatives, assembles the sparse KKT system in device memory, performs a one-time symbolic factorization on the host to determine fill-reducing ordering, and uses CUDSS.jl, a Julia wrapper for NVIDIA’s cuDSS, to factor and solve the KKT system entirely on the GPU. The solver then iterates, with inertia correction, regularization, and optional refinement, until convergence (Montoison et al., 4 Oct 2025).

A technical nuance is that “fully GPU-native” does not mean that every preprocessing step is device-resident: the symbolic factorization is explicitly performed on the host once, while numerical factorization and linear solves are delegated to cuDSS on the GPU. This distinction matters for interpreting the architecture. The front end, derivative evaluation, sparse KKT assembly, and numerical linear algebra are GPU-centered, but the ordering phase remains host-side.

The resulting stack can be summarized as a layered pipeline: OptimalControl.jl provides the continuous-time modeling DSL and transcription; ExaModels.jl provides sparse model instantiation and GPU kernel generation; MadNLP.jl provides the interior-point algorithm; and cuDSS, through CUDSS.jl, provides sparse linear system factorization and solution on NVIDIA hardware.

5. Empirical scaling on large discretizations

Benchmarks on NVIDIA A100 and H100 are reported to show consistent speed-ups once $$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$5 grows (Montoison et al., 4 Oct 2025).

Problem	Hardware	Reported behavior
Goddard rocket (≈$$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$6 total dim)	A100	GPU overtakes CPU at $$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$7; speed-up ∼2× for large $$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$8
Goddard rocket (≈$$\dot x(t)=f\bigl(x(t),u(t)\bigr),\quad b\bigl(x(0),x(t_f)\bigr)\le0,\quad c\bigl(x(t),u(t)\bigr)\le0.$$9 total dim)	H100	GPU wins for $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$0; similar ∼2× gain
Quadrotor (≈$0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$1 total dim)	A100	GPU faster beyond $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$2; up to ∼5× speed-up at large $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$3
Quadrotor (≈$0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$4 total dim)	H100	Crossover at $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$5; ∼5× gain

For the Goddard rocket case, the A100 benchmark reports that at $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$6, CPU ≈21.5 s versus GPU ≈8.75 s, and the H100 benchmark reports the largest run at $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$7 with approximately $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$8 million dimensions in approximately $0=t_0<t_1<\cdots<t_N=t_f,\quad h_i=t_{i+1}-t_i,$9 s. For the quadrotor case, the A100 benchmark reports that at $X_i\approx x(t_i),\quad U_i\approx u(t_i).$0, CPU ≈48.8 s versus GPU ≈8.3 s, and the H100 benchmark reports the largest run at $X_i\approx x(t_i),\quad U_i\approx u(t_i).$1, approximately $X_i\approx x(t_i),\quad U_i\approx u(t_i).$2 million dimensions, in approximately $X_i\approx x(t_i),\quad U_i\approx u(t_i).$3 s.

These results delimit the regime in which the package is intended to operate. The cited speed-ups are end-to-end speed-ups of $X_i\approx x(t_i),\quad U_i\approx u(t_i).$4–$X_i\approx x(t_i),\quad U_i\approx u(t_i).$5 on large grids with millions of variables, rather than uniform gains at all discretization sizes. The crossover points differ by problem class and hardware, which indicates that both problem structure and device characteristics materially affect performance.

6. Strengths, caveats, and recurrent points of confusion

The package’s stated strengths are tightly coupled to its architecture. The advantages listed for the workflow are: Expr-level DSL; fully GPU-native execution from modeling through differentiation to linear solves; SIMD over time stages for derivatives and residuals; structured sparsity leading to small memory footprint and fast factorization on GPUs; end-to-end speed-ups of $X_i\approx x(t_i),\quad U_i\approx u(t_i).$6–$X_i\approx x(t_i),\quad U_i\approx u(t_i).$7 on large grids; and portability within Julia’s GPU ecosystem through KernelAbstractions.jl and ExaModels.jl (Montoison et al., 4 Oct 2025).

Its limitations are equally explicit. The sparse linear solver is currently CUDA-only through cuDSS, so AMD/ROCm is not yet supported. For small $X_i\approx x(t_i),\quad U_i\approx u(t_i).$8, CPU still wins due to kernel-launch overheads. Multi-GPU or distributed time-parallel schemes are future work. The package is tailored to direct transcription, and other OCP methods such as multiple shooting and pseudospectral methods require extensions.

Several common misconceptions are clarified by these constraints. One is that GPU execution should dominate CPU execution at any problem size; the reported results directly reject that interpretation for small $X_i\approx x(t_i),\quad U_i\approx u(t_i).$9. Another is that the package is a general front end for all major OCP transcriptions; the available workflow is explicitly tailored to direct transcription. A third is that the package is backend-agnostic across GPU vendors; the current sparse linear algebra layer is described as CUDA-only.

7. Relation to other Julia optimal control systems

Within the Julia ecosystem, OptimalControl.jl sits alongside other packages that address OCP modeling from different architectural directions. OptControl.jl implements modeling optimal control problems with symbolic algebra system based on Julia language and generates the corresponding numerical optimization codes to solve them with packages from Julia. It does not define a data type, but generates a solution script by handling Julia strings and runs the script automatically; it also supports component-based modeling and provides an interface to save script files (2207.13229).

NLOptControl, by contrast, is an open-source modeling language for solving nonlinear OCPs via direct-collocation methods. It is written in Julia, extends JuMP to provide a natural algebraic syntax for modeling nonlinear OCPs, calls sparse NLP solvers such as IPOPT and KNITRO, and uses ReverseDiff.jl with the acyclic-coloring method to exploit Hessian sparsity. Its reported benchmarks against PROPT state that single, two, and four interval pseudospectral methods are roughly $$X_{i+1}-X_i - h_i\,f(X_i,U_i)=0,\quad i=0,\dots,N-1,$$0, $$X_{i+1}-X_i - h_i\,f(X_i,U_i)=0,\quad i=0,\dots,N-1,$$1, and $$X_{i+1}-X_i - h_i\,f(X_i,U_i)=0,\quad i=0,\dots,N-1,$$2 times faster than PROPT’s, respectively (Febbo et al., 2020).

This suggests a clear distinction in design point. OptControl.jl emphasizes symbolic-to-script generation and component-based modeling; NLOptControl emphasizes JuMP-based direct collocation and sparse derivative computation on conventional NLP backends; OptimalControl.jl, as described in the GPU-focused workflow, emphasizes direct transcription into structured sparse NLPs that are compiled into GPU kernels and solved with a GPU-aware interior-point stack. A plausible implication is that these packages represent complementary rather than interchangeable approaches within Julia-based optimal control: script generation, JuMP extension, and end-to-end GPU-native transcription correspond to different priorities in model expressiveness, solver integration, and hardware utilization.