Improve the resolution of pygame window by adding scaling factor

guided_mrmp/main.py

@@ -3,6 +3,7 @@ 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
@@ -62,14 +63,14 @@ 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")
+    # 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} ")
+        # print(f"planning path for robot {i} from {start} to {goal} ")
         rrtpath = plan_decoupled_path(env, (start[0],start[1]), (goal[0],goal[1]))
         xs = []
         ys = []
@@ -79,7 +80,10 @@ def initialize_robots(starts, goals, dynamics_models, radii, target_v, env):
 
         waypoints = [xs,ys]
 
-        print(f"waypoints = {waypoints}")
+        # 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)
         robots.append(r)
@@ -91,13 +95,16 @@ if __name__ == "__main__":
     # Load the settings
     settings = load_settings("settings.yaml")
 
-    set_python_seed(42)
+    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']
+    # sim_x_range = settings['simulator']['sim_x_range']
+    # sim_y_range = settings['simulator']['sim_y_range']
     env = Env(x_range, y_range, circle_obstacles, rectangle_obstacles)
 
     # Load the dynamics models
@@ -114,7 +121,7 @@ if __name__ == "__main__":
     target_v = settings['target_v']
     if robot_starts == []:
         robot_starts, robot_goals = create_random_starts_and_goals(env, len(robot_radii))
-    print(f"starts = {len(robot_starts)}")
+    # print(f"starts = {len(robot_starts)}")
 
     robots = initialize_robots(robot_starts, robot_goals, dynamics_models, robot_radii, target_v, env)
 
@@ -124,10 +131,18 @@ if __name__ == "__main__":
     policy = GuidedMRMP(env, robots, dynamics_models, T, DT, settings)
 
     # Create the simulator
-    pygame.init()
-    screen = pygame.display.set_mode((x_range[1], y_range[1])) 
+    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
-    sim.run(screen)
+    start = time.time()
+    sim.run(screen, show_vis)
+    end = time.time()
+    print(f"Simulation took {end-start} seconds")
 

guided_mrmp/simulator.py

@@ -23,59 +23,83 @@ class Simulator:
         self.num_robots = len(robots)
         self.dynamics_models = dynamics_models
         self.time = 0
+        self.scaling_factor = settings['simulator']['scaling_factor']
+
+    def all_robots_at_goal(self):
+        for r in self.robots:
+            if (np.sqrt((r.current_position[0] - r.goal[0]) ** 2 + (r.current_position[1] - r.goal[1]) ** 2) > 1):
+                return False
+        return True
 
     def advance(self, screen, dt):
         """
         Advance the simulation by dt seconds
         """
-        controls = self.policy.advance(screen,self.state, self.time, dt)
-        # print(controls)
-        for i in range(self.num_robots):
-            new_state = self.dynamics_models[i].next_state(self.state[i], controls[i], dt)
-            self.robots[i].current_position = new_state
-            self.state[i] = new_state
-        self.time += dt
+        conflict_detected, controls = self.policy.advance(screen,self.state, self.time, dt)
+
+        if not conflict_detected:
+            for i in range(self.num_robots):
+                new_state = self.dynamics_models[i].next_state(self.state[i], controls[i], dt)
+                self.robots[i].current_position = new_state
+                self.state[i] = new_state
+            self.time += dt
+
+        if conflict_detected:
+            print("Executing full trajectory")
+            for i in range(self.num_robots):
+                controls_list = controls[i]
+                for j in range(len(controls_list)):
+                    next_control = [controls_list[0][j], controls_list[1][j]]
+                    print(f"next control = {next_control}")
+                    new_state = self.dynamics_models[i].next_state(self.state[i], next_control, dt)
+                    self.robots[i].current_position = new_state
+                    self.state[i] = new_state
+            self.time += dt
 
     def draw_environment(self, screen):
         screen.fill((255,255,255))
         for obs in self.circ_obstacles:
-            pygame.draw.circle(screen, (0,0,0), obs[0:2], obs[2])
+            pygame.draw.circle(screen, (0,0,0), obs[0:2]*self.scaling_factor, obs[2]*self.scaling_factor)
         for obs in self.rect_obstacles:
             pygame.draw.rect(screen, (0,0,0), obs)
 
     def draw_robots(self, screen):
         for robot in self.robots:
             x,y,yaw = robot.current_position
-            pygame.draw.circle(screen, robot.color, (x,y), robot.radius)
+            pygame.draw.circle(screen, robot.color, (x*self.scaling_factor,y*self.scaling_factor), robot.radius*self.scaling_factor)
             # Plot direction marker
             dx = robot.radius * np.cos(robot.current_position[2])
             dy = robot.radius * np.sin(robot.current_position[2])
-            pygame.draw.line(screen, (0,0,0), (x,y), (x+dx,y+dy), width=2)
+            pygame.draw.line(screen, (0,0,0), (x*self.scaling_factor,y*self.scaling_factor), (x*self.scaling_factor+dx*self.scaling_factor,y*self.scaling_factor+dy*self.scaling_factor), width=2)
 
-    def run(self, screen):
+    def run(self, screen, show_vis=False):
         clock = pygame.time.Clock()
         pause = False
-        while 1:
+        while not self.all_robots_at_goal():
             clock.tick(30)
-            for event in pygame.event.get():
-                if event.type == pygame.QUIT:
-                    return
-                if event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
-                    return
-                if event.type == pygame.KEYDOWN and event.key == pygame.K_p:
-                    pause = not pause
-                if event.type == pygame.KEYDOWN and event.key == pygame.K_RIGHT and pause:
+            if show_vis:
+                for event in pygame.event.get():
+                    if event.type == pygame.QUIT:
+                        return
+                    if event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
+                        return
+                    if event.type == pygame.KEYDOWN and event.key == pygame.K_p:
+                        pause = not pause
+                    if event.type == pygame.KEYDOWN and event.key == pygame.K_RIGHT and pause:
+                        self.draw_environment(screen)
+                        self.policy.add_vis_guide_paths(screen, self.scaling_factor)
+                        self.draw_robots(screen)
+                        
+                        self.advance(screen,self.policy.DT)
+            if not pause:
+                if show_vis: 
                     self.draw_environment(screen)
-                    self.policy.add_vis_guide_paths(screen)
+                    self.policy.add_vis_guide_paths(screen, self.scaling_factor)
                     self.draw_robots(screen)
-                    
-                    self.advance(screen,.5)
-            if not pause:
-                self.draw_environment(screen)
-                self.policy.add_vis_guide_paths(screen)
-                self.draw_robots(screen)
                 
-                self.advance(screen,.5)
+                self.advance(screen,self.policy.DT)
 
-            pygame.display.flip()
+            if show_vis: 
                 
+                pygame.display.flip()
+