mindbot/scripts/environments/teleoperation/xr_teleop/xr_controller.py

# 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