-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy patharm_dynamics_base.py
129 lines (108 loc) · 4.74 KB
/
arm_dynamics_base.py
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
import numpy as np
import time
from geometry import rot, xaxis, yaxis
class ArmDynamicsBase:
""" Base class for arm dynamics """
def __init__(self, num_links, link_mass, link_length, joint_viscous_friction=0.001, gravity=True, dt=0.001):
self.num_links = num_links
self.link_length = link_length
self.link_mass = link_mass
self.joint_viscous_friction = joint_viscous_friction
self.gravity = gravity
self.dt = dt
self.link_lengths = np.full(self.num_links, float(self.link_length))
self.link_masses = np.full(self.num_links, float(self.link_mass))
self.link_inertias = np.zeros(self.num_links)
self.model_loaded = False
for i in range(0, self.num_links):
self.link_inertias[i] = self.link_masses[i] * self.link_lengths[i] * self.link_lengths[i] / 12.0
self.cpu_time_consumed = 0
self.residue_limit = 1e-4
self.residue_limit_flag = False
def get_num_links(self):
return self.num_links
def get_link_lengths(self):
return self.link_lengths
def get_state_dim(self):
return 2 * self.num_links
def get_action_dim(self):
return self.num_links
def get_pos_0(self, state):
return np.zeros((2, 1))
def get_vel_0(self, state):
return np.zeros((2, 1))
def get_q(self, state):
return state[0:self.num_links]
def get_qd(self, state):
return state[self.num_links: 2 * self.num_links]
def compute_theta(self, q):
return np.cumsum(q, axis=0)
def compute_omega(self, qd):
return np.cumsum(qd, axis=0)
def compute_pos(self, pos_0, theta):
pos_0 = pos_0.copy()
theta = theta.copy()
pos = np.zeros((2 * self.num_links, 1))
pos[0:2] = pos_0
for i in range(1, self.num_links):
pos[2 * i:2 * (i + 1)] = pos[2 * (i - 1):2 * i] + np.dot(rot(theta[i - 1]),
self.link_lengths[i - 1] * xaxis())
return pos
def compute_pos_com(self, pos, theta):
pos_com = np.zeros((2 * self.num_links, 1))
for i in range(0, self.num_links):
pos_com[2 * i:2 * (i + 1)] = pos[2 * i:2 * (i + 1)] + np.dot(rot(theta[i]),
0.5 * self.link_lengths[0] * xaxis())
return pos_com
def compute_vel(self, vel_0, omega, theta):
vel_0 = vel_0.copy()
theta = theta.copy()
omega = omega.copy()
vel = np.zeros((2 * self.num_links, 1))
vel[0:2] = vel_0
vel_world = np.zeros((2 * self.num_links, 1))
vel_world[0:2] = (np.dot(rot(theta[0]), vel_0))
for i in range(1, self.num_links):
vel_world[2 * i:2 * (i + 1)] = vel_world[2 * (i - 1):2 * i] + (
np.dot(rot(theta[i - 1]), omega[i - 1] * self.link_lengths[i - 1] * yaxis()))
vel[2 * i:2 * (i + 1)] = np.dot(rot(-1.0 * theta[i]), vel_world[2 * i:2 * (i + 1)])
return vel
def compute_vel_com(self, vel, omega):
vel = vel.copy()
omega = omega.copy()
vel_com = np.zeros((2 * self.num_links, 1))
for i in range(0, self.num_links):
vel_com[2 * i + 0] = vel[2 * i + 0]
vel_com[2 * i + 1] = vel[2 * i + 1] + omega[i] * 0.5 * self.link_lengths[i]
return vel_com
def advance(self, state, action):
""" Forward simulation to compute new state given state and action """
start_time = time.time()
dt = self.dt
new_state = self.dynamics_step(state, action, dt)
self.cpu_time_consumed += time.time() - start_time
return new_state
def compute_energy(self, state):
""" Computes lagrangian, kinetic engery, potential energy of the arm """
pos_0 = self.get_pos_0(state)
vel_0 = self.get_vel_0(state)
q = self.get_q(state)
qd = self.get_qd(state)
theta = self.compute_theta(q)
omega = self.compute_omega(qd)
pos = self.compute_pos(pos_0, theta)
pos_com = self.compute_pos_com(pos, theta)
vel = self.compute_vel(vel_0, omega, theta)
vel_com = self.compute_vel_com(vel, omega)
# Sum energy for all links
T = 0.0 # Kinetic energy
V = 0.0 # Potential energy
for i in range(0, self.num_links):
T += 0.5 * self.link_masses[i] * (np.linalg.norm(vel_com[2 * i:2 * (i + 1)]) ** 2) \
+ 0.5 * self.link_inertias[i] * (omega[i] ** 2)
V += 9.8 * pos_com[2 * i + 1, 0] * self.link_masses[i]
L = T - V # Lagrangian
return L, T, V
def dynamics_step(self, state, action, dt):
""" Advance simulation by one time step """
raise NotImplementedError