37 lines
2.4 KiB
Python
37 lines
2.4 KiB
Python
"""Check SPIDER imports, dependency versions and a two-world GPU contact simulation."""
|
|
import os
|
|
from pathlib import Path
|
|
ROOT=Path(__file__).resolve().parents[1]
|
|
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp')
|
|
os.environ['TORCHINDUCTOR_CACHE_DIR']=str(ROOT/'.spider_cache/torchinductor')
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
import json,sys,importlib,importlib.metadata as metadata,tomllib
|
|
import numpy as np
|
|
from packaging.requirements import Requirement
|
|
import torch,mujoco,mujoco_warp as mw,warp as wp
|
|
|
|
report={'python':sys.version,'compatibility_mode':'Python 3.11 with reused Torch; upstream declares Python >=3.12','versions':{},'imports':{},'requirements':{}}
|
|
reqs=tomllib.loads((ROOT/'third_party/spider/pyproject.toml').read_text())['project']['dependencies']
|
|
for text in reqs:
|
|
r=Requirement(text);v=metadata.version(r.name);report['versions'][r.name]=v;report['requirements'][r.name]=v in r.specifier
|
|
for name in ['spider.config','spider.simulators.mjwp','spider.optimizers.sampling_fast','spider.viewers','spider.preprocess.ik','open3d','pymeshlab','mink']:
|
|
importlib.import_module(name);report['imports'][name]=True
|
|
report['torch_path']=torch.__file__;report['cuda_available']=torch.cuda.is_available()
|
|
assert report['cuda_available'],'GPU test requires access outside the device-restricted sandbox'
|
|
report['gpu']=torch.cuda.get_device_name(0)
|
|
assert float((torch.ones(8,device='cuda')*2).sum())==16
|
|
wp.init()
|
|
xml='<mujoco><option timestep="0.002" iterations="5"/><worldbody><geom type="plane" size="1 1 .1"/><body pos="0 0 .105"><freejoint/><geom type="sphere" size=".1" mass=".1"/></body></worldbody></mujoco>'
|
|
m=mujoco.MjModel.from_xml_string(xml);d=mujoco.MjData(m);mujoco.mj_forward(m,d)
|
|
with wp.ScopedDevice('cuda:0'):
|
|
wm=mw.put_model(m);wd=mw.put_data(m,d,nworld=2,nconmax=32,njmax=128)
|
|
for _ in range(60):mw.step(wm,wd)
|
|
wp.synchronize()
|
|
q=wd.qpos.numpy();time=wd.time.numpy();contacts=wd.nacon.numpy()
|
|
assert np.isfinite(q).all() and np.all(q[:,2]>.085) and np.all(q[:,2]<.105),q
|
|
assert np.allclose(time,.12,atol=1e-5),time
|
|
report['gpu_physics']={'worlds':2,'steps':60,'time_s':time.tolist(),'sphere_height_m':q[:,2].tolist(),'contacts':contacts.tolist(),'passed':True}
|
|
assert all(report['requirements'].values())
|
|
out=ROOT/'output/spider_install';out.mkdir(exist_ok=True)
|
|
(out/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|