Skip to content
Snippets Groups Projects
Commit ce7fd788 authored by Louis Radtke's avatar Louis Radtke :vulcan_tone1:
Browse files

start pygame rewrite

parent 0ede09f1
No related branches found
No related tags found
No related merge requests found
import matplotlib
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.path as path
import pygame
# import matplotlib.pyplot as plt
# import matplotlib.patches as patches
# import matplotlib.path as path
from PIL import Image, ImageDraw
from typing import List, Tuple, Optional
from abc import ABC, abstractmethod
......@@ -17,6 +18,7 @@ class Environment:
"""
"""
self.bitmap = np.array([[0]])
self.pygame_surface = None
self.track = []
self.start_pose = np.array([[0.0, 0.0, 0.0, 1.0]]).T
self.target = (0.0, 0.0)
......@@ -75,6 +77,7 @@ class Environment:
# Convert the PIL image to a numpy array for compatibility with the simulator
self.bitmap = np.array(img, dtype=np.uint8)
self.pygame_surface = pygame.surfarray.make_surface(self.bitmap)
def conv(self, x: any, y: any) -> tuple[int, int]:
"""Convert position coordinates to bitmap coordinates."""
......@@ -140,11 +143,9 @@ class Environment:
tx, ty = self.target
return np.sqrt((tx - x) ** 2 + (ty - y) ** 2)
def show(self, ax: matplotlib.axes.Axes):
ax.imshow(self.bitmap, cmap='gray', extent=(
self.min_x - self.margin, self.max_x + self.margin,
self.min_y - self.margin, self.max_y + self.margin
))
def draw(self, screen):
"""Draw the environment surface to the screen."""
screen.blit(self.pygame_surface, (0, 0))
class Robot(ABC):
def __init__(self):
......@@ -168,68 +169,54 @@ class Robot(ABC):
self.tire_omega_left = 0.0 # Current angular velocity of left tire
self.tire_omega_right = 0.0 # Current angular velocity of right tire
# Colors for visualization
self.chassis_color = (0, 0, 255) # Blue
self.wheel_color = (128, 128, 128) # Gray
self.__left_sensor = CircularSensor(self.length, -self.sensor_width / 2, 40)
self.__right_sensor = CircularSensor(self.length, self.sensor_width / 2, 40)
self.__tire_patches = []
self.__chassis_patch = None
self.build_chassis_patches()
def load_params(self, params: dict):
"""Load robot-specific parameters (e.g., sensor positions, motor properties). If overwritten,
base implementation has to be called. For loading custom sensors, use the set_left_sensor and
set_right_sensor methods."""
self.sensor_width = float(params['sensor_width']) if 'sensor_width' in params else self.sensor_width
self.length = float(params['length']) if 'length' in params else self.length
self.tire_radius = float(params['tire_radius']) if 'tire_radius' in params else self.tire_radius
self.__left_sensor = CircularSensor(self.length, self.sensor_width / 2, 40)
self.__right_sensor = CircularSensor(self.length, -self.sensor_width / 2, 40)
self.build_chassis_patches()
def build_chassis_patches(self) -> None:
self.__chassis_patch = \
plt.Polygon([[0, -self.wheel_base / 2], [0, self.wheel_base / 2], [self.length, 0]], color='blue')
self.__chassis_patch.set_animated(True)
rect_w = self.tire_radius * 2
rect_h = self.wheel_base / 6
self.__tire_patches = [
plt.Rectangle(
(0, 0),
rect_w, rect_h, color='gray'),
plt.Rectangle(
(0, 0),
rect_w, rect_h, color='gray')
]
for patch in self.__tire_patches:
patch.set_animated(True)
def get_patches(self) -> List[matplotlib.patches.Patch]:
return [self.__chassis_patch] + self.__tire_patches + [
self.__left_sensor.get_patch(),
self.__right_sensor.get_patch()
]
def update_ui(self):
"""Update the robot's visualization in the GUI."""
rect_w = self.tire_radius * 2
rect_h = self.wheel_base / 6
for patch, sign in zip(self.__tire_patches, [-1, 1]):
patch_pose = rotate(float(self.pose[2])) \
@ translate(-0.5 * rect_w, 0.5 * sign * self.wheel_base - 0.5 * rect_h) \
@ np.array([[0, 0, 0, 1]]).T \
+ self.pose
patch.set_xy(patch_pose[:2].flatten())
patch.set_angle(np.rad2deg(self.pose[2]))
self.__chassis_patch.set_xy(self.pose[:2].T.tolist())
self.__left_sensor.update_ui(self.pose)
self.__right_sensor.update_ui(self.pose)
def draw(self, screen):
"""Draw the robot on the screen."""
# Robot position and orientation
x, y, theta = self.pose[0][0], self.pose[1][0], self.pose[2][0]
# Draw chassis (triangle)
p1 = (x, y)
p2 = (x + self.length * np.cos(theta), y + self.length * np.sin(theta))
p3 = (x + (self.length / 2) * np.cos(theta + np.pi / 2),
y + (self.length / 2) * np.sin(theta + np.pi / 2))
pygame.draw.polygon(screen, self.chassis_color, [p1, p2, p3])
# Draw wheels
wheel_width = self.tire_radius * 2
wheel_height = self.wheel_base / 6
# Left wheel
wheel_left_x = x - wheel_width / 2
wheel_left_y = y - self.wheel_base / 2 - wheel_height / 2
wheel_left_rect = pygame.Rect(wheel_left_x, wheel_left_y, wheel_width, wheel_height)
# Rotate the wheel
wheel_left_surf = pygame.Surface((wheel_width, wheel_height))
wheel_left_surf.fill(self.wheel_color)
wheel_left_surf = pygame.transform.rotate(wheel_left_surf, np.rad2deg(theta))
screen.blit(wheel_left_surf, wheel_left_surf.get_rect(center=(wheel_left_x + wheel_width / 2,
wheel_left_y + wheel_height / 2)))
# Right wheel
wheel_right_x = x - wheel_width / 2
wheel_right_y = y + self.wheel_base / 2 - wheel_height / 2
wheel_right_rect = pygame.Rect(wheel_right_x, wheel_right_y, wheel_width, wheel_height)
# Rotate the wheel
wheel_right_surf = pygame.Surface((wheel_width, wheel_height))
wheel_right_surf.fill(self.wheel_color)
wheel_right_surf = pygame.transform.rotate(wheel_right_surf, np.rad2deg(theta))
screen.blit(wheel_right_surf, wheel_right_surf.get_rect(center=(wheel_right_x + wheel_width / 2,
wheel_right_y + wheel_height / 2)))
# Draw sensors
self.__left_sensor.draw(screen, self.pose)
self.__right_sensor.draw(screen, self.pose)
def set_left_sensor(self, sensor):
"""Set the left sensor instance."""
......@@ -292,11 +279,8 @@ class Robot(ABC):
class Sensor(ABC):
@abstractmethod
def get_patch(self) -> matplotlib.patches.Patch:
pass
@abstractmethod
def update_ui(self, pose: np.ndarray):
def draw(self, screen, pose: np.ndarray):
"""Draw the sensor on the screen based on robot pose."""
pass
@abstractmethod
......@@ -305,110 +289,83 @@ class Sensor(ABC):
pass
class BinarySensor(Sensor):
def __init__(self, offset: np.ndarray):
"""
:param offset: Relative position of the sensor with respect to the robot's pose.
"""
self.offset = offset
self.__patch = None
def get_patch(self) -> matplotlib.patches.Patch:
if self.__patch is not None:
return self.__patch
# Create a cross-marker path
vertices = [
(-1, -1), (1, 1),
(0, 0),
(1, -1), (-1, 1),
]
codes = [path.Path.MOVETO, path.Path.LINETO,
path.Path.MOVETO,
path.Path.MOVETO, path.Path.LINETO]
marker_path = path.Path(vertices, codes)
self.__patch = patches.PathPatch(marker_path, color='red')
self.__patch.set_animated(True)
return self.__patch
def evaluate(self, pose: np.ndarray, environment: Environment) -> int:
"""Evaluate the bitmap at the sensor's position."""
# Calculate sensor's global position
rotation_matrix = np.array([
[np.cos(pose[2]), -np.sin(pose[2])],
[np.sin(pose[2]), np.cos(pose[2])]
])
sensor_global_position = pose[:2] + rotation_matrix @ self.offset
return environment.get_pixel(sensor_global_position)
# class BinarySensor(Sensor):
# def __init__(self, offset: np.ndarray):
# """
# :param offset: Relative position of the sensor with respect to the robot's pose.
# """
# self.offset = offset
# self.__patch = None
#
# def get_patch(self) -> matplotlib.patches.Patch:
# if self.__patch is not None:
# return self.__patch
#
# # Create a cross-marker path
# vertices = [
# (-1, -1), (1, 1),
# (0, 0),
# (1, -1), (-1, 1),
# ]
# codes = [path.Path.MOVETO, path.Path.LINETO,
# path.Path.MOVETO,
# path.Path.MOVETO, path.Path.LINETO]
#
# marker_path = path.Path(vertices, codes)
# self.__patch = patches.PathPatch(marker_path, color='red')
# self.__patch.set_animated(True)
# return self.__patch
#
# def evaluate(self, pose: np.ndarray, environment: Environment) -> int:
# """Evaluate the bitmap at the sensor's position."""
# # Calculate sensor's global position
# rotation_matrix = np.array([
# [np.cos(pose[2]), -np.sin(pose[2])],
# [np.sin(pose[2]), np.cos(pose[2])]
# ])
#
# sensor_global_position = pose[:2] + rotation_matrix @ self.offset
# return environment.get_pixel(sensor_global_position)
class CircularSensor(Sensor):
def get_patch(self) -> matplotlib.patches.Patch:
if self.__patch is not None:
return self.__patch
self.__patch = patches.Circle(
(float(self.offset[0,0]), float(self.offset[1,0])), radius=self.diameter / 2, color='red', fill=False)
self.__patch.set_animated(True)
return self.__patch
def __init__(self, x_offset: float, y_offset: float, diameter: float):
"""
Initializes a sensor with circular sensing area.
:param x_offset: The x-coordinate offset for the circle, relative to the robot's base_link.
:param y_offset: The y-coordinate offset for the circle, relative to the robot's base_link.
:param diameter: The diameter of the circle.
"""
"""Initialize a sensor with circular sensing area."""
self.offset = np.array([[x_offset, y_offset, 0.0, 1.0]]).T
self.diameter = diameter
self.radius = diameter / 2.0
self.color = (255, 0, 0) # Red
self.__patch = None
# Create filter for average calculation
diameter_int = int(self.diameter)
self.__filter = np.zeros((diameter_int, diameter_int))
for x_px in range(diameter_int):
for y_px in range(diameter_int):
# check if the pixel lies within the radius of the circle
if (x_px - diameter_int // 2) ** 2 + (y_px - diameter_int // 2) ** 2 > self.radius ** 2:
continue
# Check if pixel lies within the circle radius
if (x_px - diameter_int // 2) ** 2 + (y_px - diameter_int // 2) ** 2 <= self.radius ** 2:
self.__filter[x_px, y_px] = 1.0
def update_ui(self, pose: np.ndarray):
if self.__patch is not None:
def draw(self, screen, pose: np.ndarray):
"""Draw the circular sensor on the screen."""
# Calculate global position based on robot pose
center_pose = pose + rotate(float(pose[2][0])) @ self.offset
self.__patch.center = center_pose[:2]
# Convert to screen coordinates
x, y = int(center_pose[0][0]), int(center_pose[1][0])
# Draw circle outline
pygame.draw.circle(screen, self.color, (x, y), int(self.radius), 1)
def evaluate(self, pose: np.ndarray, environment: Environment) -> int:
"""
Evaluate the environment by returning the average illumination in the sensor's area.
:param pose: Robot's pose (x, y, theta).
:param environment: The environment object with the bitmap.
:returns: Average pixel intensity in the sensed circular area (0-255).
"""
"""Evaluate the environment by returning the average illumination in the sensor's area."""
assert pose.shape == (4, 1), f'pose must be (4, 1) array, got {pose.shape}'
x, y, theta = pose[:3].flatten()
sensor_global_position = translate(x, y) @ rotate(theta) @ self.offset
# Determine the sensing area in the bitmap
cx, cy = int(sensor_global_position[0]), int(sensor_global_position[1])
radius = int(self.radius)
y_min, y_max = cy - radius, cy + radius + 1
x_min, x_max = cx - radius, cx + radius + 1
avg_value = 0
pixel_count = 0
return environment.average_pixels(sensor_global_position[:2], self.__filter)
class Simulator:
def __init__(self):
self.sim_config = None
......@@ -417,7 +374,21 @@ class Simulator:
self.metrics = {'total_distance': 0.0, 'off_track_distance': 0.0}
self.gui_enabled = False
self.reached_target = False
self.loaded = False
self.paused = False
self.running = False
self.clock = None
self.screen = None
# Rendering settings
self.target_fps = 20
self.simulation_dt = 1/60 # Fixed simulation timestep
self.accumulated_time = 0.0
self.last_frame_time = 0.0
self.last_render_time = 0.0
self.frame_count = 0
self.sim_time = 0.0
def load(self, sim_cnf: dict, robot: Robot, environment: Environment):
"""Loads the simulator, robot, and environment configurations."""
......@@ -436,70 +407,100 @@ class Simulator:
if not self.loaded:
raise Exception('Simulator not loaded.')
self.gui_enabled = gui
dt = 1.0 / float(self.sim_config['ticks_per_second']) if 'ticks_per_second' in self.sim_config else 0.05
iterations = 0
max_iterations = min(max_its, int(max_time / dt))
gui_patches = []
sim_time = 0.0
frame_time = time.time()
target_fps = 1.0
# old_state = self.environment.start_pose.copy()
# old_state = np.array([[1, 300, 0.0, 1.0]]).T
old_state = self.environment.start_pose.copy()
self.robot.set_pose(old_state)
max_iterations = min(max_its, int(max_time / self.simulation_dt))
if gui:
# Initialize matplotlib for visualization
plt.ion()
fig, ax = plt.subplots()
self.environment.show(ax)
pygame.init()
self.screen = pygame.display.set_mode((self.environment.bitmap.shape[1], self.environment.bitmap.shape[0]))
pygame.display.set_caption("Robot Simulator")
self.clock = pygame.time.Clock()
gui_patches = self.robot.get_patches()
for patch in gui_patches:
ax.add_patch(patch)
self.running = True
self.last_frame_time = time.time()
for _ in tqdm(range(max_iterations)):
sim_time = iterations * dt
# Set robot's initial pose
self.robot.set_pose(self.environment.start_pose.copy())
old_state = self.robot.pose.copy()
# robot state update
state = self.robot.update_state(self.environment, dt)
# metrics update
# Main simulation loop
for iterations in tqdm(range(max_iterations)):
# Process events
if gui:
for event in pygame.event.get():
if event.type == pygame.QUIT:
self.running = False
elif event.type == pygame.KEYDOWN:
if event.key == pygame.K_ESCAPE:
self.running = False
elif event.key == pygame.K_SPACE:
self.paused = not self.paused
current_time = time.time()
elapsed = current_time - self.last_frame_time
self.last_frame_time = current_time
if not self.paused:
# Accumulate time and run fixed-step simulation
self.accumulated_time += elapsed
while self.accumulated_time >= self.simulation_dt:
# Run simulation step
state = self.robot.update_state(self.environment, self.simulation_dt)
# Update metrics
iteration_distance = np.linalg.norm(state[:2] - old_state[:2])
self.metrics['total_distance'] += iteration_distance
if self.environment.get_pixel(state[:2]) == 0:
self.metrics['off_track_distance'] += iteration_distance
distance_to_target = self.environment.distance_to_target(state)
# Check termination on leaving bounds or reaching target
# Check termination conditions
if not self.environment.is_inside_area(state[:2]):
flat = state[:2].flatten()
print(f'Terminating simulation, robot left bounds at pos {flat[0]:.1f}, {flat[1]:.1f} mm.')
self.running = False
break
distance_to_target = self.environment.distance_to_target(state)
if distance_to_target < 5.0:
self.reached_target = True
print("Target reached!")
self.running = False
break
# Visualization
if gui and frame_time + 1/target_fps > time.time():
self.robot.update_ui()
plt.draw()
plt.pause(1e-6)
frame_time += 1/target_fps
if real_time:
time.sleep(dt)
old_state = state.copy()
iterations += 1
if gui:
plt.ioff()
plt.show()
print(f'Simulation ended after {iterations} its/{sim_time:.2f} s. Metrics:', self.metrics)
\ No newline at end of file
self.sim_time += self.simulation_dt
self.accumulated_time -= self.simulation_dt
# Render at target framerate
if current_time - self.last_render_time >= 1.0 / self.target_fps:
self.render()
self.last_render_time = current_time
self.frame_count += 1
# Cap the frame rate to prevent excessive CPU usage
self.clock.tick(self.target_fps * 2)
# Cleanup when done
pygame.quit()
print(f'Simulation ended after {iterations} its/{self.sim_time:.2f} s. Metrics:', self.metrics)
return self.metrics
def render(self):
"""Render the current simulation state."""
# Clear screen
self.screen.fill((0, 0, 0))
# Draw environment
self.environment.draw(self.screen)
# Draw robot
self.robot.draw(self.screen)
# Draw simulation info
font = pygame.font.SysFont(None, 24)
fps_text = font.render(f"FPS: {self.clock.get_fps():.1f}", True, (255, 255, 255))
time_text = font.render(f"Sim Time: {self.sim_time:.2f}s", True, (255, 255, 255))
self.screen.blit(fps_text, (10, 10))
self.screen.blit(time_text, (10, 40))
# Flip the display
pygame.display.flip()
from get_line_follow_sim.RobotSimulationManager \
import Robot, Simulator, Environment, Sensor, CircularSensor, BinarySensor
import Robot, Simulator, Environment, Sensor, CircularSensor
......@@ -2,3 +2,4 @@ numpy>=2.2.5
matplotlib>=3.10.1
tqdm>=2.0.0
pillow~=11.2.1
pygame
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment