Unscented Kalman Filter (UKF)

f_h Functions

Derivation, argument reference, and worked example for the ukf block.

Overview

The UKF handles nonlinear systems without requiring Jacobians. Instead of linearising, it propagates a small set of deterministically chosen sigma points through the true nonlinear functions and reconstructs the mean and covariance from the results.

This makes it more accurate than the EKF for strongly nonlinear systems and eliminates the need to derive or compute Jacobians analytically.

State-Space Model

Same as the EKF — arbitrary nonlinear \(f\) and \(h\):

\[x_{k+1} = f(x_k, u_k) + w_k, \qquad w_k \sim \mathcal{N}(0, Q_k)\]
\[y_k = h(x_k) + v_k, \qquad v_k \sim \mathcal{N}(0, R_k)\]

Sigma Points

Given state dimension \(n\) and scaling parameter \(\lambda = \alpha^2(n + \kappa) - n\), construct \(2n + 1\) sigma points from the current estimate \(\hat{x}_{k-1}\) and covariance \(P_{k-1}\):

\[\mathcal{X}^{(0)} = \hat{x}_{k-1}\]
\[\mathcal{X}^{(i)} = \hat{x}_{k-1} + \left(\sqrt{(n+\lambda)\,P_{k-1}}\right)_i, \quad i = 1,\ldots,n\]
\[\mathcal{X}^{(i+n)} = \hat{x}_{k-1} - \left(\sqrt{(n+\lambda)\,P_{k-1}}\right)_i, \quad i = 1,\ldots,n\]

where \(\left(\sqrt{\cdot}\right)_i\) denotes the \(i\)-th column of the Cholesky factor.

Weights for the mean (\(W_m\)) and covariance (\(W_c\)):

\[W_m^{(0)} = \frac{\lambda}{n+\lambda}, \qquad W_m^{(i)} = \frac{1}{2(n+\lambda)}, \quad i > 0\]
\[W_c^{(0)} = W_m^{(0)} + 1 - \alpha^2 + \beta, \qquad W_c^{(i)} = W_m^{(i)}, \quad i > 0\]

Predict Step

Propagate each sigma point through \(f\), then compute the weighted mean and covariance:

\[\mathcal{X}^{(i)}_{k|k-1} = f\left(\mathcal{X}^{(i)}, u_k\right)\]
\[\hat{x}_{k|k-1} = \sum_i W_m^{(i)}\, \mathcal{X}^{(i)}_{k|k-1}\]
\[P_{k|k-1} = \sum_i W_c^{(i)}\, (\mathcal{X}^{(i)}_{k|k-1} - \hat{x}_{k|k-1})(\mathcal{X}^{(i)}_{k|k-1} - \hat{x}_{k|k-1})^\top + Q_k\]

Update Step

Map the predicted sigma points through \(h\), compute the innovation statistics, then apply the Kalman correction:

\[\mathcal{Y}^{(i)} = h\left(\mathcal{X}^{(i)}_{k|k-1}\right)\]
\[\hat{y}_{k|k-1} = \sum_i W_m^{(i)}\, \mathcal{Y}^{(i)}\]
\[S_k = \sum_i W_c^{(i)}\,(\mathcal{Y}^{(i)} - \hat{y})(\mathcal{Y}^{(i)} - \hat{y})^\top + R_k\]
\[P_{xy} = \sum_i W_c^{(i)}\,(\mathcal{X}^{(i)}_{k|k-1} - \hat{x}_{k|k-1})(\mathcal{Y}^{(i)} - \hat{y})^\top\]
\[K_k = P_{xy}\, S_k^{-1}\]
\[\hat{x}_k = \hat{x}_{k|k-1} + K_k\,(y_k - \hat{y}_{k|k-1})\]
\[P_k = P_{k|k-1} - K_k\, S_k\, K_k^\top\]

Argument Reference

zk — UKF state   tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]]

Field

Symbol

Shape

Description

zk[0]

\(\hat{x}_k\)

(n,)

Current state estimate

zk[1]

\(P_k\)

(n, n)

Error covariance — symmetric, positive semi-definite

Initialisation: zk = (x0, P0). Same convention as the KF and EKF.


uk — control input   float

Scalar control input at step \(k\). Passed directly to f_nl(x, uk). Pass 0.0 if uncontrolled.


yk — measurement   Float[np.ndarray, "m"]

Raw sensor reading at step \(k\), shape (m,). Must match the output space of h_nl.


f_nl — nonlinear dynamics   Callable

Signature: f_nl(x: ndarray, u: float) -> ndarray

The true nonlinear state transition. Each sigma point is propagated through this function. No Jacobian of f_nl is required.


h_nl — nonlinear observation   Callable

Signature: h_nl(x: ndarray) -> ndarray

The true nonlinear observation function. Each predicted sigma point is mapped through this to form the innovation statistics. No Jacobian of h_nl is required.


Qk — process noise covariance   Float[np.ndarray, "n n"]

Shape (n, n), symmetric positive semi-definite. Added to the predicted covariance. Same role as in the KF.


Rk — measurement noise covariance   Float[np.ndarray, "m m"]

Shape (m, m), symmetric positive definite. Added to the innovation covariance. Same role as in the KF.


alpha — sigma point spread   float (default 1e-3)

Controls how far the sigma points spread around the mean. Typical range: \(10^{-4}\) to \(1\). Smaller values keep sigma points close to the mean; larger values explore more of the nonlinear region. Affects \(\lambda = \alpha^2(n + \kappa) - n\).


beta — distribution parameter   float (default 2.0)

Incorporates prior knowledge of the state distribution. \(\beta = 2\) is optimal for Gaussian distributions and is the correct default in most robotics/control applications. Only the zeroth covariance weight \(W_c^{(0)}\) is affected.


kappa — secondary scaling   float (default 0.0)

Secondary scaling parameter. Common choices are \(\kappa = 0\) or \(\kappa = 3 - n\) (the latter ensures the fourth-order terms of the Taylor expansion are captured exactly for Gaussian distributions). Typically left at 0.0.

Implementation

ukf_f runs the full sigma-point predict→update cycle and returns the new state and covariance. ukf_h extracts \(\hat{x}_k\) for downstream use.

The **kwargs in ukf_h absorbs any extra arguments passed by _smart_call that h does not need.

import numpy as np
from typing import Callable
from jaxtyping import Float, jaxtyped
from beartype import beartype
from dynamicalnodes import DynamicalSystem
@jaxtyped(typechecker=beartype)
def ukf_f(
    zk: tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]],
    uk: float,
    yk: Float[np.ndarray, "m"],
    f_nl: Callable,
    h_nl: Callable,
    Qk: Float[np.ndarray, "n n"],
    Rk: Float[np.ndarray, "m m"],
    alpha: float = 1e-3,
    beta: float = 2.0,
    kappa: float = 0.0,
) -> tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]]:
    """UKF predict+update via sigma-point propagation; returns updated state and covariance.

    Args:
        zk:    UKF state — 2-tuple of (x̂_k, P_k), state estimate and error covariance.
        uk:    Scalar control input.
        yk:    Measurement vector (m,).
        f_nl:  Nonlinear state transition f(x, u) → x_next.
        h_nl:  Nonlinear observation function h(x) → y.
        Qk:    Process noise covariance (n, n).
        Rk:    Measurement noise covariance (m, m).
        alpha: Sigma point spread around the mean (default 1e-3).
        beta:  Distribution parameter; 2.0 optimal for Gaussian (default 2.0).
        kappa: Secondary scaling parameter, typically 0 or 3-n (default 0.0).
    """
    x, P = zk
    n = len(x)
    lam = alpha**2 * (n + kappa) - n

    sqrtP = np.linalg.cholesky((n + lam) * P)
    sigmas = np.column_stack(
        [x] + [x + sqrtP[:, i] for i in range(n)] + [x - sqrtP[:, i] for i in range(n)]
    )
    Wm = np.full(2 * n + 1, 1.0 / (2 * (n + lam)))
    Wm[0] = lam / (n + lam)
    Wc = Wm.copy()
    Wc[0] += 1 - alpha**2 + beta

    # Predict
    sf = np.column_stack([f_nl(sigmas[:, i], uk) for i in range(2 * n + 1)])
    xp = sf @ Wm
    Pp = (
        sum(Wc[i] * np.outer(sf[:, i] - xp, sf[:, i] - xp) for i in range(2 * n + 1))
        + Qk
    )

    # Update
    yk = np.atleast_1d(yk)
    sh = np.column_stack([np.atleast_1d(h_nl(sf[:, i])) for i in range(2 * n + 1)])
    yp = sh @ Wm
    S = (
        sum(Wc[i] * np.outer(sh[:, i] - yp, sh[:, i] - yp) for i in range(2 * n + 1))
        + Rk
    )
    Pxy = sum(Wc[i] * np.outer(sf[:, i] - xp, sh[:, i] - yp) for i in range(2 * n + 1))
    K = Pxy @ np.linalg.inv(S)
    return (xp + K @ (yk - yp), Pp - K @ S @ K.T)


@jaxtyped(typechecker=beartype)
def ukf_h(
    zk: tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]],
    **kwargs,
) -> Float[np.ndarray, "n"]:
    """Return the current state estimate from the UKF state tuple.

    Args:
        zk: UKF state — 2-tuple of (x̂_k, P_k), state estimate and error covariance.
    """
    return zk[0]


ukf = DynamicalSystem(f=ukf_f, h=ukf_h)

Worked Example — Pendulum

Same pendulum as the EKF example — no Jacobians required.

g, L, dt = 9.81, 1.0, 0.05


def pend_f(x, u):
    th, w = x
    return np.array([th + dt * w, w + dt * (-g / L * np.sin(th) + u)])


def pend_h(x):
    return np.array([x[0]])


Qk = np.diag([1e-4, 1e-3])
Rk = np.array([[0.05]])

zk = (np.array([0.1, 0.0]), np.eye(2))
zk_next, x_est = ukf.step(
    zk=zk, uk=0.0, yk=np.array([0.09]), f_nl=pend_f, h_nl=pend_h, Qk=Qk, Rk=Rk
)
print("estimate (pre-update):", x_est)
print("updated estimate:     ", zk_next[0])
print("updated covariance:\n",  zk_next[1])

DynamicalSystemROSNode

from dynamicalnodes import ROSNode
from dynamicalnodes.ros2py_py2ros import (
    ros2py_float64_multiarray,
    ros2py_float64,
    py2ros_float64_multiarray,
)
from std_msgs.msg import Float64MultiArray, Float64

ukf_node = ROSNode(
    dynamical_system=ukf,
    subscribes_to=[
        {
            "topic": "/measurement",
            "msg_type": Float64MultiArray,
            "arg": "yk",
            "ros2py": ros2py_float64_multiarray,
        },
        {
            "topic": "/control",
            "msg_type": Float64,
            "arg": "uk",
            "ros2py": ros2py_float64,
        },
    ],
    publishes_to=[
        {
            "topic": "/state_estimate",
            "msg_type": Float64MultiArray,
            "py2ros": py2ros_float64_multiarray,
        },
    ],
)

ROSNode.write_ROSNode_to_rclpy()

g, L, dt_ukf = 9.81, 1.0, 0.1


def pendulum_f(x, u):
    theta, omega = x
    alpha = -g / L * np.sin(theta) + u
    return np.array([theta + omega * dt_ukf, omega + alpha * dt_ukf])


def pendulum_h(x):
    return np.array([x[0]])


ukf_node.write_ROSNode_to_rclpy(
    "ukf_node.py",
    node_name="unscented_kalman_filter",
    deps=[pendulum_f, pendulum_h],
    initial_state=(
        np.array([0.0, 0.0]),  # x̂_0: [θ, ω]
        np.eye(2),             # P_0
    ),
    static_params={
        "f_nl": pendulum_f,
        "h_nl": pendulum_h,
        "Qk": np.diag([1e-4, 1e-3]),
        "Rk": np.array([[0.01]]),
        "alpha": 1e-3,
        "beta": 2.0,
        "kappa": 0.0,
    },
    initial_inputs={
        "yk": np.array([0.0]),
        "uk": 0.0,
    },
)

f_h Functions