Files
mindbot/scripts/environments/teleoperation/teleop_xr_agent.py
2026-03-06 17:31:00 +08:00

538 lines
22 KiB
Python

#!/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.",
)
parser.add_argument(
"--base_speed", type=float, default=3.0,
help="Max wheel speed (rad/s) for joystick full forward.",
)
parser.add_argument(
"--base_turn", type=float, default=2.0,
help="Max wheel differential (rad/s) for joystick full left/right.",
)
# 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 scipy.spatial.transform import Rotation as R
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
# =====================================================================
def _quat_wxyz_to_rotation(quat_wxyz: np.ndarray) -> R:
"""Convert Isaac-style wxyz quaternion to scipy Rotation."""
return R.from_quat([quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]])
def _rotation_to_quat_wxyz(rot: R) -> np.ndarray:
"""Convert scipy Rotation quaternion to Isaac-style wxyz."""
quat_xyzw = rot.as_quat()
return np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
def world_pose_to_root_frame(
pos_w: np.ndarray,
quat_wxyz: np.ndarray,
root_pos_w: np.ndarray,
root_quat_wxyz: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
"""Express a world-frame pose in the robot root frame."""
root_rot = _quat_wxyz_to_rotation(root_quat_wxyz)
pose_rot = _quat_wxyz_to_rotation(quat_wxyz)
pos_root = root_rot.inv().apply(pos_w - root_pos_w)
quat_root = _rotation_to_quat_wxyz(root_rot.inv() * pose_rot)
return pos_root, quat_root
class XrTeleopController:
"""Teleop controller for PICO XR headset."""
def __init__(self, arm="left", pos_sensitivity=1.0, rot_sensitivity=0.3, is_absolute=False):
self.xr_client = XrClient()
self.pos_sensitivity = pos_sensitivity
self.rot_sensitivity = rot_sensitivity
self.arm = arm
self.is_absolute = is_absolute
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
self.prev_xr_pos = None
self.prev_xr_quat = None
# Absolute target states
self.target_eef_pos = None
self.target_eef_quat = None # wxyz
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
# 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
self.target_eef_pos = None
self.target_eef_quat = None
# Require one grip release after reset so stale controller motion
# cannot immediately drive the robot back toward the previous pose.
self.require_grip_reengage = True
def close(self):
self.xr_client.close()
def get_zero_action(self, trigger, current_eef_pos=None, current_eef_quat=None):
if self.is_absolute:
action = torch.zeros(8)
# Stay at the current valid pose, or fallback to the cached target
if 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())
elif 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())
else:
action[3] = 1.0 # default w=1 for quat
action[7] = 1.0 if trigger > 0.5 else -1.0
else:
action = torch.zeros(7)
action[6] = 1.0 if trigger > 0.5 else -1.0
return action
def advance(self, current_eef_pos=None, current_eef_quat=None) -> torch.Tensor:
"""
Reads the XR controller.
Relative bounds return 7D action tensor: [dx, dy, dz, drx, dry, drz, gripper]
Absolute bounds return 8D action tensor: [x, y, z, qw, qx, qy, qz, gripper]
Note: in absolute mode current_eef_* and the returned target are in WORLD frame.
The caller is responsible for converting to root frame before sending to IK.
"""
# XR buttons check (e.g. A or B for reset)
try:
reset_pressed = self.xr_client.get_button("B") or self.xr_client.get_button("Y")
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 as e:
return self.get_zero_action(0.0, current_eef_pos, current_eef_quat)
# Skip transformation if quaternion is invalid
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 self.is_absolute and 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)
# Use hysteresis so noisy analog grip values do not accidentally re-enable teleop.
if self.grip_active:
grip_pressed = grip > self.grip_release_threshold
else:
grip_pressed = grip >= self.grip_engage_threshold
# 握持键作为离合器 (Clutch)
if not grip_pressed:
self.prev_xr_pos = None
self.prev_xr_quat = None
self.grip_active = False
# Ensure target tracks the current pose while we are not grabbing
if self.is_absolute and 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() # wxyz
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 self.is_absolute and 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() # wxyz
self.grip_active = True
return self.get_zero_action(trigger, current_eef_pos, current_eef_quat)
# Since OpenXR headset zero is not guaranteed to match robot zero, we compute the
# raw transformation in World Frame, but apply it relatively to the stored target.
# 1. First, find current XR pose projected into World frame
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)
# 2. Extract Delta POS in World frame
world_delta_pos = (xr_world_pos - prev_xr_world_pos) * self.pos_sensitivity
# 3. Extract Delta ROT in World frame
world_delta_rot = quat_diff_as_rotvec_xyzw(prev_xr_world_quat_xyzw, xr_world_quat_xyzw) * self.rot_sensitivity
# 4. Gripper
gripper_action = 1.0 if trigger > 0.5 else -1.0
if self.is_absolute:
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])
# Accumulate in world frame so VR direction always matches sim direction.
self.target_eef_pos += world_delta_pos
target_r = _quat_wxyz_to_rotation(self.target_eef_quat)
delta_r = R.from_rotvec(world_delta_rot)
self.target_eef_quat = _rotation_to_quat_wxyz(delta_r * target_r)
action = torch.tensor([
self.target_eef_pos[0], self.target_eef_pos[1], self.target_eef_pos[2],
self.target_eef_quat[0], self.target_eef_quat[1], self.target_eef_quat[2], self.target_eef_quat[3],
gripper_action], dtype=torch.float32)
self.prev_xr_pos = raw_pose[:3].copy()
self.prev_xr_quat = raw_pose[3:].copy()
else:
max_pos_delta = 0.05
world_pos_norm = np.linalg.norm(world_delta_pos)
if world_pos_norm > max_pos_delta:
world_delta_pos = world_delta_pos * (max_pos_delta / world_pos_norm)
max_rot_delta = 0.15
world_rot_norm = np.linalg.norm(world_delta_rot)
if world_rot_norm > max_rot_delta:
world_delta_rot = world_delta_rot * (max_rot_delta / world_rot_norm)
action = torch.tensor([
world_delta_pos[0], world_delta_pos[1], world_delta_pos[2],
world_delta_rot[0], world_delta_rot[1], world_delta_rot[2],
gripper_action], dtype=torch.float32)
self.prev_xr_pos = raw_pose[:3].copy()
self.prev_xr_quat = raw_pose[3:].copy()
self.frame_count += 1
if self.frame_count % 30 == 0:
np.set_printoptions(precision=4, suppress=True, floatmode='fixed')
print("\n====================== [VR TELEOP DEBUG] ======================")
print(f"| Task Mode: {'ABSOLUTE' if self.is_absolute else 'RELATIVE'}")
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 (world): {world_delta_pos} (norm={np.linalg.norm(world_delta_pos):.4f})")
print(f"| XR Delta Rot (world): {world_delta_rot} (norm={np.linalg.norm(world_delta_rot):.4f})")
if self.is_absolute:
print(f"| Targ Pos (world): {action[:3].numpy()}")
print(f"| Targ Quat (world, wxyz): {action[3:7].numpy()}")
else:
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
is_abs_mode = "Abs" in args_cli.task
print(f"\n[INFO] Connecting to PICO XR Headset using {args_cli.arm} controller...")
print(f"[INFO] Using IK Mode: {'ABSOLUTE' if is_abs_mode else 'RELATIVE'}")
teleop_interface = XrTeleopController(
arm=args_cli.arm,
pos_sensitivity=args_cli.sensitivity,
rot_sensitivity=args_cli.sensitivity * (1.0 if is_abs_mode else 0.3), # Absolute rotation handles 1:1 better than relative
is_absolute=is_abs_mode
)
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)
def get_arm_action_term():
return env.action_manager._terms["arm_action"]
def clear_ik_target_state():
"""Clear the internal IK target so reset does not reuse the previous pose command."""
if not is_abs_mode:
return
arm_action_term = get_arm_action_term()
ee_pos_b, ee_quat_b = arm_action_term._compute_frame_pose()
arm_action_term._raw_actions.zero_()
arm_action_term._processed_actions.zero_()
arm_action_term._ik_controller._command.zero_()
arm_action_term._ik_controller.ee_pos_des[:] = ee_pos_b
arm_action_term._ik_controller.ee_quat_des[:] = ee_quat_b
def convert_action_world_to_root(action_tensor: torch.Tensor) -> torch.Tensor:
"""Convert an absolute IK action from world frame to robot root frame."""
robot = env.unwrapped.scene["robot"]
root_pos_w = robot.data.root_pos_w[0].detach().cpu().numpy()
root_quat_w = robot.data.root_quat_w[0].detach().cpu().numpy() # wxyz
target_pos_w = action_tensor[:3].numpy()
target_quat_w = action_tensor[3:7].numpy()
pos_root, quat_root = world_pose_to_root_frame(
target_pos_w, target_quat_w, root_pos_w, root_quat_w,
)
out = action_tensor.clone()
out[:3] = torch.tensor(pos_root, dtype=torch.float32)
out[3:7] = torch.tensor(quat_root, dtype=torch.float32)
return out
def get_wheel_action() -> torch.Tensor:
"""Read left joystick and return 4-DOF wheel velocity command.
Skid-steer differential drive.
Joystick: Y-axis (+1 = forward), X-axis (+1 = right turn).
Joint order from articulation (terminal log):
[right_b, left_b, left_f, right_f]
Right/left sign convention assumes both sides' joints have the same
axis direction (positive velocity = forward). If the robot drives
backward when pushing forward, negate base_speed in the launch command.
"""
try:
joy = teleop_interface.xr_client.get_joystick("left")
jy = float(joy[1]) # forward / backward
jx = float(joy[0]) # right / left
except Exception:
return torch.zeros(4)
v = jy * args_cli.base_speed
omega = jx * args_cli.base_turn
# Positive omega = turn right → left wheels faster, right wheels slower
right_vel = v - omega
left_vel = v + omega
return torch.tensor(
[right_vel, left_vel, left_vel, right_vel], dtype=torch.float32
)
env.reset()
clear_ik_target_state()
teleop_interface.reset()
print("\n" + "=" * 50)
print(" 🚀 Teleoperation Started!")
print(" 🎮 Use the TRIGGER to open/close gripper.")
print(" ✊ Hold GRIP and move the controller to move the arm.")
print(" 🕹️ Left joystick: Y=forward/back, X=turn left/right.")
print(" 🔄 Press B or Y to reset the environment.")
print("=" * 50 + "\n")
device = env.unwrapped.device
sim_frame = 0
obs, _ = env.reset()
clear_ik_target_state()
while simulation_app.is_running():
try:
with torch.inference_mode():
if should_reset:
obs, _ = env.reset()
clear_ik_target_state()
teleop_interface.reset()
should_reset = False
sim_frame = 0
continue
# Read current EEF in world frame from observations.
policy_obs = obs["policy"]
eef_pos = policy_obs["eef_pos"][0].cpu().numpy()
eef_quat = policy_obs["eef_quat"][0].cpu().numpy()
# Get action from XR Controller (world frame for absolute mode).
action_np = teleop_interface.advance(
current_eef_pos=eef_pos, current_eef_quat=eef_quat,
)
# IK expects root-frame commands; convert just before sending.
if is_abs_mode:
action_np = convert_action_world_to_root(action_np)
# Action manager order: arm_action | wheel_action | gripper_action
# arm=7, wheel=4, gripper=1 → total 12 dims.
wheel_np = get_wheel_action()
action_np = torch.cat([action_np[:7], wheel_np, action_np[7:]])
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 arm joint indices dynamically by looking at the first 6-7 joints that aren't fingers or hands
arm_idx = [i for i, n in enumerate(jnames) if n.startswith("l_joint")]
print(f" Deduced 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]
try:
joy_dbg = teleop_interface.xr_client.get_joystick("left")
joy_str = f"[{joy_dbg[0]:+.2f}, {joy_dbg[1]:+.2f}]"
except Exception:
joy_str = "N/A"
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"| Joystick (x=turn,y=fwd): {joy_str}")
print(f"----------------------------------------------------------------")
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()