ft[sensors]: Adds occupancy sensor. Creates parent class GridSensor f…

…or SDF and occupancy sensor as most functions overlap.

examples/point_robot_occupancy_sensor.py

@@ -0,0 +1,68 @@
+import gymnasium as gym
+import numpy as np
+
+from urdfenvs.urdf_common.urdf_env import UrdfEnv
+from urdfenvs.robots.generic_urdf import GenericUrdfReacher
+from urdfenvs.sensors.occupancy_sensor import OccupancySensor
+from urdfenvs.scene_examples.obstacles import (
+    cylinder_obstacle,
+    sphereObst2,
+    sphereObst1,
+)
+
+
+def get_index_from_coordinates(point, mesh) -> tuple:
+    distances = np.linalg.norm(mesh - point, axis=3)
+    return np.unravel_index(np.argmin(distances), mesh.shape[:-1])
+
+def evaluate_occupancy(point, mesh, occupancy, resolution) -> int:
+    index = list(get_index_from_coordinates(point, mesh))
+    return occupancy[tuple(index)]
+
+
+
+
+def run_point_robot_with_occupancy_sensor(n_steps=10, render=False, obstacles=True, goal=True):
+    robots = [
+        GenericUrdfReacher(urdf="pointRobot.urdf", mode="vel"),
+    ]
+    env: UrdfEnv = gym.make(
+        "urdf-env-v0",
+        dt=0.01, robots=robots, render=render
+    )
+    env.add_obstacle(sphereObst2)
+    env.add_obstacle(cylinder_obstacle)
+    env.add_obstacle(sphereObst1)
+
+    # add sensor
+    sensor = OccupancySensor(
+        limits =  np.array([[-5, 5], [-5, 5], [0, 0]]),
+        resolution = np.array([101, 101, 1], dtype=int),
+        interval=10,
+    )
+
+    env.add_sensor(sensor, [0])
+    # Set spaces AFTER all components have been added.
+    env.set_spaces()
+    defaultAction = np.array([0.5, -0.2, 0.0])
+    pos0 = np.array([0.0, 0.0, 0.0])
+    vel0 = np.array([1.0, 0.0, 0.0])
+    initial_observations = []
+    ob, _ = env.reset(pos=pos0, vel=vel0)
+    initial_observations.append(ob)
+
+    history = []
+    for _ in range(n_steps):
+        action = defaultAction
+        ob, *_ = env.step(action)
+        point = np.append(ob['robot_0']['joint_state']['position'][0:2], 0.0)
+        occupancy = ob['robot_0']['Occupancy']
+        occupancy_eval = evaluate_occupancy(point, sensor.mesh(), occupancy, [0.2, 0.2, 1])
+        print(occupancy_eval)
+        history.append(ob)
+    env.close()
+    return history
+
+
+if __name__ == "__main__":
+    run_point_robot_with_occupancy_sensor(render=True, n_steps=300)

examples/test_examples.py

@@ -82,6 +82,14 @@ def test_point_robot_with_obstacle_sensor():
     from point_robot_obstacle_sensor import run_point_robot_with_obstacle_sensor
     blueprint_test(run_point_robot_with_obstacle_sensor)
 
+def test_point_robot_with_occupancy_sensor():
+    from point_robot_occupancy_sensor import run_point_robot_with_occupancy_sensor
+    blueprint_test(run_point_robot_with_occupancy_sensor)
+
+def test_point_robot_with_sdf_sensor():
+    from point_robot_sdf_sensor import run_point_robot_with_sdf_sensor
+    blueprint_test(run_point_robot_with_sdf_sensor)
+
 def test_tiago_robot():
     from tiago import run_tiago
     blueprint_test(run_tiago)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "urdfenvs"
-version = "0.8.5"
+version = "0.8.6"
 description = "Simple simulation environment for robots, based on the urdf files."
 authors = ["Max Spahn <[email protected]>"]
 maintainers = [

urdfenvs/sensors/grid_sensor.py

@@ -0,0 +1,64 @@
+"""Module for occupancy sensor simulation."""
+from abc import abstractmethod
+from time import perf_counter
+import numpy as np
+import gymnasium as gym
+
+from urdfenvs.sensors.sensor import Sensor
+
+
+class LinkIdNotFoundError(Exception):
+    pass
+
+
+class GridSensor(Sensor):
+    def __init__(
+        self,
+        limits: np.ndarray = np.array([[-1, -1], [-1, 1], [-1, 1]]),
+        resolution: np.ndarray = np.array([10, 10, 10], dtype=int),
+        interval: int = -1,
+        name: str = "Grid",
+    ):
+        super().__init__(name)
+        self._resolution = resolution
+        self._limits = limits
+        self._interval = interval
+        self._call_counter = 13
+        self._computed = False
+        x_values = np.linspace(limits[0][0], limits[0][1], resolution[0])
+        y_values = np.linspace(limits[1][0], limits[1][1], resolution[1])
+        z_values = np.linspace(limits[2][0], limits[2][1], resolution[2])
+        self._mesh = np.stack(
+            np.meshgrid(x_values, y_values, z_values, indexing="ij"), axis=-1
+        )
+        self._grid_values = np.zeros(shape=self._mesh.shape[0:3], dtype=int)
+
+    def mesh(self) -> np.ndarray:
+        return self._mesh
+
+    def get_observation_size(self):
+        return self._grid_values.shape
+
+    @abstractmethod
+    def get_observation_space(self, obstacles: dict, goals: dict):
+        """Create observation space, all observations should be inside the
+        observation space."""
+        pass
+
+    def distances(self, obstacles: dict) -> np.ndarray:
+        mesh_flat = self._mesh.reshape((-1, 3))
+        distances = np.min(
+            np.array(
+                [
+                    obstacle.distance(mesh_flat)
+                    for obstacle in list(obstacles.values())
+                ]
+            ),
+            axis=0,
+        )
+        return distances
+
+    @abstractmethod
+    def sense(self, robot, obstacles: dict, goals: dict, t: float) -> np.ndarray:
+        pass
+

urdfenvs/sensors/occupancy_sensor.py

@@ -0,0 +1,42 @@
+"""Module for occupancy sensor simulation."""
+from time import perf_counter
+import numpy as np
+import gymnasium as gym
+
+from urdfenvs.sensors.grid_sensor import GridSensor
+
+
+class OccupancySensor(GridSensor):
+    def __init__(
+        self,
+        limits: np.ndarray = np.array([[-1, -1], [-1, 1], [-1, 1]]),
+        resolution: np.ndarray = np.array([10, 10, 10], dtype=int),
+        interval: int = -1,
+    ):
+        super().__init__(limits=limits, resolution=resolution, interval=interval, name="Occupancy")
+
+    def get_observation_space(self, obstacles: dict, goals: dict):
+        """Create observation space, all observations should be inside the
+        observation space."""
+        observation_space = gym.spaces.Box(
+            0,
+            1,
+            shape=self.get_observation_size(),
+            dtype=int,
+        )
+        return gym.spaces.Dict({self._name: observation_space})
+
+    def sense(self, robot, obstacles: dict, goals: dict, t: float):
+        self._call_counter += 1
+        if self._computed and (
+            self._interval < 0 or self._call_counter % self._interval != 0
+        ):
+            return self._grid_values
+        start_time = perf_counter()
+        distances = self.distances(obstacles)
+        self._grid_values = np.array(distances <= 0.0, dtype=int).reshape(self._resolution)
+        end_time = perf_counter()
+
+        print(f"Computed Occupancy in {end_time-start_time} s")
+        self._computed = True
+        return self._grid_values

urdfenvs/sensors/sdf_sensor.py

@@ -4,40 +4,21 @@
 import pybullet as p
 import gymnasium as gym
 
-from urdfenvs.sensors.sensor import Sensor
+from urdfenvs.sensors.grid_sensor import GridSensor
 
 
 class LinkIdNotFoundError(Exception):
     pass
 
 
-class SDFSensor(Sensor):
+class SDFSensor(GridSensor):
     def __init__(
         self,
         limits: np.ndarray = np.array([[-1, -1], [-1, 1], [-1, 1]]),
         resolution: np.ndarray = np.array([10, 10, 10], dtype=int),
         interval: int = -1,
     ):
-        super().__init__("SDFSensor")
-        self._resolution = resolution
-        self._limits = limits
-        self._interval = interval
-        self._call_counter = 13
-        self._computed = False
-        x_values = np.linspace(limits[0][0], limits[0][1], resolution[0])
-        y_values = np.linspace(limits[1][0], limits[1][1], resolution[1])
-        z_values = np.linspace(limits[2][0], limits[2][1], resolution[2])
-        self._mesh = np.stack(
-            np.meshgrid(x_values, y_values, z_values, indexing="ij"), axis=-1
-        )
-        self._sdf = np.zeros(shape=self._mesh.shape[0:3])
-
-    def mesh(self) -> np.ndarray:
-        return self._mesh
-
-    def get_observation_size(self):
-        return self._sdf.shape
-
+        super().__init__(limits=limits, resolution=resolution, interval=interval, name="SDFSensor")
     def get_observation_space(self, obstacles: dict, goals: dict):
         """Create observation space, all observations should be inside the
         observation space."""
@@ -54,23 +35,12 @@ def sense(self, robot, obstacles: dict, goals: dict, t: float):
         if self._computed and (
             self._interval < 0 or self._call_counter % self._interval != 0
         ):
-            return self._sdf
+            return self._grid_values
         start_time = perf_counter()
-        mesh_flat = self._mesh.reshape((-1, 3))
-        self._sdf = np.maximum(
-            np.min(
-                np.array(
-                    [
-                        obstacle.distance(mesh_flat)
-                        for obstacle in list(obstacles.values())
-                    ]
-                ),
-                axis=0,
-            ),
-            0.0,
-        ).reshape(self._resolution)
+        distances = self.distances(obstacles)
+        self._grid_values = np.maximum(distances, 0.0).reshape(self._resolution)
         end_time = perf_counter()
 
         print(f"Computed SDF in {end_time-start_time} s")
         self._computed = True
-        return self._sdf
+        return self._grid_values