Skip to content
Snippets Groups Projects
Select Git revision
  • main default protected
  • updated-place-grid
2 results
Blame Permalink
main.py 4.91 KiB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151 






import yaml
import os
import random
import numpy as np
import pygame
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.planners.db_guided_mrmp import GuidedMRMP
from guided_mrmp.simulator import Simulator
from guided_mrmp.utils.library import Library
from guided_mrmp.utils.helpers import initialize_libraries

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=20, goal_sample_rate=.5, num_samples=500000, r=80):
    """
    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]

        # print(f"waypoints = {waypoints}")

        start_heading = np.arctan2(goal[1] - start[1], goal[0] - 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__":
    # Load the settings
    settings = load_settings("settings_files/settings.yaml")

    set_python_seed(1123)
    # set_python_seed(11235813)

    # Load and create the environment
    circle_obstacles = settings['environment']['circle_obstacles']
    rectangle_obstacles = settings['environment']['rectangle_obstacles']
    x_range = settings['environment']['x_range']
    y_range = settings['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 == []:
        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 = GuidedMRMP(env, robots, dynamics_models, T, DT, libs, settings)

    # Create the simulator
    scaling_factor = settings['simulator']['scaling_factor']
    show_vis = True
    if show_vis:
        pygame.init()
        screen = pygame.display.set_mode((x_range[1]*scaling_factor, y_range[1]*scaling_factor)) 

    else: screen = None

    sim = Simulator(robots, dynamics_models, env, policy, settings)

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