apply correct transforms
This commit is contained in:
Yoonyoung Cho
parent
da645b8242
commit
4f7510241d
|
|
@ -1,13 +1,27 @@
|
|||
#!/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,
|
||||
config: Config,
|
||||
ikctrl: IKCtrl):
|
||||
self.config = config
|
||||
self.ikctrl = ikctrl
|
||||
self.lim_lo_pin = self.ikctrl.robot.model.lowerPositionLimit
|
||||
|
|
@ -53,8 +67,8 @@ 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]
|
||||
|
||||
|
|
@ -67,26 +81,31 @@ class ActToDof:
|
|||
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('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
|
||||
|
|
@ -129,6 +148,7 @@ class ActToDof:
|
|||
|
||||
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,7 +175,6 @@ 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
|
||||
|
|
@ -201,5 +219,6 @@ def main():
|
|||
ax[i].legend()
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
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
|
||||
world_from_pelvis = self.tf_buffer.lookup_transform(
|
||||
'world',
|
||||
'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]
|
||||
)
|
||||
)
|
||||
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],
|
||||
|
|
@ -691,7 +693,8 @@ class Controller:
|
|||
|
||||
target_dof_pos = self.actmap(
|
||||
self.obs,
|
||||
self.action
|
||||
self.action,
|
||||
root_state_w[3:7]
|
||||
)
|
||||
# print('??',
|
||||
# target_dof_pos,
|
||||
|
|
|
|||
Loading…
Reference in New Issue