Coverage for src/pyroboplan/models/two_dof.py: 41%

1"""Utilities to load example 2-DOF manipulator.""" 

2 

3import coal 

4import numpy as np 

5import os 

6import pinocchio 

7 

8from ..core.utils import set_collisions 

9from .utils import get_example_models_folder 

10 

11 

12def load_models(): 

13 """ 

14 Gets the example 2-DOF models. 

15 

16 Returns 

17 ------- 

18 tuple[`pinocchio.Model`] 

19 A 3-tuple containing the model, collision geometry model, and visual geometry model. 

20 """ 

21 models_folder = get_example_models_folder() 

22 package_dir = os.path.join(models_folder, "two_dof_description") 

23 urdf_filename = os.path.join(package_dir, "two_dof.urdf") 

24 

25 return pinocchio.buildModelsFromUrdf(urdf_filename, package_dirs=models_folder) 

26 

27 

28def add_object_collisions(model, collision_model, visual_model): 

29 """ 

30 Adds obstacles and collisions to the 2-DOF manipulator collision model. 

31 

32 Parameters 

33 ---------- 

34 model : `pinocchio.Model` 

35 The robot model. 

36 collision_model : `pinocchio.Model` 

37 The collision geometry model. 

38 visual_model : `pinocchio.Model` 

39 The visual geometry model. 

40 """ 

41 obstacle_1 = pinocchio.GeometryObject( 

42 "obstacle_1", 

43 0, 

44 pinocchio.SE3(np.eye(3), np.array([1.0, 1.0, 0.0])), 

45 coal.Cylinder(0.3, 0.1), 

46 ) 

47 obstacle_1.meshColor = np.array([0.0, 1.0, 0.0, 0.5]) 

48 visual_model.addGeometryObject(obstacle_1) 

49 collision_model.addGeometryObject(obstacle_1) 

50 

51 obstacle_2 = pinocchio.GeometryObject( 

52 "obstacle_2", 

53 0, 

54 pinocchio.SE3(np.eye(3), np.array([-1.0, -0.75, 0.0])), 

55 coal.Box(0.5, 0.5, 0.1), 

56 ) 

57 obstacle_2.meshColor = np.array([1.0, 0.0, 0.0, 0.5]) 

58 visual_model.addGeometryObject(obstacle_2) 

59 collision_model.addGeometryObject(obstacle_2) 

60 

61 # Define the active collision pairs between the robot and obstacle links. 

62 collision_names = [ 

63 cobj.name for cobj in collision_model.geometryObjects if "arm" in cobj.name 

64 ] 

65 obstacle_names = ["obstacle_1", "obstacle_2"] 

66 for obstacle_name in obstacle_names: 

67 for collision_name in collision_names: 

68 set_collisions(model, collision_model, obstacle_name, collision_name, True)