Separating optimizer from MPC

guided_mrmp/controllers/mpc.py

 import numpy as np
 from guided_mrmp.utils import Roomba
+from guided_mrmp.optimizer import Optimizer
 
 np.seterr(divide="ignore", invalid="ignore")
 
@@ -49,6 +50,10 @@ class MPC:
         # weight for error in control
         self.R = np.diag(input_cost)
         self.P = np.diag(input_rate_cost)
+
+        # Instantiate the optimizer
+        self.optimizer = Optimizer(self.nx, self.nu, self.control_horizon, self.Q, self.Qf, self.R, self.P)
+
  
     def get_linear_model_matrices_roomba(self,x_bar,u_bar):
         """
@@ -149,75 +154,12 @@ class MPC:
         return A_lin, B_lin, C_lin
 
     def step(self, initial_state, target, prev_cmd):
-        """
+        A, B, C = self.get_linear_model_matrices_roomba(initial_state, prev_cmd)  # Use Roomba model
 
-        Args:
-            initial_state (array-like): current estimate of [x, y, heading]
-            target (ndarray): state space reference, in the same frame as the provided current state
-            prev_cmd (array-like): previous [v, delta].
+        # Use the Optimizer class to solve the optimization problem
+        x_opt, u_opt = self.optimizer.solve(initial_state, target, prev_cmd, A, B, C, self.robot_model, self.dt)
 
-        Returns:
-
-        """
-        assert len(initial_state) == self.nx
-        assert len(prev_cmd) == self.nu
-        assert target.shape == (self.nx, self.control_horizon)
-
-        # Create variables needed for setting up cvxpy problem
-        x = opt.Variable((self.nx, self.control_horizon + 1), name="states")
-        u = opt.Variable((self.nu, self.control_horizon), name="actions")
-        cost = 0
-        constr = []
-
-        # NOTE: here the state linearization is performed around the starting condition to simplify the controller.
-        # This approximation gets more inaccurate as the controller looks at the future.
-        # To improve performance we can keep track of previous optimized x, u and compute these matrices for each timestep k
-        # Ak, Bk, Ck = self.get_linear_model_matrices(x_prev[:,k], u_prev[:,k])
-        A, B, C = self.get_linear_model_matrices_roomba(initial_state, prev_cmd) # for a differential drive roomba
-
-
-        # Tracking error cost
-        # we want the difference bt our state and the target to be small
-        for k in range(self.control_horizon):
-            cost += opt.quad_form(x[:, k + 1] - target[:, k], self.Q)
-
-
-        # Final point tracking cost
-        # we want the final goals to match up
-        cost += opt.quad_form(x[:, -1] - target[:, -1], self.Qf)
-
-        # Actuation magnitude cost
-        # we want the controls to be small
-        for k in range(self.control_horizon):
-            cost += opt.quad_form(u[:, k], self.R)
-
-        # Actuation rate of change cost
-        # we want the difference in controls between time steps to be small
-        for k in range(1, self.control_horizon):
-            cost += opt.quad_form(u[:, k] - u[:, k - 1], self.P)
-
-        # Kinematics Constraints
-        # Need to obey the kinematics of the robot x_{k+1} = A*x_k + B*u_k + C
-        for k in range(self.control_horizon):
-            constr += [x[:, k + 1] == A @ x[:, k] + B @ u[:, k] + C]
-
-        # initial state
-        constr += [x[:, 0] == initial_state]
-
-        # actuation bounds
-        constr += [opt.abs(u[:, 0]) <= self.robot_model.max_acc]
-        constr += [opt.abs(u[:, 1]) <= self.robot_model.max_steer]
-
-        # Actuation rate of change bounds
-        constr += [opt.abs(u[0, 0] - prev_cmd[0]) / self.dt <= self.robot_model.max_d_acc]
-        constr += [opt.abs(u[1, 0] - prev_cmd[1]) / self.dt <= self.robot_model.max_d_steer]
-        for k in range(1, self.control_horizon):
-            constr += [opt.abs(u[0, k] - u[0, k - 1]) / self.dt <= self.robot_model.max_d_acc]
-            constr += [opt.abs(u[1, k] - u[1, k - 1]) / self.dt <= self.robot_model.max_d_steer]
-
-        prob = opt.Problem(opt.Minimize(cost), constr)
-        solution = prob.solve(solver=opt.OSQP, warm_start=True, verbose=False)
-        return x, u
+        return x_opt, u_opt
     
 if __name__ == "__main__":
     # Example usage:

guided_mrmp/optimizer.py

+import cvxpy as opt
+import numpy as np
+
+class Optimizer:
+    def __init__(self, nx, nu, control_horizon, Q, Qf, R, P):
+        self.nx = nx
+        self.nu = nu
+        self.control_horizon = control_horizon
+        self.Q = Q
+        self.Qf = Qf
+        self.R = R
+        self.P = P
+
+    def solve(self, initial_state, target, prev_cmd, A, B, C, robot_model, dt):
+        """
+        Sets up and solves the optimization problem.
+
+        Args:
+            initial_state (array-like): current estimate of [x, y, heading]
+            target (ndarray): state space reference, in the same frame as the provided current state
+            prev_cmd (array-like): previous [v, delta]
+            A, B, C: Linearized state-space matrices
+            robot_model: Robot model containing constraints
+            dt: Time step
+
+        Returns:
+            x, u: Optimal state and input trajectories
+        """
+
+        # set up variables for the optimization problem
+        x = opt.Variable((self.nx, self.control_horizon + 1), name="states")
+        u = opt.Variable((self.nu, self.control_horizon), name="actions")
+        cost = 0
+        constr = []
+
+        # Tracking error cost
+        for k in range(self.control_horizon):
+            cost += opt.quad_form(x[:, k + 1] - target[:, k], self.Q)
+
+        # Final point tracking cost
+        cost += opt.quad_form(x[:, -1] - target[:, -1], self.Qf)
+
+        # Actuation magnitude cost
+        for k in range(self.control_horizon):
+            cost += opt.quad_form(u[:, k], self.R)
+
+        # Actuation rate of change cost
+        for k in range(1, self.control_horizon):
+            cost += opt.quad_form(u[:, k] - u[:, k - 1], self.P)
+
+        # Kinematics Constraints
+        for k in range(self.control_horizon):
+            constr += [x[:, k + 1] == A @ x[:, k] + B @ u[:, k] + C]
+
+        # initial state
+        constr += [x[:, 0] == initial_state]
+
+        # actuation bounds
+        constr += [opt.abs(u[:, 0]) <= robot_model.max_acc]
+        constr += [opt.abs(u[:, 1]) <= robot_model.max_steer]
+
+        # Actuation rate of change bounds
+        constr += [opt.abs(u[0, 0] - prev_cmd[0]) / dt <= robot_model.max_d_acc]
+        constr += [opt.abs(u[1, 0] - prev_cmd[1]) / dt <= robot_model.max_d_steer]
+        for k in range(1, self.control_horizon):
+            constr += [opt.abs(u[0, k] - u[0, k - 1]) / dt <= robot_model.max_d_acc]
+            constr += [opt.abs(u[1, k] - u[1, k - 1]) / dt <= robot_model.max_d_steer]
+
+        prob = opt.Problem(opt.Minimize(cost), constr)
+        solution = prob.solve(solver=opt.OSQP, warm_start=True, verbose=False)
+        return x, u