| 43 | |
| 44 | |
| 45 | def create_cartpole(N, add_floor): |
| 46 | model = pin.Model() |
| 47 | geom_model = pin.GeometryModel() |
| 48 | |
| 49 | if add_floor: |
| 50 | # add floor |
| 51 | floor_collision_shape = hppfcl.Halfspace(0, 0, 1, 0) |
| 52 | M = pin.SE3.Identity() |
| 53 | floor_collision_object = pin.GeometryObject( |
| 54 | "floor", 0, 0, M, floor_collision_shape |
| 55 | ) |
| 56 | color = GREY |
| 57 | color[3] = 0.5 |
| 58 | floor_collision_object.meshColor = color |
| 59 | geom_floor = geom_model.addGeometryObject(floor_collision_object) |
| 60 | |
| 61 | parent_id = 0 |
| 62 | |
| 63 | cart_radius = 0.1 |
| 64 | cart_length = 5 * cart_radius |
| 65 | cart_mass = 1.0 |
| 66 | joint_name = "joint_cart" |
| 67 | |
| 68 | geometry_placement = pin.SE3.Identity() |
| 69 | geometry_placement.rotation = pin.Quaternion( |
| 70 | np.array([0.0, 0.0, 1.0]), np.array([0.0, 1.0, 0.0]) |
| 71 | ).toRotationMatrix() |
| 72 | |
| 73 | joint_id = model.addJoint( |
| 74 | parent_id, pin.JointModelPY(), pin.SE3.Identity(), joint_name |
| 75 | ) |
| 76 | |
| 77 | body_inertia = pin.Inertia.FromCylinder(cart_mass, cart_radius, cart_length) |
| 78 | body_placement = geometry_placement |
| 79 | model.appendBodyToJoint( |
| 80 | joint_id, body_inertia, body_placement |
| 81 | ) # We need to rotate the inertia as it is expressed in the LOCAL frame of the geometry |
| 82 | |
| 83 | shape_cart = hppfcl.Cylinder(cart_radius, cart_length) |
| 84 | |
| 85 | geom_cart = pin.GeometryObject( |
| 86 | "shape_cart", joint_id, geometry_placement, shape_cart |
| 87 | ) |
| 88 | geom_cart.meshColor = np.array([1.0, 0.1, 0.1, 1.0]) |
| 89 | geom_model.addGeometryObject(geom_cart) |
| 90 | |
| 91 | parent_id = joint_id |
| 92 | joint_placement = pin.SE3.Identity() |
| 93 | body_mass = 1.0 |
| 94 | body_radius = 0.1 |
| 95 | for k in range(N): |
| 96 | joint_name = "joint_" + str(k + 1) |
| 97 | joint_id = model.addJoint( |
| 98 | parent_id, pin.JointModelRX(), joint_placement, joint_name |
| 99 | ) |
| 100 | |
| 101 | body_inertia = pin.Inertia.FromSphere(body_mass, body_radius) |
| 102 | body_placement = joint_placement.copy() |