"""Offline right-hand vector retargeting and MuJoCo kinematic replay.""" import os os.environ.setdefault('MUJOCO_GL', 'osmesa') os.environ.setdefault('OMP_NUM_THREADS', '4') import argparse import json import time from pathlib import Path import xml.etree.ElementTree as ET import numpy as np import torch import mujoco import imageio.v2 as imageio from PIL import Image, ImageDraw from dex_retargeting.constants import RobotName, RetargetingType, HandType, get_default_config_path, OPERATOR2MANO_RIGHT from dex_retargeting.retargeting_config import RetargetingConfig ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / 'output/dex_2047635068' ASSETS = ROOT / 'third_party/dex-assets/robots/hands' def operator_frame(joints): # Same palm frame construction as upstream SingleHandDetector, without MediaPipe. points = joints[[0, 5, 9]] x = points[0] - points[2] points = points - points.mean(axis=0) _, _, v = np.linalg.svd(points) normal = v[2] x -= x.dot(normal) * normal x /= np.linalg.norm(x) z = np.cross(x, normal) if z.dot(points[1] - points[2]) < 0: normal *= -1 z *= -1 frame = np.stack([x, normal, z], axis=1) assert np.allclose(frame.T @ frame, np.eye(3), atol=1e-5) return frame @ OPERATOR2MANO_RIGHT def load_scene(urdf): tree = ET.parse(urdf) root = tree.getroot() ext = ET.SubElement(root, 'mujoco') ET.SubElement(ext, 'compiler', discardvisual='false', fusestatic='false', strippath='false') for mesh in root.iter('mesh'): mesh.set('filename', str(urdf.parent / mesh.get('filename'))) adapted = OUT / 'leap_replay.urdf' tree.write(adapted) model = mujoco.MjModel.from_xml_path(str(adapted)) scene = OUT / 'scene.xml' mujoco.mj_saveLastXML(str(scene), model) tree = ET.parse(scene) root = tree.getroot() visual = ET.SubElement(root, 'visual') ET.SubElement(visual, 'global', offwidth='1280', offheight='480') ET.SubElement(visual, 'headlight', ambient='0.5 0.5 0.5', diffuse='0.7 0.7 0.7') ET.SubElement(root.find('worldbody'), 'light', pos='0 -1 1', dir='0 1 -1', directional='true') tree.write(scene) model = mujoco.MjModel.from_xml_path(str(scene)) # Replay is kinematic. Retain only visual geometry in the viewer. model.geom_rgba[:, 3] = 1 # LEAP fingertip pads are supplied only as collision meshes in this asset. # Display these meshes too, while keeping collision primitives hidden. model.geom_group[model.geom_type == mujoco.mjtGeom.mjGEOM_MESH] = 1 return model def main(): parser = argparse.ArgumentParser() parser.add_argument('--frames', type=int, default=0) parser.add_argument('--no-render', action='store_true') args = parser.parse_args() torch.set_num_threads(4) human = np.load(OUT / 'human_joints.npz') joints = human['joints'][:args.frames or None] RetargetingConfig.set_default_urdf_dir(ASSETS) config_path = get_default_config_path(RobotName.leap, RetargetingType.vector, HandType.right) cfg = RetargetingConfig.load_from_file(config_path) (OUT / 'retargeting_config.yml').write_text(Path(config_path).read_text()) retarget = cfg.build() robot = retarget.optimizer.robot index = retarget.optimizer.target_link_human_indices tip_ids = [robot.get_link_index(x) for x in cfg.target_task_link_names] base_id = robot.get_link_index('base') qs, targets, actual, codes = [], [], [], [] start = time.monotonic() for i, j in enumerate(joints): local = (j - j[:1]) @ operator_frame(j) ref = local[index[1]] - local[index[0]] q = retarget.retarget(ref) robot.compute_forward_kinematics(q) base = robot.get_link_pose(base_id)[:3, 3] tips = np.array([robot.get_link_pose(k)[:3, 3] for k in tip_ids]) qs.append(q); targets.append(ref * cfg.scaling_factor + base); actual.append(tips) codes.append(retarget.optimizer.opt.last_optimize_result()) if i % 200 == 0: print(f'retarget {i}/{len(joints)}', flush=True) qpos, targets, actual = map(np.array, (qs, targets, actual)) elapsed = time.monotonic() - start limits = robot.joint_limits assert np.isfinite(qpos).all() assert np.all(qpos >= limits[:, 0] - 1e-6) and np.all(qpos <= limits[:, 1] + 1e-6) np.savez_compressed(OUT / 'leap_right_qpos.npz', qpos=qpos, joint_names=retarget.joint_names, fps=30, time=np.arange(len(qpos))/30, target_tips=targets, actual_tips=actual, joint_limits=limits, source_track=human['source_track']) np.savetxt(OUT / 'leap_right_qpos.csv', np.c_[np.arange(len(qpos))/30, qpos], delimiter=',', header=','.join(['time_s'] + retarget.joint_names), comments='') errors = np.linalg.norm(actual-targets, axis=2)*1000 stats = dict(frames=len(qpos), fps=30, dof=qpos.shape[1], finite=True, joint_limits_pass=True, retarget_seconds=elapsed, tip_error_mean_mm=float(errors.mean()), tip_error_p95_mm=float(np.percentile(errors,95)), tip_error_by_finger_mm=dict(zip(['thumb','index','middle','ring'],errors.mean(0).tolist())), max_frame_joint_step_rad=float(np.max(np.abs(np.diff(qpos,axis=0)))), optimizer_status_counts={str(k):int(v) for k,v in zip(*np.unique(codes,return_counts=True))}, scope='Fixed-wrist kinematic replay; no object, contact, dynamics, or real-robot validation.') (OUT / 'validation.json').write_text(json.dumps(stats, indent=2)) print(json.dumps(stats, indent=2), flush=True) if args.no_render: return model = load_scene(ASSETS / 'leap_hand/leap_hand_right.urdf') data = mujoco.MjData(model) addresses = [model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in retarget.joint_names] assert len(set(addresses)) == len(addresses) == model.nq tip_bodies = [mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,n) for n in cfg.target_task_link_names] assert min(tip_bodies) >= 0 fk_errors = [] for i in np.linspace(0,len(qpos)-1,min(20,len(qpos)),dtype=int): data.qpos[addresses] = qpos[i] mujoco.mj_forward(model,data) fk_errors.append(np.max(np.abs(data.xpos[tip_bodies]-actual[i]))) assert max(fk_errors) < 1e-5, f'Pinocchio/MuJoCo FK mismatch: {max(fk_errors)}' stats['pinocchio_mujoco_max_fk_error_m'] = float(max(fk_errors)) renderer = mujoco.Renderer(model, height=480, width=640) cameras=[] for azimuth in [135, 225]: cam=mujoco.MjvCamera(); cam.lookat[:]=model.stat.center;cam.distance=.65;cam.azimuth=azimuth;cam.elevation=15 cameras.append(cam) option=mujoco.MjvOption() option.geomgroup[0]=0 # URDF collision geometry; show visual group 1 only. writer=imageio.get_writer(OUT/'leap_right_replay.mp4',fps=30,codec='libx264',quality=8) try: for i,q in enumerate(qpos): data.qpos[addresses]=q; mujoco.mj_forward(model,data) panels=[] for cam in cameras: renderer.update_scene(data,camera=cam,scene_option=option) for p in targets[i]: scene=renderer.scene mujoco.mjv_initGeom(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_SPHERE, np.array([.004]*3),p,np.eye(3).ravel(),np.array([1,.25,.1,1])) scene.ngeom+=1 panels.append(renderer.render().copy()) img=Image.fromarray(np.concatenate(panels,axis=1)); draw=ImageDraw.Draw(img) draw.rectangle([0,0,1280,46],fill=(20,25,32)) draw.text((12,8),f'LEAP right | frame {i+1}/{len(qpos)} | red dots: scaled human fingertips | fixed wrist',fill='white') draw.text((12,26),f'Mean fingertip error: {errors[i].mean():.1f} mm | kinematic replay',fill='white') writer.append_data(np.asarray(img)) if i in [0,len(qpos)//2,len(qpos)-1]: img.save(OUT/f'preview_{i:04d}.png') if i % 200==0: print(f'render {i}/{len(qpos)}',flush=True) finally: writer.close(); renderer.close() stats['rendered_frames']=len(qpos) (OUT/'validation.json').write_text(json.dumps(stats,indent=2)) print('COMPLETE',flush=True) if __name__ == '__main__': main()