reachy-mini 1.2.5rc1__py3-none-any.whl

reachy_mini/daemon/backend/mujoco/__init__.py

@@ -0,0 +1,36 @@
+"""MuJoCo Backend for Reachy Mini Daemon."""
+
+from dataclasses import dataclass
+
+try:
+    import mujoco  # noqa: F401
+
+    from reachy_mini.daemon.backend.mujoco.backend import (
+        MujocoBackend,
+        MujocoBackendStatus,
+    )
+
+except ImportError:
+
+    class MujocoMockupBackend:
+        """Mockup class to avoid import errors when MuJoCo is not installed."""
+
+        def __init__(self, *args, **kwargs) -> None:  # type: ignore[no-untyped-def]
+            """Raise ImportError when trying to instantiate the class."""
+            raise ImportError(
+                "MuJoCo is not installed. MuJoCo backend is not available."
+                " To use MuJoCo backend, please install the 'mujoco' extra dependencies"
+                " with 'pip install reachy_mini[mujoco]'."
+            )
+
+    MujocoBackend = MujocoMockupBackend  # type: ignore[assignment, misc]
+
+    @dataclass
+    class MujocoMockupBackendStatus:
+        """Mockup class to avoid import errors when MuJoCo is not installed."""
+
+        pass
+
+    MujocoBackendStatus = MujocoMockupBackendStatus  # type: ignore[assignment, misc]
+
+__all__ = ["MujocoBackend", "MujocoBackendStatus"]

reachy_mini/daemon/backend/mujoco/backend.py

@@ -0,0 +1,395 @@
+"""Mujoco Backend for Reachy Mini.
+
+This module provides the MujocoBackend class for simulating the Reachy Mini robot using the MuJoCo physics engine.
+
+It includes methods for running the simulation, getting joint positions, and controlling the robot's joints.
+
+"""
+
+import json
+import time
+from dataclasses import dataclass
+from importlib.resources import files
+from threading import Thread
+from typing import Annotated, Any, Optional
+
+import cv2
+import log_throttling
+import mujoco
+import mujoco.viewer
+import numpy as np
+import numpy.typing as npt
+
+import reachy_mini
+from reachy_mini.io.video_ws import AsyncWebSocketFrameSender
+
+from ..abstract import Backend, MotorControlMode
+from .utils import (
+    get_actuator_names,
+    get_joint_addr_from_name,
+    get_joint_id_from_name,
+)
+from .video_udp import UDPJPEGFrameSender
+
+CAMERA_REACHY = "eye_camera"
+CAMERA_STUDIO_CLOSE = "studio_close"
+CAMERA_SIZES = {CAMERA_REACHY: (1280, 720), CAMERA_STUDIO_CLOSE: (640, 640)}
+
+
+class MujocoBackend(Backend):
+    """Simulated Reachy Mini using MuJoCo."""
+
+    def __init__(
+        self,
+        scene: str = "empty",
+        check_collision: bool = False,
+        kinematics_engine: str = "AnalyticalKinematics",
+        headless: bool = False,
+        use_audio: bool = False,
+        websocket_uri: Optional[str] = None,
+    ) -> None:
+        """Initialize the MujocoBackend with a specified scene.
+
+        Args:
+            scene (str): The name of the scene to load. Default is "empty".
+            check_collision (bool): If True, enable collision checking. Default is False.
+            kinematics_engine (str): Kinematics engine to use. Defaults to "AnalyticalKinematics".
+            headless (bool): If True, run Mujoco in headless mode (no GUI). Default is False.
+            use_audio (bool): If True, use audio. Default is False.
+            websocket_uri (Optional[str]): If set, allow streaming of the robot view through a WebSocket connection to the specified uri. Defaults to None.
+
+        """
+        super().__init__(
+            check_collision=check_collision,
+            kinematics_engine=kinematics_engine,
+            use_audio=use_audio,
+        )
+
+        self.headless = headless
+        self.websocket_uri = websocket_uri
+
+        from reachy_mini.reachy_mini import (
+            SLEEP_ANTENNAS_JOINT_POSITIONS,
+            SLEEP_HEAD_JOINT_POSITIONS,
+        )
+
+        # Real robot convention for the order of the antennas joints is [right, left], but in mujoco it's [left, right]
+        self._SLEEP_ANTENNAS_JOINT_POSITIONS = [
+            SLEEP_ANTENNAS_JOINT_POSITIONS[1],
+            SLEEP_ANTENNAS_JOINT_POSITIONS[0],
+        ]
+        self._SLEEP_HEAD_JOINT_POSITIONS = SLEEP_HEAD_JOINT_POSITIONS
+
+        mjcf_root_path = str(
+            files(reachy_mini).joinpath("descriptions/reachy_mini/mjcf/")
+        )
+        self.model = mujoco.MjModel.from_xml_path(
+            f"{mjcf_root_path}/scenes/{scene}.xml"
+        )
+        self.data = mujoco.MjData(self.model)
+        self.model.opt.timestep = 0.002  # s, simulation timestep, 500hz
+        self.decimation = 10  # -> 50hz control loop
+        self.rendering_timestep = 0.04  # s, rendering loop # 25Hz
+        self.streaming_timestep = 0.04  # s, streaming loop # 25Hz
+
+        self.head_site_id = mujoco.mj_name2id(
+            self.model,
+            mujoco.mjtObj.mjOBJ_SITE,
+            "head",
+        )
+
+        self.current_head_pose = np.eye(4)
+
+        self.joint_names = get_actuator_names(self.model)
+
+        self.joint_ids = [
+            get_joint_id_from_name(self.model, n) for n in self.joint_names
+        ]
+        self.joint_qpos_addr = [
+            get_joint_addr_from_name(self.model, n) for n in self.joint_names
+        ]
+
+        # Disable collisions at the beginning for smoother initialization
+        self.col_inds = []
+        for i, type in enumerate(self.model.geom_contype):
+            if type != 0:
+                geom_name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_GEOM, i)
+                # monkey-patch: the geoms in the minimal scene are named (duck_geom, table_top_collision ... ).
+                # We don't disable the collision for them so that the objects don't fall to the ground at initialization
+                if geom_name is None:
+                    self.col_inds.append(i)
+                    self.model.geom_contype[i] = 0
+                    self.model.geom_conaffinity[i] = 0
+
+    def _get_camera_id(self, camera_name: str) -> Any:
+        """Get the id of the virtual camera."""
+        return mujoco.mj_name2id(
+            self.model,
+            mujoco.mjtObj.mjOBJ_CAMERA,
+            camera_name,
+        )
+
+    def _get_renderer(self, camera_name: str) -> mujoco.Renderer:
+        """Get the renderer for the virtual camera."""
+        camera_size = CAMERA_SIZES[camera_name]
+        return mujoco.Renderer(self.model, height=camera_size[1], width=camera_size[0])
+
+    def streaming_loop(self, camera_name: str, ws_uri: str) -> None:
+        """Streaming loop for the Mujoco simulation over WebSocket.
+
+        Capture the image from the virtual camera and send it over WebSocket to the ws_uri.
+        """
+        streamer = AsyncWebSocketFrameSender(ws_uri=ws_uri + "/video_stream")
+        offscreen_renderer = self._get_renderer(camera_name)
+        camera_id = self._get_camera_id(camera_name)
+
+        while not self.should_stop.is_set():
+            start_t = time.time()
+            offscreen_renderer.update_scene(self.data, camera_id)
+
+            # OPTIMIZATION: Disable expensive rendering effects on the scene
+            offscreen_renderer.scene.flags[mujoco.mjtRndFlag.mjRND_SHADOW] = 0
+            offscreen_renderer.scene.flags[mujoco.mjtRndFlag.mjRND_REFLECTION] = 0
+
+            im = offscreen_renderer.render()
+
+            im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
+            streamer.send_frame(im)
+
+            took = time.time() - start_t
+            time.sleep(max(0, self.streaming_timestep - took))
+
+    def rendering_loop(self, camera_name: str, port: int) -> None:
+        """Offline Rendering loop for the Mujoco simulation.
+
+        Capture the image from the virtual camera_name and send it over UDP to the port or over WebSocket to the ws_uri.
+        """
+        streamer = UDPJPEGFrameSender(dest_port=port)
+        offscreen_renderer = self._get_renderer(camera_name)
+        camera_id = self._get_camera_id(camera_name)
+
+        while not self.should_stop.is_set():
+            start_t = time.time()
+            offscreen_renderer.update_scene(self.data, camera_id)
+
+            im = offscreen_renderer.render()
+            streamer.send_frame(im)
+
+            took = time.time() - start_t
+            time.sleep(max(0, self.rendering_timestep - took))
+
+    def run(self) -> None:
+        """Run the Mujoco simulation with a viewer.
+
+        This method initializes the viewer and enters the main simulation loop.
+        It updates the joint positions at a rate and publishes the joint positions.
+        """
+        step = 1
+        if self.websocket_uri:
+            robot_view_streaming_thread = Thread(
+                target=self.streaming_loop,
+                args=(CAMERA_STUDIO_CLOSE, self.websocket_uri),
+                daemon=True,
+            )
+            robot_view_streaming_thread.start()
+
+        if not self.headless:
+            viewer = mujoco.viewer.launch_passive(
+                self.model, self.data, show_left_ui=False, show_right_ui=False
+            )
+            with viewer.lock():
+                viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
+                viewer.cam.distance = 0.8  # ≃ ||pos - lookat||
+                viewer.cam.azimuth = 160  # degrees
+                viewer.cam.elevation = -20  # degrees
+                viewer.cam.lookat[:] = [0, 0, 0.15]
+
+                # force one render with your new camera
+                mujoco.mj_step(self.model, self.data)
+                viewer.sync()
+
+                # im = self.get_camera()
+                # self.streamer_udp.send_frame(im)
+
+        self.data.qpos[self.joint_qpos_addr] = np.array(
+            self._SLEEP_HEAD_JOINT_POSITIONS + self._SLEEP_ANTENNAS_JOINT_POSITIONS
+        ).reshape(-1, 1)
+        self.data.ctrl[:] = np.array(
+            self._SLEEP_HEAD_JOINT_POSITIONS + self._SLEEP_ANTENNAS_JOINT_POSITIONS
+        )
+
+        # recompute all kinematics, collisions, etc.
+        mujoco.mj_forward(self.model, self.data)
+
+        for i in range(100):
+            mujoco.mj_step(self.model, self.data)
+
+        # enable collisions
+        for i in self.col_inds:
+            self.model.geom_contype[i] = 1
+            self.model.geom_conaffinity[i] = 1
+
+        for i in range(100):
+            mujoco.mj_step(self.model, self.data)
+
+        # one more frame so the viewer shows your startup pose
+        mujoco.mj_step(self.model, self.data)
+        if not self.headless:
+            viewer.sync()
+
+            rendering_thread = Thread(
+                target=self.rendering_loop, args=(CAMERA_REACHY, 5005), daemon=True
+            )
+            rendering_thread.start()
+
+        # Update the internal states of the IK and FK to the current configuration
+        # This is important to avoid jumps when starting the robot (beore wake-up)
+        self.head_kinematics.ik(self.get_mj_present_head_pose(), no_iterations=20)
+        self.head_kinematics.fk(
+            self.get_present_head_joint_positions(), no_iterations=20
+        )
+
+        # 3) now enter your normal loop
+        while not self.should_stop.is_set():
+            start_t = time.time()
+
+            if step % self.decimation == 0:
+                # update the current states
+                self.current_head_joint_positions = (
+                    self.get_present_head_joint_positions()
+                )
+                self.current_antenna_joint_positions = (
+                    self.get_present_antenna_joint_positions()
+                )
+                # Update the Placo kinematics model to recompute passive joints
+                self.update_head_kinematics_model(
+                    self.current_head_joint_positions,
+                    self.current_antenna_joint_positions,
+                )
+                self.current_head_pose = self.get_mj_present_head_pose()
+
+                # Update the target head joint positions from IK if necessary
+                # - does nothing if the targets did not change
+                if self.ik_required:
+                    try:
+                        self.update_target_head_joints_from_ik(
+                            self.target_head_pose, self.target_body_yaw
+                        )
+                    except ValueError as e:
+                        log_throttling.by_time(self.logger, interval=0.5).warning(
+                            f"IK error: {e}"
+                        )
+
+                if self.target_head_joint_positions is not None:
+                    self.data.ctrl[:7] = self.target_head_joint_positions
+                if self.target_antenna_joint_positions is not None:
+                    self.data.ctrl[-2:] = -self.target_antenna_joint_positions
+
+                if (
+                    self.joint_positions_publisher is not None
+                    and self.pose_publisher is not None
+                ):
+                    if not self.is_shutting_down:
+                        self.joint_positions_publisher.put(
+                            json.dumps(
+                                {
+                                    "head_joint_positions": self.current_head_joint_positions.tolist(),
+                                    "antennas_joint_positions": self.current_antenna_joint_positions.tolist(),
+                                }
+                            ).encode("utf-8")
+                        )
+                        self.pose_publisher.put(
+                            json.dumps(
+                                {
+                                    "head_pose": self.get_present_head_pose().tolist(),
+                                }
+                            ).encode("utf-8")
+                        )
+                    self.ready.set()
+
+                if not self.headless:
+                    viewer.sync()
+
+            mujoco.mj_step(self.model, self.data)
+
+            took = time.time() - start_t
+            time.sleep(max(0, self.model.opt.timestep - took))
+            # print(f"Step {step}: took {took*1e6:.1f}us")
+            step += 1
+
+        if not self.headless:
+            viewer.close()
+            rendering_thread.join()
+        if self.websocket_uri:
+            robot_view_streaming_thread.join()
+
+    def get_mj_present_head_pose(self) -> Annotated[npt.NDArray[np.float64], (4, 4)]:
+        """Get the current head pose from the Mujoco simulation.
+
+        Returns:
+            np.ndarray: The current head pose as a 4x4 transformation matrix.
+
+        """
+        mj_current_head_pose = np.eye(4)
+
+        mj_current_head_pose[:3, :3] = self.data.site_xmat[self.head_site_id].reshape(
+            3, 3
+        )
+        mj_current_head_pose[:3, 3] = self.data.site_xpos[self.head_site_id]
+        mj_current_head_pose[2, 3] -= 0.177
+        return mj_current_head_pose
+
+    def close(self) -> None:
+        """Close the Mujoco backend."""
+        # TODO Do something in mujoco here ?
+        pass
+
+    def get_status(self) -> "MujocoBackendStatus":
+        """Get the status of the Mujoco backend.
+
+        Returns:
+            dict: An empty dictionary as the Mujoco backend does not have a specific status to report.
+
+        """
+        return MujocoBackendStatus(motor_control_mode=self.get_motor_control_mode())
+
+    def get_present_head_joint_positions(
+        self,
+    ) -> Annotated[npt.NDArray[np.float64], (7,)]:
+        """Get the current joint positions of the head."""
+        pos: npt.NDArray[np.float64] = self.data.qpos[
+            self.joint_qpos_addr[:7]
+        ].flatten()
+        return pos
+
+    def get_present_antenna_joint_positions(
+        self,
+    ) -> Annotated[npt.NDArray[np.float64], (2,)]:
+        """Get the current joint positions of the antennas."""
+        pos: npt.NDArray[np.float64] = self.data.qpos[
+            self.joint_qpos_addr[-2:]
+        ].flatten()
+        return -pos
+
+    def get_motor_control_mode(self) -> MotorControlMode:
+        """Get the motor control mode."""
+        return MotorControlMode.Enabled
+
+    def set_motor_control_mode(self, mode: MotorControlMode) -> None:
+        """Set the motor control mode."""
+        pass
+
+    def set_motor_torque_ids(self, ids: list[str], on: bool) -> None:
+        """Set the motor torque state for specific motor names."""
+        pass
+
+
+@dataclass
+class MujocoBackendStatus:
+    """Dataclass to represent the status of the Mujoco backend.
+
+    Empty for now, as the Mujoco backend does not have a specific status to report.
+    """
+
+    motor_control_mode: MotorControlMode
+    error: str | None = None

reachy_mini/daemon/backend/mujoco/utils.py

@@ -0,0 +1,59 @@
+"""Mujoco utilities for Reachy Mini.
+
+This module provides utility functions for working with MuJoCo models, including
+homogeneous transformation matrices, joint positions, and actuator names.
+"""
+
+from typing import Annotated
+
+import mujoco
+import numpy as np
+import numpy.typing as npt
+from mujoco._structs import MjData, MjModel
+from scipy.spatial.transform import Rotation as R
+
+
+def get_homogeneous_matrix_from_euler(
+    position: tuple[float, float, float] = (0, 0, 0),  # (x, y, z) meters
+    euler_angles: tuple[float, float, float] = (0, 0, 0),  # (roll, pitch, yaw)
+    degrees: bool = False,
+) -> Annotated[npt.NDArray[np.float64], (4, 4)]:
+    """Return a homogeneous transformation matrix from position and Euler angles."""
+    homogeneous_matrix = np.eye(4)
+    homogeneous_matrix[:3, :3] = R.from_euler(
+        "xyz", euler_angles, degrees=degrees
+    ).as_matrix()
+    homogeneous_matrix[:3, 3] = position
+    return homogeneous_matrix
+
+
+def get_joint_qpos(model: MjModel, data: MjData, joint_name: str) -> float:
+    """Return the qpos (rad) of a specified joint in the model."""
+    # Get the joint id
+    joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
+    if joint_id == -1:
+        raise ValueError(f"Joint '{joint_name}' not found.")
+
+    # Get the address of the joint's qpos in the qpos array
+    qpos_addr = model.jnt_qposadr[joint_id]
+
+    # Get the qpos value
+    qpos: float = data.qpos[qpos_addr]
+    return qpos
+
+
+def get_joint_id_from_name(model: MjModel, name: str) -> int:
+    """Return the id of a specified joint."""
+    return mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)  # type: ignore
+
+
+def get_joint_addr_from_name(model: MjModel, name: str) -> int:
+    """Return the address of a specified joint."""
+    addr: int = model.joint(name).qposadr
+    return addr
+
+
+def get_actuator_names(model: MjModel) -> list[str]:
+    """Return the list of the actuators names from the MuJoCo model."""
+    actuator_names = [model.actuator(k).name for k in range(0, model.nu)]
+    return actuator_names

reachy_mini/daemon/backend/mujoco/video_udp.py

@@ -0,0 +1,57 @@
+"""UDP JPEG Frame Sender.
+
+This module provides a class to send JPEG frames over UDP. It encodes the frames as JPEG images and splits them into chunks to fit within the maximum packet size for UDP transmission.
+"""
+
+import socket
+import struct
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+
+
+class UDPJPEGFrameSender:
+    """A class to send JPEG frames over UDP."""
+
+    def __init__(
+        self,
+        dest_ip: str = "127.0.0.1",
+        dest_port: int = 5005,
+        max_packet_size: int = 1400,
+    ) -> None:
+        """Initialize the UDPJPEGFrameSender.
+
+        Args:
+            dest_ip (str): Destination IP address.
+            dest_port (int): Destination port number.
+            max_packet_size (int): Maximum size of each UDP packet.
+
+        """
+        self.addr = (dest_ip, dest_port)
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.max_packet_size = max_packet_size
+
+    def send_frame(self, frame: npt.NDArray[np.uint8]) -> None:
+        """Send a frame as a JPEG image over UDP.
+
+        Args:
+            frame (np.ndarray): The frame to be sent, in RGB format.
+
+        """
+        frame_cvt = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        ret, jpeg_bytes = cv2.imencode(
+            ".jpg", frame_cvt, [int(cv2.IMWRITE_JPEG_QUALITY), 80]
+        )
+        data = jpeg_bytes.tobytes()
+        total_size = len(data)
+        n_chunks = (total_size + self.max_packet_size - 1) // self.max_packet_size
+        self.sock.sendto(struct.pack("!II", n_chunks, total_size), self.addr)
+        for i in range(n_chunks):
+            start = i * self.max_packet_size
+            end = min(start + self.max_packet_size, total_size)
+            self.sock.sendto(data[start:end], self.addr)
+
+    def close(self) -> None:
+        """Close the socket."""
+        self.sock.close()

reachy_mini/daemon/backend/robot/__init__.py

@@ -0,0 +1,8 @@
+"""Real robot backend for Reachy Mini."""
+
+from reachy_mini.daemon.backend.robot.backend import (
+    RobotBackend,
+    RobotBackendStatus,
+)
+
+__all__ = ["RobotBackend", "RobotBackendStatus"]