mjx_planner/visualizer.py at main · patsyuk03/mjx_planner · 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
67
68
69
70
71
import os
import numpy as np
import time


import mujoco
from mujoco import viewer


class Visualizer():
    def __init__(self, ctrl: bool=False, traj: bool=False):
        self.init_joint_state = np.array([1.5, -1.8, 1.75, -1.25, -1.6, 0])

        if ctrl:
            model_path = f"{os.path.dirname(__file__)}/ur5e_hande_mjx/scene_control.xml"
        else:
            model_path = f"{os.path.dirname(__file__)}/ur5e_hande_mjx/scene.xml"

        self.model = mujoco.MjModel.from_xml_path(model_path)
        self.model.opt.timestep = 0.05
        self.data = mujoco.MjData(self.model)
        self.data.qpos[:6] = self.init_joint_state

        if traj:
            file_path = f"{os.path.dirname(__file__)}/data/thetadot.csv"
            self.thetadot = np.genfromtxt(file_path, delimiter=',')


    def view_model(self):
        viewer.launch(self.model, self.data)

    def view_traj_mujoco(self):
        with viewer.launch_passive(self.model, self.data) as viewer_:
            viewer_.cam.distance = 4
            viewer_.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True
            viewer_.opt.sitegroup[:] = False
            viewer_.opt.sitegroup[1] = True

            i = 0
            while viewer_.is_running():
                step_start = time.time()
                self.data.qvel[:6] = self.thetadot[i]

                mujoco.mj_step(self.model, self.data)
                viewer_.sync()

                time_until_next_step = self.model.opt.timestep - (time.time() - step_start)
                if time_until_next_step > 0:
                    time.sleep(time_until_next_step)

                if i < self.thetadot.shape[0]-1:
                    i+=1
                else:
                    self.data.qvel[:6] = np.zeros(6)
                    self.data.qpos[:6] = self.init_joint_state
                    i=0


def main():
    viz = Visualizer(ctrl=False, traj=True)
    # viz.view_model()
    viz.view_traj_mujoco()


if __name__=="__main__":
    main()