"""finite_horizon_lqr.py — finite-horizon discrete-time LQR for MPC."""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import numpy.typing as npt
from numpy.linalg import LinAlgError
from scipy.linalg import solve, solve_discrete_are
FloatMatrix = npt.NDArray[np.float64]
FloatVector = npt.NDArray[np.float64]
@dataclass(frozen=True)
class LQRSolution:
"""Output of the backward Riccati recursion.
Attributes:
gains: Stacked feedback gains, shape (N, nu, nx).
gains[k] is K_k, the optimal gain at step k.
cost_to_go: Stacked cost-to-go matrices, shape (N + 1, nx, nx).
cost_to_go[k] is P_k; cost_to_go[N] is the terminal weight.
"""
gains: FloatMatrix
cost_to_go: FloatMatrix
def _as_float_matrix(
name: str,
matrix: npt.ArrayLike,
*,
expected_shape: tuple[int, int] | None = None,
) -> FloatMatrix:
"""Coerce input to a finite 2-D float64 matrix, with clear error messages."""
arr = np.asarray(matrix, dtype=np.float64)
if arr.ndim != 2:
raise ValueError(f"{name} must be a 2-D matrix, got ndim={arr.ndim}, shape={arr.shape}.")
if expected_shape is not None and arr.shape != expected_shape:
raise ValueError(f"{name} must have shape {expected_shape}, got {arr.shape}.")
if not np.all(np.isfinite(arr)):
raise ValueError(f"{name} contains non-finite entries (NaN or inf).")
return arr
def _check_symmetric(name: str, matrix: FloatMatrix, *, tol: float = 1e-8) -> None:
asymmetry = float(np.max(np.abs(matrix - matrix.T)))
if asymmetry > tol:
raise ValueError(f"{name} must be symmetric; max |M - Mᵀ| = {asymmetry:.2e} > {tol:.0e}.")
def _check_positive_semidefinite(name: str, matrix: FloatMatrix, *, tol: float = 1e-8) -> None:
min_eig = float(np.min(np.linalg.eigvalsh(matrix)))
if min_eig < -tol:
raise ValueError(f"{name} must be positive semi-definite; min eigenvalue = {min_eig:.2e}.")
def _check_positive_definite(name: str, matrix: FloatMatrix, *, tol: float = 1e-8) -> None:
min_eig = float(np.min(np.linalg.eigvalsh(matrix)))
if min_eig <= tol:
raise ValueError(f"{name} must be positive definite; min eigenvalue = {min_eig:.2e}.")
def solve_finite_horizon_lqr(
A: npt.ArrayLike,
B: npt.ArrayLike,
Q: npt.ArrayLike,
R: npt.ArrayLike,
N: int,
*,
terminal_cost: npt.ArrayLike | None = None,
) -> LQRSolution:
"""Compute the time-varying optimal gains via the backward Riccati recursion.
Args:
A: State matrix, shape (nx, nx).
B: Input matrix, shape (nx, nu).
Q: State weight, symmetric PSD, shape (nx, nx).
R: Input weight, symmetric PD, shape (nu, nu).
N: Horizon length (positive integer).
terminal_cost: Optional terminal weight Q_f, symmetric PSD, shape (nx, nx).
Defaults to Q.
Returns:
LQRSolution with the gain schedule and cost-to-go matrices.
Raises:
ValueError: On malformed or mis-shaped inputs, or violated definiteness.
LinAlgError: If (R + Bᵀ P B) is singular at some recursion step.
"""
if not isinstance(N, (int, np.integer)) or int(N) < 1:
raise ValueError(f"Horizon N must be a positive integer, got {N!r}.")
N = int(N)
A_mat = _as_float_matrix("A", A)
nx = A_mat.shape[0]
if A_mat.shape[1] != nx:
raise ValueError(f"A must be square (nx, nx); got {A_mat.shape}.")
B_mat = _as_float_matrix("B", B)
if B_mat.shape[0] != nx:
raise ValueError(f"B must have {nx} rows to match A; got {B_mat.shape}.")
nu = B_mat.shape[1]
Q_mat = _as_float_matrix("Q", Q, expected_shape=(nx, nx))
R_mat = _as_float_matrix("R", R, expected_shape=(nu, nu))
_check_symmetric("Q", Q_mat)
_check_positive_semidefinite("Q", Q_mat)
_check_symmetric("R", R_mat)
_check_positive_definite("R", R_mat)
if terminal_cost is None:
P_terminal = Q_mat.copy()
else:
P_terminal = _as_float_matrix("terminal_cost", terminal_cost, expected_shape=(nx, nx))
_check_symmetric("terminal_cost", P_terminal)
_check_positive_semidefinite("terminal_cost", P_terminal)
gains = np.zeros((N, nu, nx), dtype=np.float64)
cost_to_go = np.zeros((N + 1, nx, nx), dtype=np.float64)
cost_to_go[N] = P_terminal
for k in range(N - 1, -1, -1):
P_next = cost_to_go[k + 1]
S = R_mat + B_mat.T @ P_next @ B_mat
rhs = B_mat.T @ P_next @ A_mat
try:
K_k = solve(S, rhs, assume_a="pos")
except LinAlgError as exc:
raise LinAlgError(
f"Riccati step k={k} failed: (R + Bᵀ P B) is singular or ill-conditioned."
) from exc
closed_loop = A_mat - B_mat @ K_k
P_k = Q_mat + K_k.T @ R_mat @ K_k + closed_loop.T @ P_next @ closed_loop
P_k = 0.5 * (P_k + P_k.T)
gains[k] = K_k
cost_to_go[k] = P_k
return LQRSolution(gains=gains, cost_to_go=cost_to_go)
def first_control(solution: LQRSolution, x: npt.ArrayLike) -> FloatVector:
"""Return u_0 = -K_0 @ x, the control to apply in the current MPC step."""
K0 = solution.gains[0]
nx = K0.shape[1]
x_vec = np.asarray(x, dtype=np.float64).reshape(-1)
if x_vec.shape[0] != nx:
raise ValueError(f"State x must have length {nx}, got {x_vec.shape[0]}.")
if not np.all(np.isfinite(x_vec)):
raise ValueError("State x contains non-finite entries (NaN or inf).")
return -K0 @ x_vec
def simulate(
A: npt.ArrayLike,
B: npt.ArrayLike,
solution: LQRSolution,
x0: npt.ArrayLike,
) -> tuple[FloatMatrix, FloatMatrix]:
"""Roll the closed loop forward using the full time-varying gain schedule.
Returns:
states: shape (N + 1, nx), states[0] == x0.
controls: shape (N, nu).
"""
B_mat = _as_float_matrix("B", B)
nx, nu = B_mat.shape
A_mat = _as_float_matrix("A", A, expected_shape=(nx, nx))
x_vec = np.asarray(x0, dtype=np.float64).reshape(-1)
if x_vec.shape[0] != nx:
raise ValueError(f"x0 must have length {nx}, got {x_vec.shape[0]}.")
horizon = solution.gains.shape[0]
states = np.zeros((horizon + 1, nx), dtype=np.float64)
controls = np.zeros((horizon, nu), dtype=np.float64)
states[0] = x_vec
for k in range(horizon):
u = -solution.gains[k] @ states[k]
controls[k] = u
states[k + 1] = A_mat @ states[k] + B_mat @ u
return states, controls
def mpc_step(
A: npt.ArrayLike,
B: npt.ArrayLike,
Q: npt.ArrayLike,
R: npt.ArrayLike,
N: int,
x: npt.ArrayLike,
*,
u_min: npt.ArrayLike | None = None,
u_max: npt.ArrayLike | None = None,
terminal_cost: npt.ArrayLike | None = None,
) -> FloatVector:
"""One receding-horizon step: solve, take the first move, clamp to limits.
NOTE: clamping is a *safety clamp only*. Under active constraints it is not
optimal — use a QP-based MPC if your limits bind frequently.
"""
solution = solve_finite_horizon_lqr(A, B, Q, R, N, terminal_cost=terminal_cost)
u = first_control(solution, x)
if u_min is not None or u_max is not None:
lower = -np.inf if u_min is None else np.asarray(u_min, dtype=np.float64).reshape(-1)
upper = np.inf if u_max is None else np.asarray(u_max, dtype=np.float64).reshape(-1)
u = np.clip(u, lower, upper)
return u
def steady_state_lqr_control(
A: npt.ArrayLike,
B: npt.ArrayLike,
Q: npt.ArrayLike,
R: npt.ArrayLike,
x: npt.ArrayLike,
) -> FloatVector:
"""Infinite-horizon (steady-state) LQR control via the discrete ARE.
This is the N -> infinity limit of the finite-horizon gain, so it ignores any
horizon length: use it when the horizon is long enough that the end-of-horizon
transient is negligible. `scipy.linalg.solve_discrete_are` returns the Riccati
solution X directly — no optional dependencies required.
"""
A_mat = _as_float_matrix("A", A)
nx = A_mat.shape[0]
if A_mat.shape[1] != nx:
raise ValueError(f"A must be square (nx, nx); got {A_mat.shape}.")
B_mat = _as_float_matrix("B", B)
if B_mat.shape[0] != nx:
raise ValueError(f"B must have {nx} rows to match A; got {B_mat.shape}.")
Q_mat = _as_float_matrix("Q", Q, expected_shape=(nx, nx))
R_mat = _as_float_matrix("R", R, expected_shape=(B_mat.shape[1], B_mat.shape[1]))
_check_symmetric("Q", Q_mat)
_check_positive_semidefinite("Q", Q_mat)
_check_symmetric("R", R_mat)
_check_positive_definite("R", R_mat)
X = solve_discrete_are(A_mat, B_mat, Q_mat, R_mat)
K_inf = solve(R_mat + B_mat.T @ X @ B_mat, B_mat.T @ X @ A_mat, assume_a="pos")
x_vec = np.asarray(x, dtype=np.float64).reshape(-1)
if x_vec.shape[0] != nx:
raise ValueError(f"State x must have length {nx}, got {x_vec.shape[0]}.")
return -K_inf @ x_vec
if __name__ == "__main__":
dt = 0.1
A = np.array([[1.0, dt], [0.0, 1.0]])
B = np.array([[0.5 * dt**2], [dt]])
Q = np.diag([10.0, 1.0])
R = np.array([[0.1]])
N = 50
x0 = np.array([1.0, 0.0])
solution = solve_finite_horizon_lqr(A, B, Q, R, N)
states, controls = simulate(A, B, solution, x0)
print(f"Initial state : {x0}")
print(f"Final state : {states[-1]} (should be near the origin)")
print(f"First control : {controls[0]}")
print(f"Gain K_0 : {solution.gains[0]}")
import numpy as np
import pytest
from finite_horizon_lqr import (
solve_finite_horizon_lqr,
simulate,
steady_state_lqr_control,
)
def _double_integrator(dt: float = 0.1):
A = np.array([[1.0, dt], [0.0, 1.0]])
B = np.array([[0.5 * dt**2], [dt]])
Q = np.diag([10.0, 1.0])
R = np.array([[0.1]])
return A, B, Q, R
def test_gain_converges_to_steady_state():
A, B, Q, R = _double_integrator()
sol = solve_finite_horizon_lqr(A, B, Q, R, N=200)
x_probe = np.array([1.0, 0.0])
u_finite = -sol.gains[0] @ x_probe
u_infinite = steady_state_lqr_control(A, B, Q, R, x_probe)
np.testing.assert_allclose(u_finite, u_infinite, rtol=1e-4, atol=1e-6)
def test_closed_loop_is_stable():
A, B, Q, R = _double_integrator()
sol = solve_finite_horizon_lqr(A, B, Q, R, N=100)
states, _ = simulate(A, B, sol, x0=np.array([1.0, 0.0]))
assert np.linalg.norm(states[-1]) < 1e-2
def test_rejects_non_positive_definite_R():
A, B, Q, _ = _double_integrator()
bad_R = np.array([[0.0]])
with pytest.raises(ValueError):
solve_finite_horizon_lqr(A, B, Q, bad_R, N=10)
def test_rejects_mismatched_shapes():
A, B, Q, R = _double_integrator()
wrong_B = np.array([[1.0, 0.0]])
with pytest.raises(ValueError):
solve_finite_horizon_lqr(A, wrong_B, Q, R, N=10)
