In this example, we are solving the robust scenario optimization problem
$$ \begin{aligned} \min_{x, u} \quad & \frac{1}{N_s} \sum_{j=1}^{N_s} \left( \sum_{i=0}^{N-1} \ell_i(x_i^j,u_i^j) + \ell_N(x_N^j) \right) \\ \text {s.t.}\quad & x_0^j = \bar{x}_0, u_0^1 = u_0^j \\ & x^j_{i+1} = A^j x_i^j + B^j u_i^j, \\ & -4 \cdot \mathbf{1}_{n_x} \leq x^j_i \leq 4 \cdot \mathbf{1}_{n_x}, \\ & -0.5 \cdot \mathbf{1}_{n_u} \leq u^j_i \leq 0.5 \cdot \mathbf{1}_{n_u}, \end{aligned} $$
for a system of oscillating masses. It's an extension of the classic MPC problem. Essentially, we are rolling out $N_s$ scenarios, where each scenario $j$ has different dynamics $A^j$ and $B^j$. The first state $x_0$ is fixed and the first input $u_0$ is the same as for all scenarios. This allows us to be robust in respect to all scenario roll-outs and achieve optimal performance in the expectation.
We'll see how to formulate it using Python and how PIQP's sparse_multistage KKT solver backend can be used to achieve a considerable speedup compared to the standard sparse KKT solver backend.
import numpy as np
import scipy.sparse as sp
from scipy.signal import cont2discrete
from scipy.linalg import solve_discrete_are
import piqp
import warnings
warnings.filterwarnings('ignore')
import matplotlib.pyplot as plt
%matplotlib inline
%config InlineBackend.figure_format='retina'
np.random.seed(42)
Model¶
We start by defining the system dynamics. We are considering the system of oscillating masses as shown below:
The system consists out of $M$ masses, where each mass is described by its position $p_i$ and its velocity $v_i$. We then describe the state of the whole system as $x=[p_1,\dots,p_M,v_1,\dots,v_M]^\top$. The last $M-1$ masses can be controlled through a input force acting on them which we collect in the input vector $u$. The following class builds the system dynamics in continuous time and discretizes it with a time step of 0.5s.
class ChainMassSystem:
def __init__(self, M, m=1.0, c=0.1, k=1.0):
self.M = M
# System parameters
self.m = m # mass
self.c = c # damping
self.k = k # spring constant
self.nx_max = 4.0
self.nu_max = 0.5
# System dimensions
self.nx = 2 * M
self.nu = M - 1
# Initialize system matrices
self._setup_system()
def _setup_system(self):
# Build continuous system matrices
L = np.eye(self.M, k=-1)
self.A = np.block([
[np.zeros((self.M, self.M)),
np.eye(self.M)],
[(-2 * self.k * np.eye(self.M) + self.k * L + self.k * L.T) / self.m,
(-2 * self.c * np.eye(self.M)) / self.m]
])
self.B = np.block([[np.zeros((2 * self.M - self.nu, self.nu))],
[np.eye(self.nu)]])
self.C = np.eye(self.nx)
self.D = np.zeros((self.nx, self.nu))
# Discretize system
dt = 0.5
d_sys = cont2discrete((self.A, self.B, self.C, self.D), dt, method='zoh')
self.Ad = d_sys[0]
self.Bd = d_sys[1]
# Cost matrices
self.Q = 1e3 * np.eye(self.nx)
self.R = 1e-1 * np.eye(self.nu)
self.QN = solve_discrete_are(self.Ad, self.Bd, self.Q, self.R)
We start by setting up some parameters and sample $N_s$ systems, where we the spring constant is uniformly sampled from $[1.0,2.0]$.
M = 3 # number of masses
N = 5 # MPC horizon
Ns = 3 # number of scenarios
def setup_systems():
global systems, nx, nu
# sample Ns spring constants between 1.0 and 2.0
ks = np.linspace(1.0, 2.0, Ns)
# generate Ns systems
systems = [ChainMassSystem(M, k=k) for k in ks]
nx = systems[0].nx
nu = systems[0].nu
setup_systems()
Let's sample the initial state uniformly from $[-1.0,1.0]$, i.e., the positions and velocities of the masses are in this range.
def random_x0():
return np.random.uniform(-1.0, 1.0, nx)
x0 = random_x0()
We will transcribe the problem into a general sparse QP formulation expected by PIQP:
$$ \begin{aligned} \min_{x} \quad & \frac{1}{2} x^\top P x + c^\top x \\ \text {s.t.}\quad & Ax=b, \\ & Gx \leq h, \\ & x_{lb} \leq x \leq x_{ub} \end{aligned} $$
While the generic sparse_ldlt solver backend can accept any sparse problem formulation, we need to make sure that ordering of the variables is correct for the sparse_multistage backend. Since $x_0$ and $u_0$ are coupled with all scenarios, we put them at the end, i.e. $g = (x_0,u_0)$. Hence, we define the primal variable as $x=(x_1^1,u_1^1,\dots,x_{N}^1,x_1^2,u_1^2,\dots,x_{N}^{N_s}, x_0, u_0)$.
def setup_problem():
global n, p, P, c, A, b, x_lb, x_ub
# number of decision variables
n = nx + nu + Ns * ((N - 1) * (nx + nu) + nx)
# number of equality constraints
p = Ns * N * nx
# empty problem data matrices
P = sp.csc_matrix((n, n))
c = np.zeros(n)
A = sp.csc_matrix((p, n))
b = np.zeros(p)
x_lb = np.zeros(n)
x_ub = np.zeros(n)
setup_problem()
Let's add the constraints for the initial state $x_0$ and the bounds for $u_0$. Note that they are located at the end of the primal decision variable $x$.
def setup_x0_u0():
# initial state x0
x_lb[-(nx+nu):-nu] = x0
x_ub[-(nx+nu):-nu] = x0
# bounds for u0
x_lb[-nu:] = -systems[0].nu_max
x_ub[-nu:] = systems[0].nu_max
setup_x0_u0()
Let's fill the $P$ matrix for the cost. We will use a standard quadratic cost which regulates the state of the system to zero, i.e.
$$ \ell_i(x_i,u_i) = x_i^\top Q x_i + u_i^\top R u_i, \quad \ell_N(x_N) = x_N^\top Q_N x_N, $$
where $Q_N$ is the solution of the discrete Riccati equation.
The structure of the cost matrix P will have the following structure:
$$ P = \frac{1}{N_s} \begin{bmatrix} Q & & & & & & & & \\ & R & & & & & & & \\ & & Q & & & & & & \\ & & & R & & & & & \\ & & & & \ddots & & & & \\ & & & & & Q_N & & & \\ & & & & & & \ddots & & \\ & & & & & & & N_s Q & \\ & & & & & & & & N_s R \\ \end{bmatrix} $$
def setup_cost():
# initial state
P[-(nx+nu):-nu, -(nx+nu):-nu] = systems[0].Q
P[-nu:, -nu:] = systems[0].R
for s in range(Ns):
offset = s * ((N - 1) * (nx + nu) + nx) - (nx + nu)
for i in range(1, N):
# state cost
P[offset+i*(nx+nu):offset+i*(nx+nu)+nx,
offset+i*(nx+nu):offset+i*(nx+nu)+nx] = systems[s].Q / Ns
# input cost
P[offset+i*(nx+nu)+nx:offset+i*(nx+nu)+nx+nu,
offset+i*(nx+nu)+nx:offset+i*(nx+nu)+nx+nu] = systems[s].R / Ns
# terminal cost
P[offset+N*(nx+nu):offset+N*(nx+nu)+nx,
offset+N*(nx+nu):offset+N*(nx+nu)+nx] = systems[s].QN / Ns
setup_cost()
Looking at the sparsity pattern confirms that we populated $P$ correctly.
plt.spy(P, markersize=1.75)
plt.show()
While the order of the dynamics constraints don't matter, we model them in chronological order to be consistent with the original problem formulation. This means that our $A$ matrix should look something like:
$$ A = \begin{bmatrix} -I & & & & & & & & & & & A^1 & B^1 \\ A^1 & B^1 & -I & & & & & & & & & & \\ & & A^1 & B^1 & -I & & & & & & & & \\ & & & & \ddots & & & & & & \cdots & & \\ & & & & & -I & & & & & & A^2 & B^2 \\ & & & & & A^2 & B^2 & -I & & & & & \\ & & & & & & & A^2 & B^2 & -I & & & \\ & & & & & & & & & \ddots & & \vdots & \vdots \end{bmatrix} $$
Note that for $A^j$ and $B^j$, the $j$ indicates the $j$-th scenario and not an exponent.
def setup_dynamics():
for s in range(Ns):
offset = s * ((N - 1) * (nx + nu) + nx) - (nx + nu)
eq_offset = s * N * nx
for i in range(N):
if i == 0:
A[eq_offset+i*nx:eq_offset+(i+1)*nx, -(nx+nu):-nu] = systems[s].Ad
A[eq_offset+i*nx:eq_offset+(i+1)*nx, -nu:] = systems[s].Bd
else:
A[eq_offset+i*nx:eq_offset+(i+1)*nx,
offset+i*(nx+nu):offset+i*(nx+nu)+nx] = systems[s].Ad
A[eq_offset+i*nx:eq_offset+(i+1)*nx,
offset+i*(nx+nu)+nx:offset+i*(nx+nu)+nx+nu] = systems[s].Bd
A[eq_offset+i*nx:eq_offset+(i+1)*nx,
offset+(i+1)*(nx+nu):offset+(i+1)*(nx+nu)+nx] = -np.eye(nx)
setup_dynamics()
Again, the sparsity pattern makes looks like expected.
plt.spy(A, markersize=2)
plt.show()
As a last check, we can calculate the condensed KKT sparsity and verify that is has indeed a block-tri-diagonal-arrow structure.
def setup_bounds():
for s in range(Ns):
offset = s * ((N - 1) * (nx + nu) + nx) - (nx + nu)
for i in range(N):
if i > 0:
x_lb[offset+i*(nx+nu)+nx:offset+i*(nx+nu)+nx+nu] = -systems[s].nu_max
x_ub[offset+i*(nx+nu)+nx:offset+i*(nx+nu)+nx+nu] = systems[s].nu_max
x_lb[offset+(i+1)*(nx+nu):offset+(i+1)*(nx+nu)+nx] = -systems[s].nx_max
x_ub[offset+(i+1)*(nx+nu):offset+(i+1)*(nx+nu)+nx] = systems[s].nx_max
setup_bounds()
The last remaining part are the state and input constraints which are direct bounds on the primal decision variable.
KKT_structure = P + A.T @ A
plt.spy(KKT_structure, markersize=1.75)
plt.show()
It's time to solve the problem using PIQP. In a first step we'll use the default sparse KKT backend.
sparse_solver = piqp.SparseSolver()
sparse_solver.settings.verbose = True
sparse_solver.settings.compute_timings = True
sparse_solver.setup(P=P, c=c, A=A, b=b, x_lb=x_lb, x_ub=x_ub)
status = sparse_solver.solve()
----------------------------------------------------------
PIQP
(c) Roland Schwan
Ecole Polytechnique Federale de Lausanne (EPFL) 2025
----------------------------------------------------------
sparse backend (sparse_ldlt)
variables n = 122, nzz(P upper triangular) = 375
equality constraints p = 90, nnz(A) = 1260
inequality constraints m = 0, nnz(G) = 0
variable lower bounds n_lb = 122
variable upper bounds n_ub = 122
iter prim_obj dual_obj duality_gap prim_inf dual_inf rho delta mu p_step d_step
0 3.25459e+02 -5.73062e+06 5.73095e+06 7.50007e+01 5.37853e-04 1.000e-06 1.000e-04 5.421e+04 0.0000 0.0000
1 1.14177e+03 -1.44820e+06 1.44934e+06 1.08927e+01 5.41721e+01 4.343e-07 1.506e-05 8.164e+03 0.8548 0.9900
2 3.46065e+03 -7.90135e+04 8.24741e+04 9.64642e-01 3.66062e+01 2.437e-08 8.453e-07 4.582e+02 0.9116 0.9701
3 4.54021e+03 -2.70652e+03 7.24673e+03 1.16036e-01 6.56384e+01 9.556e-09 7.367e-08 3.993e+01 0.8803 0.9715
4 4.62977e+03 3.82282e+03 8.06949e+02 1.58108e-03 5.46240e+00 8.181e-10 6.307e-09 3.419e+00 0.9862 0.9162
5 4.47718e+03 4.39599e+03 8.11923e+01 5.54654e-05 4.26787e+00 1.000e-10 6.595e-10 3.575e-01 0.9648 0.9260
6 4.45430e+03 4.44829e+03 6.01381e+00 3.06044e-06 3.08251e-01 1.000e-10 1.000e-10 2.415e-02 0.9448 0.9609
7 4.45188e+03 4.45160e+03 2.80125e-01 4.43645e-08 4.01175e-01 1.000e-10 1.000e-10 1.375e-03 0.9856 0.9578
8 4.45174e+03 4.45172e+03 1.70776e-02 5.73521e-10 4.22093e-02 1.000e-10 1.000e-10 7.743e-05 0.9900 0.9827
9 4.45173e+03 4.45173e+03 1.98353e-03 1.02807e-10 4.22093e-04 1.000e-10 1.000e-10 8.222e-06 0.9900 0.9900
10 4.45173e+03 4.45173e+03 2.91227e-04 4.03604e-11 4.22087e-06 1.000e-10 1.000e-10 1.200e-06 0.9900 0.9900
11 4.45173e+03 4.45173e+03 4.19868e-05 1.53622e-11 4.21851e-08 1.000e-10 1.000e-10 1.728e-07 0.9900 0.9900
12 4.45173e+03 4.45173e+03 5.84887e-06 5.76911e-12 4.12971e-10 1.000e-10 1.000e-10 2.406e-08 0.9900 0.9900
status: solved
number of iterations: 12
objective: 4.45173e+03
total run time: 1.370e-03s
setup time: 3.596e-04s
solve time: 1.011e-03s
So sunrises here. PIQP is able so solve this problem to very high accuracy, but we can solve the problem faster. Let's switch to the sparse_multistage KKT solver backend. Note that the interface doesn't change and PIQP does all the heavy lifting under the hood.
multistage_solver = piqp.SparseSolver()
multistage_solver.settings.verbose = True
multistage_solver.settings.compute_timings = True
multistage_solver.settings.kkt_solver = piqp.KKTSolver.sparse_multistage
multistage_solver.setup(P=P, c=c, A=A, b=b, x_lb=x_lb, x_ub=x_ub)
status = multistage_solver.solve()
----------------------------------------------------------
PIQP
(c) Roland Schwan
Ecole Polytechnique Federale de Lausanne (EPFL) 2025
----------------------------------------------------------
sparse backend (sparse_multistage)
variables n = 122, nzz(P upper triangular) = 375
equality constraints p = 90, nnz(A) = 1260
inequality constraints m = 0, nnz(G) = 0
variable lower bounds n_lb = 122
variable upper bounds n_ub = 122
block sizes: 8,6 8,6 8,6 14,0 8,6 8,6 8,6 14,0 8,6 8,6 8,6 14,0
arrow width: 8
iter prim_obj dual_obj duality_gap prim_inf dual_inf rho delta mu p_step d_step
0 3.25459e+02 -5.73062e+06 5.73095e+06 7.50007e+01 5.37853e-04 1.000e-06 1.000e-04 5.421e+04 0.0000 0.0000
1 1.14177e+03 -1.44820e+06 1.44934e+06 1.08927e+01 5.41721e+01 4.343e-07 1.506e-05 8.164e+03 0.8548 0.9900
2 3.46065e+03 -7.90135e+04 8.24741e+04 9.64642e-01 3.66062e+01 2.437e-08 8.453e-07 4.582e+02 0.9116 0.9701
3 4.54021e+03 -2.70652e+03 7.24673e+03 1.16036e-01 6.56384e+01 9.556e-09 7.367e-08 3.993e+01 0.8803 0.9715
4 4.62977e+03 3.82282e+03 8.06949e+02 1.58108e-03 5.46240e+00 8.181e-10 6.307e-09 3.419e+00 0.9862 0.9162
5 4.47718e+03 4.39599e+03 8.11923e+01 5.54654e-05 4.26787e+00 1.000e-10 6.595e-10 3.575e-01 0.9648 0.9260
6 4.45430e+03 4.44829e+03 6.01381e+00 3.06044e-06 3.08251e-01 1.000e-10 1.000e-10 2.415e-02 0.9448 0.9609
7 4.45188e+03 4.45160e+03 2.80122e-01 4.43645e-08 4.01174e-01 1.000e-10 1.000e-10 1.375e-03 0.9856 0.9578
8 4.45174e+03 4.45172e+03 1.70773e-02 5.73521e-10 4.22089e-02 1.000e-10 1.000e-10 7.743e-05 0.9900 0.9827
9 4.45173e+03 4.45173e+03 1.98363e-03 1.02807e-10 4.22046e-04 1.000e-10 1.000e-10 8.222e-06 0.9900 0.9900
10 4.45173e+03 4.45173e+03 2.91230e-04 4.03605e-11 4.21736e-06 1.000e-10 1.000e-10 1.200e-06 0.9900 0.9900
11 4.45173e+03 4.45173e+03 4.19654e-05 1.53621e-11 3.97996e-08 1.000e-10 1.000e-10 1.728e-07 0.9900 0.9900
12 4.45173e+03 4.45173e+03 5.84894e-06 5.76911e-12 1.09918e-08 1.000e-10 1.000e-10 2.406e-08 0.9900 0.9900
status: solved
number of iterations: 12
objective: 4.45173e+03
total run time: 1.255e-03s
setup time: 5.525e-04s
solve time: 7.029e-04s
Looking at the debug information, we can see that we are now using the sparse_multistage KKT solver backend, and we also get information about the detected block sizes and the detected arrow width. We can visualize this using the sparsity pattern of the condensed KKT matrix.
sp_solve_time = sparse_solver.result.info.solve_time
ms_solve_time = multistage_solver.result.info.solve_time
speedup = sp_solve_time / ms_solve_time
print(f'Solve speedup: {speedup:0.3}x')
Solve speedup: 1.44x
For such a relatively small problem, the speedup is not particularly high with only 44% speedup. Let's rerun the previous cells with a bigger problem instance and see the speedup.
M = 15 # number of masses
N = 20 # MPC horizon
Ns = 25 # number of scenarios
setup_systems()
x0 = random_x0()
setup_problem()
setup_x0_u0()
setup_cost()
setup_dynamics()
setup_bounds()
sparse_solver = piqp.SparseSolver()
sparse_solver.settings.verbose = False
sparse_solver.settings.compute_timings = True
sparse_solver.setup(P=P, c=c, A=A, b=b, x_lb=x_lb, x_ub=x_ub)
status = sparse_solver.solve()
multistage_solver = piqp.SparseSolver()
multistage_solver.settings.verbose = False
multistage_solver.settings.compute_timings = True
multistage_solver.settings.kkt_solver = piqp.KKTSolver.sparse_multistage
multistage_solver.setup(P=P, c=c, A=A, b=b, x_lb=x_lb, x_ub=x_ub)
status = multistage_solver.solve()
sp_solve_time = sparse_solver.result.info.solve_time
ms_solve_time = multistage_solver.result.info.solve_time
speedup = sp_solve_time / ms_solve_time
print(f'Solve speedup: {speedup:0.3}x')
Solve speedup: 5.61x
Now we have considerably higher speedup. This is many due to more efficient use of vectorized instructions and cache locality.
In a last step, let's extract the solution and visualize it.
# extract solution
X = np.zeros((nx, N + 1, Ns))
U = np.zeros((nu, N, Ns))
x = multistage_solver.result.x
for s in range(Ns):
X[:, 0, s] = x[-(nx+nu):-nu]
U[:, 0, s] = x[-nu:]
offset = s * ((N - 1) * (nx + nu) + nx) - (nx + nu)
for i in range(1, N):
X[:, i, s] = x[offset+i*(nx+nu):offset+i*(nx+nu)+nx]
U[:, i, s] = x[offset+i*(nx+nu)+nx:offset+i*(nx+nu)+nx+nu]
X[:, N, s] = x[offset+N*(nx+nu):offset+N*(nx+nu)+nx]
# plot solution
time = np.linspace(0, 0.5 * N, N + 1)
for s in range(Ns):
plt.gca().set_prop_cycle(None) # Reset the color cycle to the beginning
plt.plot(time, X[:M, :, s].T, alpha=0.8)
plt.xlabel('Time')
plt.ylabel('Position')
plt.show()
for s in range(Ns):
plt.gca().set_prop_cycle(None) # Reset the color cycle to the beginning
for i in range(U.shape[0]):
plt.step(time, np.append(U[i, :, s], U[i, -1, s]), alpha=0.5, where='post')
plt.xlabel('Time')
plt.ylabel('Input')
plt.show()
