map test

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()

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))