LQR computes the cart-pole's gains from the Riccati equation

from cartpole import cartpole_draw, cartpole_rate, cartpole_step

DT = 0.01
T_END = 6.0
FRAME_EVERY = 4
START_DEG = 25.0
RICCATI_DT = 0.002
RICCATI_STEPS = 40000
# the state: cart position, cart velocity, pole angle, pole angular rate
NSTATE = 4

#| The state weights: position, velocity, angle, angular rate. The angle is
#| weighted 20 times the others, so the cart may wander a little to keep the
#| pole upright.
def weights_Q() {
    q = zeros(NSTATE, NSTATE)
    q[0, 0] = 1.0
    q[1, 1] = 1.0
    q[2, 2] = 20.0
    q[3, 3] = 1.0
    return q
}
R_COST = 0.05

#| $A = \partial f/\partial z$ at the upright equilibrium, by central differences
#| on the library's own dynamics.
def linear_A(eps) {
    a = zeros(NSTATE, NSTATE)
    for j to NSTATE {
        zp = zeros(NSTATE)
        zm = zeros(NSTATE)
        zp[j] = eps
        zm[j] = -eps
        fp = cartpole_rate(zp, 0.0)
        fm = cartpole_rate(zm, 0.0)
        for i to NSTATE {
            a[i, j] = (fp[i] - fm[i]) / (2.0 * eps)
        }
    }
    return a
}

#| $B = \partial f/\partial u$ at the same point, the same way.
def linear_B(eps) {
    b = zeros(NSTATE, 1)
    z = zeros(NSTATE)
    fp = cartpole_rate(z, eps)
    fm = cartpole_rate(z, -eps)
    for i to NSTATE {
        b[i, 0] = (fp[i] - fm[i]) / (2.0 * eps)
    }
    return b
}

#| One Euler step of $\dot P = -(A^\top P + PA - PBR^{-1}B^\top P + Q)$, taken
#| backwards in time. Once P stops changing, it solves the algebraic equation.
def riccati_step(p, a, b, q, r, dt) {
    bt = transpose(b)
    pb = p @ b
    quad = (1.0 / r) * (pb @ (bt @ p))
    dp = transpose(a) @ p + p @ a - quad + q
    return p + dt * dp
}

#| Solve the algebraic Riccati equation by taking `steps` of those Euler steps from P = 0.
def riccati_solve(a, b, q, r, dt, steps) {
    p = zeros(NSTATE, NSTATE)
    for k to steps {
        p = riccati_step(p, a, b, q, r, dt)
    }
    return p
}

#| How far $P$ is from solving the equation: the largest entry of the residual.
def riccati_residual(p, a, b, q, r) {
    bt = transpose(b)
    pb = p @ b
    res = transpose(a) @ p + p @ a - (1.0 / r) * (pb @ (bt @ p)) + q
    worst = 0.0
    for i to NSTATE {
        for j to NSTATE {
            worst = max(worst, abs(res[i, j]))
        }
    }
    return worst
}

#| The gains $K = R^{-1}B^\top P$, as a 4-vector: the force is $-K z$.
def lqr_gains(p, b, r) {
    k = (1.0 / r) * (transpose(b) @ p)
    out = zeros(NSTATE)
    for i to NSTATE {
        out[i] = k[0, i]
    }
    return out
}

A = linear_A(1.0e-5)
B = linear_B(1.0e-5)
Q = weights_Q()
P = riccati_solve(A, B, Q, R_COST, RICCATI_DT, RICCATI_STEPS)
K = lqr_gains(P, B, R_COST)
residual = riccati_residual(P, A, B, Q, R_COST)
show K
show residual

#| The closed loop $A - BK$: its eigenvalues decide whether the upright is stable.
def closed_loop(a, b, k) {
    m = zeros(NSTATE, NSTATE)
    for i to NSTATE {
        for j to NSTATE {
            m[i, j] = a[i, j] - b[i, 0] * k[j]
        }
    }
    return m
}

CL = closed_loop(A, B, K)
cl_re = zeros(NSTATE)
cl_im = zeros(NSTATE)
eig_general(CL, cl_re, cl_im)
slowest = -1.0e30
for i to NSTATE {
    slowest = max(slowest, cl_re[i])
}
show cl_re
show slowest

#| Simulate the nonlinear cart-pole under the force $-K z$ from angle `theta0`; returns
#| [final |theta|, cost, largest |force|].
def fly(gains, theta0, draw) {
    s = zeros(NSTATE)
    s[2] = theta0
    cost = 0.0
    peak = 0.0
    steps = round(T_END / DT)
    for step to steps {
        u = 0.0
        for i to NSTATE {
            u = u - gains[i] * s[i]
        }
        peak = max(peak, abs(u))
        cost = cost + DT * (s[0] * s[0] + s[1] * s[1] + 20.0 * s[2] * s[2] + s[3] * s[3] + R_COST * u * u)
        s = cartpole_step(s, u, DT)
        if draw {
            if mod(step, FRAME_EVERY) == 0.0 {
                cartpole_draw(s)
                frame(DT * FRAME_EVERY)
            }
        }
    }
    return [abs(s[2]), cost, peak]
}

#| The hand-tuned gains of the earlier post, for comparison.
def pd_gains() {
    g = zeros(NSTATE)
    g[0] = -1.5
    g[1] = -3.0
    g[2] = -60.0
    g[3] = -12.0
    return g
}

theta0 = START_DEG * 3.141592653589793 / 180.0
lqr_run = fly(K, theta0, false)
pd_run = fly(pd_gains(), theta0, false)
show lqr_run
show pd_run

#| The angle over time under each controller, from the same start.
def angle_trace(gains, theta0, n) {
    s = zeros(NSTATE)
    s[2] = theta0
    out = zeros(n)
    for step to n {
        u = 0.0
        for i to NSTATE {
            u = u - gains[i] * s[i]
        }
        out[step] = s[2] * 180.0 / 3.141592653589793
        s = cartpole_step(s, u, DT)
    }
    return out
}

N_TRACE = 300
ts = linspace(0.0, N_TRACE * DT, N_TRACE)
figure("The pole's angle, from 25 degrees", "seconds", "degrees from upright")
plot(ts, angle_trace(K, theta0, N_TRACE), "LQR, gains computed")
plot(ts, angle_trace(pd_gains(), theta0, N_TRACE), "PD, gains by hand")
style("PD, gains by hand", "dash")

test "the linear model A z + B u matches the nonlinear rate near upright" {
    z = zeros(NSTATE)
    z[0] = 0.01
    z[1] = -0.02
    z[2] = 0.005
    z[3] = 0.03
    u = 0.4
    exact = cartpole_rate(z, u)
    for i to NSTATE {
        approx = B[i, 0] * u
        for j to NSTATE {
            approx = approx + A[i, j] * z[j]
        }
        assert_near(approx, exact[i], 2.0e-3)
    }
}
test "P solves the algebraic Riccati equation: the residual is under 1e-8" {
    assert_lt(residual, 1.0e-8)
}
test "every closed-loop eigenvalue has a negative real part" {
    assert_lt(slowest, -1.0e-6)
}
test "from 25 degrees the pole ends upright, within a thousandth of a radian" {
    assert_lt(lqr_run[0], 1.0e-3)
}
test "the computed gains cost less than the hand-tuned ones on the same start" {
    assert_lt(lqr_run[1], pd_run[1])
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors