重构xr teleop

scripts/environments/teleoperation/xr_teleop/xr_controller.py

@@ -0,0 +1,176 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, The Isaac Lab Project Developers.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""XR teleoperation controller with absolute IK target accumulation and grip clutch."""
+
+from collections.abc import Callable
+
+import numpy as np
+import torch
+from scipy.spatial.transform import Rotation as R
+
+from xr_utils import XrClient, transform_xr_pose, quat_diff_as_rotvec_xyzw, is_valid_quaternion
+
+from .frame_utils import quat_wxyz_to_rotation, rotation_to_quat_wxyz
+
+
+class XrTeleopController:
+    """Teleop controller for PICO XR headset (absolute IK mode).
+
+    Accumulates an absolute EEF target in Isaac Sim world frame.
+    Grip button acts as clutch; B/Y buttons trigger environment reset.
+    """
+
+    def __init__(self, arm="left", pos_sensitivity=1.0, rot_sensitivity=0.3,
+                 max_pos_per_step=0.05, max_rot_per_step=0.08, xr_client=None):
+        self.xr_client = xr_client if xr_client is not None else XrClient()
+        self._owns_client = xr_client is None
+        self.pos_sensitivity = pos_sensitivity
+        self.rot_sensitivity = rot_sensitivity
+        self.max_pos_per_step = max_pos_per_step
+        self.max_rot_per_step = max_rot_per_step
+        self.arm = arm
+
+        self.controller_name = "left_controller" if arm == "left" else "right_controller"
+        self.grip_name = "left_grip" if arm == "left" else "right_grip"
+        self.trigger_name = "left_trigger" if arm == "left" else "right_trigger"
+
+        from xr_utils.geometry import R_HEADSET_TO_WORLD
+        self.R_headset_world = R_HEADSET_TO_WORLD
+
+        self.prev_xr_pos = None
+        self.prev_xr_quat = None
+        self.target_eef_pos = None
+        self.target_eef_quat = None
+
+        self.grip_active = False
+        self.frame_count = 0
+        self.reset_button_latched = False
+        self.require_grip_reengage = False
+        self.grip_engage_threshold = 0.8
+        self.grip_release_threshold = 0.2
+        self.callbacks = {}
+
+    def add_callback(self, name: str, func: Callable):
+        self.callbacks[name] = func
+
+    def reset(self):
+        self.prev_xr_pos = None
+        self.prev_xr_quat = None
+        self.grip_active = False
+        self.frame_count = 0
+        self.target_eef_pos = None
+        self.target_eef_quat = None
+        self.require_grip_reengage = True
+
+    def close(self):
+        if self._owns_client:
+            self.xr_client.close()
+
+    def get_zero_action(self, trigger, current_eef_pos=None, current_eef_quat=None) -> torch.Tensor:
+        """Return 8D hold action [x, y, z, qw, qx, qy, qz, gripper] at frozen target pose."""
+        action = torch.zeros(8)
+        if self.target_eef_pos is not None and self.target_eef_quat is not None:
+            action[:3] = torch.tensor(self.target_eef_pos.copy())
+            action[3:7] = torch.tensor(self.target_eef_quat.copy())
+        elif current_eef_pos is not None and current_eef_quat is not None:
+            action[:3] = torch.tensor(current_eef_pos.copy())
+            action[3:7] = torch.tensor(current_eef_quat.copy())
+        else:
+            action[3] = 1.0
+        action[7] = 1.0 if trigger > 0.5 else -1.0
+        return action
+
+    def advance(self, current_eef_pos=None, current_eef_quat=None) -> torch.Tensor:
+        """Read XR controller and return 8D absolute action [x, y, z, qw, qx, qy, qz, gripper].
+
+        Position and quaternion are in Isaac Sim world frame.
+        Caller must convert to robot root frame before sending to IK.
+        """
+        try:
+            if self.arm == "left":
+                reset_pressed = self.xr_client.get_button("Y")
+            else:
+                reset_pressed = False
+            if reset_pressed and not self.reset_button_latched:
+                if "RESET" in self.callbacks:
+                    self.callbacks["RESET"]()
+            self.reset_button_latched = reset_pressed
+        except Exception:
+            pass
+
+        try:
+            raw_pose = self.xr_client.get_pose(self.controller_name)
+            grip = self.xr_client.get_key_value(self.grip_name)
+            trigger = self.xr_client.get_key_value(self.trigger_name)
+        except Exception:
+            return self.get_zero_action(0.0, current_eef_pos, current_eef_quat)
+
+        if not is_valid_quaternion(raw_pose[3:]):
+            return self.get_zero_action(trigger, current_eef_pos, current_eef_quat)
+
+        if self.require_grip_reengage:
+            if grip <= self.grip_release_threshold:
+                self.require_grip_reengage = False
+            else:
+                if current_eef_pos is not None and current_eef_quat is not None:
+                    self.target_eef_pos = current_eef_pos.copy()
+                    self.target_eef_quat = current_eef_quat.copy()
+                return self.get_zero_action(trigger, current_eef_pos, current_eef_quat)
+
+        grip_pressed = grip > self.grip_release_threshold if self.grip_active else grip >= self.grip_engage_threshold
+
+        if not grip_pressed:
+            self.prev_xr_pos = None
+            self.prev_xr_quat = None
+            if self.grip_active or self.target_eef_pos is None:
+                if current_eef_pos is not None and current_eef_quat is not None:
+                    self.target_eef_pos = current_eef_pos.copy()
+                    self.target_eef_quat = current_eef_quat.copy()
+            self.grip_active = False
+            return self.get_zero_action(trigger, current_eef_pos, current_eef_quat)
+
+        if not self.grip_active:
+            self.prev_xr_pos = raw_pose[:3].copy()
+            self.prev_xr_quat = raw_pose[3:].copy()
+            if current_eef_pos is not None and current_eef_quat is not None:
+                self.target_eef_pos = current_eef_pos.copy()
+                self.target_eef_quat = current_eef_quat.copy()
+            self.grip_active = True
+            return self.get_zero_action(trigger, current_eef_pos, current_eef_quat)
+
+        xr_world_pos, xr_world_quat_xyzw = transform_xr_pose(raw_pose[:3], raw_pose[3:])
+        prev_xr_world_pos, prev_xr_world_quat_xyzw = transform_xr_pose(self.prev_xr_pos, self.prev_xr_quat)
+
+        world_delta_pos = (xr_world_pos - prev_xr_world_pos) * self.pos_sensitivity
+        pos_norm = np.linalg.norm(world_delta_pos)
+        if pos_norm > self.max_pos_per_step:
+            world_delta_pos *= self.max_pos_per_step / pos_norm
+
+        world_delta_rot = quat_diff_as_rotvec_xyzw(prev_xr_world_quat_xyzw, xr_world_quat_xyzw) * self.rot_sensitivity
+        rot_norm = np.linalg.norm(world_delta_rot)
+        if rot_norm > self.max_rot_per_step:
+            world_delta_rot *= self.max_rot_per_step / rot_norm
+
+        gripper_action = 1.0 if trigger > 0.5 else -1.0
+
+        if self.target_eef_pos is None:
+            self.target_eef_pos = np.zeros(3)
+            self.target_eef_quat = np.array([1.0, 0.0, 0.0, 0.0])
+
+        self.target_eef_pos += world_delta_pos
+
+        target_r = quat_wxyz_to_rotation(self.target_eef_quat)
+        self.target_eef_quat = rotation_to_quat_wxyz(R.from_rotvec(world_delta_rot) * target_r)
+
+        self.prev_xr_pos = raw_pose[:3].copy()
+        self.prev_xr_quat = raw_pose[3:].copy()
+        self.frame_count += 1
+
+        action = torch.tensor([
+            *self.target_eef_pos,
+            *self.target_eef_quat,
+            gripper_action,
+        ], dtype=torch.float32)
+
+        return action