166 lines
8.2 KiB
Python
166 lines
8.2 KiB
Python
"""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()
|