add static obstacle support to the aviary

PyFlyt/core/abstractions/__init__.py

@@ -9,4 +9,5 @@
 from .gimbals import Gimbals
 from .lifting_surfaces import LiftingSurface, LiftingSurfaces
 from .motors import Motors
+from .obstacle import Obstacle
 from .pid import PID

PyFlyt/core/abstractions/obstacle.py

@@ -0,0 +1,135 @@
+"""Obstacle class for spawning static URDF obstacles in the Aviary."""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+from pybullet_utils import bullet_client
+
+# directory containing the built-in obstacle URDFs
+_BUILTIN_OBSTACLE_DIR = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)),
+    "..",
+    "..",
+    "models",
+    "obstacles",
+)
+
+# mapping of friendly names to their built-in URDF files
+_BUILTIN_OBSTACLES: dict[str, str] = {
+    "cube": os.path.join(_BUILTIN_OBSTACLE_DIR, "cube.urdf"),
+    "cylinder": os.path.join(_BUILTIN_OBSTACLE_DIR, "cylinder.urdf"),
+    "sphere": os.path.join(_BUILTIN_OBSTACLE_DIR, "sphere.urdf"),
+}
+
+
+class Obstacle:
+    """An `Obstacle` represents a static object in the Aviary that drones may collide with.
+
+    Obstacles are loaded from URDF files. PyFlyt ships with a small number of built-in
+    primitive obstacles (`"cube"`, `"cylinder"`, `"sphere"`), accessible by name.
+    Arbitrary URDF files may also be loaded by providing a filesystem path.
+
+    Obstacles are always spawned with `useFixedBase=True`, which means PyBullet will not
+    update their positions during simulation regardless of the mass defined in the URDF.
+    This is the intended behaviour for static scene geometry.
+
+    Example:
+        >>> from PyFlyt.core import Aviary
+        >>> import numpy as np
+        >>>
+        >>> env = Aviary(
+        ...     start_pos=np.array([[0.0, 0.0, 1.0]]),
+        ...     start_orn=np.array([[0.0, 0.0, 0.0]]),
+        ...     drone_type="quadx",
+        ... )
+        >>>
+        >>> # spawn a built-in cube obstacle
+        >>> env.add_obstacle("cube", position=np.array([2.0, 0.0, 0.5]))
+        >>>
+        >>> # spawn an obstacle from a custom URDF
+        >>> env.add_obstacle("/path/to/my_obstacle.urdf", position=np.array([4.0, 0.0, 1.0]))
+
+    Args:
+        p (bullet_client.BulletClient): PyBullet physics client.
+        urdf (str): either the name of a built-in obstacle (`"cube"`, `"cylinder"`, `"sphere"`),
+            or an absolute path to a URDF file.
+        position (np.ndarray): an `(3,)` array for the X, Y, Z spawn position.
+        orientation (np.ndarray): an `(3,)` array for the spawn orientation as Euler angles
+            (roll, pitch, yaw) in radians. Defaults to zero rotation.
+        scale (float): a uniform scaling factor applied to the URDF on load. Defaults to 1.0.
+
+    """
+
+    def __init__(
+        self,
+        p: bullet_client.BulletClient,
+        urdf: str,
+        position: np.ndarray,
+        orientation: np.ndarray | None = None,
+        scale: float = 1.0,
+    ):
+        """Loads the obstacle URDF into the PyBullet client and stores its body ID.
+
+        Args:
+            p (bullet_client.BulletClient): PyBullet physics client.
+            urdf (str): name of a built-in obstacle or path to a URDF file.
+            position (np.ndarray): `(3,)` array for the X, Y, Z spawn position.
+            orientation (np.ndarray): `(3,)` array of Euler angles in radians.
+            scale (float): a uniform scaling factor applied to the URDF on load.
+
+        """
+        # resolve a built-in name to its packaged URDF path
+        urdf_path = _BUILTIN_OBSTACLES.get(urdf, urdf)
+        if not os.path.isfile(urdf_path):
+            raise FileNotFoundError(
+                f"Could not find obstacle URDF `{urdf}`. "
+                f"Expected either a built-in name from {list(_BUILTIN_OBSTACLES)} "
+                f"or a path to an existing URDF file."
+            )
+
+        # validate the position shape
+        position = np.asarray(position, dtype=np.float64)
+        if position.shape != (3,):
+            raise ValueError(
+                f"`position` must be shape (3,), got {position.shape}."
+            )
+
+        # default orientation is no rotation
+        if orientation is None:
+            orientation = np.zeros(3, dtype=np.float64)
+        orientation = np.asarray(orientation, dtype=np.float64)
+        if orientation.shape != (3,):
+            raise ValueError(
+                f"`orientation` must be shape (3,), got {orientation.shape}."
+            )
+
+        # store handles
+        self.p = p
+        self.urdf_path = urdf_path
+        self.start_pos = position
+        self.start_orn = orientation
+        self.scale = float(scale)
+
+        # convert euler to quaternion for pybullet
+        quat = self.p.getQuaternionFromEuler(orientation.tolist())
+
+        # spawn the obstacle as a fixed-base body
+        self.Id: int = self.p.loadURDF(
+            urdf_path,
+            basePosition=position.tolist(),
+            baseOrientation=quat,
+            useFixedBase=True,
+            globalScaling=self.scale,
+        )
+
+    @classmethod
+    def builtin_obstacles(cls) -> list[str]:
+        """Returns the list of built-in obstacle names recognised by `Obstacle`.
+
+        Returns:
+            list[str]: names usable as the `urdf` argument to `add_obstacle`.
+
+        """
+        return list(_BUILTIN_OBSTACLES)

PyFlyt/core/aviary.py

@@ -12,7 +12,7 @@
 import pybullet_data
 from pybullet_utils import bullet_client
 
-from PyFlyt.core.abstractions import DroneClass, WindFieldClass
+from PyFlyt.core.abstractions import DroneClass, Obstacle, WindFieldClass
 from PyFlyt.core.drones import Fixedwing, QuadX, Rocket
 
 DroneIndex = int
@@ -212,6 +212,10 @@ def __init__(
             text="RTF here", textPosition=[0, 0, 0], textColorRGB=[1, 0, 0]
         )
 
+        # obstacles registered with the aviary; respawned on every reset()
+        self._obstacle_specs: list[dict[str, Any]] = []
+        self.obstacles: list[Obstacle] = []
+
         # initialize the environment
         self.reset()
 
@@ -262,6 +266,11 @@ def reset(self) -> None:
                 )
             )
 
+        # respawn any obstacles that were registered before reset
+        self.obstacles = []
+        for spec in self._obstacle_specs:
+            self.obstacles.append(Obstacle(p=self, **spec))
+
         # initialize the wind field
         self.wind_field: None | WindFieldClass | Callable
         if self.wind_type is None:
@@ -332,6 +341,48 @@ def register_wind_field_function(self, wind_field: Callable) -> None:
         WindFieldClass._check_wind_field_validity(wind_field)
         self.wind_field = wind_field
 
+    def add_obstacle(
+        self,
+        urdf: str,
+        position: np.ndarray,
+        orientation: np.ndarray | None = None,
+        scale: float = 1.0,
+    ) -> Obstacle:
+        """Spawns a static obstacle in the simulation that drones may collide with.
+
+        The obstacle is loaded from a URDF and pinned in place with `useFixedBase=True`.
+        The obstacle is also remembered by the `Aviary`, so subsequent calls to `reset()`
+        will respawn it at the same pose. After spawning, collision tracking for the new
+        body is automatically (re-)registered via `register_all_new_bodies()`.
+
+        Args:
+            urdf (str): either the name of a built-in obstacle
+                (`"cube"`, `"cylinder"`, `"sphere"`) or a path to a URDF file.
+            position (np.ndarray): `(3,)` array for the X, Y, Z spawn position.
+            orientation (np.ndarray): `(3,)` array of Euler angles (roll, pitch, yaw)
+                in radians. Defaults to zero rotation.
+            scale (float): uniform scaling factor applied to the URDF. Defaults to 1.0.
+
+        Returns:
+            Obstacle: the spawned obstacle. Its PyBullet body ID is available as `.Id`.
+
+        """
+        spec: dict[str, Any] = dict(
+            urdf=urdf,
+            position=np.asarray(position, dtype=np.float64),
+            orientation=(
+                None if orientation is None else np.asarray(orientation, dtype=np.float64)
+            ),
+            scale=float(scale),
+        )
+        obstacle = Obstacle(p=self, **spec)
+        self.obstacles.append(obstacle)
+        self._obstacle_specs.append(spec)
+
+        # ensure the new body is tracked for collisions
+        self.register_all_new_bodies()
+        return obstacle
+
     def state(self, index: DroneIndex) -> np.ndarray:
         """Returns the state for the indexed drone.
 

PyFlyt/models/obstacles/cube.urdf

@@ -0,0 +1,26 @@
+<?xml version="1.0"?>
+<robot name="cube_obstacle">
+  <link name="base_link">
+
+    <collision>
+      <geometry>
+        <box size="1 1 1"/>
+      </geometry>
+    </collision>
+
+    <inertial>
+      <mass value="0"/>
+      <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0"/>
+    </inertial>
+
+    <visual>
+      <geometry>
+        <box size="1 1 1"/>
+      </geometry>
+      <material name="grey">
+        <color rgba="0.5 0.5 0.5 1.0"/>
+      </material>
+    </visual>
+
+  </link>
+</robot>

PyFlyt/models/obstacles/cylinder.urdf

@@ -0,0 +1,26 @@
+<?xml version="1.0"?>
+<robot name="cylinder_obstacle">
+  <link name="base_link">
+
+    <collision>
+      <geometry>
+        <cylinder radius="0.5" length="2.0"/>
+      </geometry>
+    </collision>
+
+    <inertial>
+      <mass value="0"/>
+      <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0"/>
+    </inertial>
+
+    <visual>
+      <geometry>
+        <cylinder radius="0.5" length="2.0"/>
+      </geometry>
+      <material name="grey">
+        <color rgba="0.5 0.5 0.5 1.0"/>
+      </material>
+    </visual>
+
+  </link>
+</robot>

PyFlyt/models/obstacles/sphere.urdf

@@ -0,0 +1,26 @@
+<?xml version="1.0"?>
+<robot name="sphere_obstacle">
+  <link name="base_link">
+
+    <collision>
+      <geometry>
+        <sphere radius="0.5"/>
+      </geometry>
+    </collision>
+
+    <inertial>
+      <mass value="0"/>
+      <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0"/>
+    </inertial>
+
+    <visual>
+      <geometry>
+        <sphere radius="0.5"/>
+      </geometry>
+      <material name="grey">
+        <color rgba="0.5 0.5 0.5 1.0"/>
+      </material>
+    </visual>
+
+  </link>
+</robot>

docs_source/documentation/core/abstractions.md

@@ -9,6 +9,7 @@ abstractions/camera
 abstractions/gimbals
 abstractions/lifting_surfaces
 abstractions/motors
+abstractions/obstacle
 ```
 
 ## Description

docs_source/documentation/core/abstractions/obstacle.md

@@ -0,0 +1,45 @@
+# Obstacle
+
+## Description
+
+The `Obstacle` component represents a static URDF object spawned in the `Aviary`.
+Obstacles are intended for building scene geometry (walls, boxes, cylinders, etc.) that drones can perceive and collide with, without altering the drone abstractions themselves.
+
+PyFlyt ships with a small set of primitive obstacles --- `"cube"`, `"cylinder"`, and `"sphere"` --- which can be loaded by name.
+Arbitrary URDFs may also be loaded by passing a path.
+
+Obstacles are always spawned with `useFixedBase=True`, so the URDF's mass is ignored and the body will not move under gravity or contact.
+
+## Usage
+
+Obstacles are most commonly added through the [`Aviary.add_obstacle`](aviary) helper, which constructs the `Obstacle`, tracks it for collision bookkeeping, and respawns it after every `reset()`:
+
+```python
+import numpy as np
+from PyFlyt.core import Aviary
+
+env = Aviary(
+    start_pos=np.array([[0.0, 0.0, 1.0]]),
+    start_orn=np.array([[0.0, 0.0, 0.0]]),
+    drone_type="quadx",
+)
+
+# built-in primitive obstacle
+env.add_obstacle("cube", position=np.array([2.0, 0.0, 0.5]))
+
+# custom URDF, with rotation and scaling
+env.add_obstacle(
+    "/path/to/wall.urdf",
+    position=np.array([4.0, 0.0, 1.0]),
+    orientation=np.array([0.0, 0.0, np.pi / 2]),
+    scale=1.5,
+)
+```
+
+Collisions between drones and obstacles are reflected in `env.contact_array` just like any other contact.
+
+## Class Description
+```{eval-rst}
+.. autoclass:: PyFlyt.core.abstractions.Obstacle
+    :members:
+```