mpcjax

Pure JAX implementation of Model Predictive Control (SCP MPC).

This is non-linear dynamics finite horizon MPC solver and support for arbitrary cost. It supports batching on the CPU and GPU.

Installation

Install by issuing

$ pip install mpcjax

Alternatively, installing from source

$ git clone https://github.com/rdyro/mpcjax.git
$ cd mpcjax
$ pip install .

Changelog

• version 0.4.0

  • solve - the main routine, is not jit compatible
  • optimality_fn is not a jit-compatible way of getting optimality conditions - fixed-point
  • added remote solver, allowing a persistent solver process (or a process pool)
    • redis is a hard dependency for the remote mode

• version 0.3.0

  • introduced direct_solve mode in which the rollout itself is differentiate as part of optimization, i.e., direct nonlinear optimization

Basic Usage

The solver is capable of MPC consensus optimization for several system instantiations. For the basic usage, we'll focus on a single system MPC.

A basic MPC problem is defined using the dynamics and a quadratic cost

Defining dynamics

• x0 the initial state of shape

where

np.shape(x0) == (xdim,)

• f, fx, fu = f_fx_fu_fn(xt, ut) an affine dynamics linearization, such that $$x^{(i+1)} \approx f^{(i)} + f_x^{(i)} (x^{(i)} - \tilde{x}^{(i)}) + f_u^{(i)} (u^{(i)} - \tilde{u}^{(i)}) $$ where

np.shape(xt) == (N, xdim)kkk
np.shape(ut) == (N, udim)
np.shape(f) == (N, xdim)
np.shape(fx) == (N, xdim, xdim)
np.shape(fu) == (N, xdim, udim)

Defining Cost

• X_ref, Q a reference and quadratic weight matrix for state cost
• U_ref, R a reference and quadratic weight matrix for control cost

The cost is given as

$$J = \sum_{i=0}^N \frac{1}{2} (x^{(i+1)} - x_\text{ref}^{(i+1)}) Q^{(i)} (x^{(i+1)} - x_\text{ref}^{(i+1)}) + \frac{1}{2} (u^{(i)} - u_\text{ref}^{(i)}) R^{(i)} (u^{(i)} - u_\text{ref}^{(i)})$$

Note: Initial state, x0, is assumed constant and thus does not feature in the cost.

Note: When handling controls, we'll always have np.shape(U) == (N, udim)

Note: When handling states, we'll have either np.shape(X) == (N + 1, xdim) with x0 included at the beginning or np.shape(X) == (N, xdim) with x0 NOT INCLUDED. X[:, -1] always refers to the state N, whereas U[:, -1] always refers to control N - 1.

Thus, an example call would be

>>> import mpcjax
>>> X, U, debug = mpcjax.solve(f_fx_fu_fn, Q, R, x0, X_ref, U_ref)
>>> help(mpcjax.solve)

Take a look at

• tests/simple.py for simple usage
• tests/dubins_car.py for defining dynamics

solve Method Arguments Glossary

Solver Hyperparameters

The solver has two scalar hyperparameters, the dynamics linearization deviation penalty for states and controls

$$ J_\text{deviation} = \sum_{i=0}^N \frac{1}{2} \rho_x (x^{(i+1)} - x_\text{prev}^{(i+1)})^T (x^{(i+1)} - x_\text{prev}^{(i+1)}) + \frac{1}{2} \rho_u (u^{(i)} - u_\text{prev}^{(i)})^T (u^{(i)} - u_\text{prev}^{(i)}) $$

The hyperparameters are

• reg_x - state deviation in-between SCP iterations regularization
• reg_u - control deviation in-between SCP iterations regularization

Higher values will slow evolution between SCP iterations and will require more SCP iterations to converge to a solution, but will avoid instability in the solution if the dynamics are not sufficiently smooth.

Solver Settings

• verbose - whether to print iteration status (user-facing)
• debug - whether to print very low-level debugging information (developer-facing)
• max_it - maximum number of SCP iterations to perform (can be fewer if tolerance met earlier)
• time_limit - the time limit in seconds for SCP iteration
• res_tol - deviation tolerance past which solution is accepted (measure of convergence)
• slew_rate - the quadratic penalty between time-consecutive controls (encourages smooth controls)
• u_slew - the previous action taken to align the first plan action with (useful for smooth receding horizon control)

Additional Dynamics Settings

• X_prev - previous state solution (guess), $x^{(i)} ~~ \forall i \in [1, \dots, N]$, shape = (N, xdim)
• U_prev - previous control solution (guess), $u^{(i)} ~~ \forall i \in [0, \dots, N - 1]$, shape = (N, udim)
• x_l - state lower box constraints, $x^{(i)} ~~ \forall i \in [1, \dots, N]$, shape = (N, xdim)
• x_u - state upper box constraints, $x^{(i)} ~~ \forall i \in [1, \dots, N]$, shape = (N, xdim)
• u_l - control lower box constraints, $u^{(i)} ~~ \forall i \in [0, \dots, N - 1]$, shape = (N, udim)
• u_u - control upper box constraints, $u^{(i)} ~~ \forall i \in [0, \dots, N - 1]$, shape = (N, udim)

Nonlinear Cost and Constraints

The solver supports custom arbitrary cost via each-SCP-iteration cost linearization and custom constraints via each-SCP-iteration constraint reformulation into any convex-cone constraint.

• lin_cost_fn is an optional callable which allows specifying a custom cost, it should take arbitrary X, U and return a tuple
  • cx, the linearization of the cost with respect to the state, np.shape(cx) == (N, xdim) or cx is None
  • cu, the linearization of the cost with respect to the controls, np.shape(cu) == (N, udim) or cu is None

I highly recommend using an auto-diff library to produce the linearizations to avoid unnecessary bugs.

• diff_cost_fn is an optional callable for specifying arbitrary differentiable cost
  • this includes custom penalized constraints (e.g., log-barrier or augmented-Lagrangian penalties)
  • the function can return NaNs for infeasible arguments, but a feasible guess must be provided in X_prev, U_prev

The functions must accept the call lin_cost_fn(X, U, problem) and diff_cost_fn(X, U, problem) where problem is a Python dictionary with problem data.

Misc Settings

• solver_settings - a dictionary of settings to pass to the lower-level Julia solver
• solver_state - a previous solver state to pass to the lower-level Julia solver

Advanced Usage

Non-convex Cost Example

N, xdim, udim = 10, 3, 2
X_ref = np.random.rand(N, xdim)

def lin_cost_fn(X, U):
    # cost is np.sum((X - X_ref) ** 2)
    cx = 2 * (X - X_ref)
    cu = None
    return cx, cu

Warm-start support

Warm-start in SCP MPC can refer to either

• warm-starting the SCP procedure through a good X_prev, U_prev - this is supported
• warm-starting the underlying convex solver - not supported

Warm-starting of the SCP procedure by providing a good X_prev, U_prev guess is supported and very much encouraged for good SCP performance!

Warm-starting of the underlying convex solver is currently not supported, as it does not lead to a noticeable performance improvement on problems we tested the solver on.