Coverage for src/pyroboplan/core/utils.py: 87%
131 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-03-15 22:28 -0400
1"""Core utilities for robot modeling."""
3import copy
4import numpy as np
5import pinocchio
8def check_collisions_at_state(
9 model,
10 collision_model,
11 q,
12 data=None,
13 collision_data=None,
14 distance_padding=0.0,
15):
16 """
17 Checks whether a specified joint configuration is collision-free.
19 Parameters
20 ----------
21 model : `pinocchio.Model`
22 The model to use for collision checking.
23 collision_model : `pinocchio.Model`
24 The collision model to use for collision checking.
25 q : array-like
26 The joint configuration of the model.
27 data : `pinocchio.Data`, optional
28 The model data to use for collision checking. If None, data is created automatically.
29 collision_data : `pinocchio.GeometryData`, optional
30 The collision_model data to use for collision checking. If None, data is created automatically.
31 distance_padding : float, optional
32 The padding, in meters, to use for distance to nearest collision.
34 Returns
35 -------
36 bool
37 True if there are any collisions or minimum distance violations, otherwise False.
38 """
39 if not data:
40 data = model.createData()
41 if not collision_data:
42 collision_data = collision_model.createData()
43 stop_at_first_collision = True # For faster computation
45 for elem in collision_data.collisionRequests:
46 elem.security_margin = distance_padding
48 pinocchio.computeCollisions(
49 model, data, collision_model, collision_data, q, stop_at_first_collision
50 )
51 return np.any([cr.isCollision() for cr in collision_data.collisionResults])
54def get_minimum_distance_at_state(
55 model, collision_model, q, data=None, collision_data=None
56):
57 """
58 Gets the minimum distance to collision at a specified state.
60 Parameters
61 ----------
62 model : `pinocchio.Model`
63 The model to use for collision checking.
64 collision_model : `pinocchio.Model`
65 The collision model to use for collision checking.
66 q : array-like
67 The joint configuration of the model.
68 data : `pinocchio.Data`, optional
69 The model data to use for collision checking. If None, data is created automatically.
70 collision_data : `pinocchio.GeometryData`, optional
71 The collision_model data to use for collision checking. If None, data is created automatically.
73 Returns
74 -------
75 float
76 The minimum distance to collision, in meters.
77 """
78 if len(collision_model.collisionPairs) == 0:
79 return np.inf
81 if not data:
82 data = model.createData()
83 if not collision_data:
84 collision_data = collision_model.createData()
86 pinocchio.computeDistances(model, data, collision_model, collision_data, q)
87 return np.min([dr.min_distance for dr in collision_data.distanceResults])
90def check_collisions_along_path(
91 model, collision_model, q_path, data=None, collision_data=None, distance_padding=0.0
92):
93 """
94 Checks whether a path consisting of multiple joint configurations is collision-free.
96 Parameters
97 ----------
98 model : `pinocchio.Model`
99 The model from which to generate a random state.
100 collision_model : `pinocchio.Model`
101 The model to use for collision checking.
102 q_path : list[array-like]
103 A list of joint configurations describing the path.
104 data : `pinocchio.Data`, optional
105 The model data to use for collision checking. If None, data is created automatically.
106 collision_data : `pinocchio.GeometryData`, optional
107 The collision_model data to use for collision checking. If None, data is created automatically.
108 distance_padding : float, optional
109 The padding, in meters, to use for distance to nearest collision.
111 Returns
112 -------
113 bool
114 True if there are any collisions or minimum distance violations, otherwise False.
115 """
116 if not data: 116 ↛ 118line 116 didn't jump to line 118 because the condition on line 116 was always true
117 data = model.createData()
118 if not collision_data: 118 ↛ 120line 118 didn't jump to line 120 because the condition on line 118 was always true
119 collision_data = collision_model.createData()
120 stop_at_first_collision = True # For faster computation
122 for elem in collision_data.collisionRequests:
123 elem.security_margin = distance_padding
125 for q in q_path:
126 pinocchio.computeCollisions(
127 model, data, collision_model, collision_data, q, stop_at_first_collision
128 )
129 if np.any([cr.isCollision() for cr in collision_data.collisionResults]):
130 return True
132 return False
135def configuration_distance(q_start, q_end):
136 """
137 Returns the distance between two joint configurations.
139 Parameters
140 ----------
141 q_start : array-like
142 The start joint configuration.
143 q_end : array-like
144 The end joint configuration.
146 Returns
147 -------
148 float
149 The distance between the two joint configurations.
150 """
151 return np.linalg.norm(q_end - q_start)
154def get_path_length(q_path):
155 """
156 Returns the configuration distance of a path.
158 Parameters
159 ----------
160 q_path : list[array-like]
161 A list of joint configurations describing a path.
163 Returns
164 -------
165 float
166 The total configuration distance of the entire path.
167 """
168 total_distance = 0.0
169 for idx in range(1, len(q_path)):
170 total_distance += configuration_distance(q_path[idx - 1], q_path[idx])
171 return total_distance
174def get_random_state(model, padding=0.0):
175 """
176 Returns a random state that is within the model's joint limits.
178 Parameters
179 ----------
180 model : `pinocchio.Model`
181 The model from which to generate a random state.
182 padding : float or array-like, optional
183 The padding to use around the sampled joint limits.
185 Returns
186 -------
187 array-like
188 A set of randomly generated joint variables.
189 """
190 return np.random.uniform(
191 model.lowerPositionLimit + padding, model.upperPositionLimit - padding
192 )
195def get_random_collision_free_state(
196 model, collision_model, joint_padding=0.0, distance_padding=0.0, max_tries=100
197):
198 """
199 Returns a random state that is within the model's joint limits and is collision-free according to the collision model.
201 Parameters
202 ----------
203 model : `pinocchio.Model`
204 The model from which to generate a random state.
205 collision_model : `pinocchio.Model`
206 The model to use for collision checking.
207 joint_padding : float or array-like, optional
208 The padding to use around the sampled joint limits.
209 distance_padding : float, optional
210 The padding, in meters, to use for distance to nearest collision.
211 max_tries : int, optional
212 The maximum number of tries for sampling a collision-free state.
214 Returns
215 -------
216 array-like
217 A set of randomly generated collision-free joint variables, or None if one cannot be found.
218 """
219 num_tries = 0
220 while num_tries < max_tries: 220 ↛ 228line 220 didn't jump to line 228 because the condition on line 220 was always true
221 state = get_random_state(model, padding=joint_padding)
222 if not check_collisions_at_state(
223 model, collision_model, state, distance_padding=distance_padding
224 ):
225 return state
226 num_tries += 1
228 print(f"Could not generate collision-free state after {max_tries} tries.")
229 return None
232def get_random_transform(model, target_frame, joint_padding=0.0):
233 """
234 Returns a random transform for a target frame that is within the model's joint limits.
236 Parameters
237 ----------
238 model : `pinocchio.Model`
239 The model from which to generate a random transform.
240 target_frame : str
241 The name of the frame for which to generate a random transform.
242 joint_padding : float or array-like, optional
243 The padding to use around the sampled joint limits.
245 Returns
246 -------
247 `pinocchio.SE3`
248 A randomly generated transform for the specified target frame.
249 """
250 q_target = get_random_state(model, joint_padding)
251 data = model.createData()
252 target_frame_id = model.getFrameId(target_frame)
253 pinocchio.framesForwardKinematics(model, data, q_target)
254 return data.oMf[target_frame_id]
257def get_random_collision_free_transform(
258 model,
259 collision_model,
260 target_frame,
261 joint_padding=0.0,
262 distance_padding=0.0,
263 max_tries=100,
264):
265 """
266 Returns a random transform for a target frame that is within the model's joint limits and is collision-free.
268 Parameters
269 ----------
270 model : `pinocchio.Model`
271 The model from which to generate a random transform.
272 collision_model : `pinocchio.Model`
273 The model to use for collision checking.
274 target_frame : str
275 The name of the frame for which to generate a random transform.
276 joint_padding : float or array-like, optional
277 The padding to use around the sampled joint limits.
278 distance_padding : float, optional
279 The padding, in meters, to use for distance to nearest collision.
280 max_tries : int, optional
281 The maximum number of tries for sampling a collision-free state.
283 Returns
284 -------
285 `pinocchio.SE3`
286 A randomly generated transform for the specified target frame.
287 """
288 q_target = get_random_collision_free_state(
289 model,
290 collision_model,
291 joint_padding=joint_padding,
292 distance_padding=distance_padding,
293 max_tries=max_tries,
294 )
295 data = model.createData()
296 target_frame_id = model.getFrameId(target_frame)
297 pinocchio.framesForwardKinematics(model, data, q_target)
298 return data.oMf[target_frame_id]
301def check_within_limits(model, q):
302 """
303 Checks whether a particular joint configuration is within the model's joint limits.
305 Parameters
306 ----------
307 model : `pinocchio.Model`
308 The model from which to generate a random state.
309 q : array-like
310 The joint configuration for the model.
312 Returns
313 -------
314 bool
315 True if the configuration is within joint limits, otherwise False.
316 """
317 return np.all(q >= model.lowerPositionLimit) and np.all(
318 q <= model.upperPositionLimit
319 )
322def extract_cartesian_pose(model, target_frame, q, data=None):
323 """
324 Extracts the Cartesian pose of a specified model frame given a joint configuration.
326 Parameters
327 ----------
328 model : `pinocchio.Model`
329 The model from which to perform forward kinematics.
330 target_frame : str
331 The name of the target frame.
332 q : array-like
333 The joint configuration values describing the robot state.
334 data : `pinocchio.Data`, optional
335 The model data to use. If not set, one will be created.
337 Returns
338 -------
339 `pinocchio.SE3`
340 The transform describing the Cartesian pose of the specified frame at the provided joint configuration.
341 """
342 if data is None:
343 data = model.createData()
344 target_frame_id = model.getFrameId(target_frame)
345 pinocchio.framesForwardKinematics(model, data, q)
346 return copy.deepcopy(data.oMf[target_frame_id])
349def extract_cartesian_poses(model, target_frame, q_vec, data=None):
350 """
351 Extracts the Cartesian poses of a specified model frame given a list of joint configurations.
353 Parameters
354 ----------
355 model : `pinocchio.Model`
356 The model from which to perform forward kinematics.
357 target_frame : str
358 The name of the target frame.
359 q_vec : array-like
360 A list of joint configuration values describing the path.
361 data : `pinocchio.Data`, optional
362 The model data to use. If not set, one will be created.
364 Returns
365 -------
366 list[`pinocchio.SE3`]
367 A list of transforms describing the Cartesian poses of the specified frame at the provided joint configurations.
368 """
369 if data is None: 369 ↛ 371line 369 didn't jump to line 371 because the condition on line 369 was always true
370 data = model.createData()
371 target_frame_id = model.getFrameId(target_frame)
372 tforms = []
373 for q in q_vec:
374 pinocchio.framesForwardKinematics(model, data, q)
375 tforms.append(copy.deepcopy(data.oMf[target_frame_id]))
376 return tforms
379def get_collision_geometry_ids(model, collision_model, body):
380 """
381 Gets a list of collision geometry model IDs for a specified body name.
383 Parameters
384 ----------
385 model : `pinocchio.Model`
386 The model to use for getting frame IDs.
387 collision_model : `pinocchio.Model`
388 The model to use for collision checking.
389 body : str
390 The name of the body.
391 This can be directly the name of a geometry in the collision model,
392 or it can be the name of the frame in the main model.
394 Return
395 ------
396 list[int]
397 A list of collision geometry IDs corresponding to the body name.
398 """
399 body_collision_geom_ids = []
401 # First, check if this is directly a collision geometry.
402 body_collision_geom_id = collision_model.getGeometryId(body)
403 if body_collision_geom_id < collision_model.ngeoms:
404 body_collision_geom_ids.append(body_collision_geom_id)
405 else:
406 # Otherwise, look for the frame name in the model and return its associated collision objects.
407 body_frame_id = model.getFrameId(body)
408 if body_frame_id < model.nframes:
409 for id, obj in enumerate(collision_model.geometryObjects):
410 if obj.parentFrame == body_frame_id:
411 body_collision_geom_ids.append(id)
413 return body_collision_geom_ids
416def get_collision_pair_indices_from_bodies(model, collision_model, body_list):
417 """
418 Returns a list of all the collision pair indices involving a list of objects.
420 Parameters
421 ----------
422 model : `pinocchio.Model`
423 The model to use for getting frame IDs.
424 collision_model : `pinocchio.Model`
425 The model to use for collision checking.
426 body_list : list[str]
427 A list containing the names of bodies.
428 These can be directly the name of a geometry in the collision model,
429 or they can be the name of the frame in the main model.
431 Return
432 ------
433 list[int]
434 The indices of the collision pairs list involving the bodies in the specified list.
435 """
436 collision_ids = set()
437 for obj in body_list:
438 ids = get_collision_geometry_ids(model, collision_model, obj)
439 collision_ids.update(ids)
441 pairs = []
442 for idx, p in enumerate(collision_model.collisionPairs):
443 if p.first in collision_ids or p.second in collision_ids: 443 ↛ 442line 443 didn't jump to line 442 because the condition on line 443 was always true
444 pairs.append(idx)
446 return pairs
449def set_collisions(model, collision_model, body1, body2, enable):
450 """
451 Sets collision checking between two bodies by searching for their corresponding geometry objects in the collision model.
453 Parameters
454 ----------
455 model : `pinocchio.Model`
456 The model to use for getting frame IDs.
457 collision_model : `pinocchio.Model`
458 The model to use for collision checking.
459 body1 : str
460 The name of the first body.
461 body2 : str
462 The name of the second body.
463 enable : bool
464 If True, enables collisions. If False, disables collisions.
465 """
466 body1_collision_ids = get_collision_geometry_ids(model, collision_model, body1)
467 body2_collision_ids = get_collision_geometry_ids(model, collision_model, body2)
469 for id1 in body1_collision_ids:
470 for id2 in body2_collision_ids:
471 pair = pinocchio.CollisionPair(id1, id2)
472 if enable:
473 collision_model.addCollisionPair(pair)
474 else:
475 collision_model.removeCollisionPair(pair)
478def calculate_collision_vector_and_jacobians(
479 model, collision_model, data, collision_data, pair_idx, q
480):
481 """
482 Given collision and distance results from collision model data, computes the collision vector
483 and collision Jacobians at both collision points.
485 This is useful for algorithms that perform collision avoidance, such as IK and trajectory optimization.
487 Note that forward kinematics, collision, and distance checks must be evaluated first to populate the
488 `data` and `collision_data` variables.
490 Parameters
491 ----------
492 model : `pinocchio.Model`
493 The model to use for Jacobian computation.
494 collision_model : `pinocchio.Model`
495 The model to use for collision checking.
496 data : `pinocchio.Data`
497 The model data to use for Jacobian computation.
498 collision_data : `pinocchio.GeometryData`
499 The collision_model data to use for collision checking.
500 pair_idx : int
501 The index of the collision pair from which to extract information.
502 q : array-like
503 The joint configuration of the model.
505 Return
506 ------
507 tuple(array-like, array-like, array-like)
508 A tuple containing collision distance vector from frame1 to frame2,
509 and the collision Jacobians at frame1 and frame2.
510 """
511 cp = collision_model.collisionPairs[pair_idx]
512 dr = collision_data.distanceResults[pair_idx]
514 # According to the coal documentation, the normal always points from object1 to object2.
515 coll_points = [dr.getNearestPoint1(), dr.getNearestPoint2()]
516 distance_vec = dr.normal
517 distance = dr.min_distance
519 # Calculate the Jacobians at the parent frames of both collision points.
520 parent_frame1 = collision_model.geometryObjects[cp.first].parentFrame
521 if parent_frame1 >= model.nframes: 521 ↛ 522line 521 didn't jump to line 522 because the condition on line 521 was never true
522 parent_frame1 = 0
523 Jframe1 = pinocchio.computeFrameJacobian(
524 model, data, q, parent_frame1, pinocchio.ReferenceFrame.LOCAL_WORLD_ALIGNED
525 )
526 r1 = coll_points[0] - data.oMf[parent_frame1].translation
527 Jcoll1 = pinocchio.SE3(np.eye(3), r1).toActionMatrix()[:3, :] @ Jframe1
529 parent_frame2 = collision_model.geometryObjects[cp.second].parentFrame
530 if parent_frame2 >= model.nframes: 530 ↛ 531line 530 didn't jump to line 531 because the condition on line 530 was never true
531 parent_frame2 = 0
532 Jframe2 = pinocchio.computeFrameJacobian(
533 model, data, q, parent_frame2, pinocchio.ReferenceFrame.LOCAL_WORLD_ALIGNED
534 )
535 r2 = coll_points[1] - data.oMf[parent_frame2].translation
536 Jcoll2 = pinocchio.SE3(np.eye(3), r2).toActionMatrix()[:3, :] @ Jframe2
538 return distance_vec * distance, Jcoll1, Jcoll2