177 lines
10 KiB
Python
177 lines
10 KiB
Python
"""Contact-aware L20 pose refinement with a required, explicit object specification.
|
|
|
|
Optimizes calibrated finger coordinates and all six wrist correction coordinates.
|
|
The --fixture path is an analytic sphere test, NOT the object in the source video.
|
|
"""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
import argparse,json
|
|
from pathlib import Path
|
|
import xml.etree.ElementTree as ET
|
|
import copy
|
|
import numpy as np
|
|
import mujoco
|
|
from scipy.optimize import least_squares
|
|
from scipy.spatial.transform import Rotation
|
|
from l20_calibrated import ROOT,Calibration,FINGERS,TIPS
|
|
|
|
OUT=ROOT/'output/l20_2047635068'
|
|
|
|
def scene_for_object(spec,out):
|
|
tree=ET.parse(OUT/'l20_moving.xml');root=tree.getroot()
|
|
root.find('option').set('gravity','0 0 -9.81')
|
|
root.find('option').set('timestep','0.001')
|
|
root.find('option').set('cone','elliptic')
|
|
# Resolve all mesh references relative to the existing scene.
|
|
for mesh in root.findall('asset/mesh'):
|
|
mesh.set('file',str((OUT/mesh.get('file')).resolve()))
|
|
hand=root.find("worldbody/body[@name='hand_base_link']")
|
|
for i,geom in enumerate(hand.iter('geom')):
|
|
geom.set('name',f'hand_collision_{i}')
|
|
geom.set('contype','1');geom.set('conaffinity','2')
|
|
geom.set('friction',f"{spec.get('friction',.8)} .005 .0001")
|
|
geom.set('condim','4')
|
|
obj=ET.SubElement(root.find('worldbody'),'body',name='grasp_object',pos='0 0 0')
|
|
ET.SubElement(obj,'freejoint',name='object_free')
|
|
if spec['type']=='mjcf':
|
|
source=Path(spec['mjcf_path']).resolve();bottle=ET.parse(source).getroot()
|
|
for asset in bottle.find('asset'):
|
|
item=copy.deepcopy(asset)
|
|
if item.get('file'):item.set('file',str((source.parent/item.get('file')).resolve()))
|
|
root.find('asset').append(item)
|
|
for item in bottle.find('worldbody/body'):
|
|
if item.tag=='freejoint':continue
|
|
child=copy.deepcopy(item)
|
|
if child.tag=='geom':
|
|
child.set('name','bottle_part_'+str(len(obj)))
|
|
if child.get('contype','1')!='0':
|
|
child.set('contype','2');child.set('conaffinity','1')
|
|
child.set('friction',f"{spec.get('friction',.6)} .005 .001")
|
|
child.set('solref','.004 1');child.set('solimp','.95 .99 .001')
|
|
obj.append(child)
|
|
path=out/'contact_scene.xml';tree.write(path)
|
|
return path
|
|
attrs=dict(name='object_geom',type=spec['type'],mass=str(spec.get('mass_kg',.1)),
|
|
contype='2',conaffinity='1',friction=f"{spec.get('friction',.8)} .005 .0001",
|
|
condim='4',rgba='.9 .65 .2 1',solref='.005 1')
|
|
if spec['type']=='mesh':
|
|
mesh_path=Path(spec['mesh_path']).resolve()
|
|
ET.SubElement(root.find('asset'),'mesh',name='object_mesh',file=str(mesh_path),
|
|
scale=' '.join(map(str,spec.get('mesh_scale',[1,1,1]))))
|
|
attrs['mesh']='object_mesh'
|
|
else:attrs['size']=' '.join(map(str,spec['size']))
|
|
ET.SubElement(obj,'geom',**attrs)
|
|
path=out/'contact_scene.xml';tree.write(path)
|
|
return path
|
|
|
|
class ContactRefiner:
|
|
def __init__(self,spec,out):
|
|
self.spec=spec;self.cal=Calibration()
|
|
self.model=mujoco.MjModel.from_xml_path(str(scene_for_object(spec,out)))
|
|
self.data=mujoco.MjData(self.model)
|
|
self.names=[n for n in self.cal.tables]
|
|
self.addr={n:self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in self.names}
|
|
self.wrist=self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
|
|
self.obj=self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
|
self.object_geoms=np.flatnonzero(self.model.geom_contype==2)
|
|
self.object_geom=int(self.object_geoms[0])
|
|
self.hand_geoms=np.flatnonzero(self.model.geom_contype==1)
|
|
self.tip_geoms=[]
|
|
contact_links=spec.get('contact_links',[f+'_distal' for f in spec.get('contact_fingers',['thumb','index','middle'])])
|
|
for link in contact_links:
|
|
bid=mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_BODY,link)
|
|
self.tip_geoms.append(next(g for g in self.hand_geoms if self.model.geom_bodyid[g]==bid))
|
|
self.lower=np.array([self.cal.bounds[n][0] for n in self.cal.active]+[-.06]*3+[-.5]*3)
|
|
self.upper=np.array([self.cal.bounds[n][1] for n in self.cal.active]+[.06]*3+[.5]*3)
|
|
|
|
def expand(self,active):
|
|
q=dict(zip(self.cal.active,active))
|
|
for n in self.cal.passive:q[n]=self.cal.passive_value(n,q[self.cal.curves[n][0]])[0]
|
|
return q
|
|
|
|
def set_pose(self,x,wrist_pos,wrist_quat,obj_pos,obj_quat):
|
|
q=self.expand(x[:16])
|
|
for n,v in q.items():self.data.qpos[self.addr[n]]=v
|
|
self.data.qpos[self.wrist:self.wrist+3]=wrist_pos+x[16:19]
|
|
rotation=Rotation.from_quat(wrist_quat[[1,2,3,0]])*Rotation.from_rotvec(x[19:22])
|
|
self.data.qpos[self.wrist+3:self.wrist+7]=rotation.as_quat()[[3,0,1,2]]
|
|
self.data.qpos[self.obj:self.obj+3]=obj_pos
|
|
self.data.qpos[self.obj+3:self.obj+7]=obj_quat
|
|
mujoco.mj_kinematics(self.model,self.data)
|
|
|
|
def distances(self):
|
|
return np.array([min(mujoco.mj_geomDistance(self.model,self.data,int(g),int(o),.5,None) for o in self.object_geoms) for g in self.hand_geoms])
|
|
|
|
def solve(self,reference,wrist_pos,wrist_quat,obj_pos,obj_quat,initial=None):
|
|
initial=np.r_[reference,np.zeros(6)] if initial is None else np.asarray(initial).copy()
|
|
tip_index=[list(self.hand_geoms).index(g) for g in self.tip_geoms]
|
|
def residual(x):
|
|
self.set_pose(x,wrist_pos,wrist_quat,obj_pos,obj_quat)
|
|
dist=self.distances()
|
|
compression=self.spec.get('compression_m',.0002)
|
|
return np.r_[(dist[tip_index]+compression)/.004,
|
|
np.maximum(-dist-max(.0003,compression),0)/.0015,
|
|
.15*(x[:16]-reference),.2*x[16:19]/.03,.2*x[19:22]/.25]
|
|
before=residual(initial);before_distance=self.distances()
|
|
result=least_squares(residual,np.clip(initial,self.lower+1e-9,self.upper-1e-9),
|
|
bounds=(self.lower,self.upper),diff_step=1e-4,max_nfev=100,ftol=1e-5,xtol=1e-5,gtol=1e-5)
|
|
after=residual(result.x);after_distance=self.distances()
|
|
report=dict(converged=bool(result.success),nfev=result.nfev,
|
|
cost_before=float(before@before),cost_after=float(after@after),
|
|
mean_selected_surface_gap_before_mm=float(np.abs(before_distance[tip_index]).mean()*1000),
|
|
mean_selected_surface_gap_after_mm=float(np.abs(after_distance[tip_index]).mean()*1000),
|
|
max_penetration_before_mm=float(max(0,-before_distance.min())*1000),
|
|
max_penetration_after_mm=float(max(0,-after_distance.min())*1000),
|
|
wrist_correction_m=result.x[16:19].tolist(),wrist_correction_rotvec=result.x[19:22].tolist())
|
|
return result.x,self.data.qpos.copy(),report
|
|
|
|
def dynamic_hold(self,qpos,seconds=2.):
|
|
# A diagnostic with prescribed hand coordinates and a freely simulated object.
|
|
# It does not validate actuator effort, compliant tracking, or force closure.
|
|
mujoco.mj_resetData(self.model,self.data);self.data.qpos[:]=qpos
|
|
object_initial=qpos[self.obj:self.obj+3].copy()
|
|
hand_addr=[self.addr[n] for n in self.names]+list(range(self.wrist,self.wrist+7))
|
|
object_dof=self.model.jnt_dofadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
|
controlled_dof=[i for i in range(self.model.nv) if not object_dof<=i<object_dof+6]
|
|
frames=[];contacts=[]
|
|
steps=int(seconds/self.model.opt.timestep)
|
|
for k in range(steps):
|
|
self.data.qpos[hand_addr]=qpos[hand_addr];self.data.qvel[controlled_dof]=0
|
|
mujoco.mj_step(self.model,self.data)
|
|
if not np.isfinite(self.data.qpos).all():raise RuntimeError('Non-finite physics state')
|
|
if k%33==0:
|
|
frames.append(self.data.qpos.copy())
|
|
contacts.append(sum(any(g in self.object_geoms for g in c.geom) for c in self.data.contact))
|
|
displacement=float(np.linalg.norm(self.data.qpos[self.obj:self.obj+3]-object_initial))
|
|
return np.array(frames),dict(seconds=seconds,object_displacement_m=displacement,
|
|
contact_frames=int(np.sum(np.array(contacts)>0)),total_frames=len(frames),
|
|
retained_within_2cm=displacement<.02,
|
|
boundary='Prescribed hand, free object under gravity. No object weld or pose resets. Hand actuator/arm dynamics not validated.')
|
|
|
|
def main():
|
|
parser=argparse.ArgumentParser();parser.add_argument('--spec',type=Path);parser.add_argument('--fixture',action='store_true')
|
|
parser.add_argument('--frame',type=int,default=666);args=parser.parse_args()
|
|
motion=np.load(OUT/'motion.npz');t=args.frame
|
|
if not args.fixture and args.spec is None:parser.error('Provide an object specification, or explicitly select the synthetic --fixture test.')
|
|
out=OUT/('contact_fixture' if args.fixture else 'contact_object');out.mkdir(exist_ok=True)
|
|
if args.fixture:
|
|
pts=motion['actual'][t]
|
|
center=(pts[4]+pts[12])*.5
|
|
radius=np.linalg.norm(pts[4]-pts[12])*.42
|
|
r=Rotation.from_quat(motion['wrist_quat_wxyz'][t,[1,2,3,0]]).as_matrix()
|
|
spec=dict(type='sphere',size=[float(radius)],mass_kg=.05,friction=.8,
|
|
position=(r@center+motion['wrist_pos'][t]).tolist(),quat_wxyz=[1,0,0,0],
|
|
contact_fingers=['thumb','index','middle'],purpose='Synthetic interface/physics fixture; not source-video reconstruction')
|
|
else:spec=json.loads(args.spec.read_text())
|
|
(out/'object_spec.json').write_text(json.dumps(spec,indent=2))
|
|
solver=ContactRefiner(spec,out)
|
|
active=np.array([motion['qpos'][t,list(motion['joint_names']).index(n)] for n in solver.cal.active])
|
|
x,q,report=solver.solve(active,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],np.array(spec['position']),np.array(spec['quat_wxyz']))
|
|
path,dynamic=solver.dynamic_hold(q)
|
|
report['dynamic_hold']=dynamic;report['frame']=t;report['fixture']=args.fixture
|
|
np.savez_compressed(out/'contact_result.npz',optimization_variables=x,qpos=q,dynamic_qpos=path)
|
|
(out/'validation.json').write_text(json.dumps(report,indent=2))
|
|
print(json.dumps(report,indent=2),flush=True)
|
|
|
|
if __name__=='__main__':main()
|