linear_control_problem.py

import numpy as np
import random
import copy
import scipy.linalg as sla
import scipy.sparse as spa

def generate_problem(nx: int, N: int, infeasible_start: bool, seed: int):
    '''
    Create a linear control problem specification
    '''
    options = {"infeasible_start": infeasible_start}
    plant = LinearControlPlant(
        n=nx,
        horizon=N,
        seed=seed,
        options=options
    )
    cost = QuadraticCost(Q_in=plant.Q, QF_in=plant.QN, R_in=plant.R, xg_in=np.zeros(nx))
    x_min = copy.deepcopy(plant.xmin)
    x_max = copy.deepcopy(plant.xmax)
    u_min = copy.deepcopy(plant.umin)
    u_max = copy.deepcopy(plant.umax)
    problem = LinearControlProblem(plantObj=plant, costObj=cost, x0=plant.x0, N=N)
    nq = int(nx/2)
    problem.set_joint_limits(upper_bounds=x_max[:nq], lower_bounds=x_min[:nq])
    problem.set_velocity_limits(upper_bounds=x_max[nq:], lower_bounds=x_min[nq:])
    problem.set_torque_limits(upper_bounds=u_max, lower_bounds=u_min)
    Q, q, A, l, u = problem.formQPBlocks()
    return Q, q, A, l, u

class LinearControlPlant:
    '''
    Create a Plant object for a linear control problem.
    For discrete-time dynamics x_{t+1} = (I + A) @ x_t + B @ u_t.
    '''

    nx: int
    nu: int
    N: int
    A: np.ndarray # (nx, nx)
    B: np.ndarray # (nx, nu)
    R: np.ndarray # (nu, nu)
    Q: np.ndarray # (nx, nx)
    QN: np.ndarray # (nx, nx)
    umin: np.ndarray # (nu,)
    umax: np.ndarray # (nu,)
    xmin: np.ndarray # (nx,)
    xmax: np.ndarray # (nx,)
    x0: np.ndarray # (nx)

    def forward_dynamics_gradient(self):
        return (np.eye(self.nx) + self.A), self.B

    def get_num_pos(self):
        return self.nu

    def get_num_vel(self):
        return self.nu

    def get_num_cntrl(self):
        return self.nu

    def __init__(
        self,
        n: int = None,
        horizon: int = None,
        seed: int = 1,
        options = {}
    ):
        '''
        Generate problem in QP format. Based on 
        https://github.com/osqp/osqp_benchmarks/blob/master/problem_classes/control.py.
        '''
        # Set random seed
        np.random.seed(seed)
        infeasible_start = options.get('infeasible_start', False)

        # Generate random dynamics (and traj length)
        if horizon is None:
            self.N = random.randint(5, 30)
        else:
            self.N = horizon
            
        if n is None:
            self.nu = random.randint(2, 10)
            self.nx = self.nu*2
        else:
            if n % 2 != 0:
                print("[!]ERROR: State dimension must be a mulitple of 2!")
                exit()
            self.nx = int(n)       # States
            self.nu = int(self.nx / 2)   # Inputs

        self.A = spa.eye(self.nx) + .1 * spa.random(self.nx, self.nx,
                                                    density=1.0,
                                                    data_rvs=np.random.randn)

        # Restrict eigenvalues of A to be less than 1
        lambda_values, V = np.linalg.eig(self.A.todense())
        abs_lambda_values = np.abs(lambda_values)

        # Enforce eigenvalues to be maximum norm 1
        for i in range(len(lambda_values)):
            lambda_values[i] = lambda_values[i] \
                if abs_lambda_values[i] < 1 - 1e-02 else \
                lambda_values[i] / (abs_lambda_values[i] + 1e-02)

        # Reconstruct A = V * Lambda * V^{-1}
        self.A = spa.csc_matrix(
            V.dot(np.diag(lambda_values)).dot(np.linalg.inv(V)).real
            )

        self.B = spa.random(self.nx, self.nu, density=1.0,
                            data_rvs=np.random.randn)

        # Control penalty
        self.R = .1 * spa.eye(self.nu)
        ind07 = np.random.rand(self.nx) < 0.7   # Random 30% data
        # Choose only 70% of nonzero elements
        diagQ = np.multiply(np.random.rand(self.nx), ind07)
        self.Q = spa.diags(diagQ)
        QN = sla.solve_discrete_are(self.A.todense(), self.B.todense(),
                                    self.Q.todense(), self.R.todense())
        self.QN = spa.csc_matrix(QN.dot(QN.T))

        # Input ad state bounds
        self.umax = np.random.rand(self.nu)
        self.umin = -self.umax
        self.xmax = np.random.rand(self.nx)
        self.xmin = -self.xmax

        # Initial state (constrain to be within lower and upper bound if feasible)
        w = np.random.rand(self.nx)
        if not infeasible_start:
            self.x0 = self.xmin + 0.9 * w * (self.xmax - self.xmin)
        else:
            self.x0 = (1 + w) * self.xmax

        # convert to numpy arrays
        self.A = np.array(self.A.todense())
        self.B = np.array(self.B.todense())
        self.R = np.array(self.R.todense())
        self.Q = np.array(self.Q.todense())
        self.QN = np.array(self.QN.todense())

        # convert under euler scheme and add random noise
        self.A = self.A - np.eye(self.nx) + 1e-3 * spa.random(self.nx, self.nx, density=1.0, data_rvs=np.random.randn)
        self.B = self.B + 1e-3 * spa.random(self.nx, self.nu, density=1.0, data_rvs=np.random.randn)

class QuadraticCost:
    '''
    Quadratic cost function of the form:
    cost = 0.5 * (x - xg).T @ Q @ (x - xg) + 0.5 * u.T @ R @ u
    where Q is the state cost matrix, R is the control cost matrix,
    xg is the goal state, x is the current state, and u is the control input.
    '''
    def __init__(self, Q_in: np.ndarray, QF_in: np.ndarray, R_in: np.ndarray, xg_in: np.ndarray):
        self.Q = copy.deepcopy(Q_in)
        self.QF = copy.deepcopy(QF_in)
        self.R = copy.deepcopy(R_in)
        self.xg = copy.deepcopy(xg_in)

    def get_currQ(self, u = None, timestep = None):
        use_QF = isinstance(u,type(None))
        currQ = self.QF if use_QF else self.Q
        return currQ

    def value(self, x: np.ndarray, u: np.ndarray = None, timestep: int = None):
        delta_x = x - self.xg
        currQ = self.get_currQ(u,timestep)
        cost = 0.5*np.matmul(delta_x.transpose(),np.matmul(currQ,delta_x))
        if not isinstance(u, type(None)):
            cost += 0.5*np.matmul(u.transpose(),np.matmul(self.R,u))
        return cost

    def gradient(self, x: np.ndarray, u: np.ndarray = None, timestep: int = None):
        delta_x = x - self.xg
        currQ = self.get_currQ(u,timestep)
        top = np.matmul(delta_x.transpose(),currQ)
        if u is None:
            return top
        else:
            bottom = np.matmul(u.transpose(),self.R)
            return np.hstack((top,bottom))

    def hessian(self, x: np.ndarray, u: np.ndarray = None, timestep: int = None):
        nx = self.Q.shape[0]
        nu = self.R.shape[0]
        currQ = self.get_currQ(u,timestep)
        if u is None:
            return currQ
        else:
            top = np.hstack((currQ,np.zeros((nx,nu))))
            bottom = np.hstack((np.zeros((nu,nx)),self.R))
            return np.vstack((top,bottom))

class BoxConstraint:
    '''
    Box constraint of the form:
    lower_bounds <= x <= upper_bounds
    applied at each timestep over the horizon.
    '''
    def __init__(self, constraint_size: int = 0, num_timesteps: int = 0, upper_bounds: list[float] = [], lower_bounds: list[float] = []):
        # constants
        self.constraint_size = constraint_size
        self.num_timesteps = num_timesteps
        self.num_constraints = self.constraint_size * self.num_timesteps
        # bounds
        lblen = len(lower_bounds)
        ublen = len(upper_bounds)
        if (lblen != constraint_size and lblen != 1) or (ublen != constraint_size and ublen != 1):
            print("[!]ERROR please enter bounds of the size of constraint or constant 1")
            exit()
        self.lower_bounds = np.zeros((self.constraint_size))
        self.lower_bounds[:self.constraint_size] = lower_bounds
        self.upper_bounds = np.zeros((self.constraint_size))
        self.upper_bounds[:self.constraint_size] = upper_bounds
    
    def total_constraint_violation(self, x: np.ndarray) -> float:
        violation = 0.0
        lb_mask = x < self.lower_bounds
        ub_mask = x > self.upper_bounds
        violation += np.sum(self.lower_bounds[lb_mask] - x[lb_mask])
        violation += np.sum(x[ub_mask] - self.upper_bounds[ub_mask])
        return violation

class LinearControlProblem:
    '''
    Linear control problem specification.
    '''
    def __init__(self, plantObj:LinearControlPlant, costObj:QuadraticCost, x0: np.ndarray, N: int):
        self.update_plant(plantObj)
        self.update_cost(costObj)
        self.update_problem(x0, N)

    def update_cost(self, costObj: QuadraticCost):
        self.cost = copy.deepcopy(costObj)

    def update_plant(self, plantObj: LinearControlPlant):
        self.plant = copy.deepcopy(plantObj)

    def update_problem(self, x0: np.ndarray = None, N: int = None):
        if x0 is not None:
            self.x0 = copy.deepcopy(x0)
        if N is not None:
            self.N = N

    def set_joint_limits(self, upper_bounds: list[float], lower_bounds: list[float]):
        self.joint_limits = BoxConstraint(self.plant.get_num_pos(), self.N, upper_bounds, lower_bounds)

    def set_velocity_limits(self, upper_bounds: list[float], lower_bounds: list[float]):
        self.velocity_limits = BoxConstraint(self.plant.get_num_vel(), self.N, upper_bounds, lower_bounds)

    def set_torque_limits(self, upper_bounds: list[float], lower_bounds: list[float]):
        self.torque_limits = BoxConstraint(self.plant.get_num_cntrl(), self.N - 1, upper_bounds, lower_bounds)

    def formQPBlocks(self):
        nq = self.plant.get_num_pos()
        nv = self.plant.get_num_vel()
        nu = self.plant.get_num_cntrl()
        N = self.N
        nx = nq + nv
        n = nx + nu
        # Q,q are cost hessian and gradient with n*(N-1) + nx state and control variables 
        total_states_controls = n*(N-1) + nx
        Q = np.zeros((total_states_controls, total_states_controls))
        q = np.zeros((total_states_controls))
        # A is the constraint gradient (Jacobian)
        total_constraints = nx*N
        if hasattr(self, 'joint_limits'):
            total_constraints += self.joint_limits.num_constraints
        if hasattr(self, 'velocity_limits'):
            total_constraints += self.velocity_limits.num_constraints
        if hasattr(self, 'torque_limits'):
            total_constraints += self.torque_limits.num_constraints
        A = np.zeros((total_constraints, total_states_controls))
        l = np.zeros((total_constraints))
        u = np.zeros((total_constraints))
        # start filling from the top left of the matricies (and top of vectors)
        constraint_index = 0
        state_control_index = 0
        # begin with the initial state constraint
        A[constraint_index:constraint_index + nx, state_control_index:state_control_index + nx] = np.eye(nx)
        l[constraint_index:constraint_index + nx] = self.x0
        u[constraint_index:constraint_index + nx] = self.x0
        constraint_index += nx
        for k in range(N-1):
            # first load in the cost hessian and gradient
            Q[state_control_index:state_control_index + n, \
              state_control_index:state_control_index + n] = self.cost.hessian(np.zeros(nx), np.zeros(nu), k)
            q[state_control_index:state_control_index + n] = self.cost.gradient(np.zeros(nx), np.zeros(nu), k)

            # then load in the constraints for this timestep starting with dynamics
            Ak, Bk = self.plant.forward_dynamics_gradient()
            A[constraint_index:constraint_index + nx, \
              state_control_index:state_control_index + n + nx] = np.hstack((-Ak, -Bk, np.eye(nx)))
            # l, u are both zero for dynamics constraints
            
            # update offsets
            constraint_index += nx
            state_control_index += n
        # finish with the final cost
        Q[state_control_index:state_control_index + nx, \
          state_control_index:state_control_index + nx] = self.cost.hessian(np.zeros(nx), timestep = N-1)
        q[state_control_index:state_control_index + nx] = self.cost.gradient(np.zeros(nx), timestep = N-1)

        # add in any additional constraints
        if hasattr(self, 'joint_limits'):
            for k in range(N):
                start_col = k*(nx + nu)
                A[constraint_index:constraint_index + nq, start_col:start_col + nq] = np.eye(nq)
                l[constraint_index:constraint_index + nq] = self.joint_limits.lower_bounds
                u[constraint_index:constraint_index + nq] = self.joint_limits.upper_bounds
                constraint_index += nq
        if hasattr(self, 'velocity_limits'):
            for k in range(N):
                start_col = k*(nx + nu) + nq
                A[constraint_index:constraint_index + nv, start_col:start_col + nv] = np.eye(nv)
                l[constraint_index:constraint_index + nv] = self.velocity_limits.lower_bounds
                u[constraint_index:constraint_index + nv] = self.velocity_limits.upper_bounds
                constraint_index += nv
        if hasattr(self, 'torque_limits'):
            for k in range(N-1):
                start_col = k*(nx + nu) + nx
                A[constraint_index:constraint_index + nu, start_col:start_col + nu] = np.eye(nu)
                l[constraint_index:constraint_index + nu] = self.torque_limits.lower_bounds
                u[constraint_index:constraint_index + nu] = self.torque_limits.upper_bounds
                constraint_index += nu
        return Q, q, A, l, u