52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
"""Play a hand-only L20 motion file with a free wrist in MuJoCo."""
|
|
import argparse
|
|
import time
|
|
import mujoco
|
|
import mujoco.viewer
|
|
import numpy as np
|
|
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--directory', required=True)
|
|
a = p.parse_args()
|
|
from pathlib import Path
|
|
directory = Path(a.directory)
|
|
with np.load(directory / 'motion.npz') as data:
|
|
motion = {name:data[name] for name in data.files}
|
|
model = mujoco.MjModel.from_xml_path(str(directory / 'l20_moving.xml'))
|
|
data = mujoco.MjData(model)
|
|
addr = [model.jnt_qposadr[model.joint(str(n)).id] for n in motion['joint_names']]
|
|
wrist = model.jnt_qposadr[model.joint('wrist_free').id]
|
|
fps = float(motion['fps'])
|
|
count = len(motion['qpos'])
|
|
state = dict(frame=0., paused=False)
|
|
|
|
def key(k):
|
|
if k == 32:
|
|
state['paused'] = not state['paused']
|
|
elif k in [82,114]:
|
|
state['frame'] = 0.
|
|
elif k in [262,263]:
|
|
state['frame'] = (int(state['frame']) + (1 if k == 262 else -1)) % count
|
|
state['paused'] = True
|
|
|
|
with mujoco.viewer.launch_passive(model, data, key_callback=key) as viewer:
|
|
viewer.cam.lookat[:] = motion['wrist_pos'].mean(0)
|
|
viewer.cam.distance = .9
|
|
viewer.cam.azimuth, viewer.cam.elevation = 135, 10
|
|
viewer.opt.sitegroup[:] = 0
|
|
last = time.perf_counter()
|
|
while viewer.is_running():
|
|
now = time.perf_counter()
|
|
if not state['paused']:
|
|
state['frame'] = (state['frame'] + (now-last)*fps) % count
|
|
last = now
|
|
t = int(state['frame'])
|
|
with viewer.lock():
|
|
data.qpos[addr] = motion['qpos'][t]
|
|
data.qpos[wrist:wrist+3] = motion['wrist_pos'][t]
|
|
data.qpos[wrist+3:wrist+7] = motion['wrist_quat_wxyz'][t]
|
|
data.time = t/fps
|
|
mujoco.mj_forward(model,data)
|
|
viewer.sync()
|
|
time.sleep(.005)
|