hand-motion-pipeline/scripts/smooth_project_reference.py

"""Smooth the collision-corrected L20 reference, then re-project every frame to be collision-free (boxes + table).

Smoothing is applied separately: wrist position (Gaussian), wrist rotation (SO(3) local weighted mean),
independent finger joints (Savitzky-Golay, then exact mimic expansion). Objects are untouched.
Projection: per frame, minimal weighted change s.t. linearised contact distances >= clearance and joint limits.
Env: HF_FIX_OUT (input fix dir), HF_FIX_OLD (prepare dir), HF_SMOOTH_OUT (output dir), HF_TABLE=1 to include the table,
     HF_SIGMA_POS (frames, default 2), HF_SIGMA_ROT (default 2), HF_SG_WINDOW (default 9), HF_CLEARANCE (m, default .0025).
"""
import os, json, shutil
os.environ.setdefault('MUJOCO_GL', 'osmesa'); os.environ['OPENBLAS_NUM_THREADS'] = '1'; os.environ['OMP_NUM_THREADS'] = '1'
from pathlib import Path
import numpy as np, mujoco
from scipy.optimize import minimize
from scipy.ndimage import gaussian_filter1d
from scipy.signal import savgol_filter
from scipy.spatial.transform import Rotation as R
import l20_model_source as source
ROOT = Path(__file__).resolve().parents[1]
IN = Path(os.environ['HF_FIX_OUT']); OLD = Path(os.environ['HF_FIX_OLD']); OUT = Path(os.environ['HF_SMOOTH_OUT']); OUT.mkdir(parents=True, exist_ok=True)
TABLE = os.environ.get('HF_TABLE', '0') == '1'; PAIRS = [{1, 2}, {1, 4}] if TABLE else [{1, 2}]
SIG_POS = float(os.environ.get('HF_SIGMA_POS', '2')); SIG_ROT = float(os.environ.get('HF_SIGMA_ROT', '2')); SG = int(os.environ.get('HF_SG_WINDOW', '11')); CLEAR = float(os.environ.get('HF_CLEARANCE', '.0025'))
source.REPO_ROOT = ROOT / 'third_party/l20_assets'
TIN = IN / 'datasets/processed/current/l20/bimanual/boxes'; TOUT = OUT / 'datasets/processed/current/l20/bimanual/boxes'; (TOUT / '0').mkdir(parents=True, exist_ok=True)
shutil.copy2(TIN / 'scene_act.xml', TOUT / 'scene_act.xml'); shutil.copy2(TIN / 'task_info.json', TOUT / 'task_info.json')
cfg = json.loads((IN / 'config.json').read_text()); cfg['dataset_dir'] = str(OUT / 'datasets'); (OUT / 'config.json').write_text(json.dumps(cfg, indent=2))
for name in ['collision_v2', 'hand_collision']:
    if not (OUT / name).exists(): os.symlink(os.path.relpath(IN / name, OUT), OUT / name)
m = mujoco.MjModel.from_xml_path(str(TOUT / 'scene_act.xml')); d = mujoco.MjData(m)
ref = dict(np.load(IN / 'reference_video_rate.npz')); q0 = ref['qpos'].copy(); N = len(q0); contact = ref['contact']; cp = ref['contact_pos']
meta = {}
for side in ['right', 'left']:
    k = source.HandKinematics(OLD / f'model_{side}/l20_{side}.xml', side)
    wa = np.array([m.jnt_qposadr[m.joint(f'{side}_hand_{s}').id] for s in ['pos_x', 'pos_y', 'pos_z', 'rot_x', 'rot_y', 'rot_z']]); ja = np.array([m.jnt_qposadr[m.joint(f'{side}_{n}').id] for n in k.joint_names])
    B = np.zeros((m.nv, 22)); B[wa, :6] = np.eye(6); B[ja, 6:] = k.expansion
    meta[side] = dict(k=k, wa=wa, ja=ja, B=B, sites=[m.site(f'{side}_hand_{f}_track').id for f in source.FINGERS])

def smooth_rot_euler(e, sigma):
    rot = R.from_euler('XYZ', e); out = []
    for t in range(len(e)):
        ix = np.arange(max(0, t - int(3 * sigma)), min(len(e), t + int(3 * sigma) + 1)); w = np.exp(-.5 * ((ix - t) / sigma) ** 2); out.append(rot[ix].mean(weights=w))
    es = R.concatenate(out).as_euler('XYZ'); es += 2 * np.pi * np.round((e - es) / (2 * np.pi))   # stay on the original continuous branch
    return es
def smooth(q):
    qs = q.copy()
    for side, mt in meta.items():
        k, wa, ja = mt['k'], mt['wa'], mt['ja']
        qs[:, wa[:3]] = gaussian_filter1d(q[:, wa[:3]], SIG_POS, axis=0, mode='nearest'); qs[:, wa[3:]] = smooth_rot_euler(q[:, wa[3:]], SIG_ROT)
        ind = q[:, ja][:, k.independent_indices]; ind = savgol_filter(ind, SG, 2, axis=0, mode='nearest'); ind = np.clip(ind, k.lower, k.upper)
        qs[:, ja] = np.array([k.expand(a) for a in ind])
    return qs
def contacts(side, B, derivatives=True):
    rows, dist = [], []
    for c in d.contact:
        g0, g1 = map(int, c.geom); types = {int(m.geom_contype[g0]), int(m.geom_contype[g1])}
        if types not in PAIRS: continue
        hg = g0 if m.geom_contype[g0] == 1 else g1
        if not m.geom(hg).name.startswith(side + '_'): continue
        dist.append(float(c.dist))
        if derivatives:
            jac = np.zeros((3, m.nv)); mujoco.mj_jac(m, d, jac, None, c.pos, int(m.geom_bodyid[hg])); normal = c.frame[:3] * (1 if hg == g1 else -1); rows.append(normal @ jac @ B)
    return np.asarray(rows).reshape(-1, 22), np.array(dist)
def put(side, x, base):
    mt = meta[side]; d.qpos[:] = base; d.qpos[mt['wa']] = x[:6]; d.qpos[mt['ja']] = mt['k'].expand(x[6:]); mujoco.mj_fwdPosition(m, d)
W = np.r_[[1.] * 3, [.05] * 3, [.01] * 16]     # metres, radians: prefer moving fingers, then wrist rotation, then wrist position
STEP = np.r_[[.008] * 3, [.12] * 3, [.15] * 16]   # per-frame bound relative to the previous projected frame: 8 mm, 7 deg, 8.6 deg
def solve(side, x, base, prev):
    k, B = meta[side]['k'], meta[side]['B']
    for it in range(10):
        put(side, x, base); A, ds = contacts(side, B)
        if not len(ds) or ds.min() >= CLEAR - 2e-5: break
        lo = np.r_[[-.03] * 3, [-.2] * 3, k.lower - x[6:]]; hi = np.r_[[.03] * 3, [.2] * 3, k.upper - x[6:]]
        if prev is not None: lo = np.maximum(lo, prev - STEP - x); hi = np.minimum(hi, prev + STEP - x); lo = np.minimum(lo, 0); hi = np.maximum(hi, 0)
        H = np.diag((W / np.r_[[.01] * 3, [.05] * 3, [.05] * 16]) ** 2)
        fit = minimize(lambda z: .5 * z @ H @ z, np.zeros(22), jac=lambda z: H @ z, bounds=list(zip(lo, hi)), constraints=[dict(type='ineq', fun=lambda z: A @ z + ds - CLEAR - 3e-4, jac=lambda z: A)], method='SLSQP', options={'ftol': 1e-10, 'maxiter': 100})
        x = x + fit.x
    put(side, x, base); _, ds = contacts(side, B, False); short = max(0., CLEAR - ds.min()) if len(ds) else 0.
    return x, short
def project(q):
    qp = q.copy(); moves = {s: [] for s in meta}; worst = 0.; fallback = 0; prev = {s: None for s in meta}
    for f in range(N):
        for side, mt in meta.items():
            k, wa, ja = mt['k'], mt['wa'], mt['ja']; base = qp[f].copy(); x0 = np.r_[base[wa], base[ja][k.independent_indices]]
            x, short = solve(side, x0.copy(), base, prev[side])
            if short > 2e-4 and prev[side] is not None:   # temporal bound made it infeasible: release the bound for this frame
                x, short = solve(side, x0.copy(), base, None); fallback += 1
            worst = max(worst, short); put(side, x, base); qp[f] = d.qpos.copy(); moves[side].append(float(np.linalg.norm(x[:3] - x0[:3]) * 1000)); prev[side] = x
    print('projection: frames with released temporal bound', fallback, flush=True)
    return qp, moves, worst
def stats(q, label):
    out = {}
    for side, mt in meta.items():
        wa, ja, k = mt['wa'], mt['ja'], mt['k']; p = q[:, wa[:3]]; step = np.linalg.norm(np.diff(p, axis=0), axis=1) * 1000; acc = np.linalg.norm(np.diff(p, n=2, axis=0), axis=1) * 1000
        rot = R.from_euler('XYZ', q[:, wa[3:]]); rs = (rot[:-1].inv() * rot[1:]).magnitude() * 180 / np.pi; fj = q[:, ja][:, k.independent_indices]; fs = np.abs(np.diff(fj, axis=0)).max(1) * 180 / np.pi; fa = np.sqrt((np.diff(fj, n=2, axis=0) ** 2).mean(1)) * 180 / np.pi
        gaps = []; pen = []
        for f in range(N):
            d.qpos[:] = q[f]; mujoco.mj_fwdPosition(m, d); on = contact[f, (0 if side == 'right' else 5):(5 if side == 'right' else 10)].astype(bool)
            if on.any(): gaps.extend(np.linalg.norm(d.site_xpos[mt['sites']][on] - cp[f][(0 if side == 'right' else 5):(5 if side == 'right' else 10)][on], axis=1) * 1000)
            _, ds = contacts(side, mt['B'], False); pen.append(max(0., -ds.min()) * 1000 if len(ds) else 0.)
        out[side] = dict(wrist_step_mm=dict(median=float(np.median(step)), p95=float(np.percentile(step, 95)), max=float(step.max())), wrist_2nd_diff_rms_mm=float(np.sqrt((acc ** 2).mean())), wrist_rot_step_deg=dict(p95=float(np.percentile(rs, 95)), max=float(rs.max())), finger_step_deg=dict(p95=float(np.percentile(fs, 95)), max=float(fs.max())), finger_2nd_diff_rms_deg=float(fa.mean()), contact_gap_mm=dict(median=float(np.median(gaps)), p95=float(np.percentile(gaps, 95))) if gaps else None, max_penetration_mm=float(max(pen)), frames_penetrating_over_0_5mm=int(np.sum(np.array(pen) > .5)))
    print(label, json.dumps(out), flush=True); return out
report = {'parameters': dict(sigma_pos=SIG_POS, sigma_rot=SIG_ROT, sg_window=SG, clearance_m=CLEAR, table=TABLE), 'input': stats(q0, 'input')}
q1 = smooth(q0); report['after_smoothing_only'] = stats(q1, 'smoothed')
q2, moves, worst = project(q1); report['after_projection'] = stats(q2, 'projected'); report['projection_move_mm'] = {s: dict(median=float(np.median(v)), max=float(np.max(v))) for s, v in moves.items()}; report['worst_clearance_shortfall_mm'] = worst * 1000
q3 = smooth(q2); q3[:, :] = q3; qs = q3.copy()
# second, lighter pass: blend half-way toward the re-smoothed solution, then project again
q3 = .5 * (q2 + q3); q4, moves2, worst2 = project(q3); report['after_second_pass'] = stats(q4, 'pass2'); report['projection_move_pass2_mm'] = {s: dict(median=float(np.median(v)), max=float(np.max(v))) for s, v in moves2.items()}
assert np.array_equal(q4[:, 54:], q0[:, 54:]) and np.isfinite(q4).all()
np.savez_compressed(OUT / 'reference_video_rate.npz', qpos=q4, contact=contact, contact_pos=cp, time=ref['time'], camera_world_from_cv=ref['camera_world_from_cv'])
np.savez_compressed(OUT / 'smoothing_comparison.npz', before=q0, after=q4, time=ref['time'])
times = np.load(IN / 'datasets/processed/current/l20/bimanual/boxes/0/trajectory_kinematic_act.npz')['time']; t = ref['time']
qi = np.stack([np.interp(times, t, q4[:, i]) for i in range(m.nq)], 1); vel = np.gradient(qi, .0025, axis=0); vel[0] = 0; ctrl = qi[:, m.jnt_qposadr[m.actuator_trnid[:, 0]]]
idx = np.minimum(np.searchsorted(t, times), N - 1); ci = np.stack([np.interp(times, t, cp.reshape(N, -1)[:, j]) for j in range(30)], 1).reshape(-1, 10, 3)
np.savez_compressed(TOUT / '0/trajectory_kinematic_act.npz', qpos=qi, qvel=vel, ctrl=ctrl, contact=contact[idx], contact_pos=ci, time=times)
(OUT / 'smoothing_report.json').write_text(json.dumps(report, indent=2)); print('SMOOTH_PROJECT_DONE', flush=True)