JAX-Native Workflows¶

These examples show differentiable controller design, batched linearization, compiled simulation, and nonlinear estimation workflows.

Runnable scripts:

examples/differentiable_lqr.py
examples/linearize_lqr_simulate.py
examples/continuous_lqr.py
examples/lqr_optimal_execution.py
examples/continuous_nonlinear_estimation.py
examples/structured_nonlinear_estimation.py

Related API pages:

1. Tune LQR Weights With Gradients¶

dare and lqr can sit inside an objective and be differentiated with respect to controller weights.

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 objective(q_diag, log_r):
    Q = jnp.diag(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(objective, argnums=(0, 1)))
cost, (dq, dlog_r) = objective_and_grad(jnp.ones(2), jnp.array(0.0))

The objective differentiates through controller design and simulation as one JAX computation.

Related pages: - Tune LQR with gradients - Differentiable LQR

2. Turn Nonlinear Dynamics Into State-Space Models¶

linearize (also available as linearize_ss) lets you move from nonlinear plant code to local state-space models with one JAX-compatible call.

import jax

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

import jax.numpy as jnp
import contrax as cx


def pendulum(t, x, u):
    theta, theta_dot = x
    torque = u[0]
    return jnp.array([theta_dot, -jnp.sin(theta) + torque])


def sensor(x, u):
    return x


x_eq = jnp.array([0.1, 0.0])
u_eq = jnp.array([jnp.sin(0.1)])

sys_c = cx.linearize(pendulum, x_eq, u_eq, output=sensor)
sys_d = cx.c2d(sys_c, dt=0.05)

That makes it easy to keep your plant model nonlinear and only linearize where you actually need a controller or estimator.

Related pages: - Systems API - Discretization and linearization

3. Design Controllers Over Batches Of Operating Points¶

Because the workflow is JAX-native, you can vmap over operating points or design conditions.

import jax

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

import jax.numpy as jnp
import contrax as cx


def pendulum(t, x, u):
    return jnp.array([x[1], -jnp.sin(x[0]) + u[0]])


def sensor(x, u):
    return x


def design(x_eq, u_eq):
    sys_c = cx.linearize(pendulum, x_eq, u_eq, output=sensor)
    sys_d = cx.c2d(sys_c, dt=0.05)
    return cx.lqr(sys_d, jnp.eye(2), jnp.ones((1, 1))).K


x_eqs = jnp.array([[0.0, 0.0], [0.1, 0.0], [-0.1, 0.0]])
u_eqs = jnp.zeros((3, 1))
batched_design = jax.jit(jax.vmap(design))
Ks = batched_design(x_eqs, u_eqs)

Use the batched design for gain scheduling, operating-point sweeps, or controller redesign over a grid.

4. JIT The Full Closed-Loop Path¶

Controller design and simulation can stay in one compiled pipeline.

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)


@jax.jit
def run(q_scale):
    Q = jnp.diag(jnp.array([10.0, 1.0]) * q_scale)
    R = jnp.array([[1.0]])
    K = cx.lqr(sys, Q, R).K
    return cx.simulate(sys, jnp.array([1.0, 0.0]), lambda t, x: -K @ x, num_steps=60)


ts, xs, ys = run(jnp.array(1.0))

That keeps redesign and simulation in the same JAX program instead of pushing part of the workflow into a separate control toolbox.

5. Estimate Continuous Nonlinear Systems On A Discrete Grid¶

sample_system() integrates a continuous plant over each sample interval and returns a discrete transition map. Use foh_inputs() to provide first-order-hold endpoint pairs.

Runnable scripts: - examples/continuous_nonlinear_estimation.py - examples/structured_nonlinear_estimation.py

6. Estimation And Optimization Live In The Same JAX World¶

Batch filters use scans, one-step helpers fit runtime loops, and MHE objective construction provides a differentiable cost. These workflows share model and array conventions with LTI design.

7. Cast Quadratic Optimal Execution As LQR¶

Quadratic execution models map naturally to standard control objects. This example represents remaining inventory as the state and signed inventory change as the control.

Treat remaining inventory as the state and signed inventory change as the control:

import jax

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

import jax.numpy as jnp
import contrax as cx

A = jnp.array([[1.0]])
B = jnp.array([[1.0]])
sys = cx.dss(A, B, jnp.array([[1.0]]), jnp.zeros((1, 1)), dt=1.0)

Q = jnp.array([[2.5]])   # inventory risk
R = jnp.array([[0.4]])   # temporary impact / trading cost
K = cx.lqr(sys, Q, R).K

Then the usual LQR feedback law gives a liquidation schedule. Because the whole path stays inside JAX, you can also differentiate through the Riccati solve to tune the execution urgency or vmap the design across many assets.

Runnable script: - examples/lqr_optimal_execution.py

Related page: - LQR optimal execution