simulator.py

import pygame
import numpy as np


class Simulator:
    def __init__(self, robots, dynamics_models, env, policy, settings):
        """
        robots: list of Robot objects
        dynamics_models: list of DynamicsModel objects
        circle_obstacles: list of tuples (x,y,radius)
        rectangle_obstacles: list of tuples (x,y,width,height)
        policy: The policy that gives us the controls for each robot at a given time
        state: the current state of the world. This is a list of tuples, where each tuple is the state of a robot
        time: the current time
        """
        self.robots = robots
        self.circ_obstacles = env.circle_obs
        self.rect_obstacles = env.rect_obs
        self.policy = policy 

        self.state = [robot.current_position for robot in robots]

        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) > 10):
                return False
        return True

    def advance(self, screen, dt):
        """
        Advance the simulation by dt seconds
        """
        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]*self.scaling_factor,obs[1]*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*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*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, show_vis=False):
        clock = pygame.time.Clock()
        pause = False
        while not self.all_robots_at_goal():
            clock.tick(30)
            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.scaling_factor)
                    self.draw_robots(screen)
                
                self.advance(screen,self.policy.DT)

            if show_vis: 
                
                pygame.display.flip()