Restructuring optimizer/MPC

guided_mrmp/controllers/mpc.py

 import numpy as np
 # from guided_mrmp.utils import Roomba
-from guided_mrmp.optimizer import Optimizer
+from guided_mrmp.optimizers.optimizer import Optimizer, iLQR
 
 np.seterr(divide="ignore", invalid="ignore")
 
@@ -23,17 +23,6 @@ class MPC:
         self.nx = model.state_dimension()  # number of state vars 
         self.nu = model.control_dimension()  # number of input/control vars
 
-        if len(state_cost) != self.nx:
-            raise ValueError(f"State Error cost matrix should be of size {self.nx}")
-        if len(final_state_cost) != self.nx:
-            raise ValueError(f"End State Error cost matrix should be of size {self.nx}")
-        if len(input_cost) != self.nu:
-            raise ValueError(f"Control Effort cost matrix should be of size {self.nu}")
-        if len(input_rate_cost) != self.nu:
-            raise ValueError(
-                f"Control Effort Difference cost matrix should be of size {self.nu}"
-            )
-
         self.robot_model = model
         self.dt = DT
 
@@ -52,127 +41,93 @@ class MPC:
         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)
+        self.optimizer = iLQR(self.nx, self.nu, self.control_horizon)
 
- 
-    def get_linear_model_matrices_roomba(self,x_bar,u_bar):
+    def cost_function(self, x, u, target):
         """
-        Computes the approximated LTI state space model x' = Ax + Bu + C
-
-        Args:
-            x_bar (array-like): State vector [x, y, theta]
-            u_bar (array-like): Input vector [v, omega]
-
+        Calculate the cost function for the optimizer.
+        Parameters:
+            - target: (numpy.ndarray) The target values for tracking.
+            - u: (numpy.ndarray) The control inputs.
+            - x: (numpy.ndarray) The state variables.
         Returns:
-            A_lin, B_lin, C_lin: Linearized state-space matrices
+            - cost: (float) The calculated cost.
         """
+        cost = 0
 
-        x = x_bar[0]
-        y = x_bar[1]
-        theta = x_bar[2]
+        
 
-        v = u_bar[0]
-        omega = u_bar[1]
+        # Tracking error cost
+        for k in range(self.control_horizon):
+            cost += opt.quad_form(x[:, k + 1] - target[:, k], self.Q)
 
-        ct = np.cos(theta)
-        st = np.sin(theta)
+        # Final point tracking cost
+        cost += opt.quad_form(x[:, -1] - target[:, -1], self.Qf)
 
-        # Initialize matrix A with zeros and fill in appropriate elements
-        A = np.zeros((self.nx, self.nx))
-        A[0, 2] = -v * st
-        A[1, 2] = v * ct
+        # Actuation magnitude cost
+        for k in range(self.control_horizon):
+            cost += opt.quad_form(u[:, k], self.R)
 
-        # Discrete-time state matrix A_lin
-        A_lin = np.eye(self.nx) + self.dt * A
+        # Actuation rate of change cost
+        for k in range(1, self.control_horizon):
+            cost += opt.quad_form(u[:, k] - u[:, k - 1], self.P)
 
-        # Initialize matrix B with zeros and fill in appropriate elements
-        B = np.zeros((self.nx, self.nu))
-        B[0, 0] = ct
-        B[1, 0] = st
-        B[2, 1] = 1
+        return cost
 
-        # Discrete-time input matrix B_lin
-        B_lin = self.dt * B
+    def constraints(self, initial_state, A, B,C, x, u, robot_model, dt, prev_cmd):
 
-        # Compute the non-linear state update equation f(x, u)
-        f_xu = np.array([v * ct, v * st, omega]).reshape(self.nx, 1)
+        constr = []
+        
+        # Kinematics Constraints
+        for k in range(self.control_horizon):
+            constr += [x[:, k + 1] == A @ x[:, k] + B @ u[:, k] + C]
 
-        # Compute the constant vector C_lin
-        C_lin = (self.dt * (f_xu - np.dot(A, x_bar.reshape(self.nx, 1)) - np.dot(B, u_bar.reshape(self.nu, 1))).flatten())
+        # initial state
+        constr += [x[:, 0] == initial_state]
 
-        return A_lin, B_lin, C_lin
-    
-    def get_linear_model_matrices(self, x_bar, u_bar):
-        """
-        Computes the approximated LTI state space model x' = Ax + Bu + C
+        # actuation bounds
+        constr += [opt.abs(u[:, 0]) <= robot_model.max_acc]
+        constr += [opt.abs(u[:, 1]) <= robot_model.max_steer]
 
-        Args:
-            x_bar (array-like):
-            u_bar (array-like):
+        # 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]
 
-        Returns:
+        return constr
 
+    def step(self, initial_state, target, prev_cmd, initial_guess=None):
         """
-
-        x = x_bar[0]
-        y = x_bar[1]
-        v = x_bar[2]
-        theta = x_bar[3]
-
-        a = u_bar[0]
-        delta = u_bar[1]
-
-        ct = np.cos(theta)
-        st = np.sin(theta)
-        cd = np.cos(delta)
-        td = np.tan(delta)
-
-        A = np.zeros((self.nx, self.nx))
-        A[0, 2] = ct
-        A[0, 3] = -v * st
-        A[1, 2] = st
-        A[1, 3] = v * ct
-        A[3, 2] = v * td / self.robot_model.wheelbase
-        A_lin = np.eye(self.nx) + self.dt * A
-
-        B = np.zeros((self.nx, self.nu))
-        B[2, 0] = 1
-        B[3, 1] = v / (self.robot_model.wheelbase * cd**2)
-        B_lin = self.dt * B
-
-        f_xu = np.array([v * ct, v * st, a, v * td / self.robot_model.wheelbase]).reshape(
-            self.nx, 1
-        )
-        C_lin = (
-            self.dt
-            * (
-                f_xu
-                - np.dot(A, x_bar.reshape(self.nx, 1))
-                - np.dot(B, u_bar.reshape(self.nu, 1))
-            ).flatten()
-        )
-        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
+        Sets up and solves the optimization problem.
+        
+        Args:
+            initial_state: Current estimate of [x, y, heading]
+            target: State space reference, in the same frame as the provided current state
+            prev_cmd: Previous [v, delta]
+            A, B, C: Linearized state-space matrices
+            initial_guess: Optional initial guess for the optimizer
+        
+        Returns:
+            x_opt: Optimal state trajectory
+            u_opt: Optimal control trajectory
+        """
+        A,B,C = self.robot_model.linearize(initial_state, prev_cmd, self.dt)
+        
+        # set up variables for the optimization problem
+        x = opt.Variable((self.nx, self.control_horizon + 1))
+        u = opt.Variable((self.nu, self.control_horizon))
+        
+        # Define the cost function
+        cost = self.cost_function(x, u, target)
+        
+        # Define the constraints
+        constraints = self.constraints(initial_state, A, B, C, x, u, self.robot_model, self.dt, prev_cmd)
 
         # 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)
+        x_opt, u_opt = self.optimizer.solve(x,u,cost, constraints, initial_guess)
 
         return x_opt, u_opt
     
-if __name__ == "__main__":
-    # Example usage:
-    dt = 0.1
-    roomba = Roomba()
-    Q = [20, 20, 20]  # state error cost
-    Qf = [30, 30, 30]  # state final error cost
-    R = [10, 10]  # input cost
-    P = [10, 10]  # input rate of change cost
-    mpc = MPC(roomba, 5, dt, Q, Qf, R, P)
-    x_bar = np.array([0.0, 0.0, 0.0])
-    u_bar = np.array([1.0, 0.1])
-    A_lin, B_lin, C_lin = mpc.get_linear_model_matrices_roomba(x_bar, u_bar)
-    print(A_lin)
-    print(B_lin)
-    print(C_lin)
\ No newline at end of file
+

guided_mrmp/optimizers/optimizer.py

+import cvxpy as opt
+import numpy as np
+
+class Optimizer:
+    def __init__(self):
+        pass
+
+    def solve(self, cost_function, constraints, initial_guess):
+        """
+        Solve the optimization problem.
+        
+        Parameters:
+        - cost_function: Function to minimize
+        - constraints: List of constraints
+        - initial_guess: Initial guess for the optimizer
+        
+        Returns:
+        - optimal_solution: The optimal control input found by the optimizer
+        """
+        raise NotImplementedError("This method should be implemented by a specific optimizer.")
+
+class iLQR(Optimizer):
+    def __init__(self, nx, nu, control_horizon):
+        self.nx = nx
+        self.nu = nu
+        self.control_horizon = control_horizon
+
+    def solve(self, x,u,cost, constraints, initial_guess=None):
+        """
+        Solve the optimization problem.
+        
+        Parameters:
+        - cost_function: Function to minimize
+        - constraints: List of constraints
+        - initial_guess: Optional initial guess for the optimizer
+        
+        Returns:
+        - x_opt: Optimal state trajectory
+        - u_opt: Optimal control trajectory
+        """
+        # Set up variables for the optimization problem
+        # x = opt.Variable((self.nx, self.control_horizon + 1))
+        # u = opt.Variable((self.nu, self.control_horizon))
+        
+        # If initial guess is provided, set the initial values
+        # if initial_guess:
+        #     x.value = initial_guess['x']
+        #     u.value = initial_guess['u']
+
+        
+
+        prob = opt.Problem(opt.Minimize(cost), constraints)
+
+        # Solve the problem
+        prob.solve(solver=opt.OSQP, warm_start=True, verbose=False)
+        
+        if prob.status != opt.OPTIMAL:
+            raise ValueError("The optimization problem did not solve successfully.")
+        
+        # print(f"x val = {x.value}")
+        # print(f"u val = {u.value}")
+
+        return x.value, u.value
\ No newline at end of file

guided_mrmp/simulator.py

 import pygame
 import numpy as np
-import os
-import random
-import matplotlib.pyplot as plt
-from guided_mrmp.utils import Env, Roomba, Robot
 
-from guided_mrmp.planners import RRT
-from guided_mrmp.planners import RRTStar
-from guided_mrmp.planners.multirobot.db_guided_mrmp import GuidedMRMP
 
 class Simulator:
     def __init__(self, robots, dynamics_models, env, policy, settings):

guided_mrmp/utils/control.py

 import numpy as np
+from autograd import jacobian
+from autograd.numpy import cos, sin
 
 class ControlSpace:
     """
@@ -25,6 +27,9 @@ class ControlSpace:
         """Returns a sequence of controls that connects x to y, if
         applicable.  Return None if no such sequence exists."""
         return None
+    def linearize(self, x_bar, u_bar, dt):
+        """Compute linearized dynamics around (x_bar, u_bar). Returns A_lin, B_lin, C_lin."""
+        raise NotImplementedError()
 
 class Roomba(ControlSpace):
     """
@@ -51,12 +56,14 @@ class Roomba(ControlSpace):
         State x: [x, y, theta]
         Control u: [v, steering_angle]
         """
+
         x_pos, y_pos, theta = x
         v, steering_angle = u
 
+
         # Dynamics equations
-        dx_pos = v * np.cos(theta)
-        dy_pos = v * np.sin(theta)
+        dx_pos = v * cos(theta)
+        dy_pos = v * sin(theta)
         dtheta = steering_angle
 
         return np.array([dx_pos, dy_pos, dtheta])
@@ -76,3 +83,76 @@ class Roomba(ControlSpace):
 
         # Return the predicted next state
         return np.array([new_x, new_y, new_theta])
+
+    def linearize(self, x_bar, u_bar, dt):
+        """
+        Linearize the system dynamics around (x_bar, u_bar) and discretize them.
+
+        Returns:
+            A_lin, B_lin, C_lin: Discrete-time linearized state-space matrices
+        """
+        x_bar = np.array(x_bar, dtype=np.float64)
+        u_bar = np.array(u_bar, dtype=np.float64)
+
+        # Compute continuous-time Jacobians
+        A_continuous, B_continuous = self.compute_jacobian(x_bar, u_bar)
+
+        print(f"A_continuous: {A_continuous}")
+        print(f"B_continuous: {B_continuous}")
+
+        
+        # Discretize the system using Euler approximation
+        A_lin = np.eye(self.state_dimension()) + A_continuous * dt
+        B_lin = B_continuous * dt
+
+        # Compute C_lin
+        f_xu = self.derivative(x_bar, u_bar)
+        C_lin = dt * (f_xu - A_continuous @ x_bar - B_continuous @ u_bar)
+
+        return A_lin, B_lin, C_lin
+
+    # def compute_jacobian(self, x_bar, u_bar):
+    #     """
+    #     Compute the Jacobian matrices A and B at (x_bar, u_bar).
+
+    #     Returns:
+    #         A: Jacobian of dynamics with respect to the state x
+    #         B: Jacobian of dynamics with respect to the control input u
+    #     """
+    #     # Use autograd's jacobian function
+    #     A_func = jacobian(lambda x: self.derivative(x, u_bar))
+    #     B_func = jacobian(lambda u: self.derivative(x_bar, u))
+
+    #     # Evaluate Jacobians at (x_bar, u_bar)
+    #     A = A_func(x_bar)
+    #     B = B_func(u_bar)
+
+    #     return A, B
+    
+    def compute_jacobian(self, x, u):
+        """
+        Compute the Jacobian matrices A and B based on analytical derivatives.
+        
+        Parameters:
+        - x: state vector [x, y, theta]
+        - u: control input vector [v, omega]
+        
+        Returns:
+        - A: Jacobian of dynamics with respect to the state x
+        - B: Jacobian of dynamics with respect to the control input u
+        """
+        x_pos, y_pos, theta = x
+        v, omega = u
+
+        # Compute A (Jacobian with respect to state)
+        A = np.zeros((self.state_dimension(), self.state_dimension()))
+        A[0, 2] = -v * np.sin(theta)
+        A[1, 2] = v * np.cos(theta)
+
+        # Compute B (Jacobian with respect to control)
+        B = np.zeros((self.state_dimension(), self.control_dimension()))
+        B[0, 0] = np.cos(theta)
+        B[1, 0] = np.sin(theta)
+        B[2, 1] = 1.0
+
+        return A, B
\ No newline at end of file

guided_mrmp/utils/helpers.py

@@ -19,22 +19,22 @@ def create_random_starts_and_goals(env, num_agents):
 
     for i in range(num_agents):
         # generate a random start
-        x = random.uniform(env.x_range[0], env.x_range[1])
-        y = random.uniform(env.y_range[0], env.y_range[1])
+        x = random.uniform(env.boundary[0][0], env.boundary[0][1])
+        y = random.uniform(env.boundary[1][0], env.boundary[1][1])
     
         while [x,y] in starts:
-            x = random.uniform(env.x_range[0], env.x_range[1])
-            y = random.uniform(env.y_range[0], env.y_range[1])
+            x = random.uniform(env.boundary[0][0], env.boundary[0][1])
+            y = random.uniform(env.boundary[1][0], env.boundary[1][1])
 
-        starts.append([x,y])
+        starts.append([x,y,0])
 
         # generate a random goal
-        x = random.uniform(env.x_range[0], env.x_range[1])
-        y = random.uniform(env.y_range[0], env.y_range[1])
+        x = random.uniform(env.boundary[0][0], env.boundary[0][1])
+        y = random.uniform(env.boundary[1][0], env.boundary[1][1])
     
         while [x,y] in goals:
-            x = random.uniform(env.x_range[0], env.x_range[1])
-            y = random.uniform(env.y_range[0], env.y_range[1])
+            x = random.uniform(env.boundary[0][0], env.boundary[0][1])
+            y = random.uniform(env.boundary[1][0], env.boundary[1][1])
 
         goals.append([x,y])