Skip to content

Contrax

Contrax is a JAX-native systems, control, simulation, and estimation library with MATLAB-familiar names at the API surface and JAX-first behavior underneath.

Controller design can run in the same JAX program as the model. Linearization, simulation, batching, and differentiation compose in one workflow.

  • dare is the most mature Riccati solver path in the library
  • care is supported and validated, but less benchmarked than dare
  • continuous simulate() uses Diffrax and returns values on a fixed save grid

One Shared Loop

Closed-loop overview showing controller, plant, estimator, and their shared state-space model
One shared model: use the same state-space description for controller design, simulation, and estimation.

Why Contrax

Contrax is built for workflows such as:

  • differentiating through lqr() and closed-loop simulation
  • linearizing nonlinear dynamics directly into state-space form
  • vmapping controller design across operating points
  • keeping design and validation inside one compiled objective
Gradient descent through a differentiable Contrax LQR loop, showing a pendulum trajectory improving over 50 gradient steps

50 gradient steps through ss → c2d → lqr → simulate. The pendulum settles faster as the cost weights are tuned automatically. 

import jax

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

import jax.numpy as jnp
import contrax as cx

A = jnp.array([[1.0, 0.05], [0.0, 1.0]])
B = jnp.array([[0.0], [0.05]])
SYS = cx.dss(A, B, jnp.eye(2), jnp.zeros((2, 1)), dt=0.05)
X0 = jnp.array([1.0, 0.0])


def closed_loop_cost(log_q_diag, log_r):
    Q = jnp.diag(jnp.exp(log_q_diag))
    R = jnp.exp(log_r)[None, None]
    K = cx.lqr(SYS, Q, R).K
    trajectory = cx.simulate(SYS, X0, lambda t, x: -K @ x, num_steps=80)
    return jnp.sum(trajectory.xs**2)


objective_and_grad = jax.jit(jax.value_and_grad(closed_loop_cost, argnums=(0, 1)))
cost, grads = objective_and_grad(jnp.zeros(2), jnp.array(0.0))

Control design, simulation, and estimation use JAX arrays and transformations throughout the workflow.

Contrax focuses on systems, control, simulation, and estimation workflows in JAX. The optional python-control bridge supports conversion of LTI state-space models.

Core Equations

Model Families

These equations describe common linear, nonlinear, and port-Hamiltonian models.

\[ \dot{x} = A x + B u, \qquad y = C x + D u \]
\[ \dot{x} = f(t, x, u), \qquad y = h(t, x, u) \]
\[ \dot{x} = \bigl(J - R(x)\bigr)\nabla H(x) + G(x)u \]

Design And Estimation Loops

The design and estimation side stays just as explicit.

\[ x_{k+1} = A x_k + B u_k, \qquad u_k = -K x_k \]
\[ \hat{x}_k^+ = \hat{x}_k^- + K_k\bigl(y_k - h(\hat{x}_k^-)\bigr), \qquad \hat{x}_{k+1}^- = f(\hat{x}_k^+, u_k) \]

Capabilities

  • Systems

contrax.systems for state-space models, nonlinear models, linearization, and interconnection.

  • Control

contrax.control for LQR/LQI, Riccati solvers, pole placement, and state-feedback helpers.

  • Simulation

contrax.simulation for lsim(), simulate(), response helpers, and fixed horizon rollout().

  • Estimation

contrax.estimation for Kalman-family filters, smoothers, and MHE helpers.

  • Analysis

contrax.analysis for structural checks, transfer evaluation, poles, and finite-horizon Gramians.

  • Types

contrax.types for the public result bundles such as LQRResult and KalmanResult.

  • Interoperability

contrax.compat.python_control for optional bidirectional conversion between Contrax LTI models and python-control StateSpace objects.

Where To Start

The shortest path from zero context to a working controller.

Learning-oriented end-to-end workflows for the main control and estimation paths.

Task-oriented recipes such as tuning LQR weights or handling missing measurements.

Namespace-oriented reference pages for systems, control, simulation, estimation, analysis, and types.

Short explanation pages on solver choices, discretization, estimation pipelines, and JAX transform behavior.

Runnable examples that show how batching, transforms, and multi-step workflows fit together.

Solver Guidance

Solver maturity, transform behavior, and numerical assumptions are documented alongside each workflow. The discrete Riccati path has the broadest benchmark coverage; CARE is supported and validated with less benchmark coverage.