mujoco-learning/move_ball.py at main · LitchiCheng/mujoco-learning · 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import mujoco
import mujoco.viewer
import numpy as np
from pynput import keyboard

key_states = {
    keyboard.Key.up: False,
    keyboard.Key.down: False,
    keyboard.Key.left: False,
    keyboard.Key.right: False,
    keyboard.Key.alt_l: False,
    keyboard.Key.alt_r: False,
}

def on_press(key):
    if key in key_states:
        key_states[key] = True

def on_release(key):
    if key in key_states:
        key_states[key] = False

listener = keyboard.Listener(on_press=on_press, on_release=on_release)
listener.start()

XML = """
<mujoco>
  <visual>
    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
    <rgba haze="0.15 0.25 0.35 1"/>
    <global azimuth="140" elevation="-30"/>
  </visual>

  <asset>
    <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
    <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
      markrgb="0.8 0.8 0.8" width="300" height="300"/>
    <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
  </asset>

  <worldbody>
    <geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>

    <body name="ball" pos="0 0 1">
      <freejoint name="free_joint"/>
      <geom type="sphere" size="0.2" rgba="1 0 0 1" mass="1"/>
    </body>
  </worldbody>
</mujoco>
"""

model = mujoco.MjModel.from_xml_string(XML)
model.opt.gravity = (0, 0, 0)
data = mujoco.MjData(model)

class CustomViewer:
    def __init__(self, model, data):
        self.handle = mujoco.viewer.launch_passive(model, data)
        self.pos = 0.0001

    def is_running(self):
        return self.handle.is_running()

    def sync(self):
        self.handle.sync()

    @property
    def cam(self):
        return self.handle.cam

    @property
    def viewport(self):
        return self.handle.viewport

    def run_loop(self):
        while self.is_running():
            ball_body_name = "ball"
            pos = data.body(ball_body_name).xpos
            quat = data.body(ball_body_name).xquat
            print(f"Position: {pos}, Orientation: {quat}")

            if key_states[keyboard.Key.up]:
                data.qpos[0] += self.pos
            if key_states[keyboard.Key.down]:
                data.qpos[0] -= self.pos
            if key_states[keyboard.Key.left]:
                data.qpos[1] += self.pos
            if key_states[keyboard.Key.right]:
                data.qpos[1] -= self.pos
            if key_states[keyboard.Key.alt_l]:
                data.qpos[2] += self.pos
            if key_states[keyboard.Key.alt_r]:
                data.qpos[2] -= self.pos

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


viewer = CustomViewer(model, data)
viewer.cam.distance = 3
viewer.cam.azimuth = 0
viewer.cam.elevation = -30
viewer.run_loop()