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
def do_generation_and_simulation(sim_time=10):
vis = meshcat.Visualizer(zmq_url="tcp://127.0.0.1:6000")
scene_tree, satisfied_clearance = rejection_sample_feasible_tree(num_attempts=1000)
scene_tree, satisfied_feasibility = project_tree_to_feasibility(scene_tree, num_attempts=3)
serialize_scene_tree_to_yamls_and_sdfs(scene_tree, package_name='save', package_dir=".", remove_directory=True)
# Draw generated tree in meshcat.
#draw_scene_tree_meshcat(scene_tree, alpha=1.0, node_sphere_size=0.1)
# Draw its clearance geometry for debugging.
#draw_clearance_geometry_meshcat(scene_tree, alpha=0.3)
# Simulate the resulting scene, with a PR2 for scale.
builder, mbp, scene_graph = compile_scene_tree_to_mbp_and_sg(
scene_tree, timestep=0.001)
# Add PR2 and shift it back in front of where I know the table will be.
parser = Parser(mbp)
pr2_model_path = "/home/gizatt/drake/build/install/share/drake/examples/pr2/models/pr2_description/urdf/pr2_simplified.urdf"
parser.package_map().PopulateUpstreamToDrake(pr2_model_path);
pr2_model_id = parser.AddModelFromFile(
file_name=pr2_model_path, model_name="PR2_for_scale")
# Get the tf of the robot spawn node, and put the PR2 at that xy location.
robot_spawn_tf = scene_tree.find_nodes_by_type(RobotSpawnLocation)[0].tf.numpy()
mbp.GetJointByName("x", model_instance=pr2_model_id).set_default_translation(robot_spawn_tf[0, 3])
mbp.GetJointByName("y", model_instance=pr2_model_id).set_default_translation(robot_spawn_tf[1, 3])
mbp.Finalize()
visualizer = ConnectMeshcatVisualizer(builder, scene_graph,
zmq_url="default")
diagram = builder.Build()
diag_context = diagram.CreateDefaultContext()
mbp_context = diagram.GetMutableSubsystemContext(mbp, diag_context)
# Fix input port for PR2 to zero.
actuation_port = mbp.get_actuation_input_port(model_instance=pr2_model_id)
nu = mbp.num_actuated_dofs(model_instance=pr2_model_id)
actuation_port.FixValue(mbp_context, np.zeros(nu))
sim = Simulator(diagram, diag_context)
sim.set_target_realtime_rate(1.0)
if not satisfied_clearance:
print("WARNING: SCENE TREE NOT SATISFYING CLEARANCE")
if not satisfied_feasibility:
print("WARNING: SCENE TREE NOT SATISFYING FEASIBILITY, SIM MAY FAIL")
try:
sim.AdvanceTo(sim_time)
except RuntimeError as e:
print("Encountered error in sim: ", e)
def AddWsg(plant, iiwa_model_instance, roll=np.pi / 2.0, welded=False):
parser = Parser(plant)
if welded:
parser.package_map().Add(
"wsg_50_description",
os.path.dirname(
FindResourceOrThrow(
"drake/manipulation/models/wsg_50_description/package.xml"))
)
gripper = parser.AddModelFromFile(
FindResource("models/schunk_wsg_50_welded_fingers.sdf"), "gripper")
else:
gripper = parser.AddModelFromFile(
FindResourceOrThrow(
"drake/manipulation/models/"
"wsg_50_description/sdf/schunk_wsg_50_with_tip.sdf"))
X_7G = RigidTransform(RollPitchYaw(np.pi / 2.0, 0, roll), [0, 0, 0.114])
plant.WeldFrames(plant.GetFrameByName("iiwa_link_7", iiwa_model_instance),
plant.GetFrameByName("body", gripper), X_7G)
return gripper
def AddTwoLinkIiwa(plant, q0=[0.1, -1.2]):
urdf = FindResource("models/two_link_iiwa14.urdf")
parser = Parser(plant)
parser.package_map().Add(
"iiwa_description",
os.path.dirname(
FindResourceOrThrow(
"drake/manipulation/models/iiwa_description/package.xml")))
iiwa = parser.AddModelFromFile(urdf)
plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("iiwa_link_0"))
# Set default positions:
index = 0
for joint_index in plant.GetJointIndices(iiwa):
joint = plant.get_mutable_joint(joint_index)
if isinstance(joint, RevoluteJoint):
joint.set_default_angle(q0[index])
index += 1
return iiwa
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
def make_parser(plant):
parser = Parser(plant)
parser.package_map().PopulateFromFolder("./repos/")
return parser