Robust Optimal Control on a Pendulum Using Automatic Differentiation#
In this example, we use Automatic Differentiation to solve a robust optimal control problem in the embedding space of a nonlinear pendulum. In particular, we build an objective function with robust constraints, automatically create and compile functions evaluating their gradients, Jacobians, and Hessians, and finally setup an interior point minimization problem to find a locally optimal solution using IPOPT.
This example is from the paper “\(\texttt{immrax}\): A Parallalizable and Differentiable Toolbox for Interval Analysis and Mixed Monotone Reachability in JAX”.
Preliminaries#
To run this example, you’ll need cyipopt and the MA57 linear solver from HSL. See the Installation page for more information on how to install these.
import jax
import jax.numpy as jnp
import immrax as irx
from typing import Tuple
We will be using IPOPT to solve a minimization problem. Since IPOPT is CPU-based and 64-bit, we configure JAX with these options as well. We can use the GPU here instead of the CPU for some operations, but empirically, the time to transfer data between the CPU and GPU is much larger than the amount of time the problem takes on the CPU only.
# Some configurations
jax.config.update("jax_enable_x64", True)
# We wrap the jax.jit function to set the backend to cpu by default for convenience.
device = "cpu"
def jit(*args, **kwargs):
kwargs.setdefault("backend", device)
return jax.jit(*args, **kwargs)
The pendulum’s dynamics#
Consider the dynamics of a forced, damped pendulum
with \(m=0.15\mathrm{kg}\), \(l=0.5\mathrm{m}\), \(b=0.1\mathrm{N\cdot m \cdot s}\), and \(g=9.81\mathrm{m}/\mathrm{s}^2\). The torque \(\tau := (1 + w)u\), where \(u\in\mathbb{R}\) is the desired torque input and \(w\in[\underline{w},\overline{w}] := [-0.02,0.02]\) is a bounded multiplicative disturbance on the control input. We implement this as a \(2\)-state system with \(x := (\theta, \dot{\theta})\),
This is implemented in immrax as a System, with the specified dynamics written using jax.numpy.
g = 9.81
class Pendulum(irx.System):
m: float
l: float
b: float
def __init__(self, m: float = 0.15, l: float = 0.5, b: float = 0.1) -> None:
# Tells immrax that the system is continuous
self.evolution = "continuous"
# Tells immrax the number of states
self.xlen = 2
self.m = m
self.l = l
self.b = b
def f(self, t: float, x: jax.Array, u: jax.Array, w: jax.Array) -> jax.Array:
return jnp.array(
[
x[1],
(((1 + w[0]) * u[0] - self.b * x[1]) / ((self.m) * self.l**2))
- (g / self.l) * jnp.sin(x[0]),
]
)
sys = Pendulum()
Problem Statement#
We seek to find a finite-time closed-loop optimal control policy \(\pi:[0,T]\times\mathbb{R}^n\to\mathbb{R}\) to swing up the pendulum to an a priori safe region at the top. We consider linear feedback control policies of the form \(\pi(t,x) := K(x(t) - x_{\mathrm{nom}}(t)) + u_{\mathrm{ff}}(t)\), where \(K\) is a time invariant linear closed-loop stabilizing term to help counter the disturbance, \(u_{\mathrm{ff}}:[0,T]\to\mathbb{R}\) is a feedforward control policy, and \(x_\mathrm{nom}:[0,T]\to\mathbb{R}^n\) is the nominal trajectory of the deterministic system under the feedforward control law \(u_\mathrm{ff}\) with known disturbance mapping \(w_{\mathrm{nom}}:[0,T]\to\mathbb{R}\). The closed-loop system is thus
Embedding System for the Pendulum#
Consider the following function
with \([\textsf{M}_x\ \textsf{M}_u\ \textsf{M}_w] := [\textsf{M}_\sigma^{\xi_\mathrm{nom}(t)}(\underline{x},\overline{x},u_\mathrm{ff}(t),u_\mathrm{ff}(t),\underline{w},\overline{w})]\), where \(\textsf{M}\) is defined as Proposition~\ref{prop:MixedJacobian-Based} afor the map \(\hat{f} : \mathbb{R}^{n+p+q} \to \mathbb{R}^n\) such that \(\hat{f}((x,u,w)) := f(x,u,w)\), for some \((n+p+q)\)-permutation \(\sigma\), and \(\xi_\mathrm{nom}(t) := (x_\mathrm{nom}(t),u_\mathrm{ff}(t),w_\mathrm{nom}(t))\). This is a valid inclusion function for the closed-loop dynamics \(f^\pi\)~\eqref{eq:pendulumcldyn} (proof deferred to Appendix~\ref{apax:sec:proofclpend}, which uses Proposition~\ref{prop:MixedJacobian-Based} Part~\eqref{prop:MixedJacobian-Based:p1}).
We use the mjacM transform to implement the inclusion function in in immrax. We expect the nominal values to be given as an input to the inclusion function, which will be calculated externally.
sys_mjacM = irx.mjacM(sys.f)
# sys_mjacM = irx.jacM(sys.f)
@jit
def F(
t: irx.Interval,
x: irx.Interval,
w: irx.Interval,
*,
K: jnp.ndarray,
nominal: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray],
):
tc, xc, uc, wc = nominal
Mt, Mx, Mu, Mw = sys_mjacM(t, x, irx.interval(uc), w, centers=(nominal,))[0]
# Mt, Mx, Mu, Mw = sys_mjacM(t, x, irx.interval(uc), w, centers=(nominal,))
return (
(Mx + Mu @ irx.interval(K)) @ (x - xc) + Mw @ (w - wc) + sys.f(tc, xc, uc, wc)
)
Following the treatment from the paper, now that we have an inclusion function for the dynamics of the system, we will embed the closed-loop pendulum dynamics into a new \(2n\)-dimensional embedding system that provides efficient bounds for its reachable set.
This embedding is automatically performed in immrax by the ifemb transform, which takes a System and an inclusion function for its dynamics, and returns the corresponding EmbeddingSystem.
embsys = irx.ifemb(sys, F)
Using Autodiff to Setup an IPOPT Minimization Problem#
The problem itself#
Now that we have the embedding system embsys, we will try to solve an optimization problem of the following form:
where the embedding dynamics \(\textsf{E}\) are discretized using Euler integration with step size \(\Delta t\).
The first and second terms of the objective are typical quadratic conditioning of the decision variables. The third term is a regularization factor intended to help curb the expansion of the gap between the upper and lower bound, which emperically helps the optimization problem converge to a feasible solution. Finally, in the inequality constraints, we require that the pendulum reach a target set \([\underline{x}_f,\overline{x}_f]\) and stay within these constraints for \(t\in[t_{N_e},t_N]\).
Towards an IPOPT implementation#
In the following code, we define variables setting up the minimization problem with the following definitions:
# Problem Parameters
# Time horizon and discretization
t0, te, tf, dt = 0, 3.0, 3.25, 0.05
Ne = round((te - t0) / dt)
N = round((tf - t0) / dt) + 1
tt = jnp.arange(t0, tf + dt, dt)
# Initial condition (no initial perturbation)
x0 = irx.interval(jnp.array([0.0, 0.0]))
x0ut = irx.i2ut(x0)
x0cent, x0pert = irx.i2centpert(x0)
# Final set constraint
xf = irx.icentpert(jnp.array([jnp.pi, 0.0]), jnp.array([10.0 * (jnp.pi / 360), 0.1]))
xfl, xfu = irx.i2lu(xf)
xfut = irx.i2ut(xf)
xfcent, xfpert = irx.i2centpert(xf)
# Disturbance bounds [-0.02, 0.02]
# w = irx.icentpert(0., 0.02)
w = irx.icentpert(0.0, 0.01)
We put all decision variables into one vector.
# Initial guess for control u_ff (zeros) and linear matrix K ([-1 -1])
K = jnp.array([[-1.0, -1.0]])
u0 = jnp.concatenate((jnp.zeros(N), K.reshape(-1)))
# Function to split decision vector into u_ff and K
def split_u(u: jax.Array) -> Tuple[jax.Array, jax.Array]:
return u[:-2], u[-2:].reshape(1, 2)
Next, we implement a couple of helper functions. They use jax.lax.scan to perform a simple and efficient Euler integration of the dynamics to the specified time horizon. We could use the System.compute_trajectory function instead, which uses diffrax to integrate the dynamics with many different options of integrators.
# Function to rollout the undisturbed pendulum dynamics using Euler integration
@jit
def rollout_ol_sys_undisturbed(u: jax.Array) -> jax.Array:
u, K = split_u(u)
def f_euler(xt, ut):
xtp1 = xt + dt * sys.f(0.0, xt, jnp.array([ut]), jnp.array([0.0]))
return (xtp1, xtp1)
_, x = jax.lax.scan(f_euler, x0cent, u)
return x
# Function to rollout the closed-loop embedding system dynamics using Euler integration
@jit
def rollout_cl_embsys(u: jax.Array) -> jax.Array:
u, K = split_u(u)
def f_euler(xt, ut):
xtut, xnomt = xt
xtutp1 = xtut + dt * embsys.E(
irx.interval([0.0]),
xtut,
w,
K=K,
nominal=(jnp.array([0.0]), xnomt, jnp.array([ut]), jnp.array([0.0])),
)
xnomtp1 = xnomt + dt * sys.f(0.0, xnomt, jnp.array([ut]), jnp.array([0.0]))
return ((xtutp1, xnomtp1), xtutp1)
_, x = jax.lax.scan(f_euler, (x0ut, x0cent), u)
return x
Using the rollout_cl_embsys function, we implement the desired objective function above.
# Objective Function
@jit
def obj(u: jax.Array) -> jax.Array:
x = rollout_cl_embsys(u)
return jnp.sum(u**2) + jnp.sum((x[:, 2:] - x[:, :2]) ** 2)
Next, we implement the inequality terminal set constraints as a function con_ineq(u) \(\geq 0\).
# Inequality constraints
@jit
def con_ineq(u):
x = rollout_cl_embsys(u)
return jnp.concatenate(
((x[Ne:, :2] - xfl).reshape(-1), (xfu - x[Ne:, 2:]).reshape(-1))
)
The Automatic Differentiation Step#
Next, we use JAX’s autodiff transforms to automatically create functions to compute the objective’s gradient and Hessian, as well as the Jacobian and Hessian vector product of the constraints with respect to the Lagrange multipliers. We precompile (JIT compile) these functions to separate their compilation from the rest of the stack.
obj_grad = jit(jax.grad(obj)) # Objective Gradient
obj_hess = jit(jax.jacfwd(jax.jacrev(obj))) # Objective Hessian
con_ineq_jac = jit(jax.jacfwd(con_ineq)) # Constraint Jacobian
con_ineq_hess = jit(jax.jacfwd(jax.jacrev(con_ineq))) # Constraint Hessian
# Constraint Hessian-Vector Product
@jit
def con_ineq_hessvp(u, v):
def hessvp(u):
_, hvp = jax.vjp(con_ineq, u)
return hvp(v)[0] # One tangent, one output. u^T dc_v
return jax.jacrev(hessvp)(u)
print("JIT Compiling Autodiff Functions...")
from time import time
time0 = time()
obj(u0)
print("Finished JIT for obj")
obj_grad(u0)
print("Finished JIT for obj_grad")
obj_hess(u0)
print("Finished JIT for obj_hess")
con_ineq_jac(u0)
print("Finished JIT for con_ineq_jac")
con_ineq_hessvp(u0, jnp.ones(4 * (N - Ne)))
print("Finished JIT for con_ineq_hessvp")
timef = time()
print(f"Finished JIT Compiling Autodiff Functions in {timef - time0} seconds")
JIT Compiling Autodiff Functions...
Finished JIT for obj
Finished JIT for obj_grad
Finished JIT for obj_hess
Finished JIT for con_ineq_jac
Finished JIT for con_ineq_hessvp
Finished JIT Compiling Autodiff Functions in 35.219850301742554 seconds
Minimization in IPOPT#
Finally, we solve the optimization problem from above using a call to IPOPT.
from cyipopt import minimize_ipopt
# Constraints
cons = [
{"type": "ineq", "fun": con_ineq, "jac": con_ineq_jac, "hess": con_ineq_hessvp},
]
bnds = [(-100.0, 100.0) for _ in range(u0.size)]
ipopt_opts = {
b"disp": 4,
b"linear_solver": "ma57",
b"hsllib": "libcoinhsl.so",
b"tol": 1e-3,
b"max_iter": 100,
}
# Solve the optimization problem
res = minimize_ipopt(
obj, jac=obj_grad, hess=obj_hess, x0=u0, constraints=cons, options=ipopt_opts
)
******************************************************************************
This program contains Ipopt, a library for large-scale nonlinear optimization.
Ipopt is released as open source code under the Eclipse Public License (EPL).
For more information visit https://github.com/coin-or/Ipopt
******************************************************************************
Total number of variables............................: 68
variables with only lower bounds: 0
variables with lower and upper bounds: 0
variables with only upper bounds: 0
Total number of equality constraints.................: 0
Total number of inequality constraints...............: 24
inequality constraints with only lower bounds: 24
inequality constraints with lower and upper bounds: 0
inequality constraints with only upper bounds: 0
Reallocating memory for MA57: lfact (22061)
Reallocating memory for MA57: lfact (24303)
Number of Iterations....: 100
(scaled) (unscaled)
Objective...............: 1.0210781839478321e+01 1.0210781839478321e+01
Dual infeasibility......: 1.4144511782662028e-01 1.4144511782662028e-01
Constraint violation....: 1.0006529303496162e-03 1.0006529303496162e-03
Variable bound violation: 0.0000000000000000e+00 0.0000000000000000e+00
Complementarity.........: 1.2181650600986621e-10 1.2181650600986621e-10
Overall NLP error.......: 1.4144511782662028e-01 1.4144511782662028e-01
Number of objective function evaluations = 312
Number of objective gradient evaluations = 101
Number of equality constraint evaluations = 0
Number of inequality constraint evaluations = 315
Number of equality constraint Jacobian evaluations = 0
Number of inequality constraint Jacobian evaluations = 101
Number of Lagrangian Hessian evaluations = 100
Total seconds in IPOPT = 6.251
EXIT: Maximum Number of Iterations Exceeded.
Visualizing the Results#
# Plotting imports
%matplotlib widget
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Wedge
from immrax.utils import plot_interval_t
# Use latex fonts
plt.rcParams.update({"text.usetex": True, "font.family": "Helvetica", "font.size": 14})
uu = res.x
print(f"u_ff = {split_u(uu)[0]},\nK = {split_u(uu)[1]}")
u_ff = [-5.92930569e-02 -4.26039538e-02 -1.92477872e-02 -6.60294115e-04
-2.62717792e-02 4.76887169e-03 4.71440153e-02 8.86849578e-02
1.27199566e-01 1.60448415e-01 1.86262309e-01 2.02671775e-01
2.08041820e-01 2.01199027e-01 1.81534389e-01 1.49067206e-01
1.04464849e-01 4.90271379e-02 1.95869045e-02 -4.81270685e-02
-1.24186376e-01 -2.00041944e-01 -2.71669574e-01 -3.34729699e-01
-3.84856068e-01 -4.18072137e-01 -4.31251843e-01 -4.22492002e-01
-3.91242912e-01 -3.38142072e-01 -2.64622858e-01 -1.72477162e-01
-6.35724998e-02 3.99106442e-02 1.76629662e-01 3.22030859e-01
4.70516187e-01 6.13920428e-01 7.41725346e-01 8.42580432e-01
9.07104806e-01 9.30912942e-01 9.16182929e-01 8.70652197e-01
8.04747759e-01 7.28649752e-01 6.48848759e-01 5.70874319e-01
4.97880923e-01 4.31274170e-01 3.71387162e-01 3.17903781e-01
2.70122825e-01 2.27102239e-01 1.87716840e-01 1.50651929e-01
1.14343585e-01 7.68661758e-02 3.57578338e-02 -4.95050542e-03
-6.37781565e-02 4.08752042e-02 3.16317654e-02 2.17814256e-02
1.08588158e-02 2.02389220e-03],
K = [[-0.0305142 -0.59889171]]
# Trajectory vs time
unrolledout = rollout_ol_sys_undisturbed(uu)
clrolledout = rollout_cl_embsys(uu)
if jnp.all(con_ineq(uu) >= -1e-3):
print(f"Constraints satisfied to {jnp.min(con_ineq(uu))}!")
else:
print(f"Constraints not satisfied, worst is {jnp.min(con_ineq(uu))}!")
print(con_ineq(uu))
fig1, axs = plt.subplots(1, 2, dpi=100, figsize=[8, 4])
fig1.subplots_adjust(top=0.95, bottom=0.15, left=0.075, right=0.975)
axs[0].set_xlabel("$t$")
axs[1].set_xlabel("$t$")
axs[0].set_ylabel("$\\theta$", rotation=0, labelpad=5)
axs[1].set_ylabel("$\\dot{\\theta}$", rotation=0, labelpad=5)
axs[0].plot(tt, unrolledout[:, 0], color="k")
axs[1].plot(tt, unrolledout[:, 1], color="k")
plot_interval_t(axs[0], tt, xf[0] * jnp.ones(N), color="k")
plot_interval_t(axs[1], tt, xf[1] * jnp.ones(N), color="k")
plot_interval_t(
axs[0], tt, irx.interval(clrolledout[:, 0], clrolledout[:, 2]), color="tab:blue"
)
plot_interval_t(
axs[1], tt, irx.interval(clrolledout[:, 1], clrolledout[:, 3]), color="tab:blue"
)
fig1.savefig("figures/pendulum.pdf")
plt.show()
Constraints not satisfied, worst is -0.0010006629303496162!
[-0.00043718 0.13235455 0.00118055 0.12481568 0.00242133 0.1088238
0.00286252 0.08451305 0.00208817 0.05085413 -0.00036912 0.01071328
0.10517735 -0.00100066 0.10012731 0.00380032 0.09531733 0.01661963
0.09114831 0.03761718 0.08802917 0.06783443 0.08642089 0.10431467]
If the figure did not load, please see pendulum.pdf.
# Animation of pendulum
th2deg = lambda th: (th * 180) / jnp.pi - 90
anifig, anim = plt.subplots(1, 1, dpi=100, figsize=[4, 4])
anifig.subplots_adjust()
wedge = anim.add_patch(
Wedge(
(0, 0),
sys.l,
th2deg(clrolledout[0, 0]),
th2deg(clrolledout[0, 2]),
lw=2,
ec="k",
)
)
anim.add_patch(
Wedge((0, 0), sys.l, th2deg(xfl[0]), th2deg(xfu[0]), color="k", alpha=0.25)
)
anim.set_xlim(-sys.l * 1.2, sys.l * 1.2)
anim.set_ylim(-sys.l * 1.2, sys.l * 1.2)
def animate(t):
wedge.set(theta1=th2deg(clrolledout[t, 0]), theta2=th2deg(clrolledout[t, 2]))
anifig.suptitle(f"$t = {t * dt:.2f}$")
anifig.savefig(f"figures/frames/pendulum_{t:05d}.pdf")
ani = animation.FuncAnimation(
anifig, animate, frames=N, repeat=True, interval=dt * 1000
)
FFwriter = animation.FFMpegWriter(fps=(1 / dt))
ani.save("figures/pendulum.mp4", writer=FFwriter)
plt.close()
from IPython.display import Video
Video("figures/pendulum.mp4", embed=True)
If the video did not render, please see pendulum.mp4.