Newer
Older
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
def advance(self, screen, dt):
"""
Advance the simulation by dt seconds
"""
controls = self.policy.advance(screen,self.state, self.time, dt)
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
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
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])
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)
# 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)
def run(self, screen):
clock = pygame.time.Clock()
pause = False
while 1:
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:
self.draw_environment(screen)
self.policy.add_vis_guide_paths(screen)
if not pause:
self.draw_environment(screen)
self.policy.add_vis_guide_paths(screen)