Limit-aware noise adition

urdfenvs/sensors/lidar.py

@@ -63,12 +63,18 @@ def __init__(
             + i * (angle_limits[1] - angle_limits[0]) / self._nb_rays
             for i in range(self._nb_rays)
         ]
-        self._rel_positions = np.zeros(2 * nb_rays)
+        self._rel_positions = np.zeros((nb_rays, 2))
         self._distances = np.zeros(nb_rays)
         self._sphere_ids = [
             -1,
         ] * self._nb_rays
 
+        self._observation_limits = np.repeat(
+                np.array([[-ray_length], [ray_length]]),
+                nb_rays * 2,
+                axis=1,
+                )
+
     def get_observation_size(self):
         """Getter for the dimension of the observation space."""
         if self._raw_data:
@@ -86,6 +92,8 @@ def get_observation_space(self, obstacles: dict, goals: dict):
         )
         return gym.spaces.Dict({self._name: observation_space})
 
+
+
     def sense(self, robot, obstacles: dict, goals: dict, t: float):
         """Sense the distance toward the next object with the Lidar."""
         self._call_counter += 1
@@ -105,22 +113,20 @@ def sense(self, robot, obstacles: dict, goals: dict, t: float):
                 lidar
                 * np.array([np.cos(theta + yaw), np.sin(theta + yaw)])
             )
-            noisy_rel_positions = np.random.normal(
-                true_rel_positions, self._variance
-            )
 
-            self._rel_positions[2 * i : 2 * i + 2] = noisy_rel_positions
-            self._distances[i] = np.linalg.norm(
-                self._rel_positions[2 * i : 2 * i + 2]
-            )
+            self._rel_positions[i, :] = true_rel_positions
+        noisy_rel_positions = self.add_noise(self._rel_positions.flatten()).reshape(self._nb_rays, 2)
+        self._rel_positions = noisy_rel_positions
+
+        self._distances = np.linalg.norm(self._rel_positions, axis=1)
         if (
             self._plotting_interval > 0
             and self._call_counter % self._plotting_interval == 0
         ):
             self.update_lidar_spheres(lidar_position)
         if self._raw_data:
             return self._distances
-        return self._rel_positions
+        return self._rel_positions.flatten()
 
     def init_lidar_spheres(self, lidar_position):
         """

urdfenvs/sensors/sensor.py

@@ -2,6 +2,8 @@
 from typing import List
 from abc import abstractmethod
 
+import numpy as np
+
 from urdfenvs.sensors.physics_engine_interface import PhysicsEngineInterface, PybulletInterface, MujocoInterface
 
 class Sensor():
@@ -38,6 +40,16 @@ def name(self):
     def set_data(self, data):
         self._physics_engine.set_data(data)
 
+    def add_noise(self, exact_data: np.ndarray):
+        """Add noise to the exact data."""
+        noisy_data = np.random.normal(exact_data, self._variance)
+        if np.all(exact_data >= self._observation_limits[0]) and np.all(exact_data <= self._observation_limits[1]):
+            clipped = np.clip(noisy_data, self._observation_limits[0], self._observation_limits[1])
+            return clipped
+        else:
+            return noisy_data
+
+
     @abstractmethod
     def get_observation_size(self):
         pass