Add sensor noise (#1631)

HULKs · Feb 12, 2025 · 271d8a9 · 271d8a9
1 parent 237c2a3
commit 271d8a9
Show file tree

Hide file tree

Showing 3 changed files with 97 additions and 32 deletions.
diff --git a/tools/machine-learning/mujoco/packages/nao_env/src/nao_env/nao_base_env.py b/tools/machine-learning/mujoco/packages/nao_env/src/nao_env/nao_base_env.py
@@ -92,6 +92,8 @@ def __init__(
             fsr_sensor_delay=sensor_delay,
             gyroscope_sensor_delay=sensor_delay,
             accelerometer_sensor_delay=sensor_delay,
+            gyroscope_noise=np.array([0.003889, 0.005522, 0.002381]),
+            accelerometer_noise=np.array([0.063137, 0.029709, 0.087581]),
         )
 
     def initialize_terrain(

diff --git a/tools/machine-learning/mujoco/packages/nao_interface/src/nao_interface/nao.py b/tools/machine-learning/mujoco/packages/nao_interface/src/nao_interface/nao.py
@@ -17,6 +17,14 @@ class SensorBuffers:
     accelerometer: RingBuffer
 
 
+@dataclass
+class NoiseParameters:
+    gyroscope: NDArray[np.floating]
+    accelerometer: NDArray[np.floating]
+    fsr_left: NDArray[np.floating]
+    fsr_right: NDArray[np.floating]
+
+
 class Nao:
     def __init__(
         self,
@@ -27,9 +35,29 @@ def __init__(
         fsr_sensor_delay: int = 0,
         gyroscope_sensor_delay: int = 0,
         accelerometer_sensor_delay: int = 0,
+        gyroscope_noise: NDArray[np.floating] | None = None,
+        accelerometer_noise: NDArray[np.floating] | None = None,
+        left_fsr_noise: NDArray[np.floating] | None = None,
+        right_fsr_noise: NDArray[np.floating] | None = None,
     ) -> None:
+        if gyroscope_noise is None:
+            gyroscope_noise = np.array([0.0, 0.0, 0.0])
+        if accelerometer_noise is None:
+            accelerometer_noise = np.array([0.0, 0.0, 0.0])
+        if left_fsr_noise is None:
+            left_fsr_noise = np.array([0.0, 0.0, 0.0, 0.0])
+        if right_fsr_noise is None:
+            right_fsr_noise = np.array([0.0, 0.0, 0.0, 0.0])
+
         self.model = model
         self.data = data
+        self.rng = np.random.default_rng()
+        self.noise_parameters = NoiseParameters(
+            gyroscope=gyroscope_noise,
+            accelerometer=accelerometer_noise,
+            fsr_left=left_fsr_noise,
+            fsr_right=right_fsr_noise,
+        )
 
         self.fsr_scale = fsr_scale
 
@@ -105,46 +133,70 @@ def _read_positions(self) -> NDArray[np.floating]:
         return self.position_sensors.to_numpy()
 
     def _read_left_fsr_values(self) -> NDArray[np.floating]:
-        return self.fsr_scale * np.array(
-            [
-                self.data.sensor(
-                    "force_sensitive_resistors.left.front_left",
-                ).data[0],
-                self.data.sensor(
-                    "force_sensitive_resistors.left.front_right",
-                ).data[0],
-                self.data.sensor(
-                    "force_sensitive_resistors.left.rear_left",
-                ).data[0],
-                self.data.sensor(
-                    "force_sensitive_resistors.left.rear_right",
-                ).data[0],
-            ],
+        noise = self.rng.normal(
+            loc=[0.0, 0.0, 0.0, 0.0],
+            scale=self.noise_parameters.fsr_left,
+        )
+        return (
+            self.fsr_scale
+            * np.array(
+                [
+                    self.data.sensor(
+                        "force_sensitive_resistors.left.front_left",
+                    ).data[0],
+                    self.data.sensor(
+                        "force_sensitive_resistors.left.front_right",
+                    ).data[0],
+                    self.data.sensor(
+                        "force_sensitive_resistors.left.rear_left",
+                    ).data[0],
+                    self.data.sensor(
+                        "force_sensitive_resistors.left.rear_right",
+                    ).data[0],
+                ],
+            )
+            + noise
         )
 
     def _read_right_fsr_values(self) -> NDArray[np.floating]:
-        return self.fsr_scale * np.array(
-            [
-                self.data.sensor(
-                    "force_sensitive_resistors.right.front_left",
-                ).data[0],
-                self.data.sensor(
-                    "force_sensitive_resistors.right.front_right",
-                ).data[0],
-                self.data.sensor(
-                    "force_sensitive_resistors.right.rear_left",
-                ).data[0],
-                self.data.sensor(
-                    "force_sensitive_resistors.right.rear_right",
-                ).data[0],
-            ],
+        noise = self.rng.normal(
+            loc=[0.0, 0.0, 0.0, 0.0],
+            scale=self.noise_parameters.fsr_right,
+        )
+        return (
+            self.fsr_scale
+            * np.array(
+                [
+                    self.data.sensor(
+                        "force_sensitive_resistors.right.front_left",
+                    ).data[0],
+                    self.data.sensor(
+                        "force_sensitive_resistors.right.front_right",
+                    ).data[0],
+                    self.data.sensor(
+                        "force_sensitive_resistors.right.rear_left",
+                    ).data[0],
+                    self.data.sensor(
+                        "force_sensitive_resistors.right.rear_right",
+                    ).data[0],
+                ],
+            )
+            + noise
         )
 
     def _read_gyroscope(self) -> NDArray[np.floating]:
-        return self.data.sensor("gyroscope").data
+        noise = self.rng.normal(
+            loc=[0.0, 0.0, 0.0],
+            scale=self.noise_parameters.gyroscope,
+        )
+        return self.data.sensor("gyroscope").data + noise
 
     def _read_accelerometer(self) -> NDArray[np.floating]:
-        return self.data.sensor("accelerometer").data
+        noise = self.rng.normal(
+            loc=[0.0, 0.0, 0.0],
+            scale=self.noise_parameters.accelerometer,
+        )
+        return self.data.sensor("accelerometer").data + noise
 
     def position_encoders(self) -> NDArray[np.floating]:
         return self._sensor_buffers.positions.left()

diff --git a/tools/machine-learning/mujoco/scripts/mujoco-walking.py b/tools/machine-learning/mujoco/scripts/mujoco-walking.py
@@ -52,6 +52,13 @@ def main(*, throw_tomatoes: bool, load_policy: str | None) -> None:
     fsr_figure.add_line("Left FSR")
     fsr_figure.add_line("Right FSR")
 
+    gyro_figure = viewer.figure("gyro")
+    gyro_figure.set_title("Gyroscope")
+    gyro_figure.set_x_label("Step")
+    gyro_figure.add_line("X Gyro")
+    gyro_figure.add_line("Y Gyro")
+    gyro_figure.add_line("Z Gyro")
+
     while viewer.is_alive:
         start_time = time.time()
         # viewer.track_with_camera("Nao")
@@ -62,6 +69,10 @@ def main(*, throw_tomatoes: bool, load_policy: str | None) -> None:
         fsr_figure.push_data_to_line("Left FSR", env.nao.left_fsr().sum())
         fsr_figure.push_data_to_line("Right FSR", env.nao.right_fsr().sum())
 
+        gyro_figure.push_data_to_line("X Gyro", env.nao.gyroscope()[0])
+        gyro_figure.push_data_to_line("Y Gyro", env.nao.gyroscope()[1])
+        gyro_figure.push_data_to_line("Z Gyro", env.nao.gyroscope()[2])
+
         total_reward += reward
 
         for key, value in infos.items():