multi_path_tracking.py

from guided_mrmp.planners.singlerobot.RRTStar import RRTStar

from guided_mrmp.utils import Roomba
from guided_mrmp.utils import Conflict, Robot, Env

import numpy as np
import matplotlib.pyplot as plt

from guided_mrmp.controllers.utils import compute_path_from_wp, get_ref_trajectory
from guided_mrmp.controllers.multi_mpc import MultiMPC
from guided_mrmp.conflict_resolvers.discrete_resolver import DiscreteResolver
from guided_mrmp.utils import Roomba

from guided_mrmp.conflict_resolvers.curve_path import smooth_path, calculate_headings

def initialize_libraries(library_fnames=["guided_mrmp/database/2x3_library","guided_mrmp/database/3x3_library","guided_mrmp/database/5x2_library"]):
    """
    Load the 2x3, 3x3, and 2x5 libraries. Store them in self.lib-x- 
    Inputs: 
        library_fnames - the folder containing the library files
    """
    from guided_mrmp.utils.library import Library
    # Create each of the libraries
    print(f"Loading libraries. This usually takes about 10 seconds...")
    lib_2x3 = Library(library_fnames[0])
    lib_2x3.read_library_from_file()
    
    lib_3x3 = Library(library_fnames[1])
    lib_3x3.read_library_from_file()

    
    lib_2x5 = Library(library_fnames[2])
    lib_2x5.read_library_from_file()

    return lib_2x3, lib_3x3, lib_2x5

class DiscreteRobot:
    def __init__(self, start, goal, label):
        self.start = start
        self.goal = goal
        self.current_position = start
        self.label = label

class MultiPathTracker:
    def __init__(self, env, initial_positions, dynamics, target_v, T, DT, waypoints, settings, lib_2x3, lib_3x3, lib_2x5):
        """
        Initializes the PathTracker object.
        Parameters:
        - initial_positions: List of the initial positions of the robots [x, y, heading].
        - dynamics: The dynamics model of the robots.
        - target_v: The target velocity of the robots.
        - T: The time horizon for the model predictive control (MPC).
        - DT: The time step for the MPC.
        - waypoints: A list of waypoints defining the desired path for each robot.
        """
        # State of the robot [x,y, heading]
        self.env = env

        self.states = initial_positions
        self.num_robots = len(initial_positions)
        self.dynamics = dynamics
        self.T = T
        self.DT = DT
        self.target_v = target_v

        self.radius = dynamics.radius


        self.update_ref_paths = False

        # helper variable to keep track of mpc output
        # starting condition is 0,0
        self.control = np.zeros((self.num_robots, 2))

        self.K = int(T / DT)

        # For a car model 
        # Q = [20, 20, 10, 20]  # state error cost
        # Qf = [30, 30, 30, 30]  # state final error cost
        # R = [10, 10]  # input cost
        # P = [10, 10]  # input rate of change cost
        # self.mpc = MPC(VehicleModel(), T, DT, Q, Qf, R, P)

        # libraries for the discrete solver
        self.lib_2x3 = lib_2x3
        self.lib_3x3 = lib_3x3
        self.lib_2x5 = lib_2x5


        # For a circular robot (easy dynamics)
        Q = [40, 40, 0]  # state error cost
        Qf = [20,20, 0]  # state final error cost
        R = [10, 10]  # input cost
        P = [10, 10]  # input rate of change cost
        self.mpc = MultiMPC(self.num_robots, dynamics, T, DT, Q, Qf, R, P, settings['model_predictive_controller'], settings['environment']['circle_obstacles'])

        self.circle_obs = settings['environment']['circle_obstacles']

        # Path from waypoint interpolation
        self.paths = []
        for wp in waypoints:
            self.paths.append(compute_path_from_wp(wp[0], wp[1], 0.05))

        
        print(f"paths = {len(self.paths)}")

        # Helper variables to keep track of the sim
        self.sim_time = 0
        self.x_history = [ [] for _ in range(self.num_robots) ]
        self.y_history = [ [] for _ in range(self.num_robots) ]
        self.v_history = [ [] for _ in range(self.num_robots) ]
        self.h_history = [ [] for _ in range(self.num_robots) ]
        self.a_history = [ [] for _ in range(self.num_robots) ]
        self.d_history = [ [] for _ in range(self.num_robots) ]
        self.optimized_trajectories_hist = [ [] for _ in range(self.num_robots) ]
        self.optimized_trajectory = None


    def trajectories_overlap(self, traj1, traj2, threshold):
        """
        Checks if two trajectories overlap. We only care about xy positions.

        Args:
            traj1 (3xn numpy array): First trajectory. First row is x, second row is y, third row is heading.
            traj2 (3xn numpy array): Second trajectory.
            threshold (float): Distance threshold to consider a collision.

        Returns:
            bool: True if trajectories overlap, False otherwise.
        """
        for i in range(traj1.shape[1]):
            for j in range(traj2.shape[1]):
                if np.linalg.norm(traj1[0:2, i] - traj2[0:2, j]) < 2*threshold:
                    return True
        return False
    

    def ego_to_global_roomba(self, state, mpc_out):
        """
        Transforms optimized trajectory XY points from ego (robot) reference
        into global (map) frame.

        Args:
            mpc_out (numpy array): Optimized trajectory points in ego reference frame.

        Returns:
            numpy array: Transformed trajectory points in global frame.
        """
        # Extract x, y, and theta from the state
        x = state[0]
        y = state[1]
        theta = state[2]

        # Rotation matrix to transform points from ego frame to global frame
        Rotm = np.array([
            [np.cos(theta), -np.sin(theta)],
            [np.sin(theta), np.cos(theta)]
        ])

        # Initialize the trajectory array (only considering XY points)
        trajectory = mpc_out[0:2, :]

        # Apply rotation to the trajectory points
        trajectory = Rotm.dot(trajectory)

        # Translate the points to the robot's position in the global frame
        trajectory[0, :] += x
        trajectory[1, :] += y

        return trajectory
    
    def get_next_control(self, state, show_plots=False):
        # optimization loop
        # start=time.time()

        # Get Reference_traj -> inputs are in worldframe
        # 1. Get the reference trajectory for each robot
        targets = []
        for i in range(self.num_robots):
            targets.append(get_ref_trajectory(np.array(state[i]), np.array(self.paths[i]), self.target_v, self.T, self.DT, len(self.x_history[i])+1))

        # dynamycs w.r.t robot frame
        # curr_state = np.array([0, 0, self.state[2], 0])
        curr_states = np.zeros((self.num_robots, 3))
        x_mpc, u_mpc = self.mpc.step(
            curr_states,
            targets,
            self.control
        )
        
        # only the first one is used to advance the simulation
        # self.control[:] = [u_mpc[0, 0], u_mpc[1, 0]]

        self.control = []
        for i in range(self.num_robots):
            self.control.append([u_mpc[i*2, 0], u_mpc[i*2+1, 0]])

        return x_mpc, self.control