Skip to content

How To Build An MHE Objective

This guide shows how to construct a fixed-window moving-horizon-estimation objective before choosing an optimizer.

Complete Working Code

import jax
import jax.numpy as jnp
import contrax as cx


def f(x, u, params):
    return jnp.array([params["a"] * x[0] + u[0] ** 2])


def h(x, params):
    return jnp.array([params["c"] * x[0]])


params = {"a": 0.5, "c": 2.0}
us = jnp.array([[0.0], [2.0], [1.0]])
xs = cx.rollout(f, jnp.array([1.0]), us, params)
ys = jax.vmap(lambda x: h(x, params))(xs)

cost = cx.mhe_objective(
    f,
    observation=h,
    xs=xs,
    us=us,
    ys=ys,
    x_prior=xs[0],
    P_prior=jnp.eye(1),
    Q_noise=jnp.eye(1),
    R_noise=jnp.eye(1),
    params=params,
)

Alignment Rules

The important shape rule is:

  • xs: (N, n)
  • us: (N - 1, m)
  • ys: (N, p)

That alignment matches:

  • one arrival state at the beginning of the window
  • one process residual per transition
  • one measurement residual per state in the window

Objective And Optimizer

mhe_objective() exposes the fixed-window cost directly. This lets you:

  • keep the horizon fixed-shape
  • add custom soft costs explicitly
  • choose your optimizer separately from the model definition

Two small helpers are useful in rolling-window workflows:

  • mhe_warm_start(xs, ...) shifts a previous trajectory guess forward by one step, optionally propagating the final guess with a transition model
  • soft_quadratic_penalty(residuals, weight) builds a weighted quadratic soft cost that you can add inside extra_cost

Use mhe() to minimize this objective with a configured optimizer.

Check

For an exact trajectory and measurement sequence, the objective should go to zero up to numerical tolerance.