From e825ee0320ce2b97f506a96959e7e26b507f4e38 Mon Sep 17 00:00:00 2001 From: yjinzero Date: Sat, 15 Feb 2025 11:08:02 +0900 Subject: [PATCH] map test --- deploy/deploy_real/act_to_dof.py | 72 +++++++++++++ deploy/deploy_real/common/np_math.py | 147 +++++++++++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 deploy/deploy_real/act_to_dof.py create mode 100644 deploy/deploy_real/common/np_math.py diff --git a/deploy/deploy_real/act_to_dof.py b/deploy/deploy_real/act_to_dof.py new file mode 100644 index 0000000..470e8d3 --- /dev/null +++ b/deploy/deploy_real/act_to_dof.py @@ -0,0 +1,72 @@ +import numpy as np +import torch +import pinocchio as pin +from common.np_math import (quat_from_angle_axis, quat_mul, + xyzw2wxyz, index_map) + +class ActToDof: + def __call__(self, obs, action): + hands_command = obs[..., 119:225] + non_arm_joint_pos = action[..., :22] + left_arm_residual = action[..., 22:29] + + q_pin = np.zeros_like(self.ikctrl.cfg.q) + for i_mot in range(len(self.config.motor_joint)): + i_pin = self.pin_from_mot[i_mot] + q_pin[i_pin] = self.low_state.motor_state[i_mot].q + + d_quat = quat_from_angle_axis( + torch.from_numpy(hands_command[..., 3:]) + ).detach().cpu().numpy() + + source_pose = self.ikctrl.fk(q_pin) + source_xyz = source_pose.translation + source_quat = xyzw2wxyz(pin.Quaternion(source_pose.rotation).coeffs()) + + target_xyz = source_xyz + hands_command[..., :3] + target_quat = quat_mul( + torch.from_numpy(d_quat), + torch.from_numpy(source_quat)).detach().cpu().numpy() + target = np.concatenate([target_xyz, target_quat]) + + res_q_ik = self.ikctrl( + q_pin, + target + ) + print('res_q_ik', res_q_ik) + + target_dof_pos = np.zeros(29) + for i_arm in range(len(res_q_ik)): + i_mot = self.mot_from_arm[i_arm] + i_pin = self.pin_from_mot[i_mot] + target_q = ( + self.low_state.motor_state[i_mot].q + + res_q_ik[i_arm] + + np.clip(0.3 * left_arm_residual[i_arm], + -0.2, 0.2) + ) + target_q = np.clip(target_q, + self.lim_lo_pin[i_pin], + self.lim_hi_pin[i_pin]) + target_dof_pos[i_mot] = target_q + return target_dof_pos + +def main(): + import yaml + + with open('configs/g1_eetrack.yaml', 'r') as fp: + d = yaml.safe_load(fp) + + act_to_dof = ActToDof() + obs = np.load('/tmp/eet4/obs001.npy')[0] + act = np.load('/tmp/eet4/act001.npy')[0] + dof_lab = np.load('/tmp/eet4/dof001.npy')[0] + mot_from_lab = index_map(d['motor_joint'], d['lab_joint']) + target_dof_pos = np.zeros_like(dof_lab) + target_dof_pos[mot_from_lab] = dof_lab + + + dof = act_to_dof(obs, act) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/deploy/deploy_real/common/np_math.py b/deploy/deploy_real/common/np_math.py new file mode 100644 index 0000000..6d43d77 --- /dev/null +++ b/deploy/deploy_real/common/np_math.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 + +import numpy as np +import torch + +def axis_angle_from_quat(quat: np.ndarray, eps: float = 1.0e-6) -> np.ndarray: + """Convert rotations given as quaternions to axis/angle. + + Args: + quat: The quaternion orientation in (w, x, y, z). Shape is (..., 4). + eps: The tolerance for Taylor approximation. Defaults to 1.0e-6. + + Returns: + Rotations given as a vector in axis angle form. Shape is (..., 3). + The vector's magnitude is the angle turned anti-clockwise in radians around the vector's direction. + + Reference: + https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py#L526-L554 + """ + # Modified to take in quat as [q_w, q_x, q_y, q_z] + # Quaternion is [q_w, q_x, q_y, q_z] = [cos(theta/2), n_x * sin(theta/2), n_y * sin(theta/2), n_z * sin(theta/2)] + # Axis-angle is [a_x, a_y, a_z] = [theta * n_x, theta * n_y, theta * n_z] + # Thus, axis-angle is [q_x, q_y, q_z] / (sin(theta/2) / theta) + # When theta = 0, (sin(theta/2) / theta) is undefined + # However, as theta --> 0, we can use the Taylor approximation 1/2 - + # theta^2 / 48 + quat = quat * (1.0 - 2.0 * (quat[..., 0:1] < 0.0)) + mag = np.linalg.norm(quat[..., 1:], axis=-1) + half_angle = np.arctan2(mag, quat[..., 0]) + angle = 2.0 * half_angle + # check whether to apply Taylor approximation + sin_half_angles_over_angles = np.where( + np.abs(angle) > eps, + np.sin(half_angle) / angle, + 0.5 - angle * angle / 48 + ) + return quat[..., 1:4] / sin_half_angles_over_angles[..., None] + + +def quat_from_angle_axis( + angle: torch.Tensor, + axis: torch.Tensor = None) -> torch.Tensor: + """Convert rotations given as angle-axis to quaternions. + + Args: + angle: The angle turned anti-clockwise in radians around the vector's direction. Shape is (N,). + axis: The axis of rotation. Shape is (N, 3). + + Returns: + The quaternion in (w, x, y, z). Shape is (N, 4). + """ + if axis is None: + axa = angle + angle = torch.linalg.norm(axa, dim=-1) + axis = axa / angle[..., None].clamp_min(1e-6) + theta = (angle / 2).unsqueeze(-1) + xyz = normalize(axis) * theta.sin() + w = theta.cos() + return normalize(torch.cat([w, xyz], dim=-1)) + + +def quat_mul(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: + """Multiply two quaternions together. + + Args: + q1: The first quaternion in (w, x, y, z). Shape is (..., 4). + q2: The second quaternion in (w, x, y, z). Shape is (..., 4). + + Returns: + The product of the two quaternions in (w, x, y, z). Shape is (..., 4). + + Raises: + ValueError: Input shapes of ``q1`` and ``q2`` are not matching. + """ + # check input is correct + if q1.shape != q2.shape: + msg = f"Expected input quaternion shape mismatch: {q1.shape} != {q2.shape}." + raise ValueError(msg) + # reshape to (N, 4) for multiplication + shape = q1.shape + q1 = q1.reshape(-1, 4) + q2 = q2.reshape(-1, 4) + # extract components from quaternions + w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3] + w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3] + # perform multiplication + ww = (z1 + x1) * (x2 + y2) + yy = (w1 - y1) * (w2 + z2) + zz = (w1 + y1) * (w2 - z2) + xx = ww + yy + zz + qq = 0.5 * (xx + (z1 - x1) * (x2 - y2)) + w = qq - ww + (z1 - y1) * (y2 - z2) + x = qq - xx + (x1 + w1) * (x2 + w2) + y = qq - yy + (w1 - x1) * (y2 + z2) + z = qq - zz + (z1 + y1) * (w2 - x2) + + return torch.stack([w, x, y, z], dim=-1).view(shape) + + +def quat_rotate(q: np.ndarray, v: np.ndarray) -> np.ndarray: + """Rotate a vector by a quaternion along the last dimension of q and v. + + Args: + q: The quaternion in (w, x, y, z). Shape is (..., 4). + v: The vector in (x, y, z). Shape is (..., 3). + + Returns: + The rotated vector in (x, y, z). Shape is (..., 3). + """ + q_w = q[..., 0] + q_vec = q[..., 1:] + a = v * (2.0 * q_w**2 - 1.0)[..., None] + b = np.cross(q_vec, v, axis=-1) * q_w[..., None] * 2.0 + c = q_vec * np.einsum("...i,...i->...", q_vec, v)[..., None] * 2.0 + return a + b + c + + +def quat_rotate_inverse(q: np.ndarray, v: np.ndarray) -> np.ndarray: + """Rotate a vector by the inverse of a quaternion along the last dimension of q and v. + + Args: + q: The quaternion in (w, x, y, z). Shape is (..., 4). + v: The vector in (x, y, z). Shape is (..., 3). + + Returns: + The rotated vector in (x, y, z). Shape is (..., 3). + """ + q_w = q[..., 0] + q_vec = q[..., 1:] + a = v * (2.0 * q_w**2 - 1.0)[..., None] + b = np.cross(q_vec, v, axis=-1) * q_w[..., None] * 2.0 + c = q_vec * np.einsum("...i,...i->...", q_vec, v)[..., None] * 2.0 + return a - b + c + + +def index_map(k_to, k_from): + """ + Returns an index mapping from k_from to k_to. + + Given k_to=a, k_from=b, + returns an index map "a_from_b" such that + array_a[a_from_b] = array_b + + Missing values are set to -1. + """ + index_dict = {k: i for i, k in enumerate(k_to)} # O(len(k_from)) + return [index_dict.get(k, -1) for k in k_from] # O(len(k_to)) \ No newline at end of file