SimpleMotionRetargeting/vis_mujoco.py at main · HongtuZ/SimpleMotionRetargeting · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import mujoco
import mujoco.viewer
import time
import argparse
import numpy as np
import joblib
from omegaconf import OmegaConf

def get_nonfree_joint_order(model):
    qpos_names = []
    for i in range(model.njnt):
        if model.jnt_type[i] == mujoco.mjtJoint.mjJNT_FREE:
            continue
        name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)
        qpos_names.append((model.jnt_qposadr[i], name))
    qpos_names.sort(key=lambda x: x[0])
    return [n for _, n in qpos_names]

def vis_mujoco(xml_path, motion_file_path):
    print('Loading motion data...')
    motion_data = joblib.load(motion_file_path)
    dt = 1.0 / motion_data['fps']
    print('FPS:', motion_data['fps'])

    print("Loading Mujoco model...")
    mj_model = mujoco.MjModel.from_xml_path(xml_path)
    mj_model.opt.timestep = dt
    mj_data = mujoco.MjData(mj_model)

    dof_names = get_nonfree_joint_order(mj_model)
    print('----------------mujoco dof names----------------')
    for i, name in enumerate(dof_names):
        print(i, name)
    print('----------------motion dof names----------------')
    for i, name in enumerate(motion_data['dof_list']):
        print(i, name)
    print('------------------------------------------------')

    print('Replay motion ...')

    with mujoco.viewer.launch_passive(mj_model, mj_data) as viewer:
        viewer.cam.distance = 2  # 相机距离
        # viewer.cam.azimuth = 180  # 水平旋转角度
        # 主循环
        cnt = 0
        for root_pos, root_rot, dof_pos in zip(motion_data['root_pos'], motion_data['root_rot'], motion_data['dof_pos']):
            t1 = time.perf_counter()
            mj_data.qpos[:3] = root_pos
            mj_data.qpos[3:7] = np.roll(root_rot, 1)  # xyzw -> wxyz
            mj_data.qpos[7:] = dof_pos
            mujoco.mj_forward(mj_model, mj_data)
            viewer.cam.lookat = root_pos
            viewer.sync()
            cnt += 1
            if cnt > 1000 and cnt % 100 == 0:
                print(cnt)
            time.sleep(max(0, dt - (time.perf_counter() - t1))) # 控制循环频率


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("config_path", type=str, help="Path to the config path")
    parser.add_argument("--data", type=str, help="Path to the robot motion file", default=None)
    args = parser.parse_args()
    config = OmegaConf.load(args.config_path)
    vis_mujoco(config.mujoco.xml_path, args.data)