Controllers

Controllers map a reference \(r_k\) and a state estimate \(\hat{x}_k\) to a control input \(u_k\). h always returns \(u_k\) so the output can be fed directly into a plant block.

PID Controller

Classic proportional–integral–derivative controller.See the pid notebook for derivation, API reference, and examples.

from jaxtyping import Float, jaxtyped
from beartype import beartype


@jaxtyped(typechecker=beartype)
def pid_f(
    ck: tuple[float, float],
    rk: float,
    xhatk: float,
    dt: float,
) -> tuple[float, float]:
    """Advance the PID state; returns (current error, updated integral).

    Args:
        ck:    PID state — 2-tuple of (e_{k-1}, e_int), previous error and accumulated integral.
        rk:    Reference (setpoint) at step k.
        xhatk: Current scalar state estimate.
        KP:    Proportional gain.
        KI:    Integral gain.
        KD:    Derivative gain.
        dt:    Timestep in seconds.

    Returns:
        tuple[float, float] — (e_k, e_int_k), the current error and updated integral.
        Fed back as ck on the next step; not published directly to ROS.
    """
    _, e_int = ck
    ek = rk - xhatk
    return (ek, e_int + ek * dt)


@jaxtyped(typechecker=beartype)
def pid_h(
    ck: tuple[float, float],
    rk: float,
    xhatk: float,
    KP: float,
    KI: float,
    KD: float,
    dt: float,
) -> float:
    """Compute the PID control output u_k = KP*e + KI*e_int + KD*de/dt.

    Args:
        ck:    PID state — 2-tuple of (e_{k-1}, e_int), previous error and accumulated integral.
        rk:    Reference (setpoint) at step k.
        xhatk: Current scalar state estimate.
        KP:    Proportional gain.
        KI:    Integral gain.
        KD:    Derivative gain.
        dt:    Timestep in seconds.

    Returns:
        float — control output u_k.
        Published as Float64 via py2ros_float64.
    """
    e_prev, e_int = ck
    ek = rk - xhatk
    return KP * ek + KI * e_int + KD * (ek - e_prev) / dt

LQR Controller

Linear Quadratic Regulator — optimal state-feedback for linear systems. See the lqr notebook for derivation, API reference, and examples.

import numpy as np
from typing import Callable
from jaxtyping import Float, jaxtyped
from beartype import beartype
from scipy.linalg import solve_discrete_are


@jaxtyped(typechecker=beartype)
def lqr_gain(
    A: Float[np.ndarray, "n n"],
    B: Float[np.ndarray, "n m"],
    Q: Float[np.ndarray, "n n"],
    R: Float[np.ndarray, "m m"],
) -> Float[np.ndarray, "m n"]:
    """Compute the discrete-time LQR gain K = (R + B'PB)^{-1} B'PA via DARE.

    Args:
        A: State transition matrix (n, n) of the linearised plant.
        B: Control input matrix (n, m) — maps u_k into state space.
        Q: State cost matrix (n, n), symmetric positive semi-definite.
        R: Control cost matrix (m, m), symmetric positive definite.

    Returns:
        ndarray(m, n) — optimal gain matrix K. Compute once; pass as a static
        parameter to lqr_h at every step.
    """
    P = solve_discrete_are(A, B, Q, R)
    return np.linalg.inv(R + B.T @ P @ B) @ B.T @ P @ A


@jaxtyped(typechecker=beartype)
def lqr_h(
    xhatk: Float[np.ndarray, "n"],
    xrefk: Float[np.ndarray, "n"],
    K: Float[np.ndarray, "m n"],
) -> Float[np.ndarray, "m"]:
    """Compute the LQR control output u_k = -K(x̂_k - x_ref_k).

    Args:
        xhatk: Current state estimate (n,) from an estimator or sensor.
        xrefk: Reference (target) state (n,).
        K:     Precomputed LQR gain matrix (m, n) from lqr_gain.

    Returns:
        ndarray(m,) — control input u_k.
        Published as Float64MultiArray via py2ros_float64_multiarray.
    """
    return -K @ (xhatk - xrefk)