Nonlinear Model Predictive Control based on CVXPY and jax
- Clone the repository:
git clone https://github.com/shaoanlu/pyNMPC.git
- Pip install in editable mode:
cd pyNMPC
pip install -e . # optional as long as the workspace is in the same directory
This also installs the core dependencies (cvxpy, jax, numpy).
The function should take input args x
, u
, and dt
and return a new state vector.
import jax.numpy as jnp
def dynamics(x: jnp.ndarray, u: jnp.ndarray, dt: float) -> jnp.ndarray:
"""
Unicycle dynamics model.
x_next = x + u[0] * cos(theta) * dt
y_next = y + u[0] * sin(theta) * dt
theta_next = theta + u[1] * dt
"""
theta = x[2]
return x + jnp.array([u[0] * jnp.cos(theta), u[0] * jnp.sin(theta), u[1]]) * dt
import jax.numpy as jnp
from pyNMPC.nmpc import MPCParams
params = MPCParams(
dt=0.1, # Time step
N=20, # Horizon length
n_states=3, # State vector dimension
n_controls=2, # Control input dimension
Q=jnp.diag(jnp.array([10.0, 10.0, 1.0])), # State weights
QN=jnp.diag(jnp.array([10.0, 10.0, 1.0])), # Terminal state weights
R=jnp.diag(jnp.array([1.0, 0.1])), # Control input weights
x_ref=jnp.array([0.0, 0.0, 0.0]), # Will be overridden
x_min=jnp.array([-3, -3, -float("inf")]), # Lower buond of state
x_max=jnp.array([3, 3, float("inf")]), # Upper buond of state
u_min=jnp.array([0.0, -1.0]), # Lower buond of control input
u_max=jnp.array([1.0, 1.0]), # Upper buond of control input
slack_weight=1e4, # Slack penalty for soft state constraints
max_sqp_iter=5, # Max SQP iterations
sqp_tol=1e-4, # SQP convergence tolerance
verbose=False # Verbosity flag
)
import jax.numpy as jnp
from pyNMPC.nmpc import NMPC
mpc = NMPC(dynamics_fn=dynamics, params=params, solver_opts={...}) # pass dynamics and parameter
current_state = jnp.array([...])
current_reference = jnp.array([...])
mpc_result = mpc.solve(x0=current_state, x_ref=current_reference)
# Optimized control
mpc_result.u_traj
# Predicted state
mpc_result.x_traj
- Python 3.10+
cvxpy[osqp, proxqp, piqp]
jax
- QP failed: Try different QP solvers in
solver_opts
.PROXQP
andPIQP
are usually more stable choices.