spacemouse ok
This commit is contained in:
335
scripts/environments/teleoperation/teleop_xr_agent.py
Normal file
335
scripts/environments/teleoperation/teleop_xr_agent.py
Normal file
@@ -0,0 +1,335 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""
|
||||
Script to run teleoperation with Isaac Lab manipulation environments using PICO XR Controllers.
|
||||
This script uses XRoboToolkit to fetch XR controller poses and maps them to differential IK actions.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
|
||||
# Ensure xr_utils (next to this script) is importable when running directly
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# add argparse arguments
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Teleoperation for Isaac Lab environments with PICO XR Controller."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_envs", type=int, default=1, help="Number of environments to simulate."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--task",
|
||||
type=str,
|
||||
default="Isaac-MindRobot-LeftArm-IK-Rel-v0",
|
||||
help="Name of the task.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sensitivity", type=float, default=5.0, help="Sensitivity factor for pos/rot."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arm",
|
||||
type=str,
|
||||
default="left",
|
||||
choices=["left", "right"],
|
||||
help="Which arm/controller to use.",
|
||||
)
|
||||
# append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli = parser.parse_args()
|
||||
app_launcher_args = vars(args_cli)
|
||||
|
||||
# Disable some rendering settings to speed up
|
||||
app_launcher_args["xr"] = False
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(app_launcher_args)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
|
||||
import isaaclab_tasks # noqa: F401
|
||||
import mindbot.tasks # noqa: F401
|
||||
from isaaclab_tasks.utils import parse_env_cfg
|
||||
|
||||
from xr_utils import XrClient, transform_xr_pose, quat_diff_as_rotvec_xyzw, is_valid_quaternion
|
||||
from xr_utils.geometry import R_HEADSET_TO_WORLD
|
||||
|
||||
# =====================================================================
|
||||
# Teleoperation Interface for XR
|
||||
# =====================================================================
|
||||
|
||||
class XrTeleopController:
|
||||
"""Teleop controller for PICO XR headset."""
|
||||
|
||||
def __init__(self, arm="left", pos_sensitivity=1.0, rot_sensitivity=0.3):
|
||||
self.xr_client = XrClient()
|
||||
self.pos_sensitivity = pos_sensitivity
|
||||
self.rot_sensitivity = rot_sensitivity
|
||||
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"
|
||||
|
||||
# Coordinate transform matrix
|
||||
self.R_headset_world = R_HEADSET_TO_WORLD
|
||||
|
||||
# Raw XR tracking space poses (NOT transformed)
|
||||
self.prev_xr_pos = None
|
||||
self.prev_xr_quat = None
|
||||
self.grip_active = False
|
||||
self.frame_count = 0
|
||||
|
||||
# Callbacks (like reset, etc)
|
||||
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
|
||||
|
||||
def close(self):
|
||||
self.xr_client.close()
|
||||
|
||||
def advance(self) -> torch.Tensor:
|
||||
"""
|
||||
Reads the XR controller and returns the 7D action tensor:
|
||||
[dx, dy, dz, droll, dpitch, dyaw, gripper_cmd]
|
||||
"""
|
||||
# XR buttons check (e.g. A or B for reset)
|
||||
try:
|
||||
if self.xr_client.get_button("B") or self.xr_client.get_button("Y"):
|
||||
if "RESET" in self.callbacks:
|
||||
self.callbacks["RESET"]()
|
||||
except:
|
||||
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 as e:
|
||||
return torch.zeros(7)
|
||||
|
||||
# Skip transformation if quaternion is invalid (e.g. before headset truly connects)
|
||||
if not is_valid_quaternion(raw_pose[3:]):
|
||||
action = torch.zeros(7)
|
||||
action[6] = 1.0 if trigger > 0.5 else -1.0
|
||||
return action
|
||||
|
||||
# Transform XR pose pos directly via Matrix mapping for safety check (though we can map diffs next)
|
||||
pos_w = self.R_headset_world @ raw_pose[:3]
|
||||
|
||||
# 握持键作为离合器 (Clutch) - 按下 Grip 时才移动机械臂
|
||||
if grip < 0.5:
|
||||
self.prev_xr_pos = None
|
||||
self.prev_xr_quat = None
|
||||
self.grip_active = False
|
||||
action = torch.zeros(7)
|
||||
action[6] = 1.0 if trigger > 0.5 else -1.0
|
||||
return action
|
||||
|
||||
if not self.grip_active:
|
||||
# We strictly log original XR coordinates to avoid quaternion base frame corruption
|
||||
self.prev_xr_pos = raw_pose[:3].copy()
|
||||
self.prev_xr_quat = raw_pose[3:].copy()
|
||||
self.grip_active = True
|
||||
action = torch.zeros(7)
|
||||
action[6] = 1.0 if trigger > 0.5 else -1.0
|
||||
return action
|
||||
|
||||
# ========== 1. Position Delta (calculate in XR frame, map to World) ==========
|
||||
xr_delta_pos = raw_pose[:3] - self.prev_xr_pos
|
||||
# Clamp raw position delta to prevent spikes (max ~4cm per frame)
|
||||
max_pos_delta = 0.04
|
||||
pos_norm = np.linalg.norm(xr_delta_pos)
|
||||
if pos_norm > max_pos_delta:
|
||||
xr_delta_pos = xr_delta_pos * (max_pos_delta / pos_norm)
|
||||
# PICO -> Isaac World mapping:
|
||||
# XR +X (Right) -> World -Y (Right is -Left)
|
||||
# XR +Y (Up) -> World +Z (Up is Up)
|
||||
# XR +Z (Back) -> World -X (Back is -Forward)
|
||||
delta_pos = np.array([-xr_delta_pos[2], -xr_delta_pos[0], xr_delta_pos[1]]) * self.pos_sensitivity
|
||||
|
||||
# ========== 2. Rotation Delta (calculate in XR frame, map to World) ==========
|
||||
# compute pure relative angular difference in local XR tracking space
|
||||
xr_delta_rot = quat_diff_as_rotvec_xyzw(self.prev_xr_quat, raw_pose[3:])
|
||||
# Clamp raw rotation delta to prevent spikes (max ~0.02 rad = ~1.1° per frame)
|
||||
# Keeping this small is critical: DLS IK can only solve small deltas as pure rotations.
|
||||
# Large deltas cause the IK to swing the whole arm instead of rotating in-place.
|
||||
max_rot_delta = 0.02
|
||||
rot_norm = np.linalg.norm(xr_delta_rot)
|
||||
if rot_norm > max_rot_delta:
|
||||
xr_delta_rot = xr_delta_rot * (max_rot_delta / rot_norm)
|
||||
# Same mapping rules apply to rotation axes:
|
||||
# Rotating around XR's X (Right) -> Rotating around World's -Y (Right)
|
||||
# Rotating around XR's Y (Up) -> Rotating around World's +Z (Up)
|
||||
# Rotating around XR's Z (Back) -> Rotating around World's -X (Back)
|
||||
delta_rot = np.array([-xr_delta_rot[2], -xr_delta_rot[0], xr_delta_rot[1]]) * self.rot_sensitivity
|
||||
|
||||
# Update cache
|
||||
self.prev_xr_pos = raw_pose[:3].copy()
|
||||
self.prev_xr_quat = raw_pose[3:].copy()
|
||||
self.frame_count += 1
|
||||
|
||||
# ========== 3. Gripper ==========
|
||||
gripper_action = 1.0 if trigger > 0.5 else -1.0
|
||||
|
||||
action = torch.tensor([delta_pos[0], delta_pos[1], delta_pos[2], delta_rot[0], delta_rot[1], delta_rot[2], gripper_action], dtype=torch.float32)
|
||||
|
||||
# ========== 4. Comprehensive Debug Log ==========
|
||||
if self.frame_count % 30 == 0:
|
||||
np.set_printoptions(precision=4, suppress=True, floatmode='fixed')
|
||||
print("\n====================== [VR TELEOP DEBUG] ======================")
|
||||
print(f"| Raw VR Pos (OpenXR): {np.array(raw_pose[:3])}")
|
||||
print(f"| Raw VR Quat (xyzw): {np.array(raw_pose[3:])}")
|
||||
print(f"| XR Delta Pos (raw): {xr_delta_pos} (norm={pos_norm:.4f})")
|
||||
print(f"| XR Delta Rot (raw): {xr_delta_rot} (norm={rot_norm:.4f})")
|
||||
print("|--------------------------------------------------------------")
|
||||
print(f"| Sent Action Pos (dx,dy,dz): {action[:3].numpy()}")
|
||||
print(f"| Sent Action Rot (rx,ry,rz): {action[3:6].numpy()}")
|
||||
print(f"| Gripper & Trigger: Grip={grip:.2f}, Trig={trigger:.2f}")
|
||||
print("===============================================================")
|
||||
|
||||
return action
|
||||
|
||||
# =====================================================================
|
||||
# Main Execution Loop
|
||||
# =====================================================================
|
||||
|
||||
def main() -> None:
|
||||
"""Run teleoperation with PICO XR Controller against Isaac Lab environment."""
|
||||
|
||||
# 1. Configuration parsing
|
||||
env_cfg = parse_env_cfg(args_cli.task, num_envs=args_cli.num_envs)
|
||||
env_cfg.env_name = args_cli.task
|
||||
|
||||
if not isinstance(env_cfg, ManagerBasedRLEnvCfg):
|
||||
raise ValueError(f"Teleoperation requires ManagerBasedRLEnvCfg. Got: {type(env_cfg)}")
|
||||
|
||||
env_cfg.terminations.time_out = None
|
||||
|
||||
# 2. Environment creation
|
||||
try:
|
||||
env = gym.make(args_cli.task, cfg=env_cfg).unwrapped
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create environment '{args_cli.task}': {e}")
|
||||
simulation_app.close()
|
||||
return
|
||||
|
||||
# 3. Teleoperation Interface Initialization
|
||||
print(f"\n[INFO] Connecting to PICO XR Headset using {args_cli.arm} controller...")
|
||||
teleop_interface = XrTeleopController(
|
||||
arm=args_cli.arm,
|
||||
pos_sensitivity=args_cli.sensitivity,
|
||||
rot_sensitivity=args_cli.sensitivity * 0.3, # Rotation must be much gentler for DLS IK
|
||||
)
|
||||
|
||||
should_reset = False
|
||||
def request_reset():
|
||||
nonlocal should_reset
|
||||
should_reset = True
|
||||
print("[INFO] Reset requested via XR button.")
|
||||
|
||||
teleop_interface.add_callback("RESET", request_reset)
|
||||
|
||||
env.reset()
|
||||
teleop_interface.reset()
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print(" 🚀 Teleoperation Started!")
|
||||
print(" 🎮 Use the TRIGGER to open/close gripper.")
|
||||
print(" ✊ Hold GRIP button and move the controller to move the arm.")
|
||||
print(" 🕹️ Press B or Y to reset the environment.")
|
||||
print("=" * 50 + "\n")
|
||||
|
||||
# 4. Simulation loop
|
||||
device = env.unwrapped.device
|
||||
sim_frame = 0
|
||||
while simulation_app.is_running():
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
# Get action from XR Controller [1, 7]
|
||||
action_np = teleop_interface.advance()
|
||||
actions = action_np.unsqueeze(0).repeat(env.num_envs, 1).to(device)
|
||||
|
||||
# Step environment
|
||||
obs, _, _, _, _ = env.step(actions)
|
||||
|
||||
# Print robot state every 30 frames
|
||||
sim_frame += 1
|
||||
if sim_frame % 30 == 0:
|
||||
np.set_printoptions(precision=4, suppress=True, floatmode='fixed')
|
||||
policy_obs = obs["policy"]
|
||||
joint_pos = policy_obs["joint_pos"][0].cpu().numpy()
|
||||
eef_pos = policy_obs["eef_pos"][0].cpu().numpy()
|
||||
eef_quat = policy_obs["eef_quat"][0].cpu().numpy()
|
||||
last_act = policy_obs["actions"][0].cpu().numpy()
|
||||
|
||||
# On first print, dump ALL joint names + positions to identify indices
|
||||
if sim_frame == 30:
|
||||
robot = env.unwrapped.scene["robot"]
|
||||
jnames = robot.joint_names
|
||||
print(f"\n{'='*70}")
|
||||
print(f" ALL {len(jnames)} JOINT NAMES AND POSITIONS (relative)")
|
||||
print(f"{'='*70}")
|
||||
for i, name in enumerate(jnames):
|
||||
print(f" [{i:2d}] {name:30s} = {joint_pos[i]:+.4f}")
|
||||
print(f"{'='*70}")
|
||||
# Find left arm joint indices
|
||||
arm_idx = [i for i, n in enumerate(jnames) if n.startswith("l_joint")]
|
||||
print(f" Left arm indices: {arm_idx}")
|
||||
print(f"{'='*70}\n")
|
||||
|
||||
# Get arm indices (cache-friendly: find once)
|
||||
if not hasattr(env, '_arm_idx_cache'):
|
||||
robot = env.unwrapped.scene["robot"]
|
||||
jnames = robot.joint_names
|
||||
env._arm_idx_cache = [i for i, n in enumerate(jnames) if n.startswith("l_joint")]
|
||||
arm_idx = env._arm_idx_cache
|
||||
arm_joints = joint_pos[arm_idx]
|
||||
|
||||
print(f"\n---------------- [ROBOT STATE frame={sim_frame}] ----------------")
|
||||
print(f"| Left Arm Joints (rad): {arm_joints}")
|
||||
print(f"| EEF Pos (world): {eef_pos}")
|
||||
print(f"| EEF Quat (world, wxyz): {eef_quat}")
|
||||
print(f"| Last Action Sent: {last_act}")
|
||||
print(f"----------------------------------------------------------------")
|
||||
|
||||
if should_reset:
|
||||
env.reset()
|
||||
teleop_interface.reset()
|
||||
should_reset = False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during simulation step: {e}")
|
||||
break
|
||||
|
||||
teleop_interface.close()
|
||||
env.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
simulation_app.close()
|
||||
33
scripts/environments/teleoperation/xr_utils/__init__.py
Normal file
33
scripts/environments/teleoperation/xr_utils/__init__.py
Normal file
@@ -0,0 +1,33 @@
|
||||
# Copyright (c) 2022-2026, The Mindbot Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""XR teleoperation utilities for PICO 4 Ultra / XRoboToolkit integration.
|
||||
|
||||
This package provides coordinate transformation, XR device communication,
|
||||
and teleoperation control utilities adapted from XRoboToolkit for use
|
||||
with Isaac Lab environments.
|
||||
"""
|
||||
|
||||
from .geometry import (
|
||||
R_HEADSET_TO_WORLD,
|
||||
apply_delta_pose,
|
||||
is_valid_quaternion,
|
||||
quat_diff_as_angle_axis,
|
||||
quat_diff_as_rotvec_xyzw,
|
||||
quaternion_to_angle_axis,
|
||||
transform_xr_pose,
|
||||
)
|
||||
from .xr_client import XrClient
|
||||
|
||||
__all__ = [
|
||||
"R_HEADSET_TO_WORLD",
|
||||
"apply_delta_pose",
|
||||
"is_valid_quaternion",
|
||||
"quat_diff_as_angle_axis",
|
||||
"quat_diff_as_rotvec_xyzw",
|
||||
"quaternion_to_angle_axis",
|
||||
"transform_xr_pose",
|
||||
"XrClient",
|
||||
]
|
||||
203
scripts/environments/teleoperation/xr_utils/geometry.py
Normal file
203
scripts/environments/teleoperation/xr_utils/geometry.py
Normal file
@@ -0,0 +1,203 @@
|
||||
# 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 geometry utilities by Zhigen Zhao.
|
||||
|
||||
"""Geometry utilities for XR-to-World coordinate transformations.
|
||||
|
||||
This module provides the core coordinate transformation matrices and functions
|
||||
needed to convert PICO 4 Ultra (OpenXR) data into Isaac Sim world coordinates.
|
||||
|
||||
Coordinate Systems:
|
||||
- OpenXR (PICO 4): +X Right, +Y Up, -Z Forward
|
||||
- Isaac Sim (World): +X Forward, +Y Left, +Z Up
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Core Coordinate Transform: OpenXR -> Isaac Sim World
|
||||
# =====================================================================
|
||||
|
||||
R_HEADSET_TO_WORLD = np.array(
|
||||
[
|
||||
[0, 0, -1], # XR -Z (forward) -> World +X (forward)
|
||||
[-1, 0, 0], # XR -X (left) -> World +Y (left)
|
||||
[0, 1, 0], # XR +Y (up) -> World +Z (up)
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
"""3x3 rotation matrix that maps OpenXR coordinates to Isaac Sim world coordinates."""
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Quaternion Utilities
|
||||
# =====================================================================
|
||||
|
||||
def is_valid_quaternion(quat: np.ndarray, tol: float = 1e-6) -> bool:
|
||||
"""Check if a quaternion is valid (unit norm).
|
||||
|
||||
Args:
|
||||
quat: Quaternion array of length 4 (any convention).
|
||||
tol: Tolerance for norm deviation from 1.0.
|
||||
|
||||
Returns:
|
||||
True if the quaternion is valid.
|
||||
"""
|
||||
if not isinstance(quat, (list, tuple, np.ndarray)) or len(quat) != 4:
|
||||
return False
|
||||
if not np.all(np.isfinite(quat)):
|
||||
return False
|
||||
magnitude = np.sqrt(np.sum(np.square(quat)))
|
||||
return abs(magnitude - 1.0) <= tol
|
||||
|
||||
|
||||
def quaternion_to_angle_axis(quat_wxyz: np.ndarray, eps: float = 1e-6) -> np.ndarray:
|
||||
"""Convert a quaternion (w, x, y, z) to angle-axis representation.
|
||||
|
||||
Args:
|
||||
quat_wxyz: Quaternion in (w, x, y, z) format.
|
||||
eps: Tolerance for small angles.
|
||||
|
||||
Returns:
|
||||
Angle-axis vector [ax*angle, ay*angle, az*angle], shape (3,).
|
||||
"""
|
||||
q = np.array(quat_wxyz, dtype=np.float64, copy=True)
|
||||
if q[0] < 0.0:
|
||||
q = -q
|
||||
|
||||
w = q[0]
|
||||
vec_part = q[1:]
|
||||
angle = 2.0 * np.arccos(np.clip(w, -1.0, 1.0))
|
||||
|
||||
if angle < eps:
|
||||
return np.zeros(3, dtype=np.float64)
|
||||
|
||||
sin_half_angle = np.sin(angle / 2.0)
|
||||
if sin_half_angle < eps:
|
||||
return np.zeros(3, dtype=np.float64)
|
||||
|
||||
axis = vec_part / sin_half_angle
|
||||
return axis * angle
|
||||
|
||||
|
||||
def quat_diff_as_angle_axis(
|
||||
q1_wxyz: np.ndarray, q2_wxyz: np.ndarray, eps: float = 1e-6
|
||||
) -> np.ndarray:
|
||||
"""Calculate the rotation from q1 to q2 as an angle-axis vector.
|
||||
|
||||
Computes DeltaQ such that q2 = DeltaQ * q1.
|
||||
|
||||
Args:
|
||||
q1_wxyz: Source quaternion (w, x, y, z).
|
||||
q2_wxyz: Target quaternion (w, x, y, z).
|
||||
eps: Tolerance for small rotations.
|
||||
|
||||
Returns:
|
||||
Angle-axis vector representing DeltaQ, shape (3,).
|
||||
"""
|
||||
# Use scipy for robust quaternion math (scipy uses xyzw convention)
|
||||
r1 = R.from_quat([q1_wxyz[1], q1_wxyz[2], q1_wxyz[3], q1_wxyz[0]])
|
||||
r2 = R.from_quat([q2_wxyz[1], q2_wxyz[2], q2_wxyz[3], q2_wxyz[0]])
|
||||
delta = r2 * r1.inv()
|
||||
rotvec = delta.as_rotvec()
|
||||
if np.linalg.norm(rotvec) < eps:
|
||||
return np.zeros(3, dtype=np.float64)
|
||||
return rotvec
|
||||
|
||||
|
||||
def quat_diff_as_rotvec_xyzw(
|
||||
q1_xyzw: np.ndarray, q2_xyzw: np.ndarray, eps: float = 1e-6
|
||||
) -> np.ndarray:
|
||||
"""Calculate the rotation from q1 to q2 as a rotation vector.
|
||||
|
||||
Same as quat_diff_as_angle_axis but takes (x, y, z, w) quaternions
|
||||
directly (scipy convention), avoiding extra conversions.
|
||||
|
||||
Args:
|
||||
q1_xyzw: Source quaternion (x, y, z, w).
|
||||
q2_xyzw: Target quaternion (x, y, z, w).
|
||||
eps: Tolerance for small rotations.
|
||||
|
||||
Returns:
|
||||
Rotation vector (angle * axis), shape (3,).
|
||||
"""
|
||||
r1 = R.from_quat(q1_xyzw)
|
||||
r2 = R.from_quat(q2_xyzw)
|
||||
delta = r2 * r1.inv()
|
||||
rotvec = delta.as_rotvec()
|
||||
if np.linalg.norm(rotvec) < eps:
|
||||
return np.zeros(3, dtype=np.float64)
|
||||
return rotvec
|
||||
|
||||
|
||||
def apply_delta_pose(
|
||||
source_pos: np.ndarray,
|
||||
source_quat_xyzw: np.ndarray,
|
||||
delta_pos: np.ndarray,
|
||||
delta_rot: np.ndarray,
|
||||
eps: float = 1.0e-6,
|
||||
) -> tuple:
|
||||
"""Apply a delta pose to a source pose.
|
||||
|
||||
Args:
|
||||
source_pos: Position of source frame, shape (3,).
|
||||
source_quat_xyzw: Quaternion of source frame (x, y, z, w), shape (4,).
|
||||
delta_pos: Position displacement, shape (3,).
|
||||
delta_rot: Rotation displacement in angle-axis format, shape (3,).
|
||||
eps: Tolerance to consider rotation as zero.
|
||||
|
||||
Returns:
|
||||
Tuple of (target_pos, target_quat_xyzw), both np.ndarray.
|
||||
"""
|
||||
target_pos = source_pos + delta_pos
|
||||
|
||||
angle = np.linalg.norm(delta_rot)
|
||||
if angle > eps:
|
||||
delta_r = R.from_rotvec(delta_rot)
|
||||
else:
|
||||
delta_r = R.identity()
|
||||
|
||||
source_r = R.from_quat(source_quat_xyzw)
|
||||
target_r = delta_r * source_r
|
||||
target_quat_xyzw = target_r.as_quat()
|
||||
|
||||
return target_pos, target_quat_xyzw
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# XR Pose Transform
|
||||
# =====================================================================
|
||||
|
||||
def transform_xr_pose(
|
||||
pos_xr: np.ndarray, quat_xyzw_xr: np.ndarray
|
||||
) -> tuple:
|
||||
"""Transform an XR device pose from OpenXR frame to Isaac Sim world frame.
|
||||
|
||||
Uses the R_HEADSET_TO_WORLD rotation matrix to transform both position
|
||||
and orientation, following the same approach as XRoboToolkit.
|
||||
|
||||
Args:
|
||||
pos_xr: Position in OpenXR frame, shape (3,).
|
||||
quat_xyzw_xr: Quaternion in (x, y, z, w) format from XR SDK, shape (4,).
|
||||
|
||||
Returns:
|
||||
Tuple of (pos_world, quat_xyzw_world):
|
||||
- pos_world: Position in world frame, shape (3,).
|
||||
- quat_xyzw_world: Quaternion in (x, y, z, w) for scipy, shape (4,).
|
||||
"""
|
||||
# Transform position
|
||||
pos_world = R_HEADSET_TO_WORLD @ pos_xr
|
||||
|
||||
# Transform orientation: R_world = R_transform * R_xr * R_transform^T
|
||||
r_xr = R.from_quat(quat_xyzw_xr) # scipy uses (x, y, z, w)
|
||||
r_transform = R.from_matrix(R_HEADSET_TO_WORLD)
|
||||
r_world = r_transform * r_xr * r_transform.inv()
|
||||
quat_xyzw_world = r_world.as_quat()
|
||||
|
||||
return pos_world, quat_xyzw_world
|
||||
169
scripts/environments/teleoperation/xr_utils/xr_client.py
Normal file
169
scripts/environments/teleoperation/xr_utils/xr_client.py
Normal file
@@ -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.")
|
||||
Reference in New Issue
Block a user