Implement Mujoco Simulator

Author: Samocx

bitbots_misc/bitbots_bringup/launch/simulator_teamplayer.launch

@@ -10,10 +10,15 @@
     <arg name="teamcom" default="false" description="Whether the team communication system should be started" />
     <arg name="vision" default="true" description="Whether the vision system should be started" />
     <arg name="world_model" default="true" description="Whether the world model should be started"/>
+    <arg name="mujoco" default="false" description="Whether to use the Mujoco simulator instead of Webots"/>
 
     <!-- load the general simulator -->
-    <include file="$(find-pkg-share bitbots_webots_sim)/launch/simulation.launch" />
-
+    <group if="$(var mujoco)">
+        <node pkg="bitbots_mujoco_sim" exec="sim" name="mujoco_simulation_node" output="screen"/>
+    </group>
+    <group unless="$(var mujoco)">
+        <include file="$(find-pkg-share bitbots_webots_sim)/launch/simulation.launch" />
+    </group>
     <!-- load teamplayer software stack -->
     <include file="$(find-pkg-share bitbots_bringup)/launch/teamplayer.launch">
         <arg name="behavior" value="$(var behavior)" />

bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/Camera.py

@@ -0,0 +1,59 @@
+import mujoco
+import numpy as np
+
+
+class Camera:
+    """Represents a camera in the MuJoCo simulation, providing image rendering capabilities."""
+
+    def __init__(self, model: mujoco.MjModel, data: mujoco.MjData, name: str, width: int = 800, height: int = 600):
+        self.model: mujoco.MjModel = model
+        self.data: mujoco.MjData = data
+        self.name: str = name
+        self.instance = model.camera(name)
+        self.width: int = width
+        self.height: int = height
+        self.renderer = mujoco.Renderer(self.model, height=self.height, width=self.width)
+
+    @property
+    def fov(self) -> float:
+        """Returns the camera's horizontal field of view in radians."""
+        if (hasattr(self, "_fov") and self._fov is not None):
+            return self._fov
+
+        # MuJoCo uses fovy (vertical field of view in degrees)
+        fovy_deg = (
+            self.instance.fovy[0]
+            if hasattr(self.instance.fovy, "__iter__")
+            else self.instance.fovy
+        )
+        fovy_rad = np.deg2rad(fovy_deg)
+
+        # Convert vertical FOV to horizontal FOV using aspect ratio
+        aspect_ratio = self.width / self.height
+        fovx_rad = 2 * np.arctan(np.tan(fovy_rad / 2) * aspect_ratio)
+
+        self._fov = fovx_rad
+        return self._fov
+
+    def render(self) -> bytes:
+        """
+        Renders and returns the camera image as raw bytes (BGRA format).
+        Returns Raw image data in BGRA8 format for ROS Image messages.
+        """
+        # Update renderer with current scene
+        self.renderer.update_scene(self.data, camera=self.name)
+
+        # Render the image - returns RGB by default
+        rgb_array = self.renderer.render()
+
+        # Convert RGB to BGRA (standard for ROS Image messages)
+        # MuJoCo returns RGB uint8 array of shape (height, width, 3)
+        # We need to add alpha channel and swap R and B
+        height, width, _ = rgb_array.shape
+        bgra_array = np.zeros((height, width, 4), dtype=np.uint8)
+        bgra_array[:, :, 0] = rgb_array[:, :, 2]  # B
+        bgra_array[:, :, 1] = rgb_array[:, :, 1]  # G
+        bgra_array[:, :, 2] = rgb_array[:, :, 0]  # R
+        bgra_array[:, :, 3] = 255  # A (fully opaque)
+
+        return bgra_array.tobytes()

bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/Joint.py

@@ -0,0 +1,39 @@
+import mujoco
+
+
+class Joint:
+    """Represents a single controllable joint and associated actuator in the MuJoCo simulation."""
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        data: mujoco.MjData,
+        ros_name: str,
+        name: str | None = None,
+    ):
+        self.model: mujoco.MjModel = model
+        self.data: mujoco.MjData = data
+        self.ros_name: str = ros_name
+        self.name: str = name if name is not None else ros_name
+        self.joint_instance: int = model.joint(self.name)
+        self.actuator_instance: int = model.actuator(self.name)
+
+    @property
+    def position(self) -> float:
+        """Gets the current joint position (angle) in radians."""
+        return self.data.qpos[self.joint_instance.qposadr[0]]
+    
+    @position.setter
+    def position(self, value: float) -> None:
+        """Sets the position target for the joint's position actuator."""
+        self.data.ctrl[self.actuator_instance.id] = value
+
+    @property
+    def velocity(self) -> float:
+        """Gets the current joint velocity in rad/s."""
+        return self.data.qvel[self.joint_instance.dofadr[0]]
+
+    @property
+    def effort(self) -> float:
+        """Gets the current effort (force/torque) applied by the position actuator."""
+        return self.data.actuator_force[self.actuator_instance.id]

bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/Robot.py

@@ -0,0 +1,78 @@
+import mujoco
+
+from bitbots_mujoco_sim.Camera import Camera
+from bitbots_mujoco_sim.Joint import Joint
+from bitbots_mujoco_sim.Sensor import Sensor
+
+
+class Robot:
+    """Represents the robot, holding all its components like the camera, joints, and sensors."""
+
+    def __init__(self, model: mujoco.MjModel, data: mujoco.MjData):
+        self.camera = Camera(model, data, name="head_camera")
+        self.joints = RobotJoints(
+            [
+                # --- Right Leg ---
+                Joint(model, data, ros_name="RHipYaw"),
+                Joint(model, data, ros_name="RHipRoll"),
+                Joint(model, data, ros_name="RHipPitch"),
+                Joint(model, data, ros_name="RKnee"),
+                Joint(model, data, ros_name="RAnklePitch"),
+                Joint(model, data, ros_name="RAnkleRoll"),
+                # --- Left Leg ---
+                Joint(model, data, ros_name="LHipYaw"),
+                Joint(model, data, ros_name="LHipRoll"),
+                Joint(model, data, ros_name="LHipPitch"),
+                Joint(model, data, ros_name="LKnee"),
+                Joint(model, data, ros_name="LAnklePitch"),
+                Joint(model, data, ros_name="LAnkleRoll"),
+                # --- Arms ---
+                Joint(model, data, ros_name="RShoulderPitch"),
+                Joint(model, data, ros_name="RShoulderRoll"),
+                Joint(model, data, ros_name="RElbow"),
+                Joint(model, data, ros_name="LShoulderPitch"),
+                Joint(model, data, ros_name="LShoulderRoll"),
+                Joint(model, data, ros_name="LElbow"),
+                # --- Head ---
+                Joint(model, data, ros_name="HeadPan"),
+                Joint(model, data, ros_name="HeadTilt"),
+            ]
+        )
+        self.sensors = RobotSensors(
+            [
+                # IMU Sensors
+                Sensor(model, data, name="gyro", ros_name="IMU_gyro"),
+                Sensor(model, data, name="accelerometer", ros_name="IMU_accelerometer"),
+                Sensor(model, data, name="orientation", ros_name="IMU_orientation",),  # Global orientation quaternion
+                Sensor(model, data, name="position", ros_name="IMU_position"),  # Global position vector
+                # Foot Sensors
+                Sensor(model, data, name="l_foot_pos", ros_name="left_foot_position"),
+                Sensor(model, data, name="r_foot_pos", ros_name="right_foot_position"),
+                Sensor(model, data, name="l_foot_global_linvel", ros_name="left_foot_velocity"),
+                Sensor(model, data, name="r_foot_global_linvel", ros_name="right_foot_velocity"),
+            ]
+        )
+
+
+class RobotSensors(list[Sensor]):
+    """A list of Robot Sensors with additional helper methods."""
+
+    def get(self, name: str) -> Joint:
+        """Finds and returns a specific joint by its ROS name."""
+        return next(filter(lambda joint: joint.ros_name == name, self))
+
+    @property
+    def gyro(self) -> Sensor:
+        return self.get("IMU_gyro")
+
+    @property
+    def accelerometer(self) -> Sensor:
+        return self.get("IMU_accelerometer")
+
+
+class RobotJoints(list[Joint]):
+    """A list of Robot Joints with additional helper methods."""
+
+    def get(self, name: str) -> Joint:
+        """Finds and returns a specific joint by its ROS name."""
+        return next(filter(lambda joint: joint.ros_name == name, self))

bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/Sensor.py

@@ -0,0 +1,28 @@
+import mujoco
+import numpy as np
+
+
+class Sensor:
+    """Represents a single sensor, providing a clean interface to its value."""
+
+    def __init__(self, model: mujoco.MjModel, data: mujoco.MjData, name: str, ros_name: str):
+        self._data = data
+        self.ros_name: str = ros_name
+        self.instance: int = model.sensor(name)
+        self.id: int = self.instance.id
+
+    @property
+    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 adr(self) -> int:
+        """Gets the address of the sensor data in sensordata array."""
+        return int(self.instance.adr)
+
+    @property
+    def dim(self) -> int:
+        """Gets the dimension of the sensor data."""
+        return int(self.instance.dim)
+

bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/Simulation.py

@@ -0,0 +1,160 @@
+import time
+import math
+
+import mujoco
+from ament_index_python.packages import get_package_share_directory
+from mujoco import viewer
+from rclpy.node import Node
+from rclpy.time import Time
+from rosgraph_msgs.msg import Clock
+from sensor_msgs.msg import CameraInfo, Imu, JointState, Image
+
+from bitbots_msgs.msg import JointCommand
+from bitbots_mujoco_sim.Robot import Robot
+
+
+class Simulation(Node):
+    """Manages the MuJoCo simulation state and its components."""
+
+    def __init__(self):
+        super().__init__("sim_interface")
+        package_path = get_package_share_directory("bitbots_mujoco_sim")
+        self.model: mujoco.MjModel = mujoco.MjModel.from_xml_path(package_path + "/xml/adult_field.xml")
+        self.data: mujoco.MjData = mujoco.MjData(self.model)
+        self.robot: Robot = Robot(self.model, self.data)
+
+        self.time = 0.0
+        self.time_message = Time(seconds=0, nanoseconds=0).to_msg()
+        self.timestep = self.model.opt.timestep
+        self.step_number = 0
+        self.clock_publisher = self.create_publisher(Clock, "clock", 1)
+        
+        self.js_publisher = self.create_publisher(JointState, "joint_states", 1)
+        self.create_subscription(JointCommand, "DynamixelController/command", self.joint_command_callback, 1)
+        
+        self.imu_frame_id = self.get_parameter_or("imu_frame", "imu_frame")
+        self.imu_publisher = self.create_publisher(Imu, "imu/data_raw", 1)
+        
+        self.camera_active = True
+        self.camera_optical_frame_id = self.get_parameter_or("camera_optical_frame", "camera_optical_frame")
+        self.camera_publisher = self.create_publisher(Image, "camera/image_proc", 1)
+        self.camera_info_publisher = self.create_publisher(CameraInfo, "camera/camera_info", 1) 
+
+        self.events = {
+            "clock": {"frequency": 1, "handler": self.publish_clock_event},
+            "joint_states": {"frequency": 3, "handler": self.publish_ros_joint_states_event},
+            "imu": {"frequency": 3, "handler": self.publish_imu_event},
+            "camera": {"frequency": 24, "handler": self.publish_camera_event},
+        }
+
+    def run(self) -> None:
+        print("Starting simulation viewer...")
+        with viewer.launch_passive(self.model, self.data) as view:
+            while view.is_running():
+                self.step()
+                view.sync()
+
+    def joint_command_callback(self, command: JointCommand) -> None:
+        if len(command.positions) != 0:
+            for i in range(len(command.joint_names)):
+                joint = self.robot.joints.get(command.joint_names[i])
+                joint.position = command.positions[i]
+                #if len(command.velocities) != 0:
+                #    joint.velocity = command.velocities[i]
+
+    def step(self) -> None:
+        real_start_time = time.time()
+        self.step_number += 1
+        self.time += self.timestep
+        self.time_message = Time(seconds=int(self.time), nanoseconds=int(self.time % 1 * 1e9)).to_msg()
+
+        mujoco.mj_step(self.model, self.data)
+
+        for event_name, event_config in self.events.items():
+            if self.step_number % event_config["frequency"] == 0:
+                event_config["handler"]()
+
+        real_end_time = time.time()
+        time.sleep(max(0.0, self.timestep - (real_end_time - real_start_time)))
+
+    def publish_clock_event(self) -> None:
+        clock_msg = Clock()
+        clock_msg.clock = self.time_message
+        self.clock_publisher.publish(clock_msg)
+
+    def publish_ros_joint_states_event(self) -> None:
+        js = JointState()
+        js.name = []
+        js.header.stamp = self.time_message
+        js.position = []
+        js.effort = []
+        for joint in self.robot.joints:
+            js.name.append(joint.ros_name)
+            js.position.append(joint.position)
+            js.velocity.append(joint.velocity)
+            js.effort.append(joint.effort)
+        self.js_publisher.publish(js)
+
+    def publish_imu_event(self) -> None:
+        imu = Imu()
+        imu.header.stamp = self.time_message
+        imu.header.frame_id = self.imu_frame_id
+        imu.linear_acceleration.x, imu.linear_acceleration.y, imu.linear_acceleration.z = self.robot.sensors.accelerometer.data
+
+        # make sure that acceleration is not completely zero or we will get error in filter.
+        # Happens if robot is moved manually in the simulation
+        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
+
+        self.imu_publisher.publish(imu)
+
+    def publish_camera_event(self) -> None:
+        if not self.camera_active:
+            return
+        
+        img = Image()
+        img.header.stamp = self.time_message
+        img.header.frame_id = self.camera_optical_frame_id
+        img.encoding = "bgra8"
+        img.height = self.robot.camera.height
+        img.width = self.robot.camera.width
+        img.step = 4 * self.robot.camera.width
+        img.data = self.robot.camera.render()
+        self.camera_publisher.publish(img)
+
+        cam_info = CameraInfo()
+        cam_info.header = img.header
+        cam_info.height = self.robot.camera.height
+        cam_info.width = self.robot.camera.width
+
+        @staticmethod
+        def focal_length_from_fov(fov: float, resolution: int) -> float:
+            "Calculate focal length from field of view and resolution."
+            return 0.5 * resolution * (math.cos(fov / 2) / math.sin(fov / 2))
+
+        @staticmethod
+        def h_fov_to_v_fov(h_fov: float, height: int, width: int) -> float:
+            "Convert horizontal FOV to vertical FOV based on image aspect ratio."
+            return 2 * math.atan(math.tan(h_fov * 0.5) * (height / width))
+        
+        camera = self.robot.camera
+
+        h_fov = camera.fov
+        v_fov = h_fov_to_v_fov(h_fov, camera.height, camera.width)
+        
+        f_x, f_y = focal_length_from_fov(h_fov, camera.width), focal_length_from_fov(v_fov, camera.height)
+        cx, cy = camera.width / 2.0, camera.height / 2.0
+        cam_info.k = [
+            f_x, 0.0, cx,
+            0.0, f_y, cy,
+            0.0, 0.0, 1.0
+        ]
+        cam_info.p = [
+            f_x, 0.0, cx, 0.0,
+            0.0, f_y, cy, 0.0,
+            0.0, 0.0, 1.0, 0.0
+        ]
+        
+        self.camera_info_publisher.publish(cam_info)

bitbots_simulation/bitbots_mujoco_sim/bitbots_mujoco_sim/__init__.py

@@ -0,0 +1,7 @@
+"""BitBots MuJoCo Simulation Package"""
+
+from .Simulation import Simulation
+
+__all__ = [
+    "Simulation",
+]