72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
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() |