Merge pull request #42 from ll7/38-add-an-example-for-recording-and-p…

…layback-of-a-simulation

examples/view_recording.py

@@ -0,0 +1,71 @@
+"""Simulate a random policy with a map defined in SVG format and view the recording."""
+import os
+from pathlib import Path
+
+import numpy as np
+from loguru import logger
+
+from robot_sf.nav.svg_map_parser import SvgMapConverter
+from robot_sf.nav.map_config import MapDefinition, MapDefinitionPool
+from robot_sf.gym_env.env_config import SimulationSettings, EnvSettings
+from robot_sf.gym_env.robot_env import RobotEnv
+from robot_sf.sensor.sensor_fusion import OBS_RAYS, OBS_DRIVE_STATE
+from robot_sf.robot.differential_drive import DifferentialDriveSettings
+from robot_sf.render.playback_recording import load_states_and_visualize
+
+
+
+
+logger.info("Simulate a random policy with a map defined in SVG format.")
+
+def test_simulation(map_definition: MapDefinition):
+    """Test the simulation with a random policy."""
+
+    logger.info("Creating the environment.")
+    env_config = EnvSettings(
+        map_pool=MapDefinitionPool(map_defs={"my_map": map_definition})
+        )
+    env = RobotEnv(env_config, debug=True, recording_enabled=True)  # Activate recording
+
+    env.reset()
+
+    logger.info("Simulating the random policy.")
+    for _ in range(1000):
+        action = env.action_space.sample()
+        env.step(action)
+        env.render()
+
+    env.reset()  # Save the recording
+    env.exit()
+
+def convert_map(svg_file: str):
+    """Create MapDefinition from svg file."""
+
+    logger.info("Converting SVG map to MapDefinition object.")
+    logger.info(f"SVG file: {svg_file}")
+
+    converter = SvgMapConverter(svg_file)
+    return converter.map_definition
+
+def get_file():
+    """Get the latest recorded file."""
+
+    filename = max(
+        os.listdir('recordings'), key=lambda x: os.path.getctime(os.path.join('recordings', x)))
+    return Path('recordings', filename)
+
+
+
+def main():
+    """Simulate a random policy with a map defined in SVG format and view the recording."""
+
+    # Create example recording
+    map_def = convert_map("maps/svg_maps/02_simple_maps.svg")
+    test_simulation(map_def)
+
+    # Load the states from the file and view the recording
+    load_states_and_visualize(get_file())
+
+
+if __name__ == "__main__":
+    main()

robot_sf/gym_env/robot_env.py

@@ -63,6 +63,7 @@ def __init__(
             a dictionary as input and returns a float as reward.
         - debug (bool): If True, enables debugging information such as 
             visualizations.
+        - recording_enabled (bool): If True, enables recording of the simulation
         """
 
         # Environment configuration details

robot_sf/render/playback_recording.py

@@ -37,10 +37,13 @@ def visualize_states(
     use the SimulationView to render a list of states
     on the recorded map defintion
     """
-    sim_view = SimulationView(map_def=map_def)
+    sim_view = SimulationView(map_def=map_def, caption='RobotSF Recording')
+    sim_view.show() # to activate key_events
     for state in states:
         sim_view.render(state)
 
+    sim_view.exit() # to automatically close the window
+
 def load_states_and_visualize(filename: str):
     """
     load a list of states from a file and visualize them

robot_sf/render/sim_view.py

@@ -76,6 +76,9 @@ class SimulationView:
     font: pygame.font.Font = field(init=False)
     redraw_needed: bool = field(init=False, default=False)
     offset: np.array = field(init=False, default=np.array([0, 0]))
+    caption: str='RobotSF Simulation'
+    focus_on_robot: bool = False
+    display_help: bool = False
     """The offset is already uses `scaling` as a factor."""
 
     @property
@@ -88,9 +91,8 @@ def __post_init__(self):
         pygame.font.init()
         self.screen = pygame.display.set_mode(
             (self.width, self.height), pygame.RESIZABLE)
-        pygame.display.set_caption('RobotSF Simulation')
+        pygame.display.set_caption(self.caption)
         self.font = pygame.font.SysFont('Consolas', 14)
-        self.surface_obstacles = self.preprocess_obstacles()
         self.clear()
 
     def _scale_tuple(self, tup: Tuple[float, float]) -> Tuple[float, float]:
@@ -99,39 +101,6 @@ def _scale_tuple(self, tup: Tuple[float, float]) -> Tuple[float, float]:
         y = tup[1] * self.scaling + self.offset[1]
         return (x, y)
 
-    def preprocess_obstacles(self) -> pygame.Surface:
-        # Scale the vertices of the obstacles
-        obst_vertices = [o.vertices_np * self.scaling for o in self.map_def.obstacles]
-
-        # Initialize the minimum and maximum x and y coordinates
-        min_x, max_x, min_y, max_y = np.inf, -np.inf, np.inf, -np.inf
-
-        # Find the minimum and maximum x and y coordinates among all the obstacles
-        for vertices in obst_vertices:
-            min_x = min(np.min(vertices[:, 0]), min_x)
-            max_x = max(np.max(vertices[:, 0]), max_x)
-            min_y = min(np.min(vertices[:, 1]), min_y)
-            max_y = max(np.max(vertices[:, 1]), max_y)
-
-        # Calculate the width and height of the surface needed to draw the obstacles
-        width, height = max_x - min_x, max_y - min_y
-
-        # Create a new surface with the calculated width and height
-        surface = pygame.Surface((width, height), pygame.SRCALPHA)
-
-        # Fill the surface with a transparent background color
-        surface.fill(BACKGROUND_COLOR_TRANSP)
-
-        # Draw each obstacle on the surface
-        for vertices in obst_vertices:
-            # Shift the vertices so that the minimum x and y coordinates are 0
-            shifted_vertices = vertices - [min_x, min_y]
-            # Draw the obstacle as a polygon with the shifted vertices
-            pygame.draw.polygon(surface, OBSTACLE_COLOR, [(x, y) for x, y in shifted_vertices])
-
-        # Return the surface with the drawn obstacles
-        return surface
-
     def show(self):
         """
         Starts a separate thread to process the event queue and handles SIGINT signal.
@@ -175,6 +144,10 @@ def _handle_keydown(self, e):
             pygame.K_DOWN: lambda: self.offset.__setitem__(1, self.offset[1] + 10),
             # reset the view
             pygame.K_r: lambda: self.offset.__setitem__(slice(None), (0, 0)),
+            # focus on the robot
+            pygame.K_f: lambda: setattr(self, 'focus_on_robot', not self.focus_on_robot),
+            # display help
+            pygame.K_h: lambda: setattr(self, 'display_help', not self.display_help),
         }
 
         if e.key in key_action_map:
@@ -224,11 +197,10 @@ def render(self, state: VisualizableSimState):
             self._resize_window()
             self.size_changed = False
         if self.redraw_needed:
-            self.surface_obstacles = self.preprocess_obstacles()
             self.redraw_needed = False
 
-        # TODO: is it correct to scale the ped state here?
-        state = self._scale_pedestrian_state(state)
+        # Adjust the view based on the focus
+        self._move_camera(state)
 
         self.screen.fill(BACKGROUND_COLOR)
 
@@ -250,8 +222,12 @@ def render(self, state: VisualizableSimState):
             self._augment_goal_position(state.action.robot_goal)
         self._draw_pedestrians(state.pedestrian_positions)
         self._draw_robot(state.robot_pose)
+
+        # information
         self._augment_timestep(state.timestep)
         self._add_text(state.timestep, state)
+        if self.display_help:
+            self._add_help_text()
 
         # update the display
         pygame.display.update()
@@ -263,11 +239,13 @@ def _resize_window(self):
             (self.width, self.height), pygame.RESIZABLE)
         self.screen.blit(old_surface, (0, 0))
 
-    def _scale_pedestrian_state(self, state: VisualizableSimState) \
-            -> Tuple[VisualizableSimState, Tuple[float, float]]:
-        state.pedestrian_positions *= self.scaling
-        state.ped_actions *= self.scaling
-        return state
+    def _move_camera(self, state: VisualizableSimState):
+        """ Moves the camera based on the focused object."""
+        if self.focus_on_robot:
+            r_x, r_y = state.robot_pose[0]
+            self.offset[0] = int(r_x * self.scaling - self.width / 2) * -1
+            self.offset[1] = int(r_y * self.scaling - self.height / 2) * -1
+        # TODO: implement moving for trained pedestrian
 
     def _draw_robot(self, pose: RobotPose):
         # TODO: display robot with an image instead of a circle
@@ -283,7 +261,7 @@ def _draw_pedestrians(self, ped_pos: np.ndarray):
             pygame.draw.circle(
                 self.screen,
                 PED_COLOR,
-                (ped_x+self.offset[0], ped_y+self.offset[1]),
+                self._scale_tuple((ped_x, ped_y)),
                 self.ped_radius * self.scaling
                 )
 
@@ -341,8 +319,8 @@ def _augment_ped_actions(self, ped_actions: np.ndarray):
             pygame.draw.line(
                 self.screen,
                 PED_ACTION_COLOR,
-                p1+self.offset,
-                p2+self.offset,
+                self._scale_tuple(p1),
+                self._scale_tuple(p2),
                 width=3
                 )
 
@@ -388,7 +366,8 @@ def _add_text(self, timestep: int, state: VisualizableSimState):
             f'y-offset: {self.offset[1]/self.scaling:.2f}',
             f'RobotPose: {state.robot_pose}',
             f'RobotAction: {state.action.robot_action}',
-            f'RobotGoal: {state.action.robot_goal}'
+            f'RobotGoal: {state.action.robot_goal}',
+            '(Press h for help)',
         ]
         for i, text in enumerate(text_lines):
             text_surface = self.font.render(text, False, TEXT_COLOR)
@@ -398,6 +377,30 @@ def _add_text(self, timestep: int, state: VisualizableSimState):
             )
             self.screen.blit(text_surface, pos)
 
+    def _add_help_text(self):
+        text_lines = [
+            'Move camera: arrow keys',
+            'Move fast: CTRL + arrow keys',
+            'Move slow: ALT + arrow keys',
+            'Reset view: r',
+            'Focus robot: f',
+            'Scale up: +',
+            'Scale down: -' ,
+            'Help: h',
+        ]
+
+        # Determine max width of the text
+        text_surface = self.font.render(text_lines[1], False, TEXT_COLOR)
+        width = text_surface.get_width() + 10
+
+        for i, text in enumerate(text_lines):
+            text_surface = self.font.render(text, False, TEXT_COLOR)
+            pos = (
+                self.width - width,
+                self._timestep_text_pos[1] + i * self.font.get_linesize()
+            )
+            self.screen.blit(text_surface, pos)
+
     def _draw_grid(
             self,
             grid_increment: int=50,