Merge pull request #187 from maxspahn/ft-generic-esdf

Ft generic esdf

examples/point_robot.py

@@ -4,12 +4,13 @@
 
 from urdfenvs.scene_examples.obstacles import *
 from urdfenvs.scene_examples.goal import *
+from urdfenvs.urdf_common.urdf_env import UrdfEnv
 
 def run_point_robot(n_steps=1000, render=False, goal=True, obstacles=True):
     robots = [
         GenericUrdfReacher(urdf="pointRobot.urdf", mode="vel"),
     ]
-    env = gym.make(
+    env: UrdfEnv = gym.make(
         "urdf-env-v0",
         dt=0.01, robots=robots, render=render
     )

examples/point_robot_sdf_sensor.py

@@ -0,0 +1,81 @@
+import gym
+import numpy as np
+
+from urdfenvs.urdf_common.urdf_env import UrdfEnv
+from urdfenvs.robots.generic_urdf import GenericUrdfReacher
+from urdfenvs.sensors.sdf_sensor import SDFSensor
+from urdfenvs.scene_examples.obstacles import (
+    cylinder_obstacle,
+    sphereObst2,
+)
+
+
+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_sdf(point, mesh, sdf, resolution) -> tuple:
+    index = list(get_index_from_coordinates(point, mesh))
+    value = sdf[tuple(index)]
+    gradient = np.zeros(3)
+    for dim in range(3):
+        lower_index = tuple(index[:dim] + [index[dim] - 1] + index[dim+1:])
+        upper_index = tuple(index[:dim] + [index[dim] + 1] + index[dim+1:])
+
+        if lower_index[dim] < 0 or upper_index[dim] >= sdf.shape[dim]:
+            gradient[dim] = 0.0
+        else:
+            gradient[dim] = (sdf[upper_index] - sdf[lower_index]) / resolution[dim]
+
+    return value, gradient
+
+
+
+
+def run_point_robot_with_sdf_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 = SDFSensor(
+        limits =  np.array([[-2, 2], [-2, 2], [0, 0]]),
+        resolution = np.array([101, 101, 1], dtype=int),
+        interval=100,
+    )
+
+    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)
+    #print(f"Initial observation : {ob}")
+    #assert np.array_equal(initial_observations[0], ob)
+
+    history = []
+    for _ in range(n_steps):
+        action = defaultAction
+        ob, reward, done, info = env.step(action)
+        indices = get_index_from_coordinates(np.array([2.0, -1.0, 0.0]), sensor.mesh())
+        point = np.append(ob['robot_0']['joint_state']['position'][0:2], 0.0)
+        sdf = ob['robot_0']['SDFSensor']
+        sdf_eval = evaluate_sdf(point, sensor.mesh(), sdf, [0.2, 0.2, 1])
+        print(sdf_eval)
+        history.append(ob)
+    env.close()
+    return history
+
+
+if __name__ == "__main__":
+    run_point_robot_with_sdf_sensor(render=True, n_steps=300)

urdfenvs/sensors/sdf_sensor.py

@@ -0,0 +1,64 @@
+"""Module for signed distance field sensor simulation."""
+from time import perf_counter
+import numpy as np
+import pybullet as p
+import gym
+
+from urdfenvs.sensors.sensor import Sensor
+
+class LinkIdNotFoundError(Exception):
+    pass
+
+
+class SDFSensor(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,
+                ):
+        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
+
+    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.0,
+            10.0,
+            shape=self.get_observation_size(),
+            dtype=np.float64,
+        )
+        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._sdf
+        t0 = perf_counter()
+        for i in range(self._resolution[0]):
+            for j in range(self._resolution[1]):
+                for k in range(self._resolution[2]):
+                    pos = self._mesh[i, j, k]
+                    distance = np.min(np.array([obstacle.distance(pos) for obstacle in list(obstacles.values())]))
+                    self._sdf[i, j, k] = np.maximum(0.0, distance)
+        t1 = perf_counter()
+        print(f"Computed SDF in {t1-t0}s")
+        self._computed = True
+        return self._sdf
+
+