56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
# 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
|