From 34ab3321e28cb62fec41f531445c8c12b21b9ab2 Mon Sep 17 00:00:00 2001 From: Max Spahn Date: Mon, 31 Jul 2023 15:02:20 +0200 Subject: [PATCH] ft[sensors]: Adds occupancy sensor. Creates parent class GridSensor for SDF and occupancy sensor as most functions overlap. --- examples/point_robot_occupancy_sensor.py | 68 ++++++++++++++++++++++++ examples/test_examples.py | 8 +++ pyproject.toml | 2 +- urdfenvs/sensors/grid_sensor.py | 64 ++++++++++++++++++++++ urdfenvs/sensors/occupancy_sensor.py | 42 +++++++++++++++ urdfenvs/sensors/sdf_sensor.py | 44 +++------------ 6 files changed, 190 insertions(+), 38 deletions(-) create mode 100644 examples/point_robot_occupancy_sensor.py create mode 100644 urdfenvs/sensors/grid_sensor.py create mode 100644 urdfenvs/sensors/occupancy_sensor.py diff --git a/examples/point_robot_occupancy_sensor.py b/examples/point_robot_occupancy_sensor.py new file mode 100644 index 0000000..33f933d --- /dev/null +++ b/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) diff --git a/examples/test_examples.py b/examples/test_examples.py index 9ba33b8..6a35639 100644 --- a/examples/test_examples.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 8450a65..030bdf2 100644 --- a/pyproject.toml +++ b/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 "] maintainers = [ diff --git a/urdfenvs/sensors/grid_sensor.py b/urdfenvs/sensors/grid_sensor.py new file mode 100644 index 0000000..5471908 --- /dev/null +++ b/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 + diff --git a/urdfenvs/sensors/occupancy_sensor.py b/urdfenvs/sensors/occupancy_sensor.py new file mode 100644 index 0000000..2baf2ce --- /dev/null +++ b/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 diff --git a/urdfenvs/sensors/sdf_sensor.py b/urdfenvs/sensors/sdf_sensor.py index 2a08eb8..ed9a659 100644 --- a/urdfenvs/sensors/sdf_sensor.py +++ b/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