hand-motion-pipeline/scripts/audit_collision_fix.py

"""Independent before/after full-link collision, limits, continuity, and CAD surface probes."""
import os
os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='1'
from pathlib import Path
import json,numpy as np,mujoco,trimesh,open3d as o3d
from scipy.spatial.transform import Rotation as R
from red_box_distance import signed_distance
ROOT=Path(__file__).resolve().parents[1];O=Path(os.environ.get('HF_FIX_OUT',ROOT/'output/collision_fix_20260915'));OLD=Path(os.environ.get('HF_FIX_OLD',ROOT/'output/foundationpose_spider_20260915'));T=O/'datasets/processed/current/l20/bimanual/boxes';m=mujoco.MjModel.from_xml_path(str(T/'scene_act.xml'));d=mujoco.MjData(m);ref=np.load(O/'reference_video_rate.npz');before=np.load(OLD/'reference_video_rate.npz');q=ref['qpos'];sites=[m.site(s+'_hand_'+f+'_track').id for s in ['right','left'] for f in ['thumb','index','middle','ring','pinky']]
def audit(rows):
 pen=[];gap=[];byhand=[[],[]];lim=[];mimic=[]
 for f,row in enumerate(rows):
  d.qpos[:]=row;mujoco.mj_fwdPosition(m,d);ds=[[],[]]
  for c in d.contact:
   g0,g1=map(int,c.geom)
   if {int(m.geom_contype[g0]),int(m.geom_contype[g1])} not in ([{1,2},{1,4}] if os.environ.get('HF_TABLE','0')=='1' else [{1,2}]):continue
   h=g0 if m.geom_contype[g0]==1 else g1;side=0 if m.geom(h).name.startswith('right_') else 1;ds[side].append(c.dist)
  for j in range(2):byhand[j].append(max(0,-min(ds[j]))*1000 if ds[j] else 0.)
  pen.append(max(byhand[0][-1],byhand[1][-1]));on=ref['contact'][f].astype(bool);gap.extend(np.linalg.norm(d.site_xpos[sites][on]-ref['contact_pos'][f][on],axis=1)*1000)
  ids=np.flatnonzero(m.jnt_limited);v=row[m.jnt_qposadr[ids]];lim.append(max(np.maximum(m.jnt_range[ids,0]-v,0).max(),np.maximum(v-m.jnt_range[ids,1],0).max()))
  for e in range(m.neq):
   j1,j2=m.eq_obj1id[e],m.eq_obj2id[e];poly=m.eq_data[e,:5];x=row[m.jnt_qposadr[j2]];mimic.append(abs(row[m.jnt_qposadr[j1]]-sum(poly[k]*x**k for k in range(5))))
 result=dict(frames=len(rows),all_finite=bool(np.isfinite(rows).all()),max_hand_object_penetration_mm=float(max(pen)),median_frame_max_penetration_mm=float(np.median(pen)),frames_over_0_5mm=int(np.sum(np.array(pen)>.5)),inferred_contact_gap_mean_mm=float(np.mean(gap)),inferred_contact_gap_p95_mm=float(np.percentile(gap,95)),max_joint_limit_violation_rad=float(max(lim)),max_mimic_error_rad=float(max(mimic)))
 for i,(side,st) in enumerate([('right',0),('left',27)]):
  result[side]={'max_penetration_mm':float(max(byhand[i])),'max_wrist_step_mm':float(np.linalg.norm(np.diff(rows[:,st:st+3],axis=0),axis=1).max()*1000),'max_wrist_rotation_step_deg':float((R.from_euler('XYZ',rows[:-1,st+3:st+6]).inv()*R.from_euler('XYZ',rows[1:,st+3:st+6])).magnitude().max()*180/np.pi)}
 return result,np.array(pen)
report={};report['before_same_new_collision_model'],bp=audit(before['qpos']);report['corrected_reference'],ap=audit(q)
report['object_qpos_unchanged']=bool(np.array_equal(q[:,-12:],before['qpos'][:,-12:]));report['wrist_correction_mm']={s:{'median':float(np.median(np.linalg.norm(q[:,j:j+3]-before['qpos'][:,j:j+3],axis=1))*1000),'max':float(np.max(np.linalg.norm(q[:,j:j+3]-before['qpos'][:,j:j+3],axis=1))*1000)} for s,j in [('right',0),('left',27)]}
# Sample the original visible link meshes, independently of the CoACD collision mesh.
probes=[]
for g in range(m.ngeom):
 n=m.geom(g).name or ''
 if '_visual_' not in n or not n.startswith(('right_','left_')):continue
 mid=m.geom_dataid[g];v=m.mesh_vert[m.mesh_vertadr[mid]:m.mesh_vertadr[mid]+m.mesh_vertnum[mid]];faces=m.mesh_face[m.mesh_faceadr[mid]:m.mesh_faceadr[mid]+m.mesh_facenum[mid]];cent=v[faces].mean(1);pts=np.r_[v[np.linspace(0,len(v)-1,min(64,len(v)),dtype=int)],cent[np.linspace(0,len(cent)-1,min(64,len(cent)),dtype=int)]];probes.append((g,pts))
report['cad_probe_scope']=f'{sum(len(p) for g,p in probes)} deterministic original-link surface samples per frame, all 352 frames; sampled test, not exhaustive mesh intersection.'
for name,side in [('upper','right'),('lower','left')]:
 mesh=trimesh.load(O/'collision_v2'/f'{name}_closed.ply');assert mesh.is_watertight;sc=o3d.t.geometry.RaycastingScene();sc.add_triangles(o3d.t.geometry.TriangleMesh.from_legacy(o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(mesh.vertices),o3d.utility.Vector3iVector(mesh.faces))));obj=m.body(side+'_object').id;values=[]
 for row in q:
  d.qpos[:]=row;mujoco.mj_kinematics(m,d);points=np.concatenate([v@d.geom_xmat[g].reshape(3,3).T+d.geom_xpos[g] for g,v in probes]);local=(points-d.xpos[obj])@d.xmat[obj].reshape(3,3);dist=signed_distance(sc,mesh,local);values.append(max(0,-dist.min())*1000)
 report[name+'_exact_CAD_surface_probes']={'max_inside_mm':float(max(values)),'frames_over_1mm':int(np.sum(np.array(values)>1))};print(name,'probe',report[name+'_exact_CAD_surface_probes'],flush=True)
if (T/'0/trajectory_mjwp_act.npz').exists() and np.load(T/'0/trajectory_mjwp_act.npz')['time'].max()>=ref['time'][-1]-.05:
 a=np.load(T/'0/trajectory_mjwp_act.npz');rows=a['qpos'].reshape(-1,66);times=a['time'].ravel();ix=np.argmin(abs(times[:,None]-ref['time'][None,:]),axis=0);report['spider_physics'],_=audit(rows[ix]);report['spider_physics']['duration_s']=float(times[-1]);report['spider_physics']['steps']=len(rows)
report['interpolation']=json.loads((O/'interpolation_validation.json').read_text());(O/'validation.json').write_text(json.dumps(report,indent=2));np.savez_compressed(O/'collision_comparison.npz',before=before['qpos'],after=q,time=ref['time'],before_penetration_mm=bp,after_penetration_mm=ap);print(json.dumps(report,indent=2))