Newer
Older
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
import pygame
import numpy as np
import os
import random
import matplotlib.pyplot as plt
from guided_mrmp.utils import Env, Roomba, Robot
from guided_mrmp.planners import RRT
from guided_mrmp.planners import RRTStar
from guided_mrmp.planners.multirobot.db_guided_mrmp import GuidedMRMP
class Simulator:
def __init__(self, robots, dynamics_models, circle_obstacles, rectangle_obstacles, policy):
"""
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 = circle_obstacles
self.rect_obstacles = rectangle_obstacles
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)
rachelmoan
committed
print(controls)
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
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)
self.draw_robots(screen)
self.advance(screen,.3)
if not pause:
self.draw_environment(screen)
self.policy.add_vis_guide_paths(screen)
self.draw_robots(screen)
self.advance(screen,.3)
pygame.display.flip()