apply correct transforms

deploy/deploy_real/act_to_dof.py

@@ -1,14 +1,28 @@
+#!/usr/bin/env python3
+
 import numpy as np
-import torch
 import pinocchio as pin
-from common.np_math import (quat_from_angle_axis, quat_mul,
+
-                            xyzw2wxyz, index_map)
+from common.np_math import (xyzw2wxyz, index_map)
+from math_utils import (
+    as_np,
+    quat_from_angle_axis,
+    quat_mul,
+    quat_rotate_inverse
+)
+quat_from_angle_axis = as_np(quat_from_angle_axis)
+quat_mul = as_np(quat_mul)
+quat_rotate_inverse = as_np(quat_rotate_inverse)
+
 from config import Config
 from ikctrl import IKCtrl
 
+
 class ActToDof:
-    def __init__(self, config, ikctrl):
+    def __init__(self,
-        self.config=config
+                 config: Config,
+                 ikctrl: IKCtrl):
+        self.config = config
         self.ikctrl = ikctrl
         self.lim_lo_pin = self.ikctrl.robot.model.lowerPositionLimit
         self.lim_hi_pin = self.ikctrl.robot.model.upperPositionLimit
@@ -53,11 +67,11 @@ class ActToDof:
             np.asarray(self.config.lab_joint_offsets)[self.lab_from_nonarm]
         )
 
-    def __call__(self, obs, action):
+    def __call__(self, obs, action, root_quat_wxyz=None):
-        hands_command = obs[..., 119:125]
+        hands_command_w = obs[..., 119:125]
         non_arm_joint_pos = action[..., :22]
         left_arm_residual = action[..., 22:29]
-        
+
         q_lab = obs[..., 32:61]
         q_pin = np.zeros_like(self.ikctrl.cfg.q)
         q_pin[self.pin_from_lab] = q_lab
@@ -66,27 +80,32 @@ class ActToDof:
         q_mot = np.zeros(29)
         q_mot[self.mot_from_lab] = q_lab
         q_mot[self.mot_from_lab] += np.asarray(self.config.lab_joint_offsets)
-        #print('q_mot (inside)', q_mot)
+        # print('q_mot (inside)', q_mot)
-
         # print('q0', q_mot[self.mot_from_arm])
         # q_mot[i_mot] = q_lab[ lab_from_mot[i_mot] ]
 
+        if root_quat_wxyz is None:
+            hands_command_b = hands_command_w
+        else:
+            world_from_pelvis_quat = root_quat_wxyz
+            hands_command_b = np.concatenate(
+                quat_rotate_inverse(world_from_pelvis_quat,
+                                    hands_command_w[..., :3]),
+                quat_rotate_inverse(world_from_pelvis_quat,
+                                    hands_command_w[..., 3:6])
+            )
+
         d_quat = quat_from_angle_axis(
-            torch.from_numpy(hands_command[..., 3:])
+            hands_command_b[..., 3:]
         ).detach().cpu().numpy()
-        # print('d_quat', d_quat)
 
         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_xyz = source_xyz + hands_command_b[..., :3]
-        target_quat = quat_mul(
+        target_quat = quat_mul(d_quat, source_quat)
-            torch.from_numpy(d_quat),
-            torch.from_numpy(source_quat)).detach().cpu().numpy()
         target = np.concatenate([target_xyz, target_quat])
-        # print('target', target)
-        # print('q_pin', q_pin)
         res_q_ik = self.ikctrl(
             q_pin,
             target
@@ -111,8 +130,8 @@ class ActToDof:
 
         if True:
             target_dof_pos[self.mot_from_arm] += np.clip(
-                    0.3 * left_arm_residual,
+                0.3 * left_arm_residual,
-                    -0.2, 0.2)
+                -0.2, 0.2)
 
             # print('default joint pos', self.default_nonarm)
             # print('joint order', self.config.non_arm_joint)
@@ -122,13 +141,14 @@ class ActToDof:
             )
 
             target_dof_pos[self.mot_from_arm] = np.clip(
-                    target_dof_pos[self.mot_from_arm],
+                target_dof_pos[self.mot_from_arm],
-                    self.lim_lo_pin[self.pin_from_arm],
+                self.lim_lo_pin[self.pin_from_arm],
-                    self.lim_hi_pin[self.pin_from_arm]
+                self.lim_hi_pin[self.pin_from_arm]
             )
 
         return target_dof_pos
 
+
 def main():
     from matplotlib import pyplot as plt
     import yaml
@@ -138,8 +158,7 @@ def main():
     config = Config('configs/g1_eetrack.yaml')
     ikctrl = IKCtrl(
         '../../resources/robots/g1_description/g1_29dof_with_hand_rev_1_0.urdf',
-        config.arm_joint
+        config.arm_joint)
-    )
     act_to_dof = ActToDof(config, ikctrl)
 
     exps = []
@@ -156,34 +175,33 @@ def main():
         target_dof_pos = np.zeros_like(dof_lab)
         target_dof_pos[mot_from_lab] = dof_lab
 
-
         dof = act_to_dof(obs, act)
 
         export = target_dof_pos
-        calc   = dof
+        calc = dof
 
         mot_from_arm = index_map(d['motor_joint'], d['arm_joint'])
         # print('exported', target_dof_pos[mot_from_arm],
         #       'calculated', dof[mot_from_arm])
         # print( (export - calc)[mot_from_arm] )
-        exps.append( target_dof_pos[mot_from_arm] )
+        exps.append(target_dof_pos[mot_from_arm])
-        cals.append( dof[mot_from_arm] )
+        cals.append(dof[mot_from_arm])
 
         q_lab = obs[..., 32:61]
         q_mot = q_lab[act_to_dof.lab_from_mot]
         q_arm = q_mot[mot_from_arm]
         poss.append(q_arm)
 
-    exps=np.asarray(exps, dtype=np.float32)
+    exps = np.asarray(exps, dtype=np.float32)
-    cals=np.asarray(cals, dtype=np.float32)
+    cals = np.asarray(cals, dtype=np.float32)
-    poss=np.asarray(poss, dtype=np.float32)
+    poss = np.asarray(poss, dtype=np.float32)
     print(exps.shape)
     print(cals.shape)
 
-    fig, ax=plt.subplots(7,1)
+    fig, ax = plt.subplots(7, 1)
-    
+
-    q_lo =act_to_dof.lim_lo_pin[act_to_dof.pin_from_arm]
+    q_lo = act_to_dof.lim_lo_pin[act_to_dof.pin_from_arm]
-    q_hi =act_to_dof.lim_hi_pin[act_to_dof.pin_from_arm]
+    q_hi = act_to_dof.lim_hi_pin[act_to_dof.pin_from_arm]
     RES = True
     for i in range(7):
         if RES:
@@ -192,8 +210,8 @@ def main():
             ax[i].axhline(q_lo[i], color='k', linestyle='--')
             ax[i].axhline(q_hi[i], color='k', linestyle='--')
         if RES:
-            ax[i].plot(exps[:, i]-poss[:,i], label='sim')
+            ax[i].plot(exps[:, i] - poss[:, i], label='sim')
-            ax[i].plot(cals[:, i]-poss[:,i], label='real')
+            ax[i].plot(cals[:, i] - poss[:, i], label='real')
         else:
             ax[i].plot(poss[:, i], label='pos')
             ax[i].plot(exps[:, i], label='sim')
@@ -201,5 +219,6 @@ def main():
         ax[i].legend()
     plt.show()
 
+
 if __name__ == '__main__':
     main()

deploy/deploy_real/deploy_real_ros_eetrack.py

@@ -35,6 +35,7 @@ class Mode(Enum):
     policy = 4
     null = 5
 
+
 axis_angle_from_quat = math_utils.as_np(math_utils.axis_angle_from_quat)
 quat_conjugate = math_utils.as_np(math_utils.quat_conjugate)
 quat_mul = math_utils.as_np(math_utils.quat_mul)
@@ -42,6 +43,7 @@ quat_rotate = math_utils.as_np(math_utils.quat_rotate)
 quat_rotate_inverse = math_utils.as_np(math_utils.quat_rotate_inverse)
 wrap_to_pi = math_utils.as_np(math_utils.wrap_to_pi)
 
+
 def body_pose(
         tf_buffer,
         frame: str,
@@ -274,7 +276,19 @@ class Observation:
         # base_lin_vel = np.zeros(3)
         ang_vel = np.array([low_state.imu_state.gyroscope],
                            dtype=np.float32)
-        quat = low_state.imu_state.quaternion
+
+        if True:
+            # NOTE(ycho): requires running `fake_world_tf_pub.py`.
+            world_from_pelvis = self.tf_buffer.lookup_transform(
+                'world',
+                'pelvis',
+                rp.time.Time()
+            )
+            rxn = world_from_pelvis.transform.rotation
+            quat = np.array([rxn.w, rxn.x, rxn.y, rxn.z])
+        else:
+            quat = low_state.imu_state.quaternion
+
         if self.config.imu_type == "torso":
             waist_yaw = low_state.motor_state[self.config.arm_waist_joint2motor_idx[0]].q
             waist_yaw_omega = low_state.motor_state[self.config.arm_waist_joint2motor_idx[0]].dq
@@ -283,6 +297,10 @@ class Observation:
                 waist_yaw_omega=waist_yaw_omega,
                 imu_quat=quat,
                 imu_omega=ang_vel)
+
+        # NOTE(ycho): `ang_vel` is _probably_ in the pelvis frame,
+        # since otherwise transform_imu_data() would be unnecessary for
+        # `ang_vel`.
         base_ang_vel = ang_vel.squeeze(0)
 
         # TODO(ycho): check if the convention "q_base^{-1} @ g" holds.
@@ -333,28 +351,13 @@ class Observation:
             "left_wrist_roll_link",
             "left_wrist_yaw_link"
         ])
-        if True:  # hack
+        world_from_pelvis = self.tf_buffer.lookup_transform(
-            lf_from_pelvis = self.tf_buffer.lookup_transform(
+            'world',
-                'left_ankle_roll_link',  # to
+            'pelvis',
-                'pelvis',
+            rp.time.Time()
-                rp.time.Time()
+        )
-            )
+        pelvis_height = [world_from_pelvis.translation[2]]
-            rf_from_pelvis = self.tf_buffer.lookup_transform(
+
-                'right_ankle_roll_link',  # to
-                'pelvis',
-                rp.time.Time()
-            )
-            # NOTE(ycho): we assume at least one of the feet is on the ground
-            #            and use the higher of the two as the pelvis height.
-            pelvis_height = max(lf_from_pelvis.transform.translation.z,
-                                rf_from_pelvis.transform.translation.z)
-            pelvis_height = [pelvis_height]
-        else:
-            pelvis_height = np.abs(np.dot(
-                projected_gravity,  # world frame
-                fp_l[0]
-            )
-            )
         obs = [
             base_ang_vel,
             projected_gravity,
@@ -468,7 +471,12 @@ class Controller:
         elif config.msg_type == "go":
             init_cmd_go(self.low_cmd, weak_motor=self.config.weak_motor)
 
-        self.mode = Mode.wait
+        # NOTE(ycho):
+        # if running from real robot:
+        # self.mode = Mode.wait
+        # if running from rosbag:
+        self.mode = Mode.policy
+
         self._mode_change = True
         self._timer = self._node.create_timer(
             self.config.control_dt, self.run_wrapper)
@@ -602,25 +610,19 @@ class Controller:
 
         # FIXME(ycho): implement `_hands_command_`
         # to use the output of `eetrack`.
-
         if True:
+            # NOTE(ycho): requires running `fake_world_tf_pub.py`.
+            world_from_pelvis = self.tf_buffer.lookup_transform(
+                'world',
+                'pelvis',
+                rp.time.Time()
+            )
+            xyz = world_from_pelvis.transform.translation
+            rxn = world_from_pelvis.transform.rotation
+            quat_wxyz = np.array([rxn.w, rxn.x, rxn.y, rxn.z])
             root_state_w = np.zeros(7)
-            lf_from_pelvis = self.tf_buffer.lookup_transform(
+            root_state_w[0:3] = xyz
-                'left_ankle_roll_link',  # to
+            root_state_w[3:7] = quat_wxyz
-                'pelvis',
-                rp.time.Time()
-            )
-            rf_from_pelvis = self.tf_buffer.lookup_transform(
-                'right_ankle_roll_link',  # to
-                'pelvis',
-                rp.time.Time()
-            )
-            # NOTE(ycho): we assume at least one of the feet is on the ground
-            #            and use the higher of the two as the pelvis height.
-            pelvis_height = max(lf_from_pelvis.transform.translation.z,
-                                rf_from_pelvis.transform.translation.z)
-            root_state_w[0:3] = [0, 0, pelvis_height]
-            root_state_w[3:7] = self.low_state.imu_state.quaternion
 
         if self.eetrack is None:
             self.eetrack = eetrack(torch.from_numpy(root_state_w)[None],
@@ -675,23 +677,24 @@ class Controller:
 
         # hands_command = obs[..., 119:125]
         if False:
-        world_from_body_quat = math_utils.quat_from_euler_xyz(
+            world_from_body_quat = math_utils.quat_from_euler_xyz(
-            th.as_tensor([0], dtype=th.float32),
+                th.as_tensor([0], dtype=th.float32),
-            th.as_tensor([0], dtype=th.float32),
+                th.as_tensor([0], dtype=th.float32),
-            th.as_tensor([1.57], dtype=th.float32)).reshape(4)
+                th.as_tensor([1.57], dtype=th.float32)).reshape(4)
-        obs_tensor[..., 119:122] = math_utils.quat_rotate(
+            obs_tensor[..., 119:122] = math_utils.quat_rotate(
-            world_from_body_quat[None],
+                world_from_body_quat[None],
-            obs_tensor[..., 119:122])
+                obs_tensor[..., 119:122])
-        obs_tensor[..., 122:125] = math_utils.quat_rotate(
+            obs_tensor[..., 122:125] = math_utils.quat_rotate(
-            world_from_body_quat[None],
+                world_from_body_quat[None],
-            obs_tensor[..., 122:125])
+                obs_tensor[..., 122:125])
         self.action = self.policy(obs_tensor).detach().numpy().squeeze()
         # np.save(F'{logpath}/act{self.counter:03d}.npy',
         #         self.action)
 
         target_dof_pos = self.actmap(
             self.obs,
-            self.action
+            self.action,
+            root_state_w[3:7]
         )
         # print('??',
         #         target_dof_pos,