mujoco-learning/panda_dynamics_hold.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
import src.mujoco_viewer as mujoco_viewer
import time
import mujoco
import numpy as np
from scipy.spatial.transform import Rotation as R
import matplotlib.pyplot as plt
import pinocchio as pin

class PandaEnv(mujoco_viewer.CustomViewer):
    def __init__(self, scene_xml, arm_xml):
        super().__init__(scene_xml, 3, azimuth=-45, elevation=-30)
        self.scene_xml = scene_xml
        self.arm_xml = arm_xml

        self.initial_pos = self.model.key_qpos[0]
        self.data.qpos[:7] = self.initial_pos[:7]
        self.data.qvel[:7] = np.zeros(7)

        self.step = 0
        self.step_list = []
        self.dynamics_tau_list = []
        self.damping_tau_list = []

    def runBefore(self):
        self.pin_model = pin.RobotWrapper.BuildFromMJCF(self.arm_xml).model
        self.pin_data = self.pin_model.createData()
        self.last_v = np.zeros(7)

    def runFunc(self):
        q = self.data.qpos[:7]
        v = self.data.qvel[:7]
        a = (v - self.last_v) / self.model.opt.timestep
        self.last_v = v.copy()
        # a = np.zeros(7)
        v = np.concatenate((v, np.zeros(2)))
        q = np.concatenate((q, np.zeros(2)))
        a = np.concatenate((a, np.zeros(2)))
        dynamics_tau = pin.rnea(self.pin_model, self.pin_data, q, v, a)

        IMPENDANCE_COMPENSATION = True
        if IMPENDANCE_COMPENSATION:
            self.Kp = np.diag([50] * 7)
            self.Kd = np.diag([100] * 7)
        else:
            self.Kp = np.diag([0] * 7)
            self.Kd = np.diag([0] * 7)
        error = self.initial_pos - q
        impedence_tau = self.Kp @ error[:7] - self.Kd @ v[:7]
        impedence_tau = np.concatenate((impedence_tau, np.zeros(2)))
        tau = dynamics_tau + impedence_tau

        self.data.ctrl[:7] = tau[:7]
        print(f"Total Torque: {np.round(tau[:7], 2)}")
        print(f"impedence_tau: {np.round(impedence_tau[:7], 2)}")

        self.step += 1
        self.step_list.append(self.step)
        self.dynamics_tau_list.append(dynamics_tau[:7].copy())
        self.damping_tau_list.append(impedence_tau[:7].copy())
        if self.step >= 2000:
            self.plotTorque()
    def plotTorque(self):
        steps = np.array(self.step_list)
        dynamics_tau = np.array(self.dynamics_tau_list)
        damping_tau = np.array(self.damping_tau_list)
        fig, axes = plt.subplots(7, 1, figsize=(10, 14), sharex=True)
        fig.suptitle('dynamics_tau & damping_tau', fontsize=14, fontweight='bold')
        joint_names = ['j1', 'j2', 'j4', 'j5', 'j6', 'j7', 'j8']
        for i in range(7):
            axes[i].plot(steps, dynamics_tau[:, i], label=f'dynamics_tau', color='blue', linewidth=1.5)
            axes[i].plot(steps, damping_tau[:, i], label=f'damping_tau', color='red', linewidth=1.5)
            axes[i].set_ylabel(f'{joint_names[i]} Tau (N·m)')
            axes[i].grid(True, alpha=0.3)
            axes[i].legend(loc='upper right')
        axes[-1].set_xlabel('step')
        plt.tight_layout()
        plt.show()


if __name__ == "__main__":
    SCENE_XML = "model/franka_emika_panda/scene_tau.xml"
    ARM_XML = "model/franka_emika_panda/panda_tau.xml"
    env = PandaEnv(SCENE_XML, ARM_XML)
    env.run_loop()