You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
dummy_simulation/robot/simulated_robot.py

218 lines
8.1 KiB
Python

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

import pygame
import math
import random
from submodules.robot_controller_base.robot_controller import RobotController
class Fire:
def __init__(self, x, y, radius):
self.x = x
self.y = y
self.radius = radius
class DummySimRobot(RobotController):
def __init__(self, config):
pygame.init()
self.width = config.MAP_WIDTH
self.height = config.MAP_HEIGHT
self.screen = pygame.display.set_mode((self.width, self.height))
pygame.display.set_caption("Robot Simulator Extinguish Fires")
self.clock = pygame.time.Clock()
self.font = pygame.font.SysFont(None, 28)
self.obstacles = config.MAP_OBSTACLES
self.start = config.START_POS
self.finish = config.FINISH_POS
self.robot_radius = config.ROBOT_RADIUS
self.robot_image = None
if config.ROBOT_IMAGE_PATH:
try:
self.robot_image = pygame.image.load(config.ROBOT_IMAGE_PATH)
self.robot_image = pygame.transform.scale(self.robot_image, (self.robot_radius*2, self.robot_radius*2))
except:
pass
self.max_hp = config.INITIAL_HP
self.hp = self.max_hp
self.max_charge = config.INITIAL_CHARGE
self.charge = self.max_charge
self.collision_damage = config.COLLISION_DAMAGE
self.extinguish_radius = config.EXTINGUISH_RADIUS
self.extinguish_angle = math.radians(config.EXTINGUISH_ANGLE)
self.score_per_fire = config.SCORE_PER_FIRE
self.fires = []
self.num_fires = config.NUM_FIRES
self.score = 0
self.game_over = False
self.won = False
self.x, self.y = self.start
self.angle = 0.0
self.linear_speed = 0.0
self.angular_speed = 0.0
self._generate_fires()
def _generate_fires(self):
self.fires.clear()
attempts = 0
while len(self.fires) < self.num_fires and attempts < 1000:
fx = random.randint(self.robot_radius, self.width - self.robot_radius)
fy = random.randint(self.robot_radius, self.height - self.robot_radius)
if self._point_collides_with_obstacle(fx, fy, self.robot_radius):
continue
if math.hypot(fx - self.start[0], fy - self.start[1]) < self.robot_radius * 2:
continue
if math.hypot(fx - self.finish[0], fy - self.finish[1]) < self.robot_radius * 2:
continue
self.fires.append(Fire(fx, fy, 15))
attempts += 1
def _point_collides_with_obstacle(self, x, y, radius):
for ox, oy, ow, oh in self.obstacles:
closest_x = max(ox, min(x, ox + ow))
closest_y = max(oy, min(y, oy + oh))
if math.hypot(x - closest_x, y - closest_y) < radius:
return True
return False
def send_speed_cmd(self, v, w):
self.linear_speed = max(-1.0, min(1.0, v))
self.angular_speed = max(-1.0, min(1.0, w))
def send_fire_ext_burst_cmd(self):
if self.charge <= 0:
return False
for fire in self.fires[:]:
if self._is_fire_in_front(fire):
self.fires.remove(fire)
self.charge -= 1
self.score += self.score_per_fire
return True
return False
def _is_fire_in_front(self, fire):
dx = fire.x - self.x
dy = fire.y - self.y
dist = math.hypot(dx, dy)
if dist > self.extinguish_radius:
return False
angle_to_fire = math.atan2(dy, dx)
angle_diff = abs(angle_to_fire - self.angle)
angle_diff = min(angle_diff, 2*math.pi - angle_diff)
return angle_diff < self.extinguish_angle / 2
def update(self, dt=1/30.0):
if self.game_over or self.won:
return
self.angle += self.angular_speed * dt * 2
dx = self.linear_speed * math.cos(self.angle) * dt * 100
dy = self.linear_speed * math.sin(self.angle) * dt * 100
new_x = self.x + dx
new_y = self.y + dy
new_x = max(self.robot_radius, min(self.width - self.robot_radius, new_x))
new_y = max(self.robot_radius, min(self.height - self.robot_radius, new_y))
if self._point_collides_with_obstacle(new_x, new_y, self.robot_radius):
self.hp -= self.collision_damage
else:
self.x, self.y = new_x, new_y
if math.hypot(self.x - self.finish[0], self.y - self.finish[1]) < self.robot_radius:
self.won = True
if self.hp <= 0:
self.game_over = True
def draw(self):
self.screen.fill((255,255,255))
for ox, oy, ow, oh in self.obstacles:
pygame.draw.rect(self.screen, (100,100,100), (ox, oy, ow, oh))
pygame.draw.circle(self.screen, (0,200,0), (int(self.finish[0]), int(self.finish[1])), self.robot_radius, 3)
for fire in self.fires:
pygame.draw.circle(self.screen, (255,0,0), (int(fire.x), int(fire.y)), fire.radius)
pygame.draw.circle(self.screen, (255,100,0), (int(fire.x), int(fire.y)), fire.radius-3)
if self.robot_image:
rotated = pygame.transform.rotate(self.robot_image, -math.degrees(self.angle))
rect = rotated.get_rect(center=(int(self.x), int(self.y)))
self.screen.blit(rotated, rect)
else:
nose = (self.x + self.robot_radius * math.cos(self.angle),
self.y + self.robot_radius * math.sin(self.angle))
left = (self.x + self.robot_radius * math.cos(self.angle + 2.5),
self.y + self.robot_radius * math.sin(self.angle + 2.5))
right = (self.x + self.robot_radius * math.cos(self.angle - 2.5),
self.y + self.robot_radius * math.sin(self.angle - 2.5))
pygame.draw.polygon(self.screen, (0,0,255), [nose, left, right])
hp_text = self.font.render(f"HP: {self.hp}", True, (0,0,0))
charge_text = self.font.render(f"Charge: {self.charge}", True, (0,0,0))
score_text = self.font.render(f"Score: {self.score}", True, (0,0,0))
fires_text = self.font.render(f"Fires left: {len(self.fires)}", True, (0,0,0))
self.screen.blit(hp_text, (10, 10))
self.screen.blit(charge_text, (10, 40))
self.screen.blit(score_text, (10, 70))
self.screen.blit(fires_text, (10, 100))
if self.game_over:
msg = self.font.render("GAME OVER Press R to restart", True, (255,0,0))
self.screen.blit(msg, (self.width//2 - 150, self.height//2))
elif self.won:
msg = self.font.render("YOU WIN! Press R to restart", True, (0,100,0))
self.screen.blit(msg, (self.width//2 - 120, self.height//2))
pygame.display.flip()
def handle_events(self):
for event in pygame.event.get():
if event.type == pygame.QUIT:
return False
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_r:
self.restart()
elif event.key == pygame.K_BACKSPACE:
self.reset_position()
return True
def step(self):
if not self.handle_events():
return False
self.update()
self.draw()
self.clock.tick(30)
return True
def restart(self):
self.x, self.y = self.start
self.angle = 0.0
self.linear_speed = 0.0
self.angular_speed = 0.0
self.hp = self.max_hp
self.charge = self.max_charge
self.score = 0
self.game_over = False
self.won = False
self._generate_fires()
def reset_position(self):
if self.game_over or self.won:
return
self.x, self.y = self.start
self.angle = 0.0
self.linear_speed = 0.0
self.angular_speed = 0.0
def get_score(self):
return self.score
def is_alive(self):
return not self.game_over and not self.won
def quit(self):
pygame.quit()