177 lines
7.2 KiB
Python
177 lines
7.2 KiB
Python
# 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
|