21 lines
1.8 KiB
Python
21 lines
1.8 KiB
Python
"""Fresh GPU rollout of saved controls, without optimization or virtual assistance."""
|
|
import os
|
|
from pathlib import Path
|
|
ROOT=Path(__file__).resolve().parents[1];OUT=ROOT/'output/spider_l20_contact'
|
|
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp');os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
import json,numpy as np,mujoco,mujoco_warp as mw,warp as wp
|
|
TASK=OUT/'datasets/processed/current/l20/right/bottle_hold';P=TASK/'0'
|
|
r=np.load(P/'trajectory_kinematic_act.npz');a=np.load(P/'trajectory_mjwp_act.npz');ctrl=a['ctrl'].reshape(-1,a['ctrl'].shape[-1]);expected=a['qpos'].reshape(-1,a['qpos'].shape[-1])
|
|
m=mujoco.MjModel.from_xml_path(str(TASK/'scene_act.xml'));m.opt.iterations=20;m.opt.ls_iterations=50;m.opt.o_solref[:]=[.02,1];m.opt.o_solimp[:]=[0,.95,.03,.5,2]
|
|
assert not m.actuator_gainprm[-6:].any() and not m.actuator_biasprm[-6:].any()
|
|
d=mujoco.MjData(m);d.qpos[:]=r['qpos'][0];d.qvel[:]=r['qvel'][0];d.ctrl[:]=r['ctrl'][0];mujoco.mj_step(m,d);d.time=0
|
|
wp.init();rows=[]
|
|
with wp.ScopedDevice('cuda:0'):
|
|
wm=mw.put_model(m);wd=mw.put_data(m,d,nworld=2,nconmax=128,njmax=512)
|
|
with wp.ScopedCapture() as capture:mw.step(wm,wd)
|
|
for i,u in enumerate(ctrl):
|
|
wd.ctrl.assign(np.stack([r['ctrl'][i+1],u]).astype(np.float32));wp.capture_launch(capture.graph);wp.synchronize();rows.append(wd.qpos.numpy())
|
|
rows=np.array(rows);np.savez_compressed(OUT/'independent_gpu_replay.npz',baseline=rows[:,0],optimized_controls=rows[:,1])
|
|
report=dict(all_finite=bool(np.isfinite(rows).all()),max_difference_from_native_spider_qpos=float(np.abs(rows[:,1]-expected).max()),optimized_final_object_pos=rows[-1,1,-6:-3].tolist(),native_final_object_pos=expected[-1,-6:-3].tolist(),zero_assistance=True)
|
|
(OUT/'independent_gpu_replay.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|