spacemouse ok

scripts/environments/teleoperation/xr_utils/xr_client.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2022-2026, The Mindbot Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Adapted from XRoboToolkit (https://github.com/NVIDIA-AI-IOT/XRoboToolkit-Teleop-Sample-Python)
+# Original XR client by Zhigen Zhao.
+
+"""XR device client for PICO 4 Ultra via XRoboToolkit SDK.
+
+This module wraps the ``xrobotoolkit_sdk`` C++ extension to provide a
+Pythonic interface for retrieving pose, button, joystick, and hand-tracking
+data from PICO XR headsets over the local network.
+
+Usage::
+
+    from xr_utils import XrClient
+
+    client = XrClient()
+    headset_pose = client.get_pose("headset")
+    left_trigger = client.get_key_value("left_trigger")
+    client.close()
+"""
+
+import numpy as np
+import xrobotoolkit_sdk as xrt
+
+
+class XrClient:
+    """Client for PICO XR devices via the XRoboToolkit SDK."""
+
+    def __init__(self):
+        """Initialize the SDK and start listening for XR data."""
+        xrt.init()
+        print("[XrClient] XRoboToolkit SDK initialized.")
+
+    # ------------------------------------------------------------------
+    # Pose
+    # ------------------------------------------------------------------
+
+    def get_pose(self, name: str) -> np.ndarray:
+        """Return the 7-DoF pose [x, y, z, qx, qy, qz, qw] of a device.
+
+        Args:
+            name: One of ``"left_controller"``, ``"right_controller"``, ``"headset"``.
+
+        Returns:
+            np.ndarray of shape (7,).
+        """
+        _POSE_FUNCS = {
+            "left_controller": xrt.get_left_controller_pose,
+            "right_controller": xrt.get_right_controller_pose,
+            "headset": xrt.get_headset_pose,
+        }
+        if name not in _POSE_FUNCS:
+            raise ValueError(
+                f"Invalid pose name: {name}. Choose from {list(_POSE_FUNCS)}"
+            )
+        return np.asarray(_POSE_FUNCS[name](), dtype=np.float64)
+
+    # ------------------------------------------------------------------
+    # Analog Keys (trigger / grip)
+    # ------------------------------------------------------------------
+
+    def get_key_value(self, name: str) -> float:
+        """Return an analog value (0.0~1.0) for a trigger or grip.
+
+        Args:
+            name: One of ``"left_trigger"``, ``"right_trigger"``,
+                  ``"left_grip"``, ``"right_grip"``.
+        """
+        _KEY_FUNCS = {
+            "left_trigger": xrt.get_left_trigger,
+            "right_trigger": xrt.get_right_trigger,
+            "left_grip": xrt.get_left_grip,
+            "right_grip": xrt.get_right_grip,
+        }
+        if name not in _KEY_FUNCS:
+            raise ValueError(
+                f"Invalid key name: {name}. Choose from {list(_KEY_FUNCS)}"
+            )
+        return float(_KEY_FUNCS[name]())
+
+    # ------------------------------------------------------------------
+    # Buttons
+    # ------------------------------------------------------------------
+
+    def get_button(self, name: str) -> bool:
+        """Return boolean state of a face / menu button.
+
+        Args:
+            name: One of ``"A"``, ``"B"``, ``"X"``, ``"Y"``,
+                  ``"left_menu_button"``, ``"right_menu_button"``,
+                  ``"left_axis_click"``, ``"right_axis_click"``.
+        """
+        _BTN_FUNCS = {
+            "A": xrt.get_A_button,
+            "B": xrt.get_B_button,
+            "X": xrt.get_X_button,
+            "Y": xrt.get_Y_button,
+            "left_menu_button": xrt.get_left_menu_button,
+            "right_menu_button": xrt.get_right_menu_button,
+            "left_axis_click": xrt.get_left_axis_click,
+            "right_axis_click": xrt.get_right_axis_click,
+        }
+        if name not in _BTN_FUNCS:
+            raise ValueError(
+                f"Invalid button name: {name}. Choose from {list(_BTN_FUNCS)}"
+            )
+        return bool(_BTN_FUNCS[name]())
+
+    # ------------------------------------------------------------------
+    # Joystick
+    # ------------------------------------------------------------------
+
+    def get_joystick(self, side: str) -> np.ndarray:
+        """Return joystick (x, y) for left or right controller.
+
+        Args:
+            side: ``"left"`` or ``"right"``.
+
+        Returns:
+            np.ndarray of shape (2,).
+        """
+        if side == "left":
+            return np.asarray(xrt.get_left_axis(), dtype=np.float64)
+        elif side == "right":
+            return np.asarray(xrt.get_right_axis(), dtype=np.float64)
+        raise ValueError(f"Invalid side: {side}. Choose 'left' or 'right'.")
+
+    # ------------------------------------------------------------------
+    # Hand Tracking
+    # ------------------------------------------------------------------
+
+    def get_hand_tracking(self, side: str) -> np.ndarray | None:
+        """Return hand tracking joint poses or None if inactive.
+
+        Args:
+            side: ``"left"`` or ``"right"``.
+
+        Returns:
+            np.ndarray of shape (27, 7) or None.
+        """
+        if side == "left":
+            if not xrt.get_left_hand_is_active():
+                return None
+            return np.asarray(xrt.get_left_hand_tracking_state(), dtype=np.float64)
+        elif side == "right":
+            if not xrt.get_right_hand_is_active():
+                return None
+            return np.asarray(xrt.get_right_hand_tracking_state(), dtype=np.float64)
+        raise ValueError(f"Invalid side: {side}. Choose 'left' or 'right'.")
+
+    # ------------------------------------------------------------------
+    # Timestamp
+    # ------------------------------------------------------------------
+
+    def get_timestamp_ns(self) -> int:
+        """Return the current XR system timestamp in nanoseconds."""
+        return int(xrt.get_time_stamp_ns())
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self):
+        """Shut down the SDK connection."""
+        xrt.close()
+        print("[XrClient] SDK connection closed.")