diff --git a/deploy/deploy_real/act_to_dof.py b/deploy/deploy_real/act_to_dof.py index 47e6c60..7afa0f3 100644 --- a/deploy/deploy_real/act_to_dof.py +++ b/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): - self.config=config + def __init__(self, + 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): - hands_command = obs[..., 119:125] + def __call__(self, obs, action, root_quat_wxyz=None): + 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_quat = quat_mul( - torch.from_numpy(d_quat), - torch.from_numpy(source_quat)).detach().cpu().numpy() + target_xyz = source_xyz + hands_command_b[..., :3] + target_quat = quat_mul(d_quat, source_quat) 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.2, 0.2) + 0.3 * left_arm_residual, + -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], - self.lim_lo_pin[self.pin_from_arm], - self.lim_hi_pin[self.pin_from_arm] + target_dof_pos[self.mot_from_arm], + self.lim_lo_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] ) - cals.append( dof[mot_from_arm] ) + exps.append(target_dof_pos[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) - cals=np.asarray(cals, dtype=np.float32) - poss=np.asarray(poss, dtype=np.float32) + exps = np.asarray(exps, dtype=np.float32) + cals = np.asarray(cals, dtype=np.float32) + poss = np.asarray(poss, dtype=np.float32) print(exps.shape) print(cals.shape) - fig, ax=plt.subplots(7,1) - - 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] + fig, ax = plt.subplots(7, 1) + + 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] 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(cals[:, i]-poss[:,i], label='real') + ax[i].plot(exps[:, i] - poss[:, i], label='sim') + 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() diff --git a/deploy/deploy_real/deploy_real_ros_eetrack.py b/deploy/deploy_real/deploy_real_ros_eetrack.py index 04cc9de..e5cf9c4 100644 --- a/deploy/deploy_real/deploy_real_ros_eetrack.py +++ b/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 - lf_from_pelvis = self.tf_buffer.lookup_transform( - 'left_ankle_roll_link', # to - '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) - pelvis_height = [pelvis_height] - else: - pelvis_height = np.abs(np.dot( - projected_gravity, # world frame - fp_l[0] - ) - ) + world_from_pelvis = self.tf_buffer.lookup_transform( + 'world', + 'pelvis', + rp.time.Time() + ) + pelvis_height = [world_from_pelvis.translation[2]] + 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( - 'left_ankle_roll_link', # to - '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 + root_state_w[0:3] = xyz + root_state_w[3:7] = quat_wxyz 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( - th.as_tensor([0], dtype=th.float32), - th.as_tensor([0], dtype=th.float32), - th.as_tensor([1.57], dtype=th.float32)).reshape(4) - obs_tensor[..., 119:122] = math_utils.quat_rotate( - world_from_body_quat[None], - obs_tensor[..., 119:122]) - obs_tensor[..., 122:125] = math_utils.quat_rotate( - world_from_body_quat[None], - obs_tensor[..., 122:125]) + 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([1.57], dtype=th.float32)).reshape(4) + obs_tensor[..., 119:122] = math_utils.quat_rotate( + world_from_body_quat[None], + obs_tensor[..., 119:122]) + obs_tensor[..., 122:125] = math_utils.quat_rotate( + world_from_body_quat[None], + 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,