Add Noise to the IMU sensors in the Mujoco simulation

src/bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/robot.py

@@ -1,20 +1,44 @@
 import mujoco
+import numpy as np
 
 from bitbots_mujoco_sim.camera import Camera
 from bitbots_mujoco_sim.joint import Joint
-from bitbots_mujoco_sim.sensor import Sensor
+from bitbots_mujoco_sim.sensor import NoisySensor, QuaternionNoisySensor, Sensor
 
 
 class Robot:
     """Represents the Pi Plus robot, holding all its components like the camera, joints, and sensors."""
 
-    def __init__(self, model: mujoco.MjModel, data: mujoco.MjData, index: int = 0):
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        data: mujoco.MjData,
+        index: int = 0,
+        noise_config: dict[str, float] | None = None,
+        rng: np.random.Generator | None = None,
+    ):
         self.index: int = index
+        self._rng = rng or np.random.default_rng()
+        self.noise_config = noise_config or {}
         self.camera = Camera(model, data, name=self._get_name("head_camera"))
 
         def j(ros_name: str) -> Joint:
             return Joint(model, data, ros_name=ros_name, name=self._get_name(ros_name))
 
+        def s(base_name: str, ros_name: str) -> Sensor:
+            return Sensor(model, data, name=self._get_name(base_name), ros_name=ros_name)
+
+        def n(base_name: str, ros_name: str, gaussian_key: str, bias_key: str) -> NoisySensor:
+            return NoisySensor(
+                model,
+                data,
+                name=self._get_name(base_name),
+                ros_name=ros_name,
+                gaussian_stddev=float(self.noise_config.get(gaussian_key, 0.0)),
+                bias_stddev=float(self.noise_config.get(bias_key, 0.0)),
+                rng=self._rng,
+            )
+
         self.joints = RobotJoints(
             [
                 # --- Head ---
@@ -48,14 +72,22 @@ def j(ros_name: str) -> Joint:
         )
         self.sensors = RobotSensors(
             [
-                Sensor(model, data, name=self._get_name("gyro"), ros_name="IMU_gyro"),
-                Sensor(model, data, name=self._get_name("accelerometer"), ros_name="IMU_accelerometer"),
-                Sensor(model, data, name=self._get_name("orientation"), ros_name="IMU_orientation"),
-                Sensor(model, data, name=self._get_name("position"), ros_name="IMU_position"),
-                Sensor(model, data, name=self._get_name("l_foot_pos"), ros_name="left_foot_position"),
-                Sensor(model, data, name=self._get_name("r_foot_pos"), ros_name="right_foot_position"),
-                Sensor(model, data, name=self._get_name("l_foot_global_linvel"), ros_name="left_foot_velocity"),
-                Sensor(model, data, name=self._get_name("r_foot_global_linvel"), ros_name="right_foot_velocity"),
+                n("gyro", "IMU_gyro", "gyro_gaussian_stddev", "gyro_bias_stddev"),
+                n("accelerometer", "IMU_accelerometer", "accelerometer_gaussian_stddev", "accelerometer_bias_stddev"),
+                QuaternionNoisySensor(
+                    model,
+                    data,
+                    name=self._get_name("orientation"),
+                    ros_name="IMU_orientation",
+                    gaussian_stddev=float(self.noise_config.get("orientation_gaussian_stddev", 0.0)),
+                    bias_stddev=float(self.noise_config.get("orientation_bias_stddev", 0.0)),
+                    rng=self._rng,
+                ),
+                s("position", "IMU_position"),
+                s("l_foot_pos", "left_foot_position"),
+                s("r_foot_pos", "right_foot_position"),
+                s("l_foot_global_linvel", "left_foot_velocity"),
+                s("r_foot_global_linvel", "right_foot_velocity"),
             ]
         )
 

src/bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/sensor.py

@@ -16,6 +16,11 @@ def data(self) -> np.ndarray:
         """Gets the current sensor reading from sensordata array."""
         return self._data.sensordata[self.adr : self.adr + self.dim]
 
+    @property
+    def noisy_data(self) -> np.ndarray:
+        """Gets the current sensor reading with measurement noise applied."""
+        return self.data
+
     @property
     def adr(self) -> int:
         """Gets the address of the sensor data in sensordata array."""
@@ -27,6 +32,64 @@ def dim(self) -> int:
         return int(self.instance.dim)
 
 
+class NoisySensor(Sensor):
+    """Represents a sensor with per-start fixed bias and additive Gaussian noise."""
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        data: mujoco.MjData,
+        name: str,
+        ros_name: str,
+        gaussian_stddev: float = 0.0,
+        bias_stddev: float = 0.0,
+        rng: np.random.Generator | None = None,
+        fixed_bias: np.ndarray | None = None,
+    ):
+        super().__init__(model, data, name, ros_name)
+        self.gaussian_stddev = float(gaussian_stddev)
+        self.bias_stddev = float(bias_stddev)
+        self._rng = rng if rng is not None else np.random.default_rng()
+
+        if fixed_bias is not None:
+            self._fixed_bias = np.asarray(fixed_bias, dtype=float)
+        elif self.bias_stddev > 0:
+            self._fixed_bias = self._rng.normal(scale=self.bias_stddev, size=self.dim)
+        else:
+            self._fixed_bias = np.zeros(self.dim, dtype=float)
+
+    @property
+    def fixed_bias(self) -> np.ndarray:
+        """Gets the fixed bias sampled at simulation start."""
+        return self._fixed_bias
+
+    @property
+    def noisy_data(self) -> np.ndarray:
+        """Gets the sensor reading with fixed bias and additive Gaussian noise."""
+        return self.data + self._fixed_bias + self._rng.normal(scale=self.gaussian_stddev, size=self.dim)
+
+    def reset_fixed_bias(self, fixed_bias: np.ndarray | None = None) -> None:
+        """Resets the fixed bias for this sensor."""
+        if fixed_bias is not None:
+            self._fixed_bias = np.asarray(fixed_bias, dtype=float)
+        elif self.bias_stddev > 0:
+            self._fixed_bias = self._rng.normal(scale=self.bias_stddev, size=self.dim)
+        else:
+            self._fixed_bias = np.zeros(self.dim, dtype=float)
+
+
+class QuaternionNoisySensor(NoisySensor):
+    """Represents a quaternion sensor with noise and normalization."""
+
+    @property
+    def noisy_data(self) -> np.ndarray:
+        noisy = super().noisy_data
+        norm = np.linalg.norm(noisy)
+        if norm == 0:
+            return np.array([1.0, 0.0, 0.0, 0.0], dtype=float)
+        return noisy / norm
+
+
 class FootPressureSensor(Sensor):
     """Represents a foot pressure sensor, providing access to individual pressure readings."""
 

src/bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/simulation.py

@@ -1,9 +1,11 @@
 import math
 import time
+from collections.abc import Callable
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 import mujoco
+import numpy as np
 from ament_index_python.packages import get_package_share_directory
 from mujoco import viewer
 from rclpy.node import Node
@@ -27,16 +29,58 @@ def __init__(self):
         self.declare_parameter("world_file", "")
         self.declare_parameter("use_namespace", True)
 
+        # These default value are oriented on some real world data in standstill.
+        # They are a bit worse than what was observed there.
+        # No systematic bias could be observed, so they are kept at zero as of now.
+        self.declare_parameter("imu_accelerometer_gaussian_stddev", 0.05)
+        self.declare_parameter("imu_gyro_gaussian_stddev", 0.005)
+
+        self.declare_parameter("imu_accelerometer_bias_stddev", 0.0)
+        self.declare_parameter("imu_gyro_bias_stddev", 0.0)
+
+        # The orientation noise IRL seems to be extremely close to 0 in standstill.
+        # Further measurements may reveal some larger deviations, but for now we just use 0.
+        self.declare_parameter("orientation_gaussian_stddev", 0.0)
+        self.declare_parameter("orientation_bias_stddev", 0.0)
+
+        self.declare_parameter("imu_frame", "imu_frame")
+        self.declare_parameter("camera_optical_frame", "zed_left_camera_optical_frame")
+
         self.package_path = get_package_share_directory("bitbots_mujoco_sim")
         world_file = self.get_parameter("world_file").get_parameter_value().string_value
         if not world_file:
             world_file = self._generate_default_world()
 
         self.use_namespace: bool = self.get_parameter("use_namespace").get_parameter_value().bool_value
+        self.imu_accelerometer_gaussian_stddev = self.get_parameter("imu_accelerometer_gaussian_stddev").get_parameter_value().double_value
+        self.imu_accelerometer_bias_stddev = self.get_parameter("imu_accelerometer_bias_stddev").get_parameter_value().double_value
+        self.imu_gyro_gaussian_stddev = self.get_parameter("imu_gyro_gaussian_stddev").get_parameter_value().double_value
+        self.imu_gyro_bias_stddev = self.get_parameter("imu_gyro_bias_stddev").get_parameter_value().double_value
+        self.orientation_gaussian_stddev = self.get_parameter("orientation_gaussian_stddev").get_parameter_value().double_value
+        self.orientation_bias_stddev = self.get_parameter("orientation_bias_stddev").get_parameter_value().double_value
+        self._rng = np.random.default_rng()
+
         self.model: mujoco.MjModel = mujoco.MjModel.from_xml_path(world_file)
         self.data: mujoco.MjData = mujoco.MjData(self.model)
         self.robots: list[RobotSimulation] = [
-            RobotSimulation(self, Robot(self.model, self.data, idx)) for idx in self._find_robot_indices()
+            RobotSimulation(
+                self,
+                Robot(
+                    self.model,
+                    self.data,
+                    idx,
+                    noise_config={
+                        "gyro_gaussian_stddev": self.imu_gyro_gaussian_stddev,
+                        "gyro_bias_stddev": self.imu_gyro_bias_stddev,
+                        "accelerometer_gaussian_stddev": self.imu_accelerometer_gaussian_stddev,
+                        "accelerometer_bias_stddev": self.imu_accelerometer_bias_stddev,
+                        "orientation_gaussian_stddev": self.orientation_gaussian_stddev,
+                        "orientation_bias_stddev": self.orientation_bias_stddev,
+                    },
+                    rng=self._rng,
+                ),
+            )
+            for idx in self._find_robot_indices()
         ]
 
         self.time = 0.0
@@ -48,8 +92,8 @@ def __init__(self):
         self.clock_publisher = self.create_publisher(Clock, "clock", 1)
         self.create_subscription(Float32, "real_time_factor", self.real_time_factor_callback, 1)
 
-        self.imu_frame_id = self.get_parameter_or("imu_frame", "imu_frame")
-        self.camera_optical_frame_id = self.get_parameter_or("camera_optical_frame", "zed_left_camera_optical_frame")
+        self.imu_frame_id = self.get_parameter("imu_frame").get_parameter_value().string_value
+        self.camera_optical_frame_id = self.get_parameter("camera_optical_frame").get_parameter_value().string_value
         self.camera_active = True
 
         self.events = [
@@ -180,14 +224,16 @@ def publish_imu_event(self) -> None:
         imu.header.stamp = self.simulation.time_message
         imu.header.frame_id = self.simulation.imu_frame_id
         imu.linear_acceleration.x, imu.linear_acceleration.y, imu.linear_acceleration.z = (
-            self.robot.sensors.accelerometer.data
+            self.robot.sensors.accelerometer.noisy_data
         )
 
         if imu.linear_acceleration.x == 0 and imu.linear_acceleration.y == 0 and imu.linear_acceleration.z == 0:
             imu.linear_acceleration.z = 0.001
 
-        imu.angular_velocity.x, imu.angular_velocity.y, imu.angular_velocity.z = self.robot.sensors.gyro.data
-        imu.orientation.w, imu.orientation.x, imu.orientation.y, imu.orientation.z = self.robot.sensors.orientation.data
+        imu.angular_velocity.x, imu.angular_velocity.y, imu.angular_velocity.z = self.robot.sensors.gyro.noisy_data
+        imu.orientation.w, imu.orientation.x, imu.orientation.y, imu.orientation.z = (
+            self.robot.sensors.orientation.noisy_data
+        )
         self.node_publishers["imu"].publish(imu)
 
     def publish_camera_event(self) -> None: