hand-motion-pipeline/scripts/retarget_l20_video.py

"""Retarget exported video joints to L20 using original URDF linear coupling."""
import os
os.environ.setdefault('MUJOCO_GL', 'osmesa')
import argparse
import copy
import json
from pathlib import Path
import xml.etree.ElementTree as ET

import mujoco
import numpy as np
import torch
from scipy.ndimage import gaussian_filter1d
from scipy.spatial.transform import Rotation
from dex_retargeting.retargeting_config import RetargetingConfig
from dex_retargeting.kinematics_adaptor import MimicJointKinematicAdaptor
from l20_calibrated import ASSET, build_assets, palm_basis, CHAINS, TIPS
from l20_consistent_optimizer import ConsistentPositionOptimizer


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--input', type=Path, required=True)
    p.add_argument('--output-dir', type=Path, required=True)
    p.add_argument('--sigma', type=float, default=2.5)
    p.add_argument('--side', choices=['right','left'], default='right')
    a = p.parse_args()
    torch.set_num_threads(4)
    out = a.output_dir.resolve()
    fixed, urdf = build_assets(out, a.side)
    asset = ASSET.parent/a.side.upper()
    original = ET.parse(asset / f'linkerhand_g20_{a.side}.urdf').getroot()
    joints = [j for j in original.findall('joint') if j.get('type') != 'fixed']
    active = [j.get('name') for j in joints if j.find('mimic') is None]
    limits = {j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in joints}
    bounds_by_name = {n: limits[n].copy() for n in active}
    mimic = {}
    for j in joints:
        m = j.find('mimic')
        if m is None:
            continue
        parent = m.get('joint')
        scale, offset = float(m.get('multiplier', '1')), float(m.get('offset', '0'))
        assert parent in active and scale != 0
        mimic[j.get('name')] = (parent, scale, offset)
        lo, hi = np.sort((limits[j.get('name')] - offset) / scale)
        bounds_by_name[parent] = np.array([max(bounds_by_name[parent][0], lo), min(bounds_by_name[parent][1], hi)])
    bounds = np.array([bounds_by_name[n] for n in active])
    config = dict(type='position', urdf_path=str(urdf), target_joint_names=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(copy.deepcopy(config)).build()
    retarget.optimizer.__class__ = ConsistentPositionOptimizer
    robot = retarget.optimizer.robot
    adaptor = MimicJointKinematicAdaptor(robot, active, [m[0] for m in mimic.values()],
                                        list(mimic), [m[1] for m in mimic.values()], [m[2] for m in mimic.values()])
    retarget.optimizer.set_kinematic_adaptor(adaptor)
    retarget.joint_limits = bounds
    retarget.optimizer.set_joint_limit(bounds, epsilon=0)
    neutral_q = np.zeros(robot.dof)
    neutral_q[adaptor.idx_pin2target] = np.clip(np.zeros(len(active)), bounds[:, 0], bounds[:, 1])
    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(i)[:3, 3] for i in sites])

    neutral = fk(neutral_q)
    robot_basis = palm_basis(neutral)
    human = np.load(a.input)
    hp = human['joints']
    count, fps = len(hp), float(human['fps'])
    assert hp.shape == (count, 21, 3) and np.isfinite(hp).all() and fps > 0
    targets, world_R, qraw, codes = [], [], [], []
    for t, pts in enumerate(hp):
        align = palm_basis(pts) @ robot_basis.T
        aligned = (pts - pts[0]) @ align
        target = neutral.copy()
        for chain in CHAINS:
            for i, j in zip(chain[:-1], chain[1:]):
                d = aligned[j] - aligned[i]
                assert np.linalg.norm(d) > 1e-8
                target[j] = target[i] + d / np.linalg.norm(d) * np.linalg.norm(neutral[j] - neutral[i])
        targets.append(target)
        world_R.append(Rotation.from_rotvec(human['root_orient'][t]).as_matrix() @ align)
        qraw.append(retarget.retarget(target).copy())
        codes.append(retarget.optimizer.opt.last_optimize_result())
    qraw, targets, world_R = np.array(qraw), np.array(targets), np.array(world_R)
    active_q = gaussian_filter1d(qraw[:, adaptor.idx_pin2target], a.sigma, axis=0) if a.sigma > 0 else qraw[:, adaptor.idx_pin2target]
    active_q = np.clip(active_q, bounds[:, 0], bounds[:, 1])
    qpos = []
    for row in active_q:
        q = np.zeros(robot.dof)
        q[adaptor.idx_pin2target] = row
        qpos.append(adaptor.forward_qpos(q).copy())
    qpos = np.array(qpos)
    actual = np.array([fk(q) for q in qpos])
    scene_R = world_R[0].T
    scene_t = np.array([0, 0, .2]) - scene_R @ human['wrist_world'][0]
    wrist_raw = human['wrist_world'] @ scene_R.T + scene_t
    wrist = gaussian_filter1d(wrist_raw, a.sigma, axis=0) if a.sigma > 0 else wrist_raw.copy()
    rot = Rotation.from_matrix(scene_R @ world_R)
    quat_raw = rot.as_quat()
    quat = quat_raw.copy()
    for t in range(1, count):
        if quat[t] @ quat[t-1] < 0:
            quat[t] *= -1
    if a.sigma > 0:
        quat = gaussian_filter1d(quat, a.sigma, axis=0)
        quat /= np.linalg.norm(quat, axis=1, keepdims=True)
    quat = quat[:, [3, 0, 1, 2]]
    times = np.asarray(human['time'], dtype=np.float64) if 'time' in human else np.arange(count) / fps
    assert times.shape == (count,) and np.isfinite(times).all() and np.all(np.diff(times) > 0)
    times = times - times[0]
    np.savez_compressed(out / 'motion.npz', qpos=qpos, qpos_raw=qraw, joint_names=robot.dof_joint_names,
                        active_qpos=active_q, active_joint_names=active, targets=targets, actual=actual,
                        wrist_pos=wrist, wrist_pos_unsmoothed=wrist_raw, wrist_quat_wxyz=quat,
                        wrist_quat_unsmoothed_wxyz=quat_raw[:, [3, 0, 1, 2]],
                        wrist_world=human['wrist_world'], wrist_world_R=world_R,
                        scene_rotation=scene_R, scene_translation=scene_t, time=times, fps=fps,
                        coupling='linear_urdf', offline_sigma_frames=a.sigma,
                        source=str(a.input.resolve()), detection_valid=human['detection_valid'])
    np.savetxt(out / 'trajectory.csv', np.c_[times, wrist, quat, qpos], delimiter=',',
               header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz'] + robot.dof_joint_names), comments='')
    model = mujoco.MjModel.from_xml_path(str(fixed))
    data = mujoco.MjData(model)
    addresses = [model.jnt_qposadr[model.joint(n).id] for n in robot.dof_joint_names]
    mids = [model.site(f'landmark_{i:02d}').id for i in range(21)]
    fk_error = 0.
    for t, row in enumerate(qpos):
        data.qpos[addresses] = row
        mujoco.mj_forward(model, data)
        fk_error = max(fk_error, float(np.abs(data.site_xpos[mids] - actual[t]).max()))
    lim = np.array([limits[n] for n in robot.dof_joint_names])
    violation = float(np.maximum(lim[:, 0] - qpos, 0).max() + np.maximum(qpos - lim[:, 1], 0).max())
    names = list(robot.dof_joint_names)
    mimic_error = max(float(np.abs(qpos[:, names.index(n)] - (sc*qpos[:, names.index(parent)] + off)).max()) for n,(parent,sc,off) in mimic.items())
    assert violation < 1e-8 and mimic_error < 1e-10 and fk_error < 1e-5
    assert all(np.isfinite(x).all() for x in [qpos, wrist, quat, actual])
    error = np.linalg.norm(actual[:, TIPS] - targets[:, TIPS], axis=-1) * 1000
    report = dict(frames=count, fps=fps, duration_seconds=count/fps, source=str(a.input.resolve()),
                  coupling='linear_urdf', independent_joints=len(active), urdf_joints=len(names),
                  tip_error_mean_mm=float(error.mean()), tip_error_p95_mm=float(np.percentile(error,95)),
                  tip_error_finger_mm=dict(zip(['thumb','index','middle','ring','pinky'],error.mean(0).tolist())),
                  max_joint_step_deg=float(np.rad2deg(np.abs(np.diff(qpos,axis=0)).max())),
                  joint_second_difference_rms_deg=float(np.rad2deg(np.sqrt(np.mean(np.diff(qpos,n=2,axis=0)**2)))),
                  wrist_translation_range_m=np.ptp(wrist,axis=0).tolist(),
                  wrist_smoothing_max_displacement_mm=float(np.linalg.norm(wrist-wrist_raw,axis=1).max()*1000),
                  urdf_limit_violation_rad=violation, linear_mimic_max_error_rad=mimic_error,
                  pinocchio_mujoco_fk_error_m=fk_error, fk_checked_frames=count, all_finite=True,
                  quaternion_norm_max_error=float(np.abs(np.linalg.norm(quat,axis=1)-1).max()),
                  offline_sigma_frames=a.sigma, uses_future_frames=a.sigma>0,
                  optimizer_status_counts={str(k):int(v) for k,v in zip(*np.unique(codes,return_counts=True))},
                  boundary='Hand-only kinematic retargeting; no object tracking, contact optimization, dynamics or hardware commands')
    (out / 'retarget_validation.json').write_text(json.dumps(report,indent=2))
    (out / 'dex_config.json').write_text(json.dumps(config,indent=2))
    tree = ET.parse(fixed)
    root = tree.getroot()
    ET.SubElement(root.find("worldbody/body[@name='hand_base_link']"), 'freejoint', name='wrist_free')
    root.find('option').set('gravity','0 0 0')
    # Replay supplies all angles; physical mimic equality is not needed for mj_forward.
    tree.write(out / 'l20_moving.xml')
    print(json.dumps(report,indent=2))


if __name__ == '__main__':
    main()