55 lines
5.5 KiB
Python
55 lines
5.5 KiB
Python
"""Independently evaluate SPIDER's zero-assistance rollout and a control-replay baseline."""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
from pathlib import Path
|
|
import json,re
|
|
import numpy as np
|
|
import mujoco,imageio.v2 as imageio
|
|
from PIL import Image,ImageDraw
|
|
ROOT=Path(__file__).resolve().parents[1];OUT=ROOT/'output/spider_l20_contact';TASK=OUT/'datasets/processed/current/l20/right/bottle_hold';TRIAL=TASK/'0'
|
|
DT=json.loads((OUT/'config.json').read_text())['sim_dt']
|
|
ref=np.load(TRIAL/'trajectory_kinematic_act.npz');result=np.load(TRIAL/'trajectory_mjwp_act.npz');q=result['qpos'].reshape(-1,result['qpos'].shape[-1]);ctrl=result['ctrl'].reshape(-1,result['ctrl'].shape[-1]);v=result['qvel'].reshape(-1,result['qvel'].shape[-1]);time=result['time'].ravel();N=len(q)
|
|
model=mujoco.MjModel.from_xml_path(str(TASK/'scene_act.xml'));model.opt.iterations=20;model.opt.ls_iterations=50;model.opt.o_solref[:]=[.02,1];model.opt.o_solimp[:]=[0,.95,.03,.5,2]
|
|
obj=model.body('right_object').id;hand=model.body('hand_base_link').id
|
|
assert np.all(model.actuator_gainprm[-6:]==0) and np.all(model.actuator_biasprm[-6:]==0)
|
|
def replay(controls):
|
|
d=mujoco.MjData(model);d.qpos[:]=ref['qpos'][0];d.qvel[:]=ref['qvel'][0];d.ctrl[:]=ref['ctrl'][0];mujoco.mj_step(model,d);d.time=0
|
|
rows=[]
|
|
for u in controls:
|
|
d.ctrl[:]=u;mujoco.mj_step(model,d);rows.append(d.qpos.copy())
|
|
return np.array(rows)
|
|
baseline=replay(ref['ctrl'][1:N+1]);cpu=replay(ctrl)
|
|
def metrics(rows):
|
|
d=mujoco.MjData(model);dr=mujoco.MjData(model);gaps=[];errors=[];relerrors=[];counts=[];wrist=[];objheight=[]
|
|
for t,row in enumerate(rows):
|
|
d.qpos[:]=row;mujoco.mj_forward(model,d);dr.qpos[:]=ref['qpos'][min(t+1,len(ref['qpos'])-1)];mujoco.mj_forward(model,dr)
|
|
errors.append(np.linalg.norm(d.xpos[obj]-dr.xpos[obj]));relerrors.append(np.linalg.norm((d.xpos[obj]-d.xpos[hand])-(dr.xpos[obj]-dr.xpos[hand])))
|
|
selected=[c for c in d.contact if {int(model.geom_contype[c.geom[0]]),int(model.geom_contype[c.geom[1]])}=={1,2}]
|
|
counts.append(len(selected));gaps.append(max([max(0,-c.dist) for c in selected],default=0));wrist.append(d.xpos[hand].copy());objheight.append(d.xpos[obj,2])
|
|
e=np.array(errors);r=np.array(relerrors)
|
|
return dict(finite=bool(np.isfinite(rows).all()),max_object_position_error_mm=float(e.max()*1000),final_object_position_error_mm=float(e[-1]*1000),mean_object_position_error_mm=float(e.mean()*1000),max_object_hand_relative_error_mm=float(r.max()*1000),contact_frames=int(np.sum(np.array(counts)>0)),contact_duration_s=float(np.sum(np.array(counts)>0)*DT),last_hand_contact_time_s=float((np.flatnonzero(np.array(counts)>0)[-1]+1)*DT) if np.any(np.array(counts)>0) else 0.,frames=len(rows),max_hand_object_penetration_mm=float(np.max(gaps)*1000),wrist_range_m=np.ptp(wrist,axis=0).tolist(),final_bottle_base_height_m=float(objheight[-1]),retained=bool(r.max()<.02 and counts[-1]>0)),np.array(counts),e
|
|
mr,counts,errors=metrics(q);mb,_,_=metrics(baseline);mc,_,_=metrics(cpu)
|
|
report=dict(spider=mr,reference_control_baseline=mb,independent_cpu_control_replay=mc,zero_object_actuator_gains=True,zero_object_actuator_bias=True,source_frames=[600,660],simulation_seconds=float(time[-1]),nonlinear_json_coupling=False,linear_mimic=True,source_object_pose='Estimated target, not measured video object motion',cpu_gpu_max_qpos_difference=float(np.abs(cpu-q).max()))
|
|
log=(OUT/'run.log').read_text();report['nonfinite_sample_warnings']=len(re.findall('NaNs or infs in rews',log));report['upstream_completion_marker']='SPIDER_RUN_COMPLETE' in log
|
|
if (OUT/'independent_gpu_replay.npz').exists():
|
|
independent=np.load(OUT/'independent_gpu_replay.npz')
|
|
report['independent_gpu_control_replay']=metrics(independent['optimized_controls'])[0]
|
|
report['reference_control_gpu_baseline']=metrics(independent['baseline'])[0]
|
|
(OUT/'validation.json').write_text(json.dumps(report,indent=2));np.savez_compressed(OUT/'comparison_motion.npz',spider=q,baseline=baseline,cpu_replay=cpu,reference=ref['qpos'][1:N+1],ctrl=ctrl,qvel=v,time=time,contact_counts=counts,object_error_m=errors)
|
|
np.savetxt(OUT/'optimized_trajectory.csv',np.c_[time,q,ctrl],delimiter=',',header=','.join(['time_s']+[model.joint(i).name for i in range(model.njnt)]+['ctrl_'+model.actuator(i).name for i in range(model.nu)]),comments='')
|
|
model.vis.headlight.ambient[:]=.8;model.vis.quality.offsamples=1
|
|
renderer=mujoco.Renderer(model,height=400,width=400);cam=mujoco.MjvCamera();cam.lookat[:]=[-.02,.045,.34];cam.distance=.65;cam.azimuth=50;cam.elevation=12
|
|
opt=mujoco.MjvOption();opt.geomgroup[3]=0;opt.sitegroup[:]=0;data=mujoco.MjData(model)
|
|
frames=np.minimum(np.round(np.arange(60)/30/DT).astype(int),N-1)
|
|
writer=imageio.get_writer(OUT/'spider_contact_comparison.mp4',fps=30,codec='libx264',quality=8)
|
|
for i,t in enumerate(frames):
|
|
panels=[]
|
|
for row in [ref['qpos'][t+1],baseline[t],q[t],cpu[t]]:
|
|
data.qpos[:]=row;mujoco.mj_forward(model,data);renderer.update_scene(data,camera=cam,scene_option=opt);panels.append(renderer.render().copy())
|
|
img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img);draw.rectangle([0,0,1600,54],fill=(20,25,35))
|
|
for j,label in enumerate(['Estimated reference','Reference controls, free bottle (CPU)','SPIDER contact guidance (GPU)','Same controls, independent CPU replay']):draw.text((j*400+8,8),label,fill='white')
|
|
draw.text((8,30),f't={time[t]:.2f}s | Moving wrist + URDF linear mimic | SPIDER error={errors[t]*1000:.1f} mm, contacts={counts[t]} | 20mm tracking gate={mr["retained"]}',fill='white')
|
|
writer.append_data(np.asarray(img))
|
|
if i in [0,30,59]:img.save(OUT/f'comparison_{i:02d}.png')
|
|
writer.close();renderer.close();print(json.dumps(report,indent=2))
|