128 lines
7.5 KiB
Python
128 lines
7.5 KiB
Python
"""Retarget Dyn-HaMR to the supplied L20 with JSON coupling and moving wrist."""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
os.environ.setdefault('OMP_NUM_THREADS','4')
|
|
import argparse,json,time
|
|
from pathlib import Path
|
|
import xml.etree.ElementTree as ET
|
|
import numpy as np
|
|
import torch,mujoco
|
|
from scipy.spatial.transform import Rotation
|
|
from scipy.ndimage import gaussian_filter1d
|
|
from dex_retargeting.retargeting_config import RetargetingConfig
|
|
from l20_calibrated import ROOT,Calibration,CalibratedAdaptor,build_assets,palm_basis,CHAINS,TIPS
|
|
|
|
OUT=ROOT/'output/l20_2047635068'
|
|
|
|
def main():
|
|
global OUT
|
|
parser=argparse.ArgumentParser();parser.add_argument('--frames',type=int,default=0)
|
|
parser.add_argument('--output-dir',type=Path,default=OUT)
|
|
parser.add_argument('--consistent-objective',action=argparse.BooleanOptionalAction,default=True)
|
|
args=parser.parse_args();OUT=args.output_dir.resolve()
|
|
torch.set_num_threads(4)
|
|
fixed,urdf=build_assets(OUT)
|
|
cal=Calibration()
|
|
config=dict(type='position',urdf_path=str(urdf),target_joint_names=cal.active,
|
|
target_link_names=[f'landmark_{i:02d}' for i in range(21)],
|
|
target_link_human_indices=list(range(21)),low_pass_alpha=-1,normal_delta=.002)
|
|
retarget=RetargetingConfig.from_dict(config.copy()).build()
|
|
if args.consistent_objective:
|
|
from l20_consistent_optimizer import ConsistentPositionOptimizer
|
|
retarget.optimizer.__class__=ConsistentPositionOptimizer
|
|
robot=retarget.optimizer.robot
|
|
adaptor=CalibratedAdaptor(robot,cal);retarget.optimizer.set_kinematic_adaptor(adaptor)
|
|
bounds=np.array([cal.bounds[n] for n in cal.active]);retarget.joint_limits=bounds
|
|
retarget.optimizer.set_joint_limit(bounds,epsilon=0)
|
|
neutral_q=np.zeros(robot.dof)
|
|
for i,n in enumerate(cal.active):neutral_q[robot.get_joint_index(n)]=np.clip(cal.values[n][cal.raw['baseline_command_u8'][cal.tables[n]['motor_index']]],*bounds[i])
|
|
neutral_q=adaptor.forward_qpos(neutral_q)
|
|
retarget.set_qpos(neutral_q)
|
|
sites=[robot.get_link_index(f'landmark_{i:02d}') for i in range(21)]
|
|
def fk(q):
|
|
robot.compute_forward_kinematics(q)
|
|
return np.array([robot.get_link_pose(k)[:3,3] for k in sites])
|
|
neutral=fk(neutral_q);basis_robot=palm_basis(neutral)
|
|
human=np.load(ROOT/'output/dex_2047635068/human_joints.npz')
|
|
joints=human['joints'][:args.frames or None]
|
|
targets=[];wrist_rot=[]
|
|
for t,p in enumerate(joints):
|
|
basis=palm_basis(p)
|
|
align=basis@basis_robot.T
|
|
aligned=(p-p[0])@align
|
|
target=neutral.copy()
|
|
for chain in CHAINS:
|
|
for a,b in zip(chain[:-1],chain[1:]):
|
|
direction=aligned[b]-aligned[a]
|
|
target[b]=target[a]+direction/np.linalg.norm(direction)*np.linalg.norm(neutral[b]-neutral[a])
|
|
targets.append(target)
|
|
wrist_rot.append(Rotation.from_rotvec(human['root_orient'][t]).as_matrix()@align)
|
|
targets=np.array(targets);wrist_rot=np.array(wrist_rot)
|
|
qraw=[];codes=[];start=time.monotonic()
|
|
for t,target in enumerate(targets):
|
|
q=retarget.retarget(target)
|
|
qraw.append(q.copy());codes.append(retarget.optimizer.opt.last_optimize_result())
|
|
if t%200==0: print(f'L20 dex {t}/{len(targets)}',flush=True)
|
|
qraw=np.array(qraw)
|
|
active=qraw[:,adaptor.idx_pin2target]
|
|
active=np.clip(gaussian_filter1d(active,1,axis=0),bounds[:,0],bounds[:,1])
|
|
qpos=[]
|
|
for a in active:
|
|
q=np.zeros(robot.dof);q[adaptor.idx_pin2target]=a;qpos.append(adaptor.forward_qpos(q))
|
|
qpos=np.array(qpos);actual=np.array([fk(q) for q in qpos])
|
|
commands=np.array([cal.command(dict(zip(robot.dof_joint_names,q))) for q in qpos])
|
|
decoded=np.array([cal.decode(c,robot.dof_joint_names) for c in commands])
|
|
wrist_world=human['wrist_world'][:len(qpos)]
|
|
# A single rigid transform makes the initial palm upright for visualization.
|
|
scene_rot=wrist_rot[0].T
|
|
wrist_pos=(wrist_world-wrist_world[0])@scene_rot.T+np.array([0,0,.2])
|
|
wrist_scene=np.einsum('ij,tjk->tik',scene_rot,wrist_rot)
|
|
quat=Rotation.from_matrix(wrist_scene).as_quat()[:,[3,0,1,2]]
|
|
for i in range(1,len(quat)):
|
|
if quat[i]@quat[i-1]<0:quat[i]*=-1
|
|
np.savez_compressed(OUT/'motion.npz',qpos=qpos,qpos_raw=qraw,joint_names=robot.dof_joint_names,
|
|
active_qpos=active,active_joint_names=cal.active,command_u8=commands,decoded_qpos=decoded,
|
|
wrist_pos=wrist_pos,wrist_quat_wxyz=quat,wrist_world=wrist_world,wrist_world_R=wrist_rot,
|
|
scene_rotation=scene_rot,scene_translation=np.array([0,0,.2])-scene_rot@wrist_world[0],
|
|
targets=targets,actual=actual,fps=30,time=np.arange(len(qpos))/30)
|
|
np.savetxt(OUT/'trajectory.csv',np.c_[np.arange(len(qpos))/30,wrist_pos,quat,qpos],delimiter=',',
|
|
header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz']+robot.dof_joint_names),comments='')
|
|
np.savetxt(OUT/'commands_u8.csv',np.c_[np.arange(len(qpos)),commands],delimiter=',',fmt='%d',
|
|
header=','.join(['frame']+[f'slot_{i}' for i in range(20)]),comments='')
|
|
limits=np.array([cal.urdf_limits[n] for n in robot.dof_joint_names])
|
|
violation=float(np.maximum(limits[:,0]-qpos,0).max()+np.maximum(qpos-limits[:,1],0).max())
|
|
assert violation<1e-6 and np.isfinite(qpos).all()
|
|
err=np.linalg.norm(actual[:,TIPS]-targets[:,TIPS],axis=-1)*1000
|
|
# Compare independent FK libraries, with all 21 nonlinear-coupled angles supplied.
|
|
model=mujoco.MjModel.from_xml_path(str(fixed));data=mujoco.MjData(model)
|
|
addresses=[model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in robot.dof_joint_names]
|
|
ids=[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_SITE,f'landmark_{i:02d}') for i in range(21)]
|
|
fk_error=0
|
|
for t in np.linspace(0,len(qpos)-1,min(20,len(qpos)),dtype=int):
|
|
data.qpos[addresses]=qpos[t];mujoco.mj_forward(model,data)
|
|
fk_error=max(fk_error,float(np.abs(data.site_xpos[ids]-actual[t]).max()))
|
|
assert fk_error<1e-5
|
|
report=dict(frames=len(qpos),fps=30,independent_joints=16,urdf_joints=21,
|
|
calibration_model=cal.raw['model'],source_directory_name='L20',
|
|
calibration_interpretation='q_urdf=angle_rad[command]+urdf_zero_offset_rad; passive joints share motor command; not hardware verified',
|
|
reserved_command_slots=[11,12,13,14],tip_error_mean_mm=float(err.mean()),tip_error_p95_mm=float(np.percentile(err,95)),
|
|
tip_error_finger_mm=dict(zip(['thumb','index','middle','ring','pinky'],err.mean(0).tolist())),
|
|
wrist_translation_range_m=np.ptp(wrist_pos,axis=0).tolist(),
|
|
wrist_rotation_max_from_initial_rad=float(np.linalg.norm(Rotation.from_matrix(wrist_scene).as_rotvec(),axis=1).max()),
|
|
urdf_limit_violation_rad=violation,quantization_max_error_rad=float(np.abs(decoded-qpos).max()),
|
|
pinocchio_mujoco_fk_error_m=fk_error,seconds=time.monotonic()-start,
|
|
optimizer_status_counts={str(k):int(v) for k,v in zip(*np.unique(codes,return_counts=True))},
|
|
consistent_objective=args.consistent_objective,
|
|
object_contact_stage='not applied in this hand-only retargeting run')
|
|
(OUT/'retarget_validation.json').write_text(json.dumps(report,indent=2))
|
|
(OUT/'dex_config.json').write_text(json.dumps(config,indent=2))
|
|
# Moving wrist replay model; contact scene construction is a separate stage.
|
|
tree=ET.parse(fixed);root=tree.getroot()
|
|
body=root.find("worldbody/body[@name='hand_base_link']")
|
|
ET.SubElement(body,'freejoint',name='wrist_free')
|
|
root.find('option').set('gravity','0 0 0')
|
|
tree.write(OUT/'l20_moving.xml')
|
|
print(json.dumps(report,indent=2),flush=True)
|
|
|
|
if __name__=='__main__':main()
|