def compile_scene_tree_clearance_geometry_to_mbp_and_sg(
        scene_tree, timestep=0.001, alpha=0.25):
    builder = DiagramBuilder()
    mbp, scene_graph = AddMultibodyPlantSceneGraph(
        builder, MultibodyPlant(time_step=timestep))
    parser = Parser(mbp)
    parser.package_map().PopulateFromEnvironment("ROS_PACKAGE_PATH")
    world_body = mbp.world_body()
    free_body_poses = []
    # For generating colors.
    node_class_to_color_dict = {}
    cmap = plt.cm.get_cmap('jet')
    cmap_counter = 0.
    for node in scene_tree.nodes:
        if node.tf is not None and node.physics_geometry_info is not None:
            # Don't have to do anything if this does not introduce geometry.
            sanity_check_node_tf_and_physics_geom_info(node)
            phys_geom_info = node.physics_geometry_info
            has_clearance_geometry = len(phys_geom_info.clearance_geometry) > 0
            if not has_clearance_geometry:
                continue

            # Add a body for this node and register the clearance geometry.
            # TODO(gizatt) This tree body index is built in to disambiguate names.
            # But I forsee a name-to-stuff resolution crisis when inference time comes...
            # this might get resolved by the solution to that.
            body = mbp.AddRigidBody(name=node.name,
                                    M_BBo_B=phys_geom_info.spatial_inertia)
            tf = torch_tf_to_drake_tf(node.tf)
            mbp.SetDefaultFreeBodyPose(body, tf)

            # Pick out a color for this class.
            node_type_string = node.__class__.__name__
            if node_type_string in node_class_to_color_dict.keys():
                color = node_class_to_color_dict[node_type_string]
            else:
                color = list(cmap(cmap_counter))
                color[3] = alpha
                node_class_to_color_dict[node_type_string] = color
                cmap_counter = np.fmod(cmap_counter + np.pi * 2., 1.)

            # Handle adding primitive geometry by adding it all to one
            # mbp.
            if len(phys_geom_info.clearance_geometry) > 0:
                for k, (tf, geometry) in enumerate(
                        phys_geom_info.clearance_geometry):
                    mbp.RegisterCollisionGeometry(
                        body=body,
                        X_BG=torch_tf_to_drake_tf(tf),
                        shape=geometry,
                        name=node.name + "_col_%03d" % k,
                        coulomb_friction=default_friction)
                    mbp.RegisterVisualGeometry(body=body,
                                               X_BG=torch_tf_to_drake_tf(tf),
                                               shape=geometry,
                                               name=node.name +
                                               "_vis_%03d" % k,
                                               diffuse_color=color)

    return builder, mbp, scene_graph
예제 #2
0
def build_mbp(seed=0, verts_geom=False, convex_collision_geom=True):
    # Make some random lumpy objects
    trimeshes = []
    np.random.seed(42)
    for k in range(3):
        # Make a small random number of triangles and chull it
        # to get a lumpy object
        mesh = trimesh.creation.random_soup(5)
        mesh = trimesh.convex.convex_hull(mesh)
        trimeshes.append(mesh)

    # Create Drake geometry from those objects by adding a small
    # sphere at each vertex
    sphere_rad = 0.05
    cmap = plt.cm.get_cmap('jet')

    builder = DiagramBuilder()
    mbp, scene_graph = AddMultibodyPlantSceneGraph(
        builder, MultibodyPlant(time_step=0.001))

    # Add ground
    friction = CoulombFriction(0.9, 0.8)
    g = pydrake_geom.Box(100., 100., 0.5)
    tf = RigidTransform(p=[0., 0., -0.25])
    mbp.RegisterVisualGeometry(body=mbp.world_body(),
                               X_BG=tf,
                               shape=g,
                               name="ground",
                               diffuse_color=[1.0, 1.0, 1.0, 1.0])
    mbp.RegisterCollisionGeometry(body=mbp.world_body(),
                                  X_BG=tf,
                                  shape=g,
                                  name="ground",
                                  coulomb_friction=friction)

    for i, mesh in enumerate(trimeshes):
        inertia = SpatialInertia(mass=1.0,
                                 p_PScm_E=np.zeros(3),
                                 G_SP_E=UnitInertia(0.01, 0.01, 0.01))
        body = mbp.AddRigidBody(name="body_%d" % i, M_BBo_B=inertia)
        color = cmap(np.random.random())
        if verts_geom:
            for j, vert in enumerate(mesh.vertices):
                g = pydrake_geom.Sphere(radius=sphere_rad)
                tf = RigidTransform(p=vert)
                mbp.RegisterVisualGeometry(body=body,
                                           X_BG=tf,
                                           shape=g,
                                           name="body_%d_color_%d" % (i, j),
                                           diffuse_color=color)
                mbp.RegisterCollisionGeometry(body=body,
                                              X_BG=tf,
                                              shape=g,
                                              name="body_%d_collision_%d" %
                                              (i, j),
                                              coulomb_friction=friction)
        # And add mesh itself for vis
        path = "/tmp/part_%d.obj" % i
        trimesh.exchange.export.export_mesh(mesh, path)
        g = pydrake_geom.Convex(path)
        mbp.RegisterVisualGeometry(body=body,
                                   X_BG=RigidTransform(),
                                   shape=g,
                                   name="body_%d_base" % i,
                                   diffuse_color=color)
        if convex_collision_geom:
            mbp.RegisterCollisionGeometry(body=body,
                                          X_BG=RigidTransform(),
                                          shape=g,
                                          name="body_%d_base_col" % i,
                                          coulomb_friction=friction)
        mbp.SetDefaultFreeBodyPose(body, RigidTransform(p=[i % 3, i / 3., 1.]))
    mbp.Finalize()
    return builder, mbp, scene_graph
def compile_scene_tree_to_mbp_and_sg(scene_tree, timestep=0.001):
    builder = DiagramBuilder()
    mbp, scene_graph = AddMultibodyPlantSceneGraph(
        builder, MultibodyPlant(time_step=timestep))
    parser = Parser(mbp)
    parser.package_map().PopulateFromEnvironment("ROS_PACKAGE_PATH")
    world_body = mbp.world_body()

    node_to_free_body_ids_map = {}
    body_id_to_node_map = {}

    free_body_poses = []
    for node in scene_tree.nodes:
        node_to_free_body_ids_map[node] = []
        if node.tf is not None and node.physics_geometry_info is not None:
            # Don't have to do anything if this does not introduce geometry.
            sanity_check_node_tf_and_physics_geom_info(node)
            phys_geom_info = node.physics_geometry_info

            # Don't have to do anything if this does not introduce geometry.
            has_models = len(phys_geom_info.model_paths) > 0
            has_prim_geometry = (len(phys_geom_info.visual_geometry) +
                                 len(phys_geom_info.collision_geometry)) > 0
            if not has_models and not has_prim_geometry:
                continue

            node_model_ids = []
            # Handle adding primitive geometry by adding it all to one
            # mbp.
            if has_prim_geometry:
                # Contain this primitive geometry in a model instance.
                model_id = mbp.AddModelInstance(node.name + "::model_%d" %
                                                len(node_model_ids))
                node_model_ids.append(model_id)
                # Add a body for this node, and register any of the
                # visual and collision geometry available.
                # TODO(gizatt) This tree body index is built in to disambiguate names.
                # But I forsee a name-to-stuff resolution crisis when inference time comes...
                # this might get resolved by the solution to that.
                body = mbp.AddRigidBody(name=node.name,
                                        model_instance=model_id,
                                        M_BBo_B=phys_geom_info.spatial_inertia)
                body_id_to_node_map[body.index()] = node
                tf = torch_tf_to_drake_tf(node.tf)
                if phys_geom_info.fixed:
                    weld = mbp.WeldFrames(world_body.body_frame(),
                                          body.body_frame(), tf)
                else:
                    node_to_free_body_ids_map[node].append(body.index())
                    mbp.SetDefaultFreeBodyPose(body, tf)
                for k, (tf, geometry,
                        color) in enumerate(phys_geom_info.visual_geometry):
                    mbp.RegisterVisualGeometry(body=body,
                                               X_BG=torch_tf_to_drake_tf(tf),
                                               shape=geometry,
                                               name=node.name +
                                               "_vis_%03d" % k,
                                               diffuse_color=color)
                for k, (tf, geometry, friction) in enumerate(
                        phys_geom_info.collision_geometry):
                    mbp.RegisterCollisionGeometry(
                        body=body,
                        X_BG=torch_tf_to_drake_tf(tf),
                        shape=geometry,
                        name=node.name + "_col_%03d" % k,
                        coulomb_friction=friction)

            # Handle adding each model from sdf/urdf.
            if has_models:
                for local_tf, model_path, root_body_name, q0_dict in phys_geom_info.model_paths:
                    model_id = parser.AddModelFromFile(
                        resolve_catkin_package_path(parser.package_map(),
                                                    model_path),
                        node.name + "::"
                        "model_%d" % len(node_model_ids))
                    node_model_ids.append(model_id)
                    if root_body_name is None:
                        root_body_ind_possibilities = mbp.GetBodyIndices(
                            model_id)
                        assert len(root_body_ind_possibilities) == 1, \
                            "Please supply root_body_name for model with path %s" % model_path
                        root_body = mbp.get_body(
                            root_body_ind_possibilities[0])
                    else:
                        root_body = mbp.GetBodyByName(name=root_body_name,
                                                      model_instance=model_id)
                    body_id_to_node_map[root_body.index()] = node
                    node_tf = torch_tf_to_drake_tf(node.tf)
                    full_model_tf = node_tf.multiply(
                        torch_tf_to_drake_tf(local_tf))
                    if phys_geom_info.fixed:
                        mbp.WeldFrames(world_body.body_frame(),
                                       root_body.body_frame(), full_model_tf)
                    else:
                        node_to_free_body_ids_map[node].append(
                            root_body.index())
                        mbp.SetDefaultFreeBodyPose(root_body, full_model_tf)
                    # Handle initial joint state
                    if q0_dict is not None:
                        for joint_name in list(q0_dict.keys()):
                            q0_this = q0_dict[joint_name]
                            joint = mbp.GetMutableJointByName(
                                joint_name, model_instance=model_id)
                            # Reshape to make Drake happy.
                            q0_this = q0_this.reshape(joint.num_positions(), 1)
                            joint.set_default_positions(q0_this)

    return builder, mbp, scene_graph, node_to_free_body_ids_map, body_id_to_node_map