Skip to content

Nonlinear MHE with Constrained Parameter Estimation

This example estimates both state and a physical parameter from noisy measurements using Moving Horizon Estimation. The parameter (damping coefficient) is kept positive throughout optimization by representing it in an unconstrained raw form and applying positive_softplus() inside the dynamics.

Runnable script: examples/nonlinear_mhe.py

Problem Setup

A damped pendulum is driven by a small random torque:

\[ \ddot{\theta} = -\frac{g}{l} \sin\theta - b\,\dot{\theta} + \tau, \qquad y_k = \theta(t_k) + v_k. \]

The damping coefficient \(b > 0\) is unknown. The estimator must recover both the state \((\theta, \dot\theta)\) and \(b\) jointly from angle measurements alone.

Augmented State and Constrained Parameterization

import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp

import contrax as cx

DT = 0.05  # sample period (s)
WINDOW = 12  # MHE window length
G_ACCEL = 9.81  # gravitational acceleration (m/s²)
L = 1.0  # pendulum length (m), assumed known
DAMPING_TRUE = 0.3  # true damping coefficient (kg·m²/s)
N_AUG = 3  # augmented state dimension: [theta, omega, b_raw]

def aug_dynamics(x_aug, u):
    """Euler step for augmented state z = [theta, omega, b_raw].

    The damping parameter b_raw is treated as a slowly-varying state:
    its dynamics are an identity (random-walk model). The physical damping
    b = positive_softplus(b_raw) is always positive.
    """
    theta, omega, b_raw = x_aug[0], x_aug[1], x_aug[2]
    b = cx.positive_softplus(b_raw)
    torque = u[0]
    alpha = -(G_ACCEL / L) * jnp.sin(theta) - b * omega + torque
    theta_next = theta + DT * omega
    omega_next = omega + DT * alpha
    return jnp.array([theta_next, omega_next, b_raw])  # b_raw unchanged


def aug_observation(x_aug):
    """Observe angle only from augmented state."""
    return x_aug[:1]

The model augments the state with \(b_\text{raw}\) using identity dynamics (random-walk model), then recovers the physical parameter with

\[ b = \text{softplus}(b_\text{raw}) > 0. \]

The softplus map keeps \(b > 0\) for every raw parameter value, so the optimizer works over unconstrained \(b_\text{raw}\) values.

Solving MHE

def run_mhe(xs_phys, ys, us):
    """Solve MHE over the first window and report results."""
    # Noise covariances in augmented state space
    Q_noise = jnp.diag(jnp.array([1e-4, 1e-3, 1e-5]))  # small drift on b_raw
    R_noise = jnp.array([[0.04]])

    # Prior: start near the first measurement; b_raw=0 → b≈0.69 (soft overestimate)
    x_prior = jnp.array([float(ys[0, 0]), 0.0, 0.0])
    P_prior = jnp.diag(jnp.array([0.1, 0.5, 2.0]))

    ys_win = ys[: WINDOW + 1]  # (WINDOW+1, 1)
    us_win = us[:WINDOW]  # (WINDOW, 1)

    # Warm-start: replicate prior across window
    xs_init = jnp.tile(x_prior, (WINDOW + 1, 1))
    model = cx.nonlinear_system(
        lambda t, x, u: aug_dynamics(x, u),
        output=lambda t, x, u: aug_observation(x),
        dt=DT,
    )

    result = cx.mhe(
        model,
        xs_init,
        us_win,
        ys_win,
        x_prior,
        P_prior,
        Q_noise,
        R_noise,
        max_steps=512,
    )
    return result

mhe() minimizes mhe_objective() over the full window trajectory using LBFGS and accepts the nonlinear system model directly. The window has N states and measurements, with N−1 input transitions. The Kalman-style prior term (x_prior, P_prior) regularizes the start of the window, and Q_noise controls how much the parameter is allowed to drift across the horizon.

What the Script Prints

The script reports convergence, status, gradient norm, final cost, estimated damping and terminal angle alongside the ground truth. It also checks positive damping and terminal-angle accuracy.

The angle estimate is close to ground truth. The damping estimate is biased: a 12-step window with a slowly varying random-walk prior gives a short identification horizon for one scalar parameter. Longer windows or a tighter prior on \(b_\text{raw}\) dynamics may improve the estimate. Here solver_converged=False; inspect status, gradient norm, and cost reduction before using the estimate or interpreting implicit gradients.