main.py

import yaml
import os
import random
import numpy as np
import time
from guided_mrmp.utils import Env, Roomba, Robot, create_random_starts_and_goals
from guided_mrmp.planners import RRT
from guided_mrmp.planners import RRTStar

from guided_mrmp.simulator import Simulator
from guided_mrmp.utils.helpers import initialize_libraries
from guided_mrmp.controllers import MultiPathTracker, MultiPathTrackerDB

class_function_names_dict = {
    'Roomba': Roomba
    # 'RRT': RRT(),
    # 'RRTStar': RRTStar(),
    }

# robot = class_names_dist[config["robot_naame"]](**config["robot_params"])

def load_settings(file_path):
    with open(file_path, 'r') as file:
        settings = yaml.safe_load(file)
    return settings

def set_python_seed(seed):
    print(f"***Setting Python Seed {seed}***")
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    random.seed(seed)

def plan_decoupled_path(env, start, goal, solver="RRT*", 
                        step_length=.5, goal_sample_rate=.05, num_samples=500, r=2):
    """
    Plan decoupled path from a given start to a given goal, using a single-agent solver.

    inputs:
        - start (tuple): (x,y) location of start 
        - goal (tuple): (x,y) location of goal 
        - solver (string): Name of single-agent solver to be used
        - step_length (float): 
        - goal_sample_rate (float):
        - num_samples (int):
        - r (float):
    output:
        - path (list): list of nodes in path 
    """
    if solver == "RRT":
        rrt = RRT(env, start, goal, step_length, goal_sample_rate, num_samples)
        path, tree = rrt.run()
    elif solver == "RRT*":
        rrtstar = RRTStar(env, start, goal, step_length, goal_sample_rate, num_samples,r)
        path, tree = rrtstar.run()
    else:
        print(f"Solver {solver} is not yet implemented. Choose something else.")
        return None

    return list(reversed(path)), tree

def initialize_robots(starts, goals, dynamics_models, radii, target_v, env):
    """
    NOTE This function (and the plan_decoupled_paths function could just exist as 
    helper functions elsewhere to make this class easier to understand)
    """
    # print("initializing robots")

    robots = []

    colors = [list(np.random.choice(range(256), size=3)) for i in range(len(starts))]

    for i, (start, goal, dynamics_model, radius, color) in enumerate(zip(starts, goals, dynamics_models, radii, colors)):
        # print(f"planning path for robot {i} from {start} to {goal} ")
        rrtpath, tree = plan_decoupled_path(env, (start[0],start[1]), (goal[0],goal[1]))
        xs = []
        ys = []
        for node in rrtpath:
            xs.append(node[0])
            ys.append(node[1])

        waypoints = [xs,ys]

        start_heading = np.arctan2(ys[1] - start[1], xs[1] - start[0])
        start = [start[0], start[1], start_heading]
        
        r = Robot(i,color,radius,start,goal,dynamics_model,target_v,rrtpath,waypoints, tree)
        robots.append(r)

    return robots

if __name__ == "__main__":
    
    # get the name of the settings file from the command line
    import sys
    if len(sys.argv) < 3:
        print("Using default settings file")
        settings_file = "settings_files/settings.yaml"
        environment_file = "settings_files/env.yaml"

    else: 
        settings_file = sys.argv[1]
        environment_file = sys.argv[2]

    # Load the settings
    settings = load_settings(settings_file)
    environment = load_settings(environment_file)

    set_python_seed(1123)

    # Load and create the environment
    circle_obstacles = environment['circle_obstacles']
    rectangle_obstacles = environment['rectangle_obstacles']
    x_range = environment['x_range']
    y_range = environment['y_range']
    env = Env(x_range, y_range, circle_obstacles, rectangle_obstacles)

    # Load the dynamics models
    dynamics_models_st = settings['dynamics_models']
    dynamics_models = []
    for model in dynamics_models_st:
        dynamics_models.append(class_function_names_dict[model](settings))

    # Load and create the robots
    robot_starts = settings['robot_starts']
    robot_goals = settings['robot_goals']
    robot_radii = settings['robot_radii']
    target_v = settings['target_v']
    if robot_starts == []: # if no starts and goals are provided, create random ones
        robot_starts, robot_goals = create_random_starts_and_goals(env, len(robot_radii))
    robots = initialize_robots(robot_starts, robot_goals, dynamics_models, robot_radii, target_v, env)

    # Load the libraries
    libs = initialize_libraries()

    # Create the Guided MRMP policy
    T = settings['prediction_horizon']
    DT = settings['discretization_step']
    policy = MultiPathTrackerDB(env=env,
                                 initial_positions=robot_starts, 
                                 dynamics=dynamics_models[0], # NOTE: Using the same dynamics model for all robots for now
                                 target_v=target_v, 
                                 T=T, 
                                 DT=DT, 
                                 waypoints=[robot.waypoints for robot in robots], 
                                 settings=settings, 
                                 lib_2x3=libs[0], 
                                 lib_3x3=libs[1], 
                                 lib_2x5=libs[2]
                                )

    # Create the simulator
    show_vis = settings['simulator']['show_plots']
    sim = Simulator(robots, dynamics_models, env, policy, settings)

    # Run the simulation
    start = time.time()
    sim.run(show_vis)
    end = time.time()
    print(f"Simulation took {end-start} seconds")