Merge pull request #44 from jjshoots/staggered_render

Staggered camera capture for drones

Author: jjshoots

PyFlyt/core/abstractions/base_drone.py

@@ -155,11 +155,15 @@ def reset(self) -> None:
         raise NotImplementedError
 
     @abstractmethod
-    def update_control(self) -> None:
+    def update_control(self, physics_step: int) -> None:
         """Updates onboard flight control laws at a rate specified by `control_hz`.
 
         Example Implementation:
-            >>> def update_control(self):
+            >>> def update_control(self, physics_step: int) -> None:
+            >>>     # skip control if we have not enough physics steps
+            >>>     if physics_step % self.physics_control_ratio != 0:
+            >>>         return
+            >>>
             >>>     # depending on the flight control mode, do things
             >>>     if self.mode == 0:
             >>>         # assign the actuator commands to be some function of the setpoint
@@ -182,7 +186,7 @@ def update_physics(self) -> None:
         """Updates all physics on the vehicle at a rate specified by `physics_hz`.
 
         Example Implementation:
-            >>> def update_physics(self):
+            >>> def update_physics(self) -> None:
             >>>     self.lifting_surfaces.physics_update(self.cmd[...])
             >>>     self.boosters.physics_update(self.cmd[...])
             >>>     ...
@@ -194,7 +198,7 @@ def update_state(self) -> None:
         """Updates the vehicle's state values at a rate specified by `phyiscs_hz`.
 
         Example Implementation:
-            >>> def update_state(self):
+            >>> def update_state(self) -> None:
             >>>     # get all relevant information from PyBullet backend
             >>>     lin_pos, ang_pos = self.p.getBasePositionAndOrientation(self.Id)
             >>>     lin_vel, ang_vel = self.p.getBaseVelocity(self.Id)
@@ -222,12 +226,13 @@ def update_state(self) -> None:
         raise NotImplementedError
 
     @abstractmethod
-    def update_last(self) -> None:
+    def update_last(self, physics_step: int) -> None:
         """Update that happens at the end of `Aviary.step()`, usually reserved for camera updates.
 
         Example Implementation:
-            >>> def update_last(self):
-            >>>     if self.use_camera:
+            >>> def update_last(self, physics_step: int) -> None:
+            >>>
+            >>>     if self.use_camera and (physics_step % self.physics_camera_ratio == 0):
             >>>         self.rgbaImg, self.depthImg, self.segImg = self.camera.capture_image()
         """
         raise NotImplementedError

PyFlyt/core/aviary.py

@@ -17,6 +17,15 @@
 DroneIndex = int
 
 
+class AviaryInitException(Exception):
+    def __init__(self, message: str) -> None:
+        self.message = message
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        return f"Aviary Error: {self.message}"
+
+
 class Aviary(bullet_client.BulletClient):
     """Aviary class, the core of how PyFlyt handles UAVs in the PyBullet simulation environment.
 
@@ -73,15 +82,18 @@ def __init__(
         print("\033[A                             \033[A")
 
         # check for starting position and orientation shapes
-        assert (
-            len(start_pos.shape) == 2
-        ), f"start_pos must be shape (n, 3), currently {start_pos.shape}."
-        assert (
-            start_pos.shape[-1] == 3
-        ), f"start_pos must be shape (n, 3), currently {start_pos.shape}."
-        assert (
-            start_orn.shape == start_pos.shape
-        ), f"start_orn must be same shape as start_pos, currently {start_orn.shape}."
+        if len(start_pos.shape) != 2:
+            raise AviaryInitException(
+                f"start_pos must be shape (n, 3), currently {start_pos.shape}."
+            )
+        if start_pos.shape[-1] != 3:
+            raise AviaryInitException(
+                f"start_pos must be shape (n, 3), currently {start_pos.shape}."
+            )
+        if start_orn.shape != start_pos.shape:
+            raise AviaryInitException(
+                f"start_orn must be same shape as start_pos, currently {start_orn.shape}."
+            )
 
         # check the physics hz
         if physics_hz != 240.0:
@@ -91,27 +103,33 @@ def __init__(
 
         # check to ensure drone type has same number as drones if is list/tuple
         if isinstance(drone_type, (tuple, list)):
-            assert (
-                len(drone_type) == start_pos.shape[0]
-            ), f"If multiple `drone_types` are used, must have same number of `drone_types` ({len(drone_type)}) as number of drones ({start_pos.shape[0]})."
+            if len(drone_type) != start_pos.shape[0]:
+                raise AviaryInitException(
+                    f"If multiple `drone_types` are used, must have same number of `drone_types` ({len(drone_type)}) as number of drones ({start_pos.shape[0]})."
+                )
+            if len(drone_type) != start_pos.shape[0]:
+                raise AviaryInitException(
+                    f"If multiple `drone_types` are used, must have same number of `drone_types` ({len(drone_type)}) as number of drones ({start_pos.shape[0]})."
+                )
         # check to ensure drone type has same number as drones if is list/tuple
         if isinstance(drone_options, (tuple, list)):
-            assert (
-                len(drone_options) == start_pos.shape[0]
-            ), f"If multiple `drone_options` ({len(drone_options)}) are used, must have same number of `drone_options` as number of drones ({start_pos.shape[0]})."
+            if len(drone_options) != start_pos.shape[0]:
+                raise AviaryInitException(
+                    f"If multiple `drone_options` ({len(drone_options)}) are used, must have same number of `drone_options` as number of drones ({start_pos.shape[0]})."
+                )
 
         # constants
-        self.num_drones = start_pos.shape[0]
-        self.start_pos = start_pos
-        self.start_orn = start_orn
+        self.num_drones: int = start_pos.shape[0]
+        self.start_pos: np.ndarray = start_pos
+        self.start_orn: np.ndarray = start_orn
 
         # do not change because pybullet doesn't like it
         # default physics looprate is 240 Hz
-        self.physics_hz = physics_hz
-        self.physics_period = 1.0 / physics_hz
+        self.physics_hz: int = physics_hz
+        self.physics_period: float = 1.0 / physics_hz
 
         # mapping of drone type string to the constructors
-        self.drone_type_mappings = dict()
+        self.drone_type_mappings: dict[str, type[DroneClass]] = dict()
         self.drone_type_mappings["quadx"] = QuadX
         self.drone_type_mappings["fixedwing"] = Fixedwing
         self.drone_type_mappings["rocket"] = Rocket
@@ -123,18 +141,24 @@ def __init__(
 
         # store all drone types
         if isinstance(drone_type, (tuple, list)):
-            assert all(
-                dt in self.drone_type_mappings for dt in drone_type
-            ), f"One of types in `drone_type` {drone_type} is not amongst known types {self.drone_type_mappings.keys()}."
+            if not all(dt in self.drone_type_mappings for dt in drone_type):
+                raise AviaryInitException(
+                    f"One of types in `drone_type` {drone_type} is not amongst known types {self.drone_type_mappings.keys()}."
+                )
+            if not all(dt in self.drone_type_mappings for dt in drone_type):
+                raise AviaryInitException(
+                    f"One of types in `drone_type` {drone_type} is not amongst known types {self.drone_type_mappings.keys()}."
+                )
             self.drone_type = drone_type
         else:
-            assert (
-                drone_type in self.drone_type_mappings
-            ), f"Can't find `drone_type` {drone_type} amongst known types {self.drone_type_mappings.keys()}."
+            if drone_type not in self.drone_type_mappings:
+                raise AviaryInitException(
+                    f"Can't find `drone_type` {drone_type} amongst known types {self.drone_type_mappings.keys()}."
+                )
             self.drone_type = repeat(drone_type)
 
         # store the drone options
-        if isinstance(drone_options, (tuple, list)):
+        if isinstance(drone_options, Sequence):
             self.drone_options = drone_options
         else:
             self.drone_options = repeat(drone_options)
@@ -251,7 +275,7 @@ def reset(self, seed: None | int = None) -> None:
         # reset all drones and initialize required states
         [drone.reset() for drone in self.drones]
         [drone.update_state() for drone in self.drones]
-        [drone.update_last() for drone in self.drones]
+        [drone.update_last(0) for drone in self.drones]
 
     def register_all_new_bodies(self) -> None:
         """Registers all new bodies in the environment to be able to handle collisions later.
@@ -441,21 +465,18 @@ def step(self) -> None:
         self.contact_array &= False
 
         # step the environment enough times for one control loop of the slowest controller
-        for step in range(self.updates_per_step):
-            # update onboard avionics conditionally
-            [
-                drone.update_control()
-                for drone in self.armed_drones
-                if step % drone.physics_control_ratio == 0
-            ]
-
-            # update physics and state
+        for _ in range(self.updates_per_step):
+            # update control and physics
+            [drone.update_control(self.physics_steps) for drone in self.armed_drones]
             [drone.update_physics() for drone in self.armed_drones]
-            [drone.update_state() for drone in self.armed_drones]
 
             # advance pybullet
             self.stepSimulation()
 
+            # update states and camera
+            [drone.update_state() for drone in self.armed_drones]
+            [drone.update_last(self.physics_steps) for drone in self.armed_drones]
+
             # splice out collisions
             for collision in self.getContactPoints():
                 self.contact_array[collision[1], collision[2]] = True
@@ -465,7 +486,4 @@ def step(self) -> None:
             self.physics_steps += 1
             self.elapsed_time = self.physics_steps / self.physics_hz
 
-        # update the last components of the drones, this is usually limited to cameras only
-        [drone.update_last() for drone in self.armed_drones]
-
         self.aviary_steps += 1

PyFlyt/core/drones/fixedwing.py

@@ -30,6 +30,7 @@ def __init__(
         camera_FOV_degrees: int = 90,
         camera_resolution: tuple[int, int] = (128, 128),
         camera_position_offset: np.ndarray = np.array([-3.0, 0.0, 1.0]),
+        camera_fps: None | int = None,
         starting_velocity: np.ndarray = np.array([20.0, 0.0, 0.0]),
     ):
         """Creates a Fixedwing UAV and handles all relevant control and physics.
@@ -49,6 +50,7 @@ def __init__(
             camera_FOV_degrees (int): camera_FOV_degrees
             camera_resolution (tuple[int, int]): camera_resolution
             camera_position_offset (np.ndarray): offset position of the camera
+            camera_fps (None | int): camera_fps
             starting_velocity (np.ndarray): vector representing the velocity at spawn
         """
         super().__init__(
@@ -178,6 +180,15 @@ def __init__(
                 is_tracking_camera=True,
             )
 
+        # compute camera fps parameters
+        if camera_fps:
+            assert (
+                (physics_hz / camera_fps) % 1 == 0
+            ), f"Expected `camera_fps` to roundly divide `physics_hz`, got {camera_fps=} and {physics_hz=}."
+            self.physics_camera_ratio = int(physics_hz / camera_fps)
+        else:
+            self.physics_camera_ratio = 1
+
     def reset(self) -> None:
         """Resets the vehicle to the initial state."""
         self.set_mode(0)
@@ -212,8 +223,16 @@ def set_mode(self, mode) -> None:
         elif mode == 0:
             self.setpoint = np.zeros(4)
 
-    def update_control(self) -> None:
-        """Runs through controllers."""
+    def update_control(self, physics_step: int) -> None:
+        """Runs through controllers.
+
+        Args:
+            physics_step (int): the current physics step
+        """
+        # skip control if we don't have enough physics steps
+        if physics_step % self.physics_control_ratio != 0:
+            return
+
         # full control over all surfaces
         if self.mode == -1:
             self.cmd = self.setpoint
@@ -267,7 +286,11 @@ def update_state(self) -> None:
             (self.lifting_surfaces.get_states(), self.motors.get_states())
         )
 
-    def update_last(self) -> None:
-        """Updates things only at the end of `Aviary.step()`."""
-        if self.use_camera:
+    def update_last(self, physics_step: int) -> None:
+        """Updates things only at the end of `Aviary.step()`.
+
+        Args:
+            physics_step (int): the current physics step
+        """
+        if self.use_camera and (physics_step % self.physics_camera_ratio == 0):
             self.rgbaImg, self.depthImg, self.segImg = self.camera.capture_image()

PyFlyt/core/drones/quadx.py

@@ -33,6 +33,7 @@ def __init__(
         camera_angle_degrees: int = 20,
         camera_FOV_degrees: int = 90,
         camera_resolution: tuple[int, int] = (128, 128),
+        camera_fps: None | int = None,
     ):
         """Creates a drone in the QuadX configuration and handles all relevant control and physics.
 
@@ -50,6 +51,7 @@ def __init__(
             camera_angle_degrees (int): camera_angle_degrees
             camera_FOV_degrees (int): camera_FOV_degrees
             camera_resolution (tuple[int, int]): camera_resolution
+            camera_fps (None | int): camera_fps
         """
         super().__init__(
             p=p,
@@ -203,6 +205,15 @@ def __init__(
                 camera_resolution=camera_resolution,
             )
 
+        # compute camera fps parameters
+        if camera_fps:
+            assert (
+                (physics_hz / camera_fps) % 1 == 0
+            ), f"Expected `camera_fps` to roundly divide `physics_hz`, got {camera_fps=} and {physics_hz=}."
+            self.physics_camera_ratio = int(physics_hz / camera_fps)
+        else:
+            self.physics_camera_ratio = 1
+
     def reset(self) -> None:
         """Resets the vehicle to the initial state."""
         self.set_mode(0)
@@ -380,8 +391,16 @@ def register_controller(
         self.registered_controllers[controller_id] = controller_constructor
         self.registered_base_modes[controller_id] = base_mode
 
-    def update_control(self) -> None:
-        """Runs through controllers."""
+    def update_control(self, physics_step: int) -> None:
+        """Runs through controllers.
+
+        Args:
+            physics_step (int): the current physics step
+        """
+        # skip control if we don't have enough physics steps
+        if physics_step % self.physics_control_ratio != 0:
+            return
+
         # this is the thing we cascade down controllers
         a_output = self.setpoint[:3].copy()
         z_output = self.setpoint[-1].copy()
@@ -507,7 +526,11 @@ def update_state(self) -> None:
         # update auxiliary information
         self.aux_state = self.motors.get_states()
 
-    def update_last(self) -> None:
-        """Updates things only at the end of `Aviary.step()`."""
-        if self.use_camera:
+    def update_last(self, physics_step: int) -> None:
+        """Updates things only at the end of `Aviary.step()`.
+
+        Args:
+            physics_step (int): the current physics step
+        """
+        if self.use_camera and (physics_step % self.physics_camera_ratio == 0):
             self.rgbaImg, self.depthImg, self.segImg = self.camera.capture_image()