mujoco-learning/pickup_cube.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
import mujoco
import numpy as np
import src.mujoco_viewer as mujoco_viewer
import src.pinocchio_kinematic as pinocchio_kinematic
import src.utils as utils
import src.key_listener as key_listener
from pynput import keyboard

# 按键状态字典
key_states = {
    keyboard.Key.up: False,
    keyboard.Key.down: False,
    keyboard.Key.left: False,
    keyboard.Key.right: False,
    keyboard.Key.page_up: False,
    keyboard.Key.page_down: False,
    keyboard.Key.ctrl_l: False,  # 夹爪开
    keyboard.Key.ctrl_r: False   # 夹爪收
}

class RobotController(mujoco_viewer.CustomViewer):
    def __init__(self, scene_path, arm_path):
        super().__init__(scene_path, distance=2.0, azimuth=-135, elevation=-20)
        self.scene_path = scene_path
        self.arm_path = arm_path

        self.key_listener = key_listener.KeyListener(key_states)
        self.key_listener.start()

        self.arm = pinocchio_kinematic.Kinematics("ee_center_body")
        self.arm.buildFromMJCF(arm_path)
        self.last_dof = None

    def runBefore(self):
        self.model.opt.timestep = 0.005
        self.ee_cmd_x = 0.1
        self.ee_cmd_y = 0.0
        self.ee_cmd_z = 0.4
        self.grip_value = 0
        self.last_dof = self.data.qpos[:9].copy()

    def runFunc(self):
        if key_states[keyboard.Key.up]:
            self.ee_cmd_z += 0.001
        if key_states[keyboard.Key.down]:
            self.ee_cmd_z -= 0.001
        if key_states[keyboard.Key.left]:
            self.ee_cmd_x -= 0.001
        if key_states[keyboard.Key.right]:
            self.ee_cmd_x += 0.001
        if key_states[keyboard.Key.page_up]:
            self.ee_cmd_y += 0.001
        if key_states[keyboard.Key.page_down]:
            self.ee_cmd_y -= 0.001
        if key_states[keyboard.Key.ctrl_l]:
            self.grip_value += 5
        if key_states[keyboard.Key.ctrl_r]:
            self.grip_value -= 5

        tf = utils.transform2mat(self.ee_cmd_x, self.ee_cmd_y, self.ee_cmd_z, np.pi, 0, 0)
        self.dof, info = self.arm.ik(tf, current_arm_motor_q=self.last_dof)
        self.last_dof = self.dof
        self.data.ctrl[:7] = self.dof[:7]
        self.data.ctrl[7] = self.grip_value

        geom_id = self.getGeomIdByName("cube_geom")
        friction = self.model.geom_friction[geom_id]
        condim = self.model.geom_condim[geom_id]
        print(f"condim: {condim} sliding: {friction[0]} torsional: {friction[1]} rolling: {friction[2]}")

if __name__ == "__main__":
    SCENE_XML_PATH = 'model/franka_emika_panda/scene_with_cube.xml'
    ARM_XML_PATH = 'model/franka_emika_panda/panda.xml'
    robot = RobotController(SCENE_XML_PATH, ARM_XML_PATH)
    robot.run_loop()