liyang
ae28d55f81
Source: RGB-D -> Dyn-HaMR -> L20 retargeting -> FoundationPose -> reference repair -> SPIDER, documented in docs/PIPELINE_LATEST.md and docs/SETUP_AND_WEIGHTS.md. Adds FoundationPose and nvdiffrast upstream snapshots, requirements/pipeline_venv.txt and the FoundationPose weight manifest/downloader. Assets (Git LFS): weights/ (WiLoR detector, HandFlow denoiser, UniDepth-L), FoundationPose checkpoints, HaMeR checkpoint, Dyn-HaMR HMP model and BMC constraints, L20 URDF/meshes, the 20260915_171525 D405 recording and the two box CADs. MANO models are not redistributed (third_party/hamer/_DATA/data/mano/README.txt). Environments, caches and run outputs excluded. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
119 lines
10 KiB
Python
119 lines
10 KiB
Python
"""Collision-constrained L20 retargeting: analytic contact Jacobians, exact mimic, verified steps."""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='1'
|
|
from pathlib import Path
|
|
import json,copy,argparse,xml.etree.ElementTree as E
|
|
import numpy as np,mujoco,trimesh
|
|
from scipy.optimize import minimize
|
|
from scipy.spatial.transform import Rotation as Rotation
|
|
import l20_model_source as source
|
|
R=Path(__file__).resolve().parents[1];OLD=Path(os.environ.get('HF_FIX_OLD',R/'output/foundationpose_spider_20260915'));OUT=Path(os.environ.get('HF_FIX_OUT',R/'output/collision_fix_20260915'));T=OUT/'datasets/processed/current/l20/bimanual/boxes';(T/'0').mkdir(parents=True,exist_ok=True)
|
|
p=argparse.ArgumentParser();p.add_argument('--legacy-geoms',action='store_true');p.add_argument('--pilot',action='store_true');args=p.parse_args()
|
|
root=E.parse(OLD/'datasets/processed/current/l20/bimanual/boxes/scene_act.xml').getroot();assets=root.find('asset')
|
|
if not args.legacy_geoms:
|
|
manifest=json.loads((OUT/'collision_v2/manifest.json').read_text());handmanifest=json.loads((OUT/'hand_collision/manifest.json').read_text());assert len(manifest)==2 and len(handmanifest)==4
|
|
for name,side in [('upper','right'),('lower','left')]:
|
|
body=root.find(f"worldbody/body[@name='{side}_object']")
|
|
for g in list(body):
|
|
if g.tag=='geom' and '_collision_' in g.get('name',''):body.remove(g)
|
|
for a in list(assets):
|
|
if a.get('name','').startswith(name+'_collision_'):assets.remove(a)
|
|
for i,path in enumerate(sorted((OUT/'collision_v2').glob(name+'_*.obj'))):
|
|
part=trimesh.load(path)
|
|
if abs(part.volume)<1e-12 or np.linalg.matrix_rank(part.vertices-part.vertices.mean(0),tol=1e-9)<3:continue
|
|
n=f'{name}_collision_{i}';E.SubElement(assets,'mesh',name=n,file=str(path),maxhullvert='128');E.SubElement(body,'geom',name=n,type='mesh',mesh=n,contype='2',conaffinity='7',friction='.8 .005 .001',condim='3',solref='.006 1',solimp='.95 .99 .001',mass='0',group='3')
|
|
for body in root.findall('.//body'):
|
|
for g in list(body.findall('geom')):
|
|
mesh=g.get('mesh')
|
|
if mesh not in handmanifest:continue
|
|
g.set('contype','0');g.set('conaffinity','0')
|
|
for i,path in enumerate(sorted((OUT/'hand_collision').glob(mesh+'_*.obj'))):
|
|
n=f'{mesh}_part_{i}';E.SubElement(assets,'mesh',name=n,file=str(path),maxhullvert='128');attrs=dict(g.attrib);attrs.update(name=n,mesh=n,contype='1',conaffinity='2',group='3',mass='0',solref='.006 1',solimp='.95 .99 .001');E.SubElement(body,'geom',**attrs)
|
|
TABLE=os.environ.get('HF_TABLE','0')=='1';PAIRS=[{1,2},{1,4}] if TABLE else [{1,2}]
|
|
if TABLE:
|
|
for fl in root.findall(".//geom[@name='right_floor']"):fl.set('conaffinity','3') # hands (contype 1) now collide with the table
|
|
# Keep contacts within 4 mm available to the solver, without early physical forces.
|
|
for g in root.findall('.//geom'):
|
|
if int(g.get('contype','0')) in ([1,2,4] if TABLE else [1,2]):g.attrib.update(margin='.004',gap='.004',solref='.005 1',solimp='.95 .99 .001')
|
|
scene=T/('scene_legacy.xml' if args.legacy_geoms else 'scene_act.xml');E.ElementTree(root).write(scene)
|
|
m=mujoco.MjModel.from_xml_path(str(scene));d=mujoco.MjData(m);source.REPO_ROOT=R/'third_party/l20_assets'
|
|
oldref=np.load(OLD/'reference_video_rate.npz');original=np.load(OLD/'reference_before_contact_fit.npz')['qpos'];q=oldref['qpos'].copy();cp=oldref['contact_pos'];contact=oldref['contact'];N=len(q);meta={}
|
|
for side in ['right','left']:
|
|
k=source.HandKinematics(OLD/('model_'+side)/('l20_'+side+'.xml'),side);wn=[side+'_hand_'+s for s in ['pos_x','pos_y','pos_z','rot_x','rot_y','rot_z']];wa=np.array([m.jnt_qposadr[m.joint(n).id] for n in wn]);ja=np.array([m.jnt_qposadr[m.joint(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]=(k,wa,ja,B,[m.site(side+'_hand_'+f+'_track').id for f in source.FINGERS])
|
|
# All hand/object geometry is used, rather than only tip points.
|
|
def collision_rows(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(x,base,k,wa,ja):
|
|
d.qpos[:]=base;d.qpos[wa]=x[:6];d.qpos[ja]=k.expand(x[6:]);mujoco.mj_fwdPosition(m,d)
|
|
frames=[0,36,46,80,160,240,351] if args.pilot else range(N);report=[];previous={};margin=float(os.environ.get('HF_MARGIN','.0005'))
|
|
for f in frames:
|
|
for si,side in enumerate(['right','left']):
|
|
k,wa,ja,B,sites=meta[side];base=q[f].copy();xorig=np.r_[original[f,wa],original[f,ja][k.independent_indices]];x=np.r_[q[f,wa],q[f,ja][k.independent_indices]]
|
|
low=np.r_[xorig[:3]-.15,xorig[3:6]-.65,k.lower];high=np.r_[xorig[:3]+.15,xorig[3:6]+.65,k.upper];x=np.clip(x,low,high)
|
|
target=cp[f,si*5:si*5+5];on=contact[f,si*5:si*5+5].astype(bool);prior=xorig.copy()
|
|
if side in previous and not args.pilot:
|
|
prior=np.clip(xorig+previous[side]*(1. if on.any() else .95),low,high);x=prior.copy()
|
|
prev=np.r_[q[f-1,wa],q[f-1,ja][k.independent_indices]]
|
|
temporal=np.r_[[.012]*3,[.18]*3,[.25]*16];low=np.maximum(low,prev-temporal);high=np.minimum(high,prev+temporal)
|
|
# Transport the previous hand with its currently contacted object to stay on the same side.
|
|
if on.any():
|
|
d.qpos[:]=base;mujoco.mj_kinematics(m,d);center=target[on].mean(0);choices=[54] if f>=177 else [54,60]
|
|
start=min(choices,key=lambda j:np.linalg.norm(center-d.site_xpos[m.site(('right' if j==54 else 'left')+'_object_track').id]))
|
|
old_obj=q[f-1,start:start+6];new_obj=q[f,start:start+6];delta=Rotation.from_euler('XYZ',new_obj[3:])*Rotation.from_euler('XYZ',old_obj[3:]).inv()
|
|
x[:3]=delta.apply(prev[:3]-old_obj[:3])+new_obj[:3]
|
|
e=(delta*Rotation.from_euler('XYZ',prev[3:6])).as_euler('XYZ');candidates=np.array([e,[e[0]+np.pi,np.pi-e[1],e[2]+np.pi]]);candidates+=2*np.pi*np.round((prev[3:6]-candidates)/(2*np.pi));x[3:6]=candidates[np.argmin(np.linalg.norm(candidates-prev[3:6],axis=1))];x[6:]=prev[6:]
|
|
x=np.clip(x,low,high)
|
|
put(x,base,k,wa,ja);_,dist=collision_rows(side,B,False);initial=max(0,-dist.min()) if len(dist) else 0
|
|
# Exit deep or contradictory overlap first; no object pose changes or mesh shrinking.
|
|
if initial>.004 and (f==0 or args.pilot):
|
|
seeds=[x.copy()];directions=np.array([[i,j,z] for i in [-1,0,1] for j in [-1,0,1] for z in [-1,0,1] if (i,j,z)!=(0,0,0)],float);directions/=np.linalg.norm(directions,axis=1)[:,None]
|
|
for radius in [.025,.05,.08,.12]:
|
|
for direction in directions:
|
|
seed=x.copy();seed[:3]=np.clip(x[:3]+direction*radius,low[:3],high[:3]);seeds.append(seed)
|
|
best=None
|
|
for seed in seeds:
|
|
put(seed,base,k,wa,ja);_,ds=collision_rows(side,B,False);pen=max(0,-ds.min()) if len(ds) else 0;gap=np.linalg.norm(d.site_xpos[sites][on]-target[on],axis=1).mean() if on.any() else 0
|
|
score=500*pen+np.linalg.norm(seed[:3]-xorig[:3])+gap*.3
|
|
if best is None or score<best[0]:best=(score,seed.copy())
|
|
x=best[1]
|
|
weights=np.r_[[float(os.environ.get('HF_WRIST_WEIGHT','25'))]*3,[1.5]*3,[.22]*16];bounds_step=np.r_[[.015]*3,[.10]*3,[.18]*16]
|
|
for it in range(45):
|
|
put(x,base,k,wa,ja);A,ds=collision_rows(side,B);J=[];res=[]
|
|
for tip in np.flatnonzero(on):
|
|
jac=np.zeros((3,m.nv));mujoco.mj_jacSite(m,d,jac,None,int(sites[tip]));J.extend((jac@B)/.012);res.extend((d.site_xpos[sites[tip]]-target[tip])/.012)
|
|
J=np.asarray(J).reshape(-1,22);res=np.array(res);H=np.diag(weights**2)+J.T@J;g=weights**2*(x-prior)+J.T@res
|
|
lo=np.maximum(low-x,-bounds_step);hi=np.minimum(high-x,bounds_step)
|
|
constraints=[] if not len(ds) else [dict(type='ineq',fun=lambda z:A@z+ds-margin,jac=lambda z:A)]
|
|
fit=minimize(lambda z:.5*z@H@z+g@z,np.zeros(22),jac=lambda z:H@z+g,bounds=list(zip(lo,hi)),constraints=constraints,method='SLSQP',options={'maxiter':80,'ftol':1e-9})
|
|
dx=fit.x;pen=max(0.,-ds.min()) if len(ds) else 0.;oldcost=.5*np.sum((weights*(x-prior))**2)+.5*(res@res)+1e5*pen
|
|
chosen=None
|
|
for alpha in [1.,.5,.25,.1]:
|
|
trial=np.clip(x+alpha*dx,low,high);put(trial,base,k,wa,ja);_,td=collision_rows(side,B,False);tp=max(0.,-td.min()) if len(td) else 0.;rr=(d.site_xpos[sites][on]-target[on]).ravel()/.012;cost=.5*np.sum((weights*(trial-prior))**2)+.5*(rr@rr)+1e5*tp
|
|
if cost<oldcost-1e-7 or (pen>.0001 and tp<pen*.8):chosen=trial;break
|
|
if chosen is None:break
|
|
x=chosen
|
|
if np.linalg.norm(alpha*dx)<1e-5:break
|
|
put(x,base,k,wa,ja);_,ds=collision_rows(side,B,False);final=max(0.,-ds.min()) if len(ds) else 0
|
|
q[f]=d.qpos.copy();previous[side]=x-xorig
|
|
report.append(dict(frame=int(f),side=side,initial_penetration_mm=float(initial*1000),final_penetration_mm=float(final*1000),wrist_shift_mm=float(np.linalg.norm(x[:3]-xorig[:3])*1000),contact_gap_mm=float(np.linalg.norm(d.site_xpos[sites][on]-target[on],axis=1).mean()*1000) if on.any() else None,iterations=it+1))
|
|
if f%20==0 or args.pilot:print('frame',f,report[-2:],flush=True)
|
|
(OUT/('pilot_fit.json' if args.pilot else 'fit_progress.json')).write_text(json.dumps(report,indent=2))
|
|
if args.pilot:
|
|
print('PILOT_FINISHED',max(x['final_penetration_mm'] for x in report),flush=True);raise SystemExit
|
|
np.savez_compressed(OUT/'reference_video_rate.npz',qpos=q,contact=contact,contact_pos=cp,time=oldref['time'],camera_world_from_cv=oldref['camera_world_from_cv'])
|
|
# Interpolated states are separately collision-checked before physical optimization.
|
|
times=np.load(OLD/'datasets/processed/current/l20/bimanual/boxes/0/trajectory_kinematic_act.npz')['time'];qi=np.stack([np.interp(times,oldref['time'],q[:,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(oldref['time'],times),N-1);ci=np.stack([np.interp(times,oldref['time'],cp.reshape(N,-1)[:,j]) for j in range(30)],1).reshape(-1,10,3)
|
|
np.savez_compressed(T/'0/trajectory_kinematic_act.npz',qpos=qi,qvel=vel,ctrl=ctrl,contact=contact[idx],contact_pos=ci,time=times)
|
|
(T/'task_info.json').write_text((OLD/'datasets/processed/current/l20/bimanual/boxes/task_info.json').read_text());config=json.loads((OLD/'config.json').read_text());config.update(dataset_dir=str(OUT/'datasets'),num_samples=64,nconmax_per_env=1024,njmax_per_env=3072);(OUT/'config.json').write_text(json.dumps(config,indent=2));print('COLLISION_FIT_DONE',flush=True)
|