29 lines
1.9 KiB
Python
29 lines
1.9 KiB
Python
"""Retry residual collision failures from the preceding valid pose."""
|
|
import json
|
|
import numpy as np
|
|
from l20_bottle_sequence import BASE,OUT
|
|
from l20_contact import ContactRefiner
|
|
|
|
src=np.load(BASE/'motion.npz');m=dict(np.load(OUT/'registered_motion.npz'))
|
|
s=ContactRefiner(json.loads((OUT/'object_spec.json').read_text()),OUT)
|
|
names=list(m['joint_names']);reports=[]
|
|
for t in np.where(m['penetration_m']>.0015)[0]:
|
|
active=np.array([src['qpos'][t,names.index(n)] for n in s.cal.active])
|
|
x,_,r=s.solve(active,src['wrist_pos'][t],src['wrist_quat_wxyz'][t],m['bottle_pos'][t],m['bottle_quat_wxyz'][t],initial=m['optimization_variables'][t-1])
|
|
if r['max_penetration_after_mm']<m['penetration_m'][t]*1000:
|
|
m['optimization_variables'][t]=x
|
|
m['qpos'][t]=[s.expand(x[:16])[str(n)] for n in names]
|
|
m['wrist_pos'][t]=s.data.qpos[s.wrist:s.wrist+3]
|
|
m['wrist_quat_wxyz'][t]=s.data.qpos[s.wrist+3:s.wrist+7]
|
|
m['penetration_m'][t]=r['max_penetration_after_mm']/1000
|
|
m['command_u8'][t]=s.cal.command(dict(zip(names,m['qpos'][t])))
|
|
reports.append(dict(frame=int(t),**r))
|
|
np.savez_compressed(OUT/'registered_motion.npz',**m)
|
|
header=(OUT/'registered_trajectory.csv').open().readline().strip()
|
|
np.savetxt(OUT/'registered_trajectory.csv',np.c_[np.arange(len(m['qpos']))/30,m['wrist_pos'],m['wrist_quat_wxyz'],m['qpos'],m['bottle_pos'],m['bottle_quat_wxyz']],delimiter=',',header=header,comments='')
|
|
r=json.loads((OUT/'sequence_validation.json').read_text());r['neighbor_retry_reports']=reports
|
|
r['max_all_frame_penetration_mm']=float(m['penetration_m'].max()*1000)
|
|
r['max_holding_penetration_mm']=float(m['penetration_m'][160:].max()*1000)
|
|
(OUT/'sequence_validation.json').write_text(json.dumps(r,indent=2))
|
|
print({'retries':len(reports),'max_penetration_mm':r['max_all_frame_penetration_mm'],'remaining_bad_frames':np.where(m['penetration_m']>.0015)[0].tolist()})
|