8000 Fix numpy contiguous order · hzyjerry/robotics-toolbox-python@2b92a92 · GitHub
[go: up one dir, main page]
Skip to content

Commit 2b92a92

Browse files
hzyjerryhzyjerry
committed
Fix numpy contiguous order
1 parent 21a36ec commit 2b92a92

File tree

2 files changed

+145
-111
lines changed

2 files changed

+145
-111
lines changed

‎roboticstoolbox/cuda/fknm.cu

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,16 +44,16 @@ __global__ void _jacob0(double *T,
4444
double *etool_i;
4545
double *invU;
4646
double *link_iA; // TODO: =ret ?
47-
4847
U = (double*) malloc(sizeof(double) * 16);
4948
invU = (double*) malloc(sizeof(double) * 16);
5049
temp = (double*) malloc(sizeof(double) * 16);
5150
int j = 0;
5251

53-
T_i = &T[tid * 16];
5452
tool_i = &tool[tid * 16];
5553
etool_i = &etool[tid * 16];
54+
5655
_eye(U);
56+
T_i = &T[tid * 16];
5757

5858
if (tid >= N) {
5959
return;
@@ -329,8 +329,10 @@ __device__ int _inv(double *m, double *invOut)
329329

330330
det = m[0] * inv[0] + m[1] * inv[4] + m[2] * inv[8] + m[3] * inv[12];
331331

332-
if (det == 0)
332+
if (det == 0) {
333+
free(inv);
333334
return 0;
335+
}
334336

335337
det = 1.0 / det;
336338

‎roboticstoolbox/examples/chomp.py

Lines changed: 140 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,92 @@
1515
visualize = True
1616
use_mesh = True
1717
parallel = False
18+
fknm_ = None
1819

1920

20-
def chomp():
21+
def jacob0_loop(robot, link, pts, qt=None):
22+
"""
23+
Non-parallel, use for loop
24+
:param pts: (num, 3) xyz positions of pts
25+
:param q: (cdim,) joint configurations
26+
"""
27+
if qt is None:
28+
qt = robot.q
29+
jacob_vec = []
30+
for pt in pts:
31+
JJ = robot.jacob0(qt, end=link, tool=SE3(pt)) # (k, 6)
32+
jacob_vec.append(JJ)
33+
return np.array(jacob_vec)
34+
35+
def jacob0_vec(robot, link, pts, qt=None, verbose=False):
36+
"""
37+
Parallel, use CUDA
38+
:param pts: (num, 3) xyz positions of pts
39+
:param q: (cdim,) joint configurations
40+
"""
41+
import ctypes as ct
42+
global fknm_
43+
if qt is None:
44+
qt = robot.q
45+
if fknm_ is None:
46+
fknm_=np.ctypeslib.load_library('roboticstoolbox/cuda/fknm','.')
47+
# Parallel, use cuda
48+
t_start = time.time()
49+
num_pts = len(pts)
50+
se3_pts = SE3(pts)
51+
pts_tool = np.array(se3_pts.A)
52+
link_base = robot.fkine(qt, end=link)
53+
print(f"1: {time.time() - t_start:.3f}\n")
54+
t_start = time.time()
55+
# pts_mat = np.array((link_base @ se3_pts).A)
56+
pts_mat = np.array(link_base.A.dot(np.array(se3_pts.A)).swapaxes(0, 1), order='C')
57+
e_pts = np.zeros((num_pts, 3))
58+
pts_etool = np.array(SE3(e_pts).A)
59+
print(f"2: {time.time() - t_start:.3f}\n")
60+
t_start = time.time()
61+
link_As = []
62+
link_axes = []
63+
link_isjoint = []
64+
path, njoints, _ = robot.get_path(end=link)
65+
nlinks = len(path)
66+
for il, link in enumerate(path):
67+
axis = get_axis(link)
68+
link_As.append(link.A(qt[link.jindex]).A)
69+
link_axes.append(axis)
70+
link_isjoint.append(link.isjoint)
71+
link_As = np.array(link_As)
72+
link_axes = np.array(link_axes, dtype=int)
73+
link_isjoint = np.array(link_isjoint, dtype=int)
74+
jacob_vec = np.zeros((num_pts, 6, njoints))
75+
76+
if verbose:
77+
print(f"pts shape {pts_mat.shape}")
78+
print(f"pts_tool shape {pts_tool.shape}")
79+
print(f"pts_etool shape {pts_etool.shape}")
80+
print(f"link_As shape {link_As.shape}")
81+
print(f"link_axes shape {link_axes.shape}")
82+
print(f"link_isjoint shape {link_isjoint.shape}")
83+
print(f"jacob_vec shape {jacob_vec.shape}")
84+
print(f"nlinks {nlinks}")
85+
print(f"njoints {njoints}")
86+
print(f"num_pts {num_pts}")
87+
88+
fknm_.jacob0(pts_mat.ctypes.data_as(ct.c_void_p),
89+
pts_tool.ctypes.data_as(ct.c_void_p),
90+
pts_etool.ctypes.data_as(ct.c_void_p),
91+
# link_As.ctypes.data_as(ct.c_void_p),
92+
ct.c_void_p(link_As.ctypes.data),
93+
link_axes.ctypes.data_as(ct.c_void_p),
94+
link_isjoint.ctypes.data_as(ct.c_void_p),
95+
ct.c_int(num_pts),
96+
ct.c_int(nlinks),
97+
ct.c_int(njoints),
98+
jacob_vec.ctypes.data_as(ct.c_void_p))
99+
return jacob_vec
100+
101+
def chomp(seed=0):
102+
np.random.seed(seed)
103+
21104
robot = rtb.models.DH.Panda() # load Mesh version (for collisions)
22105
# robot = rtb.models.URDF.Panda() # load URDF version of the Panda (for visuals)
23106
# print(robot)
@@ -87,10 +170,6 @@ def chomp():
87170
AA /= dt * dt * (nq + 1)
88171
Ainv = np.linalg.pinv(AA)
89172

90-
fknm_ = None
91-
if parallel:
92-
fknm_ = np.ctypeslib.load_library('roboticstoolbox/cuda/fknm','.')
93-
94173
for t in range(iters):
95174
nabla_smooth = AA.dot(xi) + bb
96175
nabla_obs = np.zeros(xidim)
@@ -107,56 +186,54 @@ def chomp():
107186
else:
108187
qd = 0.5 * (xi[cdim * (i+1): cdim * (i+2)] - xi[cdim * (i-1): cdim * (i)])
109188

110-
import pdb; pdb.set_trace()
189+
# import pdb; pdb.set_trace()
111190
for link in robot.links: # bodypart
112191
k = link.jindex
113192
if k is None:
114193
continue
115194

116195
link_base = robot.fkine(robot.q, end=link) # x_current, (4, 4)
117196
delta_nabla_obs = np.zeros(k + 1)
197+
mesh = meshes[link.name]
198+
idxs = np.random.choice(np.arange(len(mesh.vertices)), num_pts, replace=False)
199+
pts = mesh.vertices[idxs]
200+
118201
if not parallel:
119-
# Non-parallel, use for loop
120-
mesh = meshes[link.name]
121-
for j in range(num_pts):
122-
# For each point: compute Jacobian, compute cost, compute cost gradient
123-
124-
pt_rel = mesh.vertices[j]
125-
pt_tool = link_base @ SE3(pt_rel)
126-
pt_pos = pt_tool.t
127-
128-
JJ = robot.jacob0(qt, end=link, tool=SE3(pt_rel)) # (k, 6)
129-
import pdb; pdb.set_trace()
130-
xd = JJ.dot(qd[:k+1]) # x_delta
131-
vel = np.linalg.norm(xd)
132-
133-
xdn = xd / (vel + 1e-3) # speed normalized
134-
xdd = JJ.dot(xidd[cdim * i: cdim * i + k + 1]) # second derivative of xi
135-
prj = np.eye(6) - xdn[:, None].dot(xdn[:, None].T) # curvature vector (6, 6)
136-
kappa = (prj.dot(xdd) / (vel ** 2 + 1e-3)) # (6,)
137-
138-
cost = np.sum(pt_pos)
139-
total_cost += cost / num_pts
140-
# delta := negative gradient of obstacle cost in work space, (6, cdim)
141-
delta = -1 * np.concatenate([[1, 1, 0], np.zeros(3)])
142-
# for link in robot.links:
143-
# cost = get_link_cost(robot, meshes, link)
144-
delta_nabla_obs += JJ.T.dot(vel).dot(prj.dot(delta) - cost * kappa)
145-
146-
nabla_obs[cdim * i: cdim * i + k + 1] += (delta_nabla_obs / num_pts)
202+
jacobs = jacob0_loop(robot, link, pts, qt)
147203
else:
148-
# Parallel, use cuda
149-
import ctypes as ct
150-
pt_rel = np.array(mesh.vertices[j])
204+
jacobs = jacob0_vec(robot, link, pts, qt)
205+
206+
# TODO: vectorize cost calculation
207+
for j in range(num_pts):
208+
# For each point: compute Jacobian, compute cost, compute cost gradient
209+
pt_rel = mesh.vertices[j]
210+
pt_tool = link_base @ SE3(pt_rel)
211+
pt_pos = pt_tool.t
212+
213+
JJ = robot.jacob0(qt, end=link, tool=SE3(pt_rel)) # (k, 6)
214+
xd = JJ.dot(qd[:k+1]) # x_delta
215+
vel = np.linalg.norm(xd)
216+
xdn = xd / (vel + 1e-3) # speed normalized
217+
xdd = JJ.dot(xidd[cdim * i: cdim * i + k + 1]) # second derivative of xi
218+
prj = np.eye(6) - xdn[:, None].dot(xdn[:, None].T) # curvature vector (6, 6)
219+
kappa = (prj.dot(xdd) / (vel ** 2 + 1e-3)) # (6,)
220+
221+
cost = np.sum(pt_pos)
222+
total_cost += cost / num_pts
223+
# delta := negative gradient of obstacle cost in work space, (6, cdim)
224+
delta = -1 * np.concatenate([[1, 1, 0], np.zeros(3)])
225+
# for link in robot.links:
226+
# cost = get_link_cost(robot, meshes, link)
227+
delta_nabla_obs += JJ.T.dot(vel).dot(prj.dot(delta) - cost * kappa)
228+
229+
nabla_obs[cdim * i: cdim * i + k + 1] += (delta_nabla_obs / num_pts)
151230

152231

153232
# dxi = Ainv.dot(lmbda * nabla_smooth)
154233
dxi = Ainv.dot(nabla_obs + lmbda * nabla_smooth)
155234
xi -= dxi / eta
156235
print(f"Iteration {t} total cost {total_cost}")
157236

158-
# @vectorize(['float32(float32)'], target='cuda')
159-
# def vec_jacob0()
160237

161238
if visualize:
162239
from roboticstoolbox.backends.swift import Swift # instantiate 3D browser-based visualizer
@@ -245,8 +322,8 @@ def test_parallel():
245322
T = SE3(0.7, 0.2, 0.1) * SE3.OA([0, 1, 0], [0, 0, -1])
246323
sol = robot.ikine_LM(T) # solve IK
247324
q_pickup = sol.q
248-
qtraj = rtb.jtraj(robot.qz, q_pickup, 50)
249325
robot = rtb.models.URDF.Panda() # load URDF version of the Panda
326+
robot.q = q_pickup
250327
meshes = {}
251328
for link in robot.links:
252329
if len(link.geometry) != 1:
@@ -264,86 +341,41 @@ def test_parallel():
264341
lmbda = 1000
265342
eta = 1000
266343
iters = 4
267-
num_pts = 5000
344+
num_pts = 2000
268345

269346

270347
# Make cost field, starting & end points
271-
cdim = len(qtraj.q[0])
348+
cdim = len(robot.q)
272349
xidim = cdim * nq
273350
AA = np.zeros((xidim, xidim))
274351
xi = np.zeros(xidim)
275-
robot._set_link_fk(qtraj.q[1])
276-
277352

278-
link = robot.links[-2]
279-
k = link.jindex
280-
link_base = robot.fkine(robot.q, end=link)
281353

354+
for idx in [2, 3, 4, 5, 6, 7]:
355+
print(f"Iidx {idx}")
356+
link = robot.links[idx]
357+
k = link.jindex
282358

283-
# Non-parallel, use for loop
284-
mesh = meshes[link.name]
285-
jacob_loop = []
286-
t_start = time.time()
287-
for j in range(num_pts):
288-
pt_rel = mesh.vertices[j]
289-
JJ = robot.jacob0(robot.q, end=link, tool=SE3(pt_rel)) # (
290-
jacob_loop.append(JJ)
291-
jacob_loop = np.array(jacob_loop)
292-
print(f"\nNon-parallel time: {time.time() - t_start:.3f}\n")
293-
# print(jacob_loop)
294359

360+
# Non-parallel, use for loop
361+
mesh = meshes[link.name]
362+
pts = mesh.vertices[:num_pts]
295363

296-
# Parallel, use cuda
297-
import ctypes as ct
298-
fknm_=np.ctypeslib.load_library('roboticstoolbox/cuda/fknm','.')
299-
pts = np.array(meshes[link.name].vertices[:num_pts])
300-
pts_mat = np.array((link_base @ SE3(pts)).A)
301-
e_pts = np.zeros((num_pts, 3))
302-
pts_tool = np.array(SE3(pts).A)
303-
pts_etool = np.array(SE3(e_pts).A)
304-
link_As = []
305-
link_axes = []
306-
link_isjoint = []
307-
path, njoints, _ = robot.get_path(end=link)
308-
nlinks = len(path)
364+
t_start = time.time()
365+
jacobs1 = jacob0_loop(robot, link, pts)
366+
print(f"\nNon-parallel time: {time.time() - t_start:.3f}\n")
367+
# print(jacobs1)
368+
# Parallel, use cuda
369+
t_start = time.time()
370+
jacobs2 = jacob0_vec(robot, link, pts, verbose=True)
371+
print(f"\nParallel time: {time.time() - t_start:.3f}\n")
372+
# print(jacobs2)
373+
# print(f"Max diff 0 {np.max(np.abs(jacobs1[0] - jacobs2[0]))}")
374+
# print(f"Max diff 0 {np.argmax(np.abs(jacobs1 - jacobs2))}")
375+
# print(f"Max diff {np.max(np.abs(jacobs1 - jacobs2))}")
376+
# import pdb; pdb.set_trace()
377+
assert np.all(np.isclose(jacobs1, jacobs2))
309378

310-
for il, link in enumerate(path):
311-
axis = get_axis(link)
312-
link_As.append(link.A().A)
313-
link_axes.append(axis)
314-
link_isjoint.append(link.isjoint)
315-
link_As = np.array(link_As)
316-
link_axes = np.array(link_axes, dtype=int)
317-
link_isjoint = np.array(link_isjoint, dtype=int)
318-
jacob_vec = np.zeros((num_pts, 6, njoints))
319-
320-
321-
print(f"pts shape {pts_mat.shape}")
322-
print(f"pts_tool shape {pts_tool.shape}")
323-
print(f"pts_etool shape {pts_etool.shape}")
324-
print(f"link_As shape {link_As.shape}")
325-
print(f"link_axes shape {link_axes.shape}")
326-
print(f"link_isjoint shape {link_isjoint.shape}")
327-
print(f"jacob_vec shape {jacob_vec.shape}")
328-
print(f"nlinks {nlinks}")
329-
print(f"njoints {njoints}")
330-
print(f"num_pts {num_pts}")
331-
332-
t_start = time.time()
333-
fknm_.jacob0(pts_mat.ctypes.data_as(ct.c_void_p),
334-
pts_tool.ctypes.data_as(ct.c_void_p),
335-
pts_etool.ctypes.data_as(ct.c_void_p),
336-
# link_As.ctypes.data_as(ct.c_void_p),
337-
ct.c_void_p(link_As.ctypes.data),
338-
link_axes.ctypes.data_as(ct.c_void_p),
339-
link_isjoint.ctypes.data_as(ct.c_void_p),
340-
ct.c_int(num_pts),
341-
ct.c_int(nlinks),
342-
ct.c_int(njoints),
343-
jacob_vec.ctypes.data_as(ct.c_void_p))
344-
print(f"\nParallel time: {time.time() - t_start:.3f}\n")
345-
# print(jacob_vec)
346-
assert np.all(np.isclose(jacob_loop, jacob_vec))
347379

348380
def get_axis(link):
349381
axis = None

0 commit comments

Comments
 (0)
0