# Copyright (c) 2025, Yutang Li, SIAT (yt.li2@siat.ac.cn) # 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 = {} self._root_yaw_rad = 0.0 # current chassis heading, updated each frame 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 set_root_yaw(self, yaw_rad: float): """Update the chassis heading so XR deltas are in body frame.""" self._root_yaw_rad = yaw_rad 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 # In VR teleop, user "inhabits" the robot: their physical forward (XR +X) # maps directly to robot forward (body +X). No rotation needed — the VR view # already rotates with the robot, so XR deltas naturally align with body frame. # Diagnostic if self.frame_count % 30 == 0 and self.target_eef_pos is not None: print(f"[XR {self.arm}] delta=({world_delta_pos[0]:+.4f},{world_delta_pos[1]:+.4f},{world_delta_pos[2]:+.4f}) " f"target=({self.target_eef_pos[0]:.3f},{self.target_eef_pos[1]:.3f},{self.target_eef_pos[2]:.3f})") 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