重构xr teleop

scripts/environments/teleoperation/xr_teleop/frame_utils.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, The Isaac Lab Project Developers.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Coordinate frame utilities for XR teleoperation."""
+
+import numpy as np
+import torch
+from scipy.spatial.transform import Rotation as R
+
+
+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 to Isaac-style wxyz quaternion."""
+    q = rot.as_quat()
+    return np.array([q[3], q[0], q[1], q[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
+
+
+def convert_action_world_to_root(action_tensor: torch.Tensor, robot) -> torch.Tensor:
+    """Convert absolute IK action from world frame to robot root frame.
+
+    Args:
+        action_tensor: 8D tensor [x, y, z, qw, qx, qy, qz, gripper] in world frame.
+        robot: Isaac Lab Articulation object (env.unwrapped.scene["robot"]).
+
+    Returns:
+        8D tensor with pos/quat converted to robot root frame.
+    """
+    root_pos_w = robot.data.root_pos_w[0].detach().cpu().numpy()
+    root_quat_w = robot.data.root_quat_w[0].detach().cpu().numpy()
+    pos_root, quat_root = world_pose_to_root_frame(
+        action_tensor[:3].numpy(), action_tensor[3:7].numpy(),
+        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