reachy-mini 1.2.5rc1__py3-none-any.whl

reachy_mini/io/__init__.py

@@ -0,0 +1,15 @@
+"""IO module."""
+
+from .audio_ws import AsyncWebSocketAudioStreamer
+from .video_ws import AsyncWebSocketFrameSender
+from .ws_controller import AsyncWebSocketController
+from .zenoh_client import ZenohClient
+from .zenoh_server import ZenohServer
+
+__all__ = [
+    "AsyncWebSocketAudioStreamer",
+    "AsyncWebSocketFrameSender",
+    "AsyncWebSocketController",
+    "ZenohClient",
+    "ZenohServer",
+]

reachy_mini/io/abstract.py

@@ -0,0 +1,70 @@
+"""Base classes for server and client implementations.
+
+These abstract classes define the interface for server and client components
+in the Reachy Mini project. They provide methods for starting and stopping
+the server, handling commands, and managing client connections.
+"""
+
+from abc import ABC, abstractmethod
+from threading import Event
+from uuid import UUID
+
+from reachy_mini.io.protocol import AnyTaskRequest
+
+
+class AbstractServer(ABC):
+    """Base class for server implementations."""
+
+    @abstractmethod
+    def start(self) -> None:
+        """Start the server."""
+        pass
+
+    @abstractmethod
+    def stop(self) -> None:
+        """Stop the server."""
+        pass
+
+    @abstractmethod
+    def command_received_event(self) -> Event:
+        """Wait for a new command and return it."""
+        pass
+
+
+class AbstractClient(ABC):
+    """Base class for client implementations."""
+
+    @abstractmethod
+    def wait_for_connection(self) -> None:
+        """Wait for the client to connect to the server."""
+        pass
+
+    @abstractmethod
+    def is_connected(self) -> bool:
+        """Check if the client is connected to the server."""
+        pass
+
+    @abstractmethod
+    def disconnect(self) -> None:
+        """Disconnect the client from the server."""
+        pass
+
+    @abstractmethod
+    def send_command(self, command: str) -> None:
+        """Send a command to the server."""
+        pass
+
+    @abstractmethod
+    def get_current_joints(self) -> tuple[list[float], list[float]]:
+        """Get the current joint positions."""
+        pass
+
+    @abstractmethod
+    def send_task_request(self, task_req: AnyTaskRequest) -> UUID:
+        """Send a task request to the server and return a unique task identifier."""
+        pass
+
+    @abstractmethod
+    def wait_for_task_completion(self, task_uid: UUID, timeout: float = 5.0) -> None:
+        """Wait for the specified task to complete."""
+        pass

reachy_mini/io/audio_ws.py

@@ -0,0 +1,242 @@
+"""Async WebSocket Audio Streamer."""
+
+import asyncio
+import logging
+import threading
+import time
+from queue import Empty, Queue
+from typing import Optional, Union
+
+import numpy as np
+import numpy.typing as npt
+from websockets.asyncio.client import ClientConnection, connect
+
+logger = logging.getLogger("reachy_mini.io.audio_ws")
+
+
+class AsyncWebSocketAudioStreamer:
+    """Async WebSocket audio streamer with send and receive support."""
+
+    ws_uri: str
+    send_queue: "Queue[bytes]"
+    recv_queue: "Queue[bytes]"
+    loop: asyncio.AbstractEventLoop
+    thread: threading.Thread
+    connected: threading.Event
+    stop_flag: bool
+    keep_alive_interval: float
+
+    # --- CONFIGURATION ---
+    # Target ~2048 samples per packet (approx 128ms)
+    # 2048 samples * 2 bytes (int16) = 4096 bytes
+    BATCH_SIZE_BYTES = 4096
+    # Don't hold audio longer than 200ms even if buffer isn't full
+    BATCH_TIMEOUT = 0.2
+
+    def __init__(self, ws_uri: str, keep_alive_interval: float = 2.0) -> None:
+        """Initialize the WebSocket audio streamer.
+
+        Args:
+            ws_uri: WebSocket URI to connect to.
+            keep_alive_interval: Interval in seconds to send keep-alive pings
+                when no audio is flowing.
+
+        """
+        self.ws_uri = ws_uri
+        self.send_queue = Queue()
+        self.recv_queue = Queue()
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(target=self._run_loop, daemon=True)
+        self.connected = threading.Event()
+        self.stop_flag = False
+        self.keep_alive_interval = keep_alive_interval
+        self.thread.start()
+
+    def _run_loop(self) -> None:
+        """Run the WebSocket streamer loop in a background thread."""
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_until_complete(self._run())
+
+    async def _run(self) -> None:
+        """Run the main reconnect loop."""
+        while not self.stop_flag:
+            try:
+                async with connect(self.ws_uri) as ws:
+                    logger.info("[WS-AUDIO] Connected to Space")
+                    self.connected.set()
+
+                    send_task = asyncio.create_task(self._send_loop(ws))
+                    recv_task = asyncio.create_task(self._recv_loop(ws))
+
+                    done, pending = await asyncio.wait(
+                        {send_task, recv_task},
+                        return_when=asyncio.FIRST_EXCEPTION,
+                    )
+
+                    # Cancel the other task if one fails or finishes
+                    for task in pending:
+                        task.cancel()
+                        try:
+                            await task
+                        except Exception:
+                            pass
+
+            except Exception as e:
+                logger.info(f"[WS-AUDIO] Connection failed: {e}")
+                await asyncio.sleep(1.0)
+
+            self.connected.clear()
+
+    async def _send_loop(self, ws: ClientConnection) -> None:
+        """Send outgoing audio chunks and keep-alive pings.
+
+        To avoid audible artifacts, this method aggregates small chunks into larger batches before sending.
+        """
+        last_activity = time.time()
+
+        # Buffer to hold small chunks
+        batch_buffer = bytearray()
+        # Track when we started filling this specific batch
+        batch_start_time = time.time()
+
+        while not self.stop_flag:
+            try:
+                # 1. Try to pull data from the queue
+                # Use a short timeout so we can check time-based conditions frequently
+                chunk = self.send_queue.get(timeout=0.01)
+
+                # If this is the first chunk in the buffer, reset timer
+                if len(batch_buffer) == 0:
+                    batch_start_time = time.time()
+
+                batch_buffer.extend(chunk)
+
+            except Empty:
+                pass
+            except Exception as e:
+                logger.info(f"[WS-AUDIO] Queue error: {e}")
+                break
+
+            # 2. Check if we should send the batch
+            now = time.time()
+
+            # Condition A: Buffer is full enough (Size based)
+            is_full = len(batch_buffer) >= self.BATCH_SIZE_BYTES
+
+            # Condition B: Buffer has data, but it's getting too old (Time based)
+            # This ensures that if the robot says a short word, it sends it
+            # after 100ms instead of waiting forever for more data.
+            is_timed_out = (len(batch_buffer) > 0) and (
+                (now - batch_start_time) > self.BATCH_TIMEOUT
+            )
+            if is_full or is_timed_out:
+                try:
+                    # Send the aggregated buffer
+                    await ws.send(batch_buffer)  # type: ignore
+
+                    # Reset
+                    batch_buffer = bytearray()
+                    last_activity = now
+                except Exception as e:
+                    logger.info(f"[WS-AUDIO] Send error: {e}")
+                    break
+
+            # 3. Keep-Alive Ping (Only if completely idle)
+            # We only ping if the buffer is empty AND we haven't sent anything recently
+            if (
+                len(batch_buffer) == 0
+                and (now - last_activity) > self.keep_alive_interval
+            ):
+                try:
+                    await ws.send("ping")
+                    last_activity = now
+                    logger.debug("[WS-AUDIO] Sent keep-alive ping")
+                except Exception as e:
+                    logger.info(f"[WS-AUDIO] Ping failed: {e}")
+                    break
+
+            # Tiny sleep to yield control if we are just spinning
+            if len(batch_buffer) == 0:
+                await asyncio.sleep(0.001)
+
+    async def _recv_loop(self, ws: ClientConnection) -> None:
+        """Receive incoming audio chunks."""
+        while not self.stop_flag:
+            try:
+                msg = await ws.recv()
+            except Exception as e:
+                logger.info(f"[WS-AUDIO] Receive error: {e}")
+                break
+
+            if isinstance(msg, bytes):
+                try:
+                    self.recv_queue.put_nowait(msg)
+                except Exception as e:
+                    logger.debug(f"[WS-AUDIO] Failed to enqueue received audio: {e}")
+            else:
+                logger.debug(f"[WS-AUDIO] Received non-binary message: {msg}")
+
+    # ------------------------
+    # Public API
+    # ------------------------
+
+    def send_audio_chunk(
+        self,
+        audio: Union[bytes, npt.NDArray[np.int16], npt.NDArray[np.float32]],
+    ) -> None:
+        """Queue an audio chunk to be sent.
+
+        Args:
+            audio: Either raw bytes or a numpy array of int16 or float32.
+                   Float32 arrays are assumed to be in [-1, 1] and will
+                   be converted to int16 PCM.
+
+        """
+        if self.stop_flag:
+            return
+
+        if isinstance(audio, bytes):
+            data = audio
+        else:
+            # Convert only if needed
+            arr = np.asarray(audio)
+            # Handle stereo if accidentally passed (take channel 0)
+            if arr.ndim > 1:
+                arr = arr[:, 0]
+
+            if arr.dtype == np.float32 or arr.dtype == np.float64:
+                # If any value is above 1 or below -1, scale the entire array so the max abs value is 1 or less
+                max_abs = np.max(np.abs(arr))
+                if max_abs > 1.0:
+                    arr = arr / max_abs
+                # Convert float audio [-1,1] to int16 PCM
+                arr = np.clip(arr, -1.0, 1.0)
+                arr = (arr * 32767.0).astype(np.int16)  # type: ignore
+            elif arr.dtype != np.int16:
+                arr = arr.astype(np.int16)
+
+            data = arr.tobytes()
+
+        self.send_queue.put(data)
+
+    def get_audio_chunk(
+        self, timeout: Optional[float] = 0.01
+    ) -> Optional[npt.NDArray[np.float32]]:
+        """Retrieve a received audio chunk, if any."""
+        try:
+            if timeout == 0:
+                audio_bytes = self.recv_queue.get_nowait()
+            else:
+                audio_bytes = self.recv_queue.get(timeout=timeout)
+            # bytes -> int16 -> float32 in [-1, 1]
+            int16_arr = np.frombuffer(audio_bytes, dtype=np.int16)
+            float_arr = int16_arr.astype(np.float32) / 32767.0
+            return float_arr
+        except Empty:
+            return None
+
+    def close(self) -> None:
+        """Close the WebSocket audio streamer."""
+        self.stop_flag = True
+        if self.loop.is_running():
+            self.loop.call_soon_threadsafe(self.loop.stop)

reachy_mini/io/protocol.py

@@ -0,0 +1,44 @@
+"""Protocol definitions for Reachy Mini client/server communication."""
+
+from datetime import datetime
+from uuid import UUID
+
+from pydantic import BaseModel
+
+from reachy_mini.utils.interpolation import InterpolationTechnique
+
+
+class GotoTaskRequest(BaseModel):
+    """Class to represent a goto target task."""
+
+    head: list[float] | None  # 4x4 flatten pose matrix
+    antennas: list[float] | None  # [right_angle, left_angle] (in rads)
+    duration: float
+    method: InterpolationTechnique
+    body_yaw: float | None
+
+
+class PlayMoveTaskRequest(BaseModel):
+    """Class to represent a play move task."""
+
+    move_name: str
+
+
+AnyTaskRequest = GotoTaskRequest | PlayMoveTaskRequest
+
+
+class TaskRequest(BaseModel):
+    """Class to represent any task request."""
+
+    uuid: UUID
+    req: AnyTaskRequest
+    timestamp: datetime
+
+
+class TaskProgress(BaseModel):
+    """Class to represent task progress."""
+
+    uuid: UUID
+    finished: bool = False
+    error: str | None = None
+    timestamp: datetime

reachy_mini/io/video_ws.py

@@ -0,0 +1,135 @@
+"""Async WebSocket Frame Sender."""
+
+import asyncio
+import logging
+import threading
+from queue import Empty, Full, Queue
+from typing import Optional
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+from websockets.asyncio.client import ClientConnection, connect
+from websockets.exceptions import ConnectionClosed
+
+logger = logging.getLogger("reachy_mini.io.video_ws")
+
+
+class AsyncWebSocketFrameSender:
+    """Async WebSocket frame sender."""
+
+    ws_uri: str
+    queue: "Queue[bytes]"
+    loop: asyncio.AbstractEventLoop
+    thread: threading.Thread
+    connected: threading.Event
+    stop_flag: bool
+    _last_frame: Optional[npt.NDArray[np.uint8]]
+
+    def __init__(self, ws_uri: str) -> None:
+        """Initialize the WebSocket frame sender."""
+        self.ws_uri = ws_uri
+        self.queue = Queue(maxsize=2)
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(target=self._run_loop, daemon=True)
+        self.connected = threading.Event()
+        self.stop_flag = False
+        self._last_frame = None
+        self.thread.start()
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_until_complete(self._run())
+
+    def _clear_queue(self) -> None:
+        """Empty the queue so we don't send 10 seconds of old video on reconnect."""
+        while not self.queue.empty():
+            try:
+                self.queue.get_nowait()
+            except Empty:
+                break
+
+    async def _run(self) -> None:
+        """Run the WebSocket frame sender loop."""
+        while not self.stop_flag:
+            try:
+                ws: ClientConnection
+                async with connect(
+                    self.ws_uri,
+                    ping_interval=5,  # Every 5 seconds is plenty
+                    ping_timeout=10,  # Give it 10s to respond
+                    close_timeout=1,  # Don't wait long for polite closes
+                ) as ws:
+                    logger.info("[WS Video] Connected to Space")
+                    self.connected.set()
+                    self._clear_queue()  # Ensure we start fresh
+
+                    while not self.stop_flag:
+                        try:
+                            frame = self.queue.get_nowait()
+
+                            await ws.send(frame)
+
+                        except Empty:
+                            # Queue is empty, just yield to event loop
+                            await asyncio.sleep(0.05)
+                            continue
+
+                        except Exception as e:
+                            logger.error(f"[WS Video] Send error: {e}")
+                            break
+
+            except ConnectionClosed as e:
+                # Common network errors, retry quickly
+                logger.error(
+                    f"[WS Video] Connection lost ({type(e).__name__}). Retrying..."
+                )
+                self.connected.clear()
+                self._last_frame = None
+                await asyncio.sleep(0.5)  # Wait briefly before reconnecting
+
+            except Exception as e:
+                logger.error(f"[WS Video] Unexpected error: {e}")
+                await asyncio.sleep(1)
+
+            self.connected.clear()
+            self._last_frame = None
+
+    def send_frame(self, frame: npt.NDArray[np.uint8]) -> None:
+        """Send a frame to the WebSocket (Non-blocking)."""
+        if not self.stop_flag:
+            # 1. Frame Deduplication
+            if self._last_frame is not None:
+                if np.array_equal(frame, self._last_frame):
+                    return
+            self._last_frame = frame.copy()
+
+            # 2. Encode
+            ok, jpeg_bytes = cv2.imencode(
+                ".jpg",
+                frame,
+                [int(cv2.IMWRITE_JPEG_QUALITY), 80],
+            )
+            if not ok:
+                return
+
+            data = jpeg_bytes.tobytes()
+
+            # If queue is full (network lagging), remove the old frame and put the new one.
+            try:
+                self.queue.put_nowait(data)
+            except Full:
+                # Queue is full, network is likely slower than camera.
+                # Drop the OLDEST frame to make room for the NEWEST.
+                try:
+                    self.queue.get_nowait()  # Pop old frame
+                    self.queue.put_nowait(data)  # Push new frame
+                except Empty:
+                    logger.error(
+                        "[WS Video] Queue is full and empty, this should not happen."
+                    )
+
+    def close(self) -> None:
+        """Close the WebSocket frame sender."""
+        self.stop_flag = True
+        self.loop.call_soon_threadsafe(self.loop.stop)

reachy_mini/io/ws_controller.py

@@ -0,0 +1,120 @@
+"""Async WebSocket Controller for remote control and streaming of the robot."""
+
+import asyncio
+import json
+import logging
+import threading
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import numpy as np
+from websockets.asyncio.client import ClientConnection, connect
+
+from reachy_mini.daemon.backend.abstract import Backend
+
+logger = logging.getLogger("reachy_mini.ws_controller")
+
+
+@dataclass
+class Movement:
+    """Movement data for the WebSocket controller."""
+
+    name: str
+    x: float = 0
+    y: float = 0
+    z: float = 0
+    roll: float = 0
+    pitch: float = 0
+    yaw: float = 0
+    body_yaw: float = 0
+    left_antenna: Optional[float] = None
+    right_antenna: Optional[float] = None
+    duration: float = 1.0
+
+
+class AsyncWebSocketController:
+    """WebSocket controller for remote control and streaming of the robot."""
+
+    ws_uri: str
+    backend: Backend
+    loop: asyncio.AbstractEventLoop
+    thread: threading.Thread
+    stop_flag: bool
+
+    def __init__(self, ws_uri: str, backend: Backend) -> None:
+        """Initialize the WebSocket controller."""
+        self.ws_uri = ws_uri
+        self.backend = backend
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(target=self._run_loop, daemon=True)
+        self.stop_flag = False
+        self.thread.start()
+
+    def _run_loop(self) -> None:
+        """Run the WebSocket controller loop."""
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_until_complete(self._run())
+
+    async def on_command(self, cmd: Dict[str, Any]) -> None:
+        """Handle a command from the WebSocket."""
+        typ = cmd.get("type")
+
+        if typ == "movement":
+            logger.debug("[Daemon] Movement command received")
+            mov = cmd.get("movement", {})
+            logger.debug("[Daemon] Movement command: %s", mov)
+
+            head = mov.get("head")
+            if head is not None:
+                head_arr = np.array(head, dtype=float).reshape(4, 4)
+            else:
+                head_arr = None
+
+            antennas = mov.get("antennas")
+            if antennas is not None:
+                antennas_arr = np.array(antennas, dtype=float)
+            else:
+                antennas_arr = None
+
+            try:
+                await self.backend.goto_target(
+                    head=head_arr,
+                    antennas=antennas_arr,
+                    duration=mov.get("duration", 1.0),
+                    body_yaw=mov.get("body_yaw", 0.0),
+                )
+            except Exception as e:
+                logger.debug("[Daemon] Error in goto_target: %s", e)
+        elif typ == "ping":
+            logger.debug("[Daemon] Ping received")
+            return
+        else:
+            logger.debug("[Daemon] Unknown command type: %s", typ)
+
+    async def _run(self) -> None:
+        """Run the WebSocket controller loop."""
+        while not self.stop_flag:
+            try:
+                ws: ClientConnection
+                async with connect(self.ws_uri, ping_interval=5, ping_timeout=10) as ws:
+                    logger.info("[WS] Connected to Space")
+                    async for msg in ws:
+                        try:
+                            data = json.loads(msg)
+                        except Exception as e:
+                            logger.debug("[WS] Bad JSON: %s raw: %s", e, msg)
+                            continue
+
+                        # Now this is awaited inside the same loop
+                        await self.on_command(data)
+
+            except Exception as e:
+                logger.info("[WS] Connection failed: %s", e)
+                # small backoff before reconnect
+                await asyncio.sleep(1)
+
+    def stop(self) -> None:
+        """Stop the WebSocket controller."""
+        self.stop_flag = True
+        if self.loop.is_running():
+            self.loop.call_soon_threadsafe(lambda: None)