96 lines
7.8 KiB
Python
96 lines
7.8 KiB
Python
"""Build full video kinematic contact playback with linear URDF coupling."""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL', 'osmesa')
|
|
import json, shutil, xml.etree.ElementTree as E
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import mujoco
|
|
from scipy.spatial.transform import Rotation
|
|
from scipy.ndimage import gaussian_filter1d
|
|
from scipy.optimize import least_squares
|
|
ROOT=Path(__file__).resolve().parents[1]
|
|
OUT=ROOT/'output/l20_full_replay';OUT.mkdir(exist_ok=True)
|
|
TASK=ROOT/'output/spider_l20_contact/datasets/processed/current/l20/right/bottle_hold'
|
|
tree=E.parse(TASK/'scene_act.xml');root=tree.getroot();world=root.find('worldbody');asset=root.find('asset')
|
|
# A playback scene has no object assistance or hidden controller.
|
|
root.remove(root.find('actuator'))
|
|
for item in list(world.findall('light')):world.remove(item)
|
|
for item in list(asset.findall("texture[@type='skybox']")):asset.remove(item)
|
|
E.SubElement(asset,'texture',name='daylight',type='skybox',builtin='gradient',rgb1='.55 .67 .8',rgb2='.88 .91 .94',width='512',height='3072')
|
|
E.SubElement(asset,'texture',name='floor_grid',type='2d',builtin='checker',rgb1='.54 .59 .62',rgb2='.64 .69 .72',width='512',height='512')
|
|
E.SubElement(asset,'material',name='floor_mat',texture='floor_grid',texrepeat='8 8',texuniform='true',reflectance='.08',rgba='1 1 1 1')
|
|
world.find("geom[@name='floor']").attrib.update(material='floor_mat',rgba='1 1 1 1',size='2 2 .05',pos='0 0 0')
|
|
E.SubElement(world,'light',name='key',pos='-.5 -.5 2',dir='.3 .3 -1',directional='true',diffuse='.45 .42 .38',ambient='.04 .05 .06',castshadow='false')
|
|
E.SubElement(world,'light',name='fill',pos='1 1 1.5',dir='-.5 -.5 -1',directional='true',diffuse='.15 .18 .21',castshadow='false')
|
|
root.find('visual/headlight').attrib.update(ambient='.28 .28 .28',diffuse='.25 .25 .25',specular='.15 .15 .15')
|
|
root.find('visual/rgba').set('haze','.8 .86 .91 1')
|
|
root.find('visual/quality').attrib.update(shadowsize='2048',offsamples='2')
|
|
E.SubElement(root.find('visual'),'map',shadowclip='1',shadowscale='.8',zfar='10',znear='.01')
|
|
(OUT/'assets').mkdir(exist_ok=True)
|
|
for item in asset:
|
|
if item.get('file'):
|
|
src=Path(item.get('file'));dest=OUT/'assets'/(item.get('name')+src.suffix);shutil.copy2(src,dest);item.set('file',str(dest.relative_to(OUT)))
|
|
scene=OUT/'scene.xml';tree.write(scene)
|
|
model=mujoco.MjModel.from_xml_path(str(scene));data=mujoco.MjData(model)
|
|
addr={model.joint(j).name:int(model.jnt_qposadr[j]) for j in range(model.njnt)}
|
|
wn=['right_hand_'+s+'_'+a for s in ['pos','rot'] for a in 'xyz'];on=['right_object_'+s+'_'+a for s in ['pos','rot'] for a in 'xyz']
|
|
wa=[addr[n] for n in wn];oa=[addr[n] for n in on]
|
|
urdf=E.parse(ROOT/'third_party/l20_assets/L20/RIGHT/linkerhand_g20_right.urdf').getroot();js=[j for j in urdf.findall('joint') if j.get('type')!='fixed']
|
|
active=[j.get('name') for j in js if j.find('mimic') is None];mimic={}
|
|
limits={j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in js}
|
|
for j in js:
|
|
mi=j.find('mimic')
|
|
if mi is not None:
|
|
n=j.get('name');parent=mi.get('joint');s=float(mi.get('multiplier','1'));off=float(mi.get('offset','0'));mimic[n]=(parent,s,off)
|
|
lo,hi=np.sort((limits[n]-off)/s);limits[parent]=np.array([max(limits[parent][0],lo),min(limits[parent][1],hi)])
|
|
low=np.array([limits[n][0] for n in active]);high=np.array([limits[n][1] for n in active]);aa=[addr[n] for n in active]
|
|
src=np.load(ROOT/'output/l20_2047635068/jitter_audit/temporal/motion.npz');N=len(src['qpos']);sn=list(src['joint_names']);base=np.clip(np.stack([src['qpos'][:,sn.index(n)] for n in active],1),low,high)
|
|
wr0=Rotation.from_quat(src['wrist_quat_wxyz'][:,[1,2,3,0]]);wp0=src['wrist_pos'];rel=Rotation.from_euler('x',-np.pi/2);orr0=wr0*rel;op0=wp0+wr0.apply(np.tile([.12,-.065,.15],(N,1)))
|
|
align=orr0[600].inv();shift=np.array([0,0,.3])-align.apply(op0[600]);wp=align.apply(wp0)+shift;wr=align*wr0;op=align.apply(op0)+shift;orr=align*orr0
|
|
ref=np.load(TASK/'0/trajectory_kinematic_act.npz')['qpos'][0]
|
|
offset=ref[aa]-base[600];dp=ref[wa[:3]]-wp[600];dr=(wr[600].inv()*Rotation.from_euler('XYZ',ref[wa[3:]])).as_rotvec()
|
|
onset=160;weight=np.clip((np.arange(N)-130)/30,0,1);weight=weight*weight*(3-2*weight)
|
|
# Estimated acquisition bracket, not a measured bottle track.
|
|
op[:onset]=op[onset];oq=orr.as_quat();oq[:onset]=oq[onset];orr=Rotation.from_quat(oq)
|
|
handgeoms=np.flatnonzero(model.geom_contype==1);objgeoms=np.flatnonzero(model.geom_contype==2)
|
|
fingers=['thumb','index','middle','ring','pinky'];tips=[next(i for i,g in enumerate(handgeoms) if model.geom_bodyid[g]==model.body(f+'_distal').id) for f in fingers]
|
|
def setpose(t,x):
|
|
vals=dict(zip(active,np.clip(base[t]+x[:16],low,high)))
|
|
vals.update({n:s*vals[p]+o for n,(p,s,o) in mimic.items()})
|
|
for n,v in vals.items():data.qpos[addr[n]]=v
|
|
data.qpos[wa]=np.r_[wp[t]+x[16:19],(wr[t]*Rotation.from_rotvec(x[19:])).as_euler('XYZ')]
|
|
data.qpos[oa]=np.r_[op[t],orr[t].as_euler('XYZ')]
|
|
mujoco.mj_kinematics(model,data)
|
|
def distances():return np.array([min(mujoco.mj_geomDistance(model,data,int(g),int(o),.3,None) for o in objgeoms) for g in handgeoms])
|
|
prev=np.r_[offset,dp,dr];fits=[];frames=np.unique(np.r_[np.arange(onset,N,2),N-1])
|
|
for k,t in enumerate(frames):
|
|
previous=prev.copy()
|
|
def fun(x):
|
|
setpose(t,x);d=distances()
|
|
return np.r_[(d[tips]+.0002)/.004,np.maximum(-d-.0007,0)/.001,.18*x[:16],.8*(x[:16]-previous[:16]),.2*x[16:19]/.03,.15*x[19:]/.3,.5*(x[16:19]-previous[16:19])/.01,.5*(x[19:]-previous[19:])/.1]
|
|
lb=np.r_[low-base[t],[-.065]*3,[-.55]*3];ub=np.r_[high-base[t],[.065]*3,[.55]*3]
|
|
fit=least_squares(fun,np.clip(prev,lb+1e-9,ub-1e-9),bounds=(lb,ub),max_nfev=65,diff_step=1e-4)
|
|
prev=fit.x;fits.append(prev.copy())
|
|
if k%20==0:print('Contact keyframes',k,len(frames),flush=True)
|
|
fit=np.array(fits);correction=np.stack([np.interp(np.arange(N),frames,fit[:,j]) for j in range(22)],1)
|
|
correction=gaussian_filter1d(correction,1.5,axis=0)*weight[:,None]
|
|
rows=[];penetration=[];gaps=[];minz=1e9
|
|
for t in range(N):
|
|
setpose(t,correction[t]);rows.append(data.qpos.copy());d=distances();penetration.append(max(0,-d.min()));gaps.append(d[tips])
|
|
for g in np.r_[handgeoms,objgeoms]:
|
|
mesh=model.geom_dataid[g]
|
|
if model.geom_type[g]==mujoco.mjtGeom.mjGEOM_MESH:
|
|
vertices=model.mesh_vert[model.mesh_vertadr[mesh]:model.mesh_vertadr[mesh]+model.mesh_vertnum[mesh]]
|
|
minz=min(minz,float((vertices@data.geom_xmat[g].reshape(3,3).T+data.geom_xpos[g])[:,2].min()))
|
|
minz=min(minz,float(op[t,2]-.16))
|
|
q=np.array(rows);zshift=max(0,.025-minz);q[:,wa[2]]+=zshift;q[:,oa[2]]+=zshift
|
|
q[:,wa[3:]]=np.unwrap(q[:,wa[3:]],axis=0);q[:,oa[3:]]=np.unwrap(q[:,oa[3:]],axis=0)
|
|
center=np.mean(np.r_[q[:,wa[:3]],q[:,oa[:3]]],axis=0);center[2]+=.04
|
|
root.find('statistic').attrib.update(center=' '.join(map(str,center)),extent='1')
|
|
root.find('visual/global').attrib.update(azimuth='135',elevation='-22')
|
|
E.SubElement(E.SubElement(root,'keyframe'),'key',name='video_start',qpos=' '.join(map(str,q[0])))
|
|
tree.write(scene)
|
|
np.savez_compressed(OUT/'motion.npz',qpos=q,fps=30,time=np.arange(N)/30,wrist_pos=q[:,wa[:3]],bottle_pos=q[:,oa[:3]],camera_center=center,correction=correction,contact_gaps_m=gaps,penetration_m=penetration)
|
|
report=dict(frames=N,fps=30,duration_seconds=N/30,mode='kinematic prescribed-pose playback; not a full SPIDER dynamic rollout',bottle_pose='estimated registration; acquisition bracket 130-160, not video object tracking',nonlinear_coupling=False,linear_mimic=True,wrist_range_m=np.ptp(q[:,wa[:3]],axis=0).tolist(),finite=bool(np.isfinite(q).all()),max_finger_step_deg=float(np.rad2deg(np.abs(np.diff(q[:,aa],axis=0))).max()),mean_holding_fingertip_gap_mm=float(np.abs(gaps[onset:]).mean()*1000),max_holding_penetration_mm=float(np.max(penetration[onset:])*1000),max_all_penetration_mm=float(np.max(penetration)*1000),bottle_dimensions_m=[.05411,.05411,.157856],floor_z=0,ambient_lighting=True)
|
|
(OUT/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|