bird/mujoco.py at main · ericjenn/bird · 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

# Don't forget to do: export MUJOCO_GL=eglD
# install : pip install mujoco
# (in a virtual env :  python -m venv Venvs/mujoco

# Other imports and helper functions

import mujoco
import time
import itertools
import numpy as np
from mujoco import viewer

# Graphics and plotting.
import mediapy as media
import matplotlib.pyplot as plt

class LevenbegMarquardtIK:

    def __init__(self, model, data, step_size, tol, alpha, jacp, jacr, damping):
        self.model = model
        self.data = data
        self.step_size = step_size
        self.tol = tol
        self.alpha = alpha
        self.jacp = jacp
        self.jacr = jacr
        self.damping = damping

    def check_joint_limits(self, q):
        """Check if the joints is under or above its limits"""
        for i in range(len(q)):
            q[i] = max(self.model.jnt_range[i][0],
                       min(q[i], self.model.jnt_range[i][1]))

    #Levenberg-Marquardt pseudocode implementation
    def calculate(self, goal, init_q, body_id):
        """Calculate the desire joints angles for goal"""
        self.data.qpos = init_q
        mujoco.mj_forward(self.model, self.data)
        current_pose = self.data.body(body_id).xpos
        error = np.subtract(goal, current_pose)

        while (np.linalg.norm(error) >= self.tol):
            #calculate jacobian
            mujoco.mj_jac(self.model, self.data, self.jacp, self.jacr, goal, body_id)
            #calculate delta of joint q
            n = self.jacp.shape[1]
            I = np.identity(n)
            product = self.jacp.T @ self.jacp + self.damping * I

            if np.isclose(np.linalg.det(product), 0):
                j_inv = np.linalg.pinv(product) @ self.jacp.T
            else:
                j_inv = np.linalg.inv(product) @ self.jacp.T

            delta_q = j_inv @ error
            #compute next step
            self.data.qpos += self.step_size * delta_q
            #check limits
            self.check_joint_limits(self.data.qpos)
            #compute forward kinematics
            mujoco.mj_forward(self.model, self.data)
            #calculate new error
            error = np.subtract(goal, self.data.body(body_id).xpos)


# More legible printing from numpy.
np.set_printoptions(precision=3, suppress=True, linewidth=100)

m = mujoco.MjModel.from_xml_path('./mujoco/mjmodel.xml')
d = mujoco.MjData(m)

print(m.ngeom)
print(m.joint('joint1'))

d.geom_xpos


mujoco.mj_kinematics(m, d)
d.qpos[1]=0
print('raw access:\n', d.qpos)

paused = False
deltax=0

def key_callback(keycode):
    global paused, deltax
    print("key pressded")
    if chr(keycode) == ' ':
        paused = not paused
    if chr(keycode) == '1':
        deltax+= 10
    if chr(keycode) == '2':
        deltax-= 10


print('Total number of DoFs in the model:', m.nv)
print('Generalized positions:', d.qpos)
print('Generalized velocities:', d.qvel)

pi = 3.141592
d.qpos = [3*pi/2, -pi/2, pi/2, 3*pi/2, 3*pi/2, 0, 0, 0]

#Initial joint position
qpos0 = d.qpos.copy()


body_id = m.body('gripper').id
jacp = np.zeros((3, m.nv)) #translation jacobian
jacr = np.zeros((3, m.nv)) #rotational jacobian
step_size = 0.5
tol = 0.01
alpha = 0.5
init_q = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
damping = 0.15

ik = LevenbegMarquardtIK(m, d, step_size, tol, alpha, jacp, jacr, damping)


deltax_old = deltax


with mujoco.viewer.launch_passive(m, d, key_callback=key_callback) as viewer:
  # Close the viewer automatically after 30 wall-seconds.
  start = time.time()
  while viewer.is_running():
  #and time.time() - start < 30:
    step_start = time.time()


    #Get desire point
    if  deltax != deltax_old:
        goal = [510.49+deltax, 310.13, 110.59+deltax]
        mujoco.mj_resetDataKeyframe(m, d, 1) #reset qpos to initial value
        ik.calculate(goal, init_q, body_id) #calculate the qpos
        result = d.qpos.copy()
        #Plot results
        d.qpos = result
        mujoco.mj_forward(m,d)
        result_point = d.body('gripper').xpos


    # mj_step can be replaced with code that also evaluates
    # a policy and applies a control signal before stepping the physics.
    if not paused:
      mujoco.mj_step(m, d)
      viewer.sync()