Stream Mujoco's view of robot, useful for gradio

src/reachy_mini/daemon/app/main.py

@@ -31,6 +31,7 @@ class Args:
     """Arguments for configuring the Reachy Mini daemon."""
 
     log_level: str = "INFO"
+    log_file: str | None = None
 
     wireless_version: bool = False
 
@@ -39,6 +40,8 @@ class Args:
     sim: bool = False
     scene: str = "empty"
     headless: bool = False
+    stream_robot_view: bool = False
+    use_audio: bool = True
 
     kinematics_engine: str = "AnalyticalKinematics"
     check_collision: bool = False
@@ -73,6 +76,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
                 sim=args.sim,
                 scene=args.scene,
                 headless=args.headless,
+                stream_robot_view=args.stream_robot_view,
+                use_audio=args.use_audio,
                 kinematics_engine=args.kinematics_engine,
                 check_collision=args.check_collision,
                 wake_up_on_start=args.wake_up_on_start,
@@ -185,6 +190,19 @@ def main() -> None:
         default=default_args.headless,
         help="Run the daemon in headless mode (default: False).",
     )
+    parser.add_argument(
+        "--stream-robot-view",
+        action="store_true",
+        default=default_args.stream_robot_view,
+        help="Stream the robot view to the port 5010 (default: False).",
+    )
+    parser.add_argument(
+        "--deactivate-audio",
+        action="store_false",
+        dest="use_audio",
+        default=default_args.use_audio,
+        help="Deactivate audio (default: True).",
+    )
     # Daemon options
     parser.add_argument(
         "--autostart",
@@ -269,8 +287,23 @@ def main() -> None:
         choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
         help="Set the logging level (default: INFO).",
     )
+    parser.add_argument(
+        "--log-file",
+        type=str,
+        default=default_args.log_file,
+        help="Path to a file to write logs to.",
+    )
 
     args = parser.parse_args()
+
+    if args.log_file:
+        file_handler = logging.FileHandler(args.log_file, mode="a")
+        file_handler.setFormatter(
+            logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        )
+        logging.getLogger().addHandler(file_handler)
+        logging.getLogger().setLevel(args.log_level)
+
     run_app(Args(**vars(args)))
 
 

src/reachy_mini/daemon/app/routers/daemon.py

@@ -37,6 +37,7 @@ async def start(logger: logging.Logger) -> None:
                 check_collision=request.app.state.args.check_collision,
                 kinematics_engine=request.app.state.args.kinematics_engine,
                 headless=request.app.state.args.headless,
+                stream_robot_view=request.app.state.args.stream_robot_view,
             )
 
     job_id = bg_job_register.run_command("daemon-start", start)

src/reachy_mini/daemon/backend/abstract.py

@@ -55,11 +55,14 @@ def __init__(
         log_level: str = "INFO",
         check_collision: bool = False,
         kinematics_engine: str = "AnalyticalKinematics",
+        use_audio: bool = True,
     ) -> None:
         """Initialize the backend."""
         self.logger = logging.getLogger(__name__)
         self.logger.setLevel(log_level)
 
+        self.use_audio = use_audio
+
         self.should_stop = threading.Event()
         self.ready = threading.Event()
 
@@ -159,7 +162,10 @@ def __init__(
         # Recording lock to guard buffer swaps and appends
         self._rec_lock = threading.Lock()
 
-        self.audio = SoundDeviceAudio(log_level=log_level)
+        if self.use_audio:
+            self.audio = SoundDeviceAudio(log_level=log_level)
+        else:
+            self.audio = None  # type: ignore
 
     # Life cycle methods
     def wrapped_run(self) -> None:
@@ -612,7 +618,8 @@ def play_sound(self, sound_file: str) -> None:
             sound_file (str): The name of the sound file to play (e.g., "wake_up.wav").
 
         """
-        self.audio.play_sound(sound_file, autoclean=True)
+        if self.audio:
+            self.audio.play_sound(sound_file, autoclean=True)
 
     # Basic move definitions
     INIT_HEAD_POSE = np.eye(4)

src/reachy_mini/daemon/backend/mujoco/backend.py

@@ -6,39 +6,46 @@
 
 """
 
 import json
 import time
 from dataclasses import dataclass
 from importlib.resources import files
 from threading import Thread
 from typing import Annotated
 
 import log_throttling
 import mujoco
 import mujoco.viewer
 import numpy as np
 import numpy.typing as npt
 
 import reachy_mini
 
 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
+from .video_tcp import TCPJPEGFrameServer
+
+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,
+        stream_robot_view: bool = False,
+        use_audio: bool = False,
     ) -> None:
         """Initialize the MujocoBackend with a specified scene.
 
@@ -47,13 +54,17 @@
             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.
+            stream_robot_view (bool): If True, stream the robot view to the port 5010. Default is False.
 
         """
         super().__init__(
-            check_collision=check_collision, kinematics_engine=kinematics_engine
+            check_collision=check_collision,
+            kinematics_engine=kinematics_engine,
+            use_audio=use_audio,
         )
 
         self.headless = headless
+        self.stream_robot_view = stream_robot_view
 
         from reachy_mini.reachy_mini import (
             SLEEP_ANTENNAS_JOINT_POSITIONS,
@@ -78,12 +89,6 @@
         self.decimation = 10  # -> 50hz control loop
         self.rendering_timestep = 0.04  # s, rendering loop # 25Hz
 
-        self.camera_id = mujoco.mj_name2id(
-            self.model,
-            mujoco.mjtObj.mjOBJ_CAMERA,
-            "eye_camera",
-        )
-
         self.head_site_id = mujoco.mj_name2id(
             self.model,
             mujoco.mjtObj.mjOBJ_SITE,
@@ -106,21 +111,35 @@
             get_joint_addr_from_name(self.model, n) for n in self.joint_names
         ]
 
-    def rendering_loop(self) -> None:
+    def rendering_loop(self, camera_name: str, port: int) -> None:
         """Offline Rendering loop for the Mujoco simulation.
 
-        Capture the image from the virtual Reachy's camera and send it over UDP.
+        Capture the image from the virtual camera_name and send it over UDP to the port.
         """
-        streamer_udp = UDPJPEGFrameSender()
-        camera_size = (1280, 720)
+        if port == 5010:
+            streamer = TCPJPEGFrameServer(listen_port=port)
+        else:
+            streamer = UDPJPEGFrameSender(dest_port=port)
+        camera_id = mujoco.mj_name2id(
+            self.model,
+            mujoco.mjtObj.mjOBJ_CAMERA,
+            camera_name,
+        )
+        camera_size = CAMERA_SIZES[camera_name]
         offscreen_renderer = mujoco.Renderer(
             self.model, height=camera_size[1], width=camera_size[0]
         )
         while not self.should_stop.is_set():
             start_t = time.time()
-            offscreen_renderer.update_scene(self.data, self.camera_id)
+            offscreen_renderer.update_scene(self.data, camera_id)
+
+            if camera_name == CAMERA_STUDIO_CLOSE:
+                # 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()
-            streamer_udp.send_frame(im)
+            streamer.send_frame(im)
 
             took = time.time() - start_t
             time.sleep(max(0, self.rendering_timestep - took))
@@ -132,6 +151,10 @@
         It updates the joint positions at a rate and publishes the joint positions.
         """
         step = 1
+        if self.stream_robot_view:
+            robot_view_rendering_thread = Thread(target=self.rendering_loop, args=(CAMERA_STUDIO_CLOSE, 5010), daemon=True)
+            robot_view_rendering_thread.start()
+
         if not self.headless:
             viewer = mujoco.viewer.launch_passive(
                 self.model, self.data, show_left_ui=False, show_right_ui=False
@@ -150,24 +173,24 @@
                 # 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
-                )
+        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)
+        # recompute all kinematics, collisions, etc.
+        mujoco.mj_forward(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, daemon=True)
+            rendering_thread = Thread(target=self.rendering_loop, args=(CAMERA_REACHY, 5005), daemon=True)
             rendering_thread.start()
 
         # 3) now enter your normal loop
@@ -235,6 +258,9 @@
 
         if not self.headless:
             viewer.close()
+            rendering_thread.join()
+        if self.stream_robot_view:
+            robot_view_rendering_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.