59 lines
3.6 KiB
Python
59 lines
3.6 KiB
Python
"""Interactive MuJoCo full replay, or export the exact same poses as video."""
|
|
import argparse,os,time,json
|
|
from pathlib import Path
|
|
p=argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument('--directory',type=Path,default=Path(__file__).resolve().parents[1]/'output/l20_full_replay')
|
|
p.add_argument('--start',type=int,default=0);p.add_argument('--render',action='store_true');p.add_argument('--frames',type=int,default=0);p.add_argument('--check',action='store_true')
|
|
a=p.parse_args()
|
|
if a.render or a.check:os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
import numpy as np
|
|
import mujoco
|
|
motion=np.load(a.directory/'motion.npz');q=motion['qpos'];fps=float(motion['fps']);N=len(q)
|
|
m=mujoco.MjModel.from_xml_path(str(a.directory/'scene.xml'));d=mujoco.MjData(m)
|
|
option=mujoco.MjvOption();option.sitegroup[:]=0;option.geomgroup[3:]=0
|
|
cam=mujoco.MjvCamera();cam.lookat[:]=motion['camera_center'];cam.distance=1.05;cam.azimuth=135;cam.elevation=-23
|
|
if a.check:
|
|
for row in q:d.qpos[:]=row;mujoco.mj_forward(m,d)
|
|
print(json.dumps(dict(frames=N,nq=m.nq,actuators=m.nu,finite=bool(np.isfinite(q).all()),mode='kinematic')))
|
|
elif a.render:
|
|
import imageio.v2 as imageio
|
|
from PIL import Image,ImageDraw,ImageFont
|
|
total=min(a.frames or N,N-a.start)
|
|
renderer=mujoco.Renderer(m,height=720,width=640)
|
|
close=mujoco.MjvCamera();close.distance=.62;close.azimuth=35;close.elevation=-12
|
|
font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',16)
|
|
path=a.directory/('preview.mp4' if a.frames else 'full_replay.mp4')
|
|
writer=imageio.get_writer(path,fps=fps,codec='libx264',quality=8,macro_block_size=1)
|
|
try:
|
|
for t in range(a.start,a.start+total):
|
|
d.qpos[:]=q[t];mujoco.mj_forward(m,d)
|
|
close.lookat[:]=(d.xpos[m.body('hand_base_link').id]+d.xpos[m.body('right_object').id])/2;close.lookat[2]+=.06
|
|
panels=[]
|
|
for c in [cam,close]:renderer.update_scene(d,camera=c,scene_option=option);panels.append(renderer.render().copy())
|
|
img=Image.fromarray(np.concatenate(panels,1));draw=ImageDraw.Draw(img)
|
|
draw.rectangle((0,0,1280,64),fill=(28,38,49))
|
|
draw.text((18,9),f'L20 + bottle | full video motion | {t/fps:05.2f} / {N/fps:.2f} s',font=font,fill='white')
|
|
draw.text((18,36),'World view - moving wrist',font=font,fill=(195,220,235));draw.text((660,36),'Grasp detail - linear URDF coupling',font=font,fill=(195,220,235))
|
|
draw.rectangle((0,690,1280,720),fill=(28,38,49));draw.text((18,696),'Kinematic contact replay | estimated bottle trajectory | floor + ambient lighting',font=font,fill=(210,220,230))
|
|
writer.append_data(np.asarray(img))
|
|
if t in [a.start,0,160,600,N//2,N-1]:img.save(a.directory/f'preview_{t:04d}.png')
|
|
if t%100==0:print(f'Render {t}/{total}',flush=True)
|
|
finally:writer.close();renderer.close()
|
|
print('RENDER_COMPLETE',path,flush=True)
|
|
else:
|
|
import mujoco.viewer
|
|
state={'paused':False,'frame':0.,'last':time.perf_counter()}
|
|
def key(k):
|
|
if k==32:state['paused']=not state['paused']
|
|
elif k in [82,114]:state['frame']=0.
|
|
elif k==262:state['frame']=(int(state['frame'])+1)%N;state['paused']=True
|
|
elif k==263:state['frame']=(int(state['frame'])-1)%N;state['paused']=True
|
|
with mujoco.viewer.launch_passive(m,d,key_callback=key) as viewer:
|
|
viewer.cam.lookat[:]=cam.lookat;viewer.cam.distance=cam.distance;viewer.cam.azimuth=cam.azimuth;viewer.cam.elevation=cam.elevation
|
|
viewer.opt.sitegroup[:]=0;viewer.opt.geomgroup[3:]=0
|
|
while viewer.is_running():
|
|
now=time.perf_counter();elapsed=now-state['last'];state['last']=now
|
|
if not state['paused']:state['frame']=(state['frame']+elapsed*fps)%N
|
|
with viewer.lock():d.qpos[:]=q[int(state['frame'])];d.time=int(state['frame'])/fps;mujoco.mj_forward(m,d)
|
|
viewer.sync();time.sleep(.005)
|