def draw_clearance_geometry_meshcat(scene_tree, zmq_url=None, alpha=0.25):
vis = meshcat.Visualizer(zmq_url=zmq_url or "tcp://127.0.0.1:6000")
vis["clearance_geom"].delete()
builder, mbp, scene_graph = compile_scene_tree_clearance_geometry_to_mbp_and_sg(scene_tree)
mbp.Finalize()
vis = ConnectMeshcatVisualizer(builder, scene_graph,
zmq_url="default", prefix="clearance")
diagram = builder.Build()
context = diagram.CreateDefaultContext()
vis.load(vis.GetMyContextFromRoot(context))
diagram.Publish(context)
def __init__(self,
builder,
dt=5e-4,
N=150,
params=None,
trj_decay=0.7,
x_w_cov=1e-5,
door_angle_ref=1.0,
visualize=False):
self.plant_derivs = MultibodyPlant(time_step=dt)
parser = Parser(self.plant_derivs)
self.derivs_iiwa, _, _ = self.add_models(self.plant_derivs,
parser,
params=params)
self.plant_derivs.Finalize()
self.plant_derivs_context = self.plant_derivs.CreateDefaultContext()
self.plant_derivs.get_actuation_input_port().FixValue(
self.plant_derivs_context, [0., 0., 0., 0., 0., 0., 0.])
null_force = BasicVector([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
self.plant_derivs.GetInputPort("applied_generalized_force").FixValue(
self.plant_derivs_context, null_force)
self.plant, scene_graph = AddMultibodyPlantSceneGraph(builder,
time_step=dt)
parser = Parser(self.plant, scene_graph)
self.iiwa, self.hinge, self.bushing = self.add_models(self.plant,
parser,
params=params)
self.plant.Finalize(
) # Finalize will assign ports for compatibility w/ the scene_graph; could be cause of the issue w/ first order taylor.
self.meshcat = ConnectMeshcatVisualizer(builder,
scene_graph,
zmq_url=zmq_url)
self.sym_derivs = False # If system should use symbolic derivatives; if false, autodiff
self.custom_sim = False # If rollouts should be gathered with sys.dyn() calls
nq = self.plant.num_positions()
nv = self.plant.num_velocities()
self.n_x = nq + nv
self.n_u = self.plant.num_actuators()
self.n_y = self.plant.get_state_output_port(self.iiwa).size()
self.N = N
self.dt = dt
self.decay = trj_decay
self.V = 1e-2 * np.ones(self.n_y)
self.W = np.concatenate((1e-7 * np.ones(nq), 1e-4 * np.ones(nv)))
self.W0 = np.concatenate((1e-9 * np.ones(nq), 1e-6 * np.ones(nv)))
self.x_w_cov = x_w_cov
self.door_angle_ref = door_angle_ref
self.q0 = np.array(
[-3.12, -0.17, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55])
#self.q0 = np.array([-3.12, -0.27, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55])
self.x0 = np.concatenate((self.q0, np.zeros(nv)))
self.door_index = None
self.phi = {}
def setup_manipulation_station(T_world_objectInitial, zmq_url, T_world_targetBin, manipuland="foam", include_bin=True, include_hoop=False):
builder = DiagramBuilder()
station = builder.AddSystem(ManipulationStation(time_step=1e-3))
station.SetupClutterClearingStation()
if manipuland is "foam":
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/061_foam_brick.sdf",
T_world_objectInitial)
elif manipuland is "ball":
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/sphere.sdf",
T_world_objectInitial)
elif manipuland is "bball":
station.AddManipulandFromFile(
"drake/../../../../../../manipulation/sdfs/bball.sdf", # this is some path hackery
T_world_objectInitial)
elif manipuland is "rod":
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/rod.sdf",
T_world_objectInitial)
station_plant = station.get_multibody_plant()
parser = Parser(station_plant)
if include_bin:
parser.AddModelFromFile("extra_bin.sdf")
station_plant.WeldFrames(station_plant.world_frame(), station_plant.GetFrameByName("extra_bin_base"), T_world_targetBin)
if include_hoop:
parser.AddModelFromFile("sdfs/hoop.sdf")
station_plant.WeldFrames(station_plant.world_frame(), station_plant.GetFrameByName("base_link_hoop"), T_world_targetBin)
station.Finalize()
frames_to_draw = {"gripper": {"body"}}
meshcat = None
if zmq_url is not None:
meshcat = ConnectMeshcatVisualizer(builder,
station.get_scene_graph(),
output_port=station.GetOutputPort("pose_bundle"),
delete_prefix_on_load=False,
frames_to_draw=frames_to_draw,
zmq_url=zmq_url)
diagram = builder.Build()
plant = station.get_multibody_plant()
context = plant.CreateDefaultContext()
gripper = plant.GetBodyByName("body")
initial_pose = plant.EvalBodyPoseInWorld(context, gripper)
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
simulator.AdvanceTo(0.01)
return initial_pose, meshcat
def simulate_scene_tree(scene_tree, T, timestep=0.001, with_meshcat=False):
builder, mbp, scene_graph = compile_scene_tree_to_mbp_and_sg(
scene_tree, timestep=timestep)
mbp.Finalize()
if with_meshcat:
visualizer = ConnectMeshcatVisualizer(builder, scene_graph,
zmq_url="default")
diagram = builder.Build()
diag_context = diagram.CreateDefaultContext()
sim = Simulator(diagram)
sim.set_target_realtime_rate(1.0)
sim.AdvanceTo(T)
def draw_scene_tree_contents_meshcat(scene_tree,
prefix="scene",
zmq_url=None,
alpha=0.25,
draw_clearance_geom=False,
quiet=True):
''' Given a scene tree, draws it in meshcat at the requested ZMQ url.
Can be configured to draw the tree geometry or the clearance geometry. '''
if draw_clearance_geom:
builder, mbp, scene_graph = compile_scene_tree_clearance_geometry_to_mbp_and_sg(
scene_tree)
else:
builder, mbp, scene_graph, _, _, = compile_scene_tree_to_mbp_and_sg(
scene_tree)
mbp.Finalize()
if quiet:
with open(os.devnull, 'w') as devnull:
with contextlib.redirect_stdout(devnull):
vis = ConnectMeshcatVisualizer(builder,
scene_graph,
zmq_url=zmq_url or "default",
prefix=prefix)
else:
vis = ConnectMeshcatVisualizer(builder,
scene_graph,
zmq_url=zmq_url or "default",
prefix=prefix)
vis.delete_prefix()
diagram = builder.Build()
context = diagram.CreateDefaultContext()
vis.load(vis.GetMyContextFromRoot(context))
diagram.Publish(context)
# Necessary to manually remove this meshcat visualizer now that we're
# done with it, as a lot of Drake systems (that are involved with the
# diagram builder) don't get properly garbage collected. See Drake issue #14387.
# Meshcat collects sockets, so deleting this avoids a file descriptor
# leak.
del vis.vis
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("-t1", default=0.05, help="Extend leg")
parser.add_argument("-t2", default=0.5, help="Dwell at top")
parser.add_argument("-t3", default=0.5, help="Contract leg")
parser.add_argument("-t4", default=0.1, help="Wait at bottom")
setup_argparse_for_setup_dot_diagram(parser)
MeshcatVisualizer.add_argparse_argument(parser)
args = parser.parse_args()
t1 = float(args.t1)
t2 = float(args.t2)
t3 = float(args.t3)
t4 = float(args.t4)
q_crouch = np.array([
1600, 2100, 2000, 1600, 2100, 2000, 1400, 2100, 2000, 1400, 2100, 2000
])
q_extend = np.array([
1600, 1600, 2400, 1600, 1600, 2400, 1400, 1600, 2400, 1400, 1600, 2400
])
breaks = np.cumsum([0., t1, t2, t3, t4])
samples = np.stack([q_crouch, q_extend, q_extend, q_crouch, q_crouch]).T
trajectory = PiecewisePolynomial.FirstOrderHold(breaks, samples)
builder = DiagramBuilder()
plant, scene_graph, servo_controller = setup_dot_diagram(builder, args)
trajectory_source = builder.AddSystem(TrajectoryLooper(trajectory))
builder.Connect(trajectory_source.get_output_port(0),
servo_controller.get_input_port(0))
if args.meshcat:
meshcat = ConnectMeshcatVisualizer(
builder,
output_port=scene_graph.get_query_output_port(),
zmq_url=args.meshcat,
open_browser=args.open_browser)
diagram = builder.Build()
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
simulator.AdvanceTo(1E6)
def gripper_forward_kinematics_example():
builder = DiagramBuilder()
station = builder.AddSystem(ManipulationStation())
station.SetupClutterClearingStation()
station.Finalize()
frames_to_draw = {
"iiwa": {
"iiwa_link_1", "iiwa_link_2", "iiwa_link_3", "iiwa_link_4",
"iiwa_link_5", "iiwa_link_6", "iiwa_link_7"
},
"gripper": {"body"}
}
meshcat = ConnectMeshcatVisualizer(
builder,
station.get_scene_graph(),
output_port=station.GetOutputPort("pose_bundle"),
frames_to_draw=frames_to_draw,
axis_length=0.3,
axis_radius=0.01)
diagram = builder.Build()
context = diagram.CreateDefaultContext()
# xyz = Text(value="", description="gripper position (m): ", layout=Layout(width='500px'), style={'description_width':'initial'})
# rpy = Text(value="", description="gripper roll-pitch-yaw (rad): ", layout=Layout(width='500px'), style={'description_width':'initial'})
plant = station.get_multibody_plant()
gripper = plant.GetBodyByName("body")
def pose_callback(context):
pose = plant.EvalBodyPoseInWorld(context, gripper) # Important
# xyz.value = np.array2string(pose.translation(), formatter={'float': lambda x: "{:3.2f}".format(x)})
# rpy.value = np.array2string(RollPitchYaw(pose.rotation()).vector(), formatter={'float': lambda x: "{:3.2f}".format(x)})
meshcat.load()
MakeJointSlidersThatPublishOnCallback(station.get_multibody_plant(),
meshcat,
context,
my_callback=pose_callback)
def grasp_poses_example():
builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
parser = Parser(plant, scene_graph)
grasp = parser.AddModelFromFile(
FindResourceOrThrow(
"drake/manipulation/models/wsg_50_description/sdf/schunk_wsg_50_no_tip.sdf"
), "grasp")
brick = parser.AddModelFromFile(
FindResourceOrThrow(
"drake/examples/manipulation_station/models/061_foam_brick.sdf"))
plant.Finalize()
meshcat = ConnectMeshcatVisualizer(builder, scene_graph)
diagram = builder.Build()
context = diagram.CreateDefaultContext()
plant_context = plant.GetMyContextFromRoot(context)
B_O = plant.GetBodyByName("base_link", brick)
X_WO = plant.EvalBodyPoseInWorld(plant_context, B_O)
B_Ggrasp = plant.GetBodyByName("body", grasp)
p_GgraspO = [0, 0.12, 0]
R_GgraspO = RotationMatrix.MakeXRotation(np.pi / 2.0).multiply(
RotationMatrix.MakeZRotation(np.pi / 2.0))
X_GgraspO = RigidTransform(R_GgraspO, p_GgraspO)
X_OGgrasp = X_GgraspO.inverse()
X_WGgrasp = X_WO.multiply(X_OGgrasp)
plant.SetFreBodyPose(plant_context, B_Ggrasp, X_WGgrasp)
# Open the fingers, too
plant.GetJointByName("left_finger_sliding_joint",
grasp).set_translation(plant_context, -0.054)
plant.GetJointByName("right_finger_sliding_joint",
grasp).set_translation(plant_context, 0.054)
meshcat.load()
diagram.Publish(context)
def jacobian_controller_example():
builder = DiagramBuilder()
station = builder.AddSystem(ManipulationStation())
station.SetupClutterClearingStation()
station.Finalize()
controller = builder.AddSystem(
PseudoInverseController(station.get_multibody_plant()))
integrator = builder.AddSystem(Integrator(7))
builder.Connect(controller.get_output_port(), integrator.get_input_port())
builder.Connect(integrator.get_output_port(),
station.GetInputPort("iiwa_position"))
builder.Connect(station.GetOutputPort("iiwa_position_measured"),
controller.get_input_port())
meshcat = ConnectMeshcatVisualizer(
builder,
station.get_scene_graph(),
output_port=station.GetOutputPort("pose_bundle"))
diagram = builder.Build()
simulator = Simulator(diagram)
station_context = station.GetMyContextFromRoot(
simulator.get_mutable_context())
station.GetInputPort("iiwa_feedforward_torque").FixValue(
station_context, np.zeros((7, 1)))
station.GetInputPort("wsg_position").FixValue(station_context, [0.1])
integrator.GetMyContextFromRoot(simulator.get_mutable_context(
)).get_mutable_continuous_state_vector().SetFromVector(
station.GetIiwaPosition(station_context))
simulator.set_target_realtime_rate(1.0)
simulator.AdvanceTo(0.01)
return simulator
def main():
parser = argparse.ArgumentParser(description=__doc__)
setup_argparse_for_setup_dot_diagram(parser)
parser.add_argument("--interactive", action='store_true')
MeshcatVisualizer.add_argparse_argument(parser)
args = parser.parse_args()
q_init = np.array([
1700, 2000, 2000, 1700, 2000, 2000, 1300, 2000, 2000, 1300, 2000, 2000
])
builder = DiagramBuilder()
plant, scene_graph, servo_controller = setup_dot_diagram(builder, args)
if args.interactive:
# Add sliders to set positions of the joints.
sliders = builder.AddSystem(ServoSliders(servo_controller))
sliders.set_position(q_init)
builder.Connect(sliders.get_output_port(0),
servo_controller.get_input_port(0))
else:
source = builder.AddSystem(
ConstantVectorSource(np.zeros(servo_controller.nu)))
builder.Connect(source.get_output_port(0),
servo_controller.get_input_port(0))
if args.meshcat:
meshcat = ConnectMeshcatVisualizer(
builder,
output_port=scene_graph.get_query_output_port(),
zmq_url=args.meshcat,
open_browser=args.open_browser)
diagram = builder.Build()
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
simulator.AdvanceTo(1E6)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("sdf_path", help="path to sdf")
parser.add_argument("--interactive", action='store_true')
MeshcatVisualizer.add_argparse_argument(parser)
args = parser.parse_args()
builder = DiagramBuilder()
scene_graph = builder.AddSystem(SceneGraph())
plant = MultibodyPlant(time_step=0.01)
plant.RegisterAsSourceForSceneGraph(scene_graph)
parser = Parser(plant)
model = parser.AddModelFromFile(args.sdf_path)
plant.Finalize()
if args.meshcat:
meshcat = ConnectMeshcatVisualizer(
builder, output_port=scene_graph.get_query_output_port(),
zmq_url=args.meshcat, open_browser=args.open_browser)
if args.interactive:
# Add sliders to set positions of the joints.
sliders = builder.AddSystem(JointSliders(robot=plant))
else:
q0 = plant.GetPositions(plant.CreateDefaultContext())
sliders = builder.AddSystem(ConstantVectorSource(q0))
to_pose = builder.AddSystem(MultibodyPositionToGeometryPose(plant))
builder.Connect(sliders.get_output_port(0), to_pose.get_input_port())
builder.Connect(
to_pose.get_output_port(),
scene_graph.get_source_pose_port(plant.get_source_id()))
diagram = builder.Build()
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
simulator.AdvanceTo(1E6)
#prog = opt.get_mutable_prog()
#AddUnitQuaternionConstraintOnPlant(
# mbp_ad, q_vars, prog)
#q_targ = mbp.GetPositions(mbp.CreateDefaultContext())
## Penalize deviation from target configuration
#prog.AddQuadraticErrorCost(np.eye(q_targ.shape[0]), q_targ, q_vars)
#prog.SetInitialGuess(q_vars, q_targ)
#solver = SnoptSolver()
#result = solver.Solve(opt.prog())
#print(result, result.is_success())
#q_0 = result.GetSolution(q_vars)
#print("Q found: ", q_0)
builder, mbp, scene_graph = build_mbp(seed=seed)
q_des = mbp.GetPositions(mbp.CreateDefaultContext())
forcer = builder.AddSystem(DecayingForceToDesiredConfigSystem(mbp, q_des))
builder.Connect(mbp.get_state_output_port(), forcer.get_input_port(0))
builder.Connect(forcer.get_output_port(0),
mbp.get_applied_spatial_force_input_port())
visualizer = ConnectMeshcatVisualizer(builder,
scene_graph,
zmq_url="default")
diagram = builder.Build()
diag_context = diagram.CreateDefaultContext()
mbp_context = diagram.GetMutableSubsystemContext(mbp, diag_context)
mbp.SetPositions(mbp_context, np.random.random(q_des.shape) * 10.0)
sim = Simulator(diagram, diag_context)
#sim.set_target_realtime_rate(1.0)
sim.AdvanceTo(10.0)
builder.Connect(controller.get_output_port(), integrator.get_input_port())
builder.Connect(integrator.get_output_port(),
station.GetInputPort("iiwa_position"))
builder.Connect(station.GetOutputPort("iiwa_position_measured"),
controller.GetInputPort("iiwa_position"))
traj_wsg_command = make_wsg_command_trajectory(times)
wsg_source = builder.AddSystem(TrajectorySource(traj_wsg_command))
wsg_source.set_name("wsg_command")
builder.Connect(wsg_source.get_output_port(),
station.GetInputPort("wsg_position"))
meshcat = ConnectMeshcatVisualizer(
builder,
station.get_scene_graph(),
output_port=station.GetOutputPort("pose_bundle"),
# delete_prefix_on_load=False, # Use this if downloading is a pain.
zmq_url=zmq_url,
)
diagram = builder.Build()
diagram.set_name("pick_and_place")
simulator = Simulator(diagram)
station_context = station.GetMyContextFromRoot(simulator.get_mutable_context())
# TODO(russt): Add this missing python binding
#integrator.set_integral_value(
# integrator.GetMyContextFromRoot(simulator.get_mutable_context()),
# station.GetIiwaPosition(station_context))
integrator.GetMyContextFromRoot(simulator.get_mutable_context(
)).get_mutable_continuous_state_vector().SetFromVector(
def project_tree_to_feasibility_via_sim(tree,
constraints=[],
zmq_url=None,
prefix="projection",
timestep=0.0005,
T=10.):
# Mutates tree into tree with bodies in closest
# nonpenetrating configuration.
builder, mbp, sg, node_to_free_body_ids_map, body_id_to_node_map = \
compile_scene_tree_to_mbp_and_sg(tree, timestep=timestep)
mbp.Finalize()
# Connect visualizer if requested. Wrap carefully to keep it
# from spamming the console.
if zmq_url is not None:
with open(os.devnull, 'w') as devnull:
with contextlib.redirect_stdout(devnull):
visualizer = ConnectMeshcatVisualizer(builder,
sg,
zmq_url=zmq_url,
prefix=prefix)
# Forward sim under langevin forces
force_source = builder.AddSystem(
DecayingForceToDesiredConfigSystem(
mbp, mbp.GetPositions(mbp.CreateDefaultContext())))
builder.Connect(mbp.get_state_output_port(),
force_source.get_input_port(0))
builder.Connect(force_source.get_output_port(0),
mbp.get_applied_spatial_force_input_port())
diagram = builder.Build()
diagram_context = diagram.CreateDefaultContext()
mbp_context = diagram.GetMutableSubsystemContext(mbp, diagram_context)
q0 = mbp.GetPositions(mbp_context)
nq = len(q0)
if nq == 0:
logging.warn("Generated MBP had no positions.")
return tree
# Make 'safe' initial configuration that randomly arranges objects vertically
#k = 0
#all_pos = []
#for node in tree:
# for body_id in node_to_free_body_ids_map[node]:
# body = mbp.get_body(body_id)
# tf = mbp.GetFreeBodyPose(mbp_context, body)
# tf = RigidTransform(p=tf.translation() + np.array([0., 0., 1. + k*0.5]), R=tf.rotation())
# mbp.SetFreeBodyPose(mbp_context, body, tf)
# k += 1
sim = Simulator(diagram, diagram_context)
sim.set_target_realtime_rate(1000)
sim.AdvanceTo(T)
# Reload poses back into tree
free_bodies = mbp.GetFloatingBaseBodies()
for body_id, node in body_id_to_node_map.items():
if body_id not in free_bodies:
continue
node.tf = drake_tf_to_torch_tf(
mbp.GetFreeBodyPose(mbp_context, mbp.get_body(body_id)))
return tree
def project_tree_to_feasibility(tree,
constraints=[],
jitter_q=None,
do_forward_sim=False,
zmq_url=None,
prefix="projection",
timestep=0.001,
T=1.):
# Mutates tree into tree with bodies in closest
# nonpenetrating configuration.
builder, mbp, sg, node_to_free_body_ids_map, body_id_to_node_map = \
compile_scene_tree_to_mbp_and_sg(tree, timestep=timestep)
mbp.Finalize()
# Connect visualizer if requested. Wrap carefully to keep it
# from spamming the console.
if zmq_url is not None:
with open(os.devnull, 'w') as devnull:
with contextlib.redirect_stdout(devnull):
visualizer = ConnectMeshcatVisualizer(builder,
sg,
zmq_url=zmq_url,
prefix=prefix)
diagram = builder.Build()
diagram_context = diagram.CreateDefaultContext()
mbp_context = diagram.GetMutableSubsystemContext(mbp, diagram_context)
q0 = mbp.GetPositions(mbp_context)
nq = len(q0)
if nq == 0:
logging.warn("Generated MBP had no positions.")
return tree
# Set up projection NLP.
ik = InverseKinematics(mbp, mbp_context)
q_dec = ik.q()
prog = ik.prog()
# It's always a projection, so we always have this
# Euclidean norm error between the optimized q and
# q0.
prog.AddQuadraticErrorCost(np.eye(nq), q0, q_dec)
# Nonpenetration constraint.
ik.AddMinimumDistanceConstraint(0.001)
# Other requested constraints.
for constraint in constraints:
constraint.add_to_ik_prog(tree, ik, mbp, mbp_context,
node_to_free_body_ids_map)
# Initial guess, which can be slightly randomized by request.
q_guess = q0
if jitter_q:
q_guess = q_guess + np.random.normal(0., jitter_q, size=q_guess.size)
prog.SetInitialGuess(q_dec, q_guess)
# Solve.
solver = SnoptSolver()
options = SolverOptions()
logfile = "/tmp/snopt.log"
os.system("rm %s" % logfile)
options.SetOption(solver.id(), "Print file", logfile)
options.SetOption(solver.id(), "Major feasibility tolerance", 1E-3)
options.SetOption(solver.id(), "Major optimality tolerance", 1E-3)
options.SetOption(solver.id(), "Major iterations limit", 300)
result = solver.Solve(prog, None, options)
if not result.is_success():
logging.warn("Projection failed.")
print("Logfile: ")
with open(logfile) as f:
print(f.read())
qf = result.GetSolution(q_dec)
mbp.SetPositions(mbp_context, qf)
# If forward sim is requested, do a quick forward sim to get to
# a statically stable config.
if do_forward_sim:
sim = Simulator(diagram, diagram_context)
sim.set_target_realtime_rate(1000.)
sim.AdvanceTo(T)
# Reload poses back into tree
free_bodies = mbp.GetFloatingBaseBodies()
for body_id, node in body_id_to_node_map.items():
if body_id not in free_bodies:
continue
node.tf = drake_tf_to_torch_tf(
mbp.GetFreeBodyPose(mbp_context, mbp.get_body(body_id)))
return tree
def BuildAndSimulateTrajectory(
q_traj,
g_traj,
get_gripper_controller,
T_world_objectInitial,
T_world_targetBin,
zmq_url,
time_step,
include_target_bin=True,
include_hoop=False,
manipuland="foam"
):
"""Simulate trajectory for manipulation station.
@param q_traj: Trajectory class used to initialize TrajectorySource for joints.
@param g_traj: Trajectory class used to initialize TrajectorySource for gripper.
"""
builder = DiagramBuilder()
station = builder.AddSystem(ManipulationStation(time_step=time_step)) #1e-3 or 1e-4 probably
station.SetupClutterClearingStation()
if manipuland is "foam":
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/061_foam_brick.sdf",
T_world_objectInitial)
elif manipuland is "ball":
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/sphere.sdf",
T_world_objectInitial)
elif manipuland is "bball":
station.AddManipulandFromFile(
"drake/../../../../../../manipulation/sdfs/bball.sdf", # this is some path hackery
T_world_objectInitial)
elif manipuland is "rod":
station.AddManipulandFromFile(
"drake/examples/manipulation_station/models/rod.sdf",
T_world_objectInitial)
station_plant = station.get_multibody_plant()
parser = Parser(station_plant)
if include_target_bin:
parser.AddModelFromFile("sdfs/extra_bin.sdf")
station_plant.WeldFrames(station_plant.world_frame(), station_plant.GetFrameByName("extra_bin_base"), T_world_targetBin)
if include_hoop:
parser.AddModelFromFile("sdfs/hoop.sdf")
station_plant.WeldFrames(station_plant.world_frame(), station_plant.GetFrameByName("base_link_hoop"), T_world_targetBin)
station.Finalize()
# iiwa joint trajectory - predetermined trajectory
q_traj_system = builder.AddSystem(TrajectorySource(q_traj))
builder.Connect(q_traj_system.get_output_port(),
station.GetInputPort("iiwa_position"))
# gripper - closed loop controller
gctlr = builder.AddSystem(get_gripper_controller(station_plant))
gctlr.set_name("GripperControllerUsingIiwaState")
builder.Connect(station.GetOutputPort("iiwa_position_measured"), gctlr.GetInputPort("iiwa_position"))
builder.Connect(station.GetOutputPort("iiwa_velocity_estimated"), gctlr.GetInputPort("iiwa_velocity"))
builder.Connect(gctlr.get_output_port(), station.GetInputPort("wsg_position"))
loggers = dict(
state=builder.AddSystem(SignalLogger(31)),
v_est=builder.AddSystem(SignalLogger(7))
)
builder.Connect(station.GetOutputPort("plant_continuous_state"),
loggers["state"].get_input_port())
builder.Connect(station.GetOutputPort("iiwa_velocity_estimated"),
loggers["v_est"].get_input_port())
meshcat = None
if zmq_url is not None:
meshcat = ConnectMeshcatVisualizer(builder,
station.get_scene_graph(),
output_port=station.GetOutputPort("pose_bundle"),
delete_prefix_on_load=True,
frames_to_draw={"gripper":{"body"}},
zmq_url=zmq_url)
diagram = builder.Build()
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
return simulator, station_plant, meshcat, loggers
models_object = {}
for name in model_names:
models_object[name] = []
for i in range(3):
models_object[name].append(
parser.AddModelFromFile(
os.path.join("cad_files", name + '_simplified.sdf'),
name + str(i)))
station.Finalize()
builder = DiagramBuilder()
builder.AddSystem(station)
ConnectMeshcatVisualizer(builder,
scene_graph=station.get_scene_graph(),
output_port=station.GetOutputPort("query_object"))
diagram = builder.Build()
#%%
context = diagram.CreateDefaultContext()
context_station = diagram.GetSubsystemContext(station, context)
context_plant = station.GetSubsystemContext(station.get_multibody_plant(),
context_station)
#%%
plant = station.get_multibody_plant()
model_iiwa = plant.GetModelInstanceByName("iiwa")
model_gripper = plant.GetModelInstanceByName("gripper")
model_bin1 = plant.GetModelInstanceByName("bin1")
def __init__(self, kp=100, ki=1, kd=20, time_step=0.001, show_axis=False):
"""
Robotic Kuka IIWA juggler with paddle end effector
Args:
kp (int, optional): proportional gain. Defaults to 100.
ki (int, optional): integral gain. Defaults to 1.
kd (int, optional): derivative gain. Defaults to 20.
time_step (float, optional): time step for internal manipulation station controller. Defaults to 0.001.
show_axis (boolean, optional): whether or not to show the axis of reflection.
"""
self.time = 0
juggler_station = JugglerStation(kp=kp,
ki=ki,
kd=kd,
time_step=time_step,
show_axis=show_axis)
self.builder = DiagramBuilder()
self.station = self.builder.AddSystem(juggler_station.get_diagram())
self.position_log = []
self.velocity_log = []
self.visualizer = ConnectMeshcatVisualizer(
self.builder,
output_port=self.station.GetOutputPort("geometry_query"),
zmq_url=zmq_url)
self.plant = juggler_station.get_multibody_plant()
# ---------------------
ik_sys = self.builder.AddSystem(InverseKinematics(self.plant))
ik_sys.set_name("ik_sys")
v_mirror = self.builder.AddSystem(VelocityMirror(self.plant))
v_mirror.set_name("v_mirror")
w_tilt = self.builder.AddSystem(AngularVelocityTilt(self.plant))
w_tilt.set_name("w_tilt")
integrator = self.builder.AddSystem(Integrator(7))
self.builder.Connect(
self.station.GetOutputPort("iiwa_position_measured"),
ik_sys.GetInputPort("iiwa_pos_measured"))
self.builder.Connect(ik_sys.get_output_port(),
integrator.get_input_port())
self.builder.Connect(integrator.get_output_port(),
self.station.GetInputPort("iiwa_position"))
self.builder.Connect(
self.station.GetOutputPort("iiwa_position_measured"),
w_tilt.GetInputPort("iiwa_pos_measured"))
self.builder.Connect(
self.station.GetOutputPort("iiwa_position_measured"),
v_mirror.GetInputPort("iiwa_pos_measured"))
self.builder.Connect(
self.station.GetOutputPort("iiwa_velocity_estimated"),
v_mirror.GetInputPort("iiwa_velocity_estimated"))
self.builder.Connect(
w_tilt.get_output_port(),
ik_sys.GetInputPort("paddle_desired_angular_velocity"))
self.builder.Connect(v_mirror.get_output_port(),
ik_sys.GetInputPort("paddle_desired_velocity"))
self.builder.ExportInput(v_mirror.GetInputPort("ball_pose"),
"v_ball_pose")
self.builder.ExportInput(v_mirror.GetInputPort("ball_velocity"),
"v_ball_velocity")
self.builder.ExportInput(w_tilt.GetInputPort("ball_pose"),
"w_ball_pose")
self.builder.ExportInput(w_tilt.GetInputPort("ball_velocity"),
"w_ball_velocity")
# Useful for debugging
# desired_vel = self.builder.AddSystem(ConstantVectorSource([0, 0, 0]))
# self.builder.Connect(desired_vel.get_output_port(), ik_sys.GetInputPort("paddle_desired_angular_velocity"))
# ---------------------
self.diagram = self.builder.Build()
self.simulator = Simulator(self.diagram)
self.simulator.set_target_realtime_rate(1.0)
self.context = self.simulator.get_context()
self.station_context = self.station.GetMyContextFromRoot(self.context)
self.plant_context = self.plant.GetMyContextFromRoot(self.context)
self.plant.SetPositions(
self.plant_context, self.plant.GetModelInstanceByName("iiwa7"),
[0, np.pi / 3, 0, -np.pi / 2, 0, -np.pi / 3, 0])
self.station.GetInputPort("iiwa_feedforward_torque").FixValue(
self.station_context, np.zeros((7, 1)))
iiwa_model_instance = self.plant.GetModelInstanceByName("iiwa7")
iiwa_q = self.plant.GetPositions(self.plant_context,
iiwa_model_instance)
integrator.GetMyContextFromRoot(
self.context).get_mutable_continuous_state_vector().SetFromVector(
iiwa_q)
class Juggler:
def __init__(self, kp=100, ki=1, kd=20, time_step=0.001, show_axis=False):
"""
Robotic Kuka IIWA juggler with paddle end effector
Args:
kp (int, optional): proportional gain. Defaults to 100.
ki (int, optional): integral gain. Defaults to 1.
kd (int, optional): derivative gain. Defaults to 20.
time_step (float, optional): time step for internal manipulation station controller. Defaults to 0.001.
show_axis (boolean, optional): whether or not to show the axis of reflection.
"""
self.time = 0
juggler_station = JugglerStation(kp=kp,
ki=ki,
kd=kd,
time_step=time_step,
show_axis=show_axis)
self.builder = DiagramBuilder()
self.station = self.builder.AddSystem(juggler_station.get_diagram())
self.position_log = []
self.velocity_log = []
self.visualizer = ConnectMeshcatVisualizer(
self.builder,
output_port=self.station.GetOutputPort("geometry_query"),
zmq_url=zmq_url)
self.plant = juggler_station.get_multibody_plant()
# ---------------------
ik_sys = self.builder.AddSystem(InverseKinematics(self.plant))
ik_sys.set_name("ik_sys")
v_mirror = self.builder.AddSystem(VelocityMirror(self.plant))
v_mirror.set_name("v_mirror")
w_tilt = self.builder.AddSystem(AngularVelocityTilt(self.plant))
w_tilt.set_name("w_tilt")
integrator = self.builder.AddSystem(Integrator(7))
self.builder.Connect(
self.station.GetOutputPort("iiwa_position_measured"),
ik_sys.GetInputPort("iiwa_pos_measured"))
self.builder.Connect(ik_sys.get_output_port(),
integrator.get_input_port())
self.builder.Connect(integrator.get_output_port(),
self.station.GetInputPort("iiwa_position"))
self.builder.Connect(
self.station.GetOutputPort("iiwa_position_measured"),
w_tilt.GetInputPort("iiwa_pos_measured"))
self.builder.Connect(
self.station.GetOutputPort("iiwa_position_measured"),
v_mirror.GetInputPort("iiwa_pos_measured"))
self.builder.Connect(
self.station.GetOutputPort("iiwa_velocity_estimated"),
v_mirror.GetInputPort("iiwa_velocity_estimated"))
self.builder.Connect(
w_tilt.get_output_port(),
ik_sys.GetInputPort("paddle_desired_angular_velocity"))
self.builder.Connect(v_mirror.get_output_port(),
ik_sys.GetInputPort("paddle_desired_velocity"))
self.builder.ExportInput(v_mirror.GetInputPort("ball_pose"),
"v_ball_pose")
self.builder.ExportInput(v_mirror.GetInputPort("ball_velocity"),
"v_ball_velocity")
self.builder.ExportInput(w_tilt.GetInputPort("ball_pose"),
"w_ball_pose")
self.builder.ExportInput(w_tilt.GetInputPort("ball_velocity"),
"w_ball_velocity")
# Useful for debugging
# desired_vel = self.builder.AddSystem(ConstantVectorSource([0, 0, 0]))
# self.builder.Connect(desired_vel.get_output_port(), ik_sys.GetInputPort("paddle_desired_angular_velocity"))
# ---------------------
self.diagram = self.builder.Build()
self.simulator = Simulator(self.diagram)
self.simulator.set_target_realtime_rate(1.0)
self.context = self.simulator.get_context()
self.station_context = self.station.GetMyContextFromRoot(self.context)
self.plant_context = self.plant.GetMyContextFromRoot(self.context)
self.plant.SetPositions(
self.plant_context, self.plant.GetModelInstanceByName("iiwa7"),
[0, np.pi / 3, 0, -np.pi / 2, 0, -np.pi / 3, 0])
self.station.GetInputPort("iiwa_feedforward_torque").FixValue(
self.station_context, np.zeros((7, 1)))
iiwa_model_instance = self.plant.GetModelInstanceByName("iiwa7")
iiwa_q = self.plant.GetPositions(self.plant_context,
iiwa_model_instance)
integrator.GetMyContextFromRoot(
self.context).get_mutable_continuous_state_vector().SetFromVector(
iiwa_q)
def step(self,
simulate=True,
duration=0.1,
final=True,
verbose=False,
display_pose=False):
"""
step the closed loop system
Args:
simulate (bool, optional): whether or not to visualize the command. Defaults to True.
duration (float, optional): duration to complete command in simulation. Defaults to 0.1.
final (bool, optional): whether or not this is the final command in the sequence; relevant for recording. Defaults to True.
verbose (bool, optional): whether or not to print measured position change. Defaults to False.
display_pose (bool, optional): whether or not to show the pose of the paddle in simulation. Defaults to False.
"""
ball_pose = self.plant.EvalBodyPoseInWorld(
self.plant_context,
self.plant.GetBodyByName("ball")).translation()
ball_velocity = self.plant.EvalBodySpatialVelocityInWorld(
self.plant_context,
self.plant.GetBodyByName("ball")).translational()
self.diagram.GetInputPort("v_ball_pose").FixValue(
self.context, ball_pose)
self.diagram.GetInputPort("v_ball_velocity").FixValue(
self.context, ball_velocity)
self.diagram.GetInputPort("w_ball_pose").FixValue(
self.context, ball_pose)
self.diagram.GetInputPort("w_ball_velocity").FixValue(
self.context, ball_velocity)
if display_pose:
transform = self.plant.EvalBodyPoseInWorld(
self.plant_context,
self.plant.GetBodyByName("base_link")).GetAsMatrix4()
AddTriad(self.visualizer.vis,
name=f"paddle_{round(self.time, 1)}",
prefix='',
length=0.15,
radius=.006)
self.visualizer.vis[''][
f"paddle_{round(self.time, 1)}"].set_transform(transform)
if simulate:
self.visualizer.start_recording()
self.simulator.AdvanceTo(self.time + duration)
self.visualizer.stop_recording()
self.position_log.append(
self.station.GetOutputPort("iiwa_position_measured").Eval(
self.station_context))
self.velocity_log.append(
self.station.GetOutputPort("iiwa_velocity_estimated").Eval(
self.station_context))
if verbose:
print("Time: {}\nMeasured Position: {}\n\n".format(
round(self.time, 3),
np.around(
self.station.GetOutputPort(
"iiwa_position_measured").Eval(
self.station_context), 3)))
if len(self.position_log) >= 2 and np.equal(
np.around(self.position_log[-1], 3),
np.around(self.position_log[-2], 3)).all():
print("something went wrong, simulation terminated")
final = True
if final:
self.visualizer.publish_recording()
return self.time + duration
self.time += duration
def do_fixed_structure_hmc_with_constraint_penalties(
grammar,
original_tree,
num_samples=100,
subsample_step=5,
verbose=0,
kernel_type="NUTS",
fix_observeds=False,
with_nonpenetration=False,
with_static_stability=False,
zmq_url=None,
prefix="hmc_sample",
constraints=[],
structure_vis_kwargs={},
**kwargs):
''' Given a scene tree, resample its continuous variables
(i.e. the node poses) while keeping the root and observed
node poses fixed, and trying to keep the constraints implied
by the tree and grammar satisfied.. Returns a population of trees sampled
from the joint distribution over node poses given the fixed structure.
Verbose = 0: Print nothing
Verbose = 1: Print updates about accept rate and steps
Verbose = 2: Print NLP output and scoring info
'''
# Strategy sketch:
# - Initialize at the current scene tree, asserting that it's a feasible configuration.
# - Form a probabilistic model that samples all of the node poses,
# and uses observe() statements to implement constraint factors as tightly peaked
# energy terms.
# - Use Pyro HMC to sample from the model.
# Make a bookkeeping copy of the tree
scene_tree = deepcopy(original_tree)
# Do steps of random-walk MCMC on those variables.
initial_score = scene_tree.score()
assert torch.isfinite(initial_score), "Bad initialization for MCMC."
if with_nonpenetration:
constraints.append(NonpenetrationConstraint(0.00))
# Form probabilistic model
root = scene_tree.get_root()
def model():
# Resample the continuous structure of the tree.
node_queue = [root]
while len(node_queue) > 0:
parent = node_queue.pop(0)
children, rules = scene_tree.get_children_and_rules(parent)
for child, rule in zip(children, rules):
with scope(prefix=parent.name):
rule.sample_child(parent, child)
node_queue.append(child)
# Implement observation constraints
if fix_observeds:
xyz_observed_variance = 1E-2
rot_observed_variance = 1E-2
for node, original_node in zip(scene_tree.nodes,
original_tree.nodes):
if node.observed:
xyz_observed_dist = dist.Normal(original_node.translation,
xyz_observed_variance)
rot_observed_dist = dist.Normal(original_node.rotation,
rot_observed_variance)
pyro.sample("%s_xyz_observed" % node.name,
xyz_observed_dist,
obs=node.translation)
pyro.sample("%s_rotation_observed" % node.name,
rot_observed_dist,
obs=node.rotation)
for k, constraint in enumerate(constraints):
clamped_error_distribution = dist.Normal(0., 0.001)
violation, _, _ = constraint.eval_violation(scene_tree)
positive_violations = torch.clamp(violation, 0., np.inf)
pyro.sample("%s_%d_err" % (type(constraint).__name__, k),
clamped_error_distribution,
obs=positive_violations)
# Ugh, I shouldn't need to manually reproduce the site names here.
# Can I rearrange how traces get assembled to extract these?
initial_values = {}
for parent in original_tree.nodes:
children, rules = original_tree.get_children_and_rules(parent)
for child, rule in zip(children, rules):
for key, site_value in rule.get_site_values(parent, child).items():
initial_values["%s/%s/%s" % (parent.name, child.name,
key)] = site_value.value
trace = pyro.poutine.trace(model).get_trace()
for key in initial_values.keys():
if key not in trace.nodes.keys():
print("Trace keys: ", trace.nodes.keys())
print("Initial values keys: ", initial_values.keys())
raise ValueError("%s not in trace keys" % key)
print("Initial trace log prob: ", trace.log_prob_sum())
# If I let MCMC auto-tune its step size, it seems to do well,
# but sometimes seems to get lost, and then gets stuck with big step size and
# zero acceptances.
if kernel_type == "NUTS":
kernel = NUTS(model,
init_strategy=pyro.infer.autoguide.init_to_value(
values=initial_values),
**kwargs)
elif kernel_type is "HMC":
kernel = HMC(model,
init_strategy=pyro.infer.autoguide.init_to_value(
values=initial_values),
**kwargs)
else:
raise NotImplementedError(kernel_type)
# Get MBP for viz
if zmq_url is not None:
builder, mbp, sg, node_to_free_body_ids_map, body_id_to_node_map = compile_scene_tree_to_mbp_and_sg(
scene_tree)
mbp.Finalize()
visualizer = ConnectMeshcatVisualizer(builder,
sg,
zmq_url=zmq_url,
prefix=prefix)
diagram = builder.Build()
diagram_context = diagram.CreateDefaultContext()
mbp_context = diagram.GetMutableSubsystemContext(mbp, diagram_context)
vis_context = diagram.GetMutableSubsystemContext(
visualizer, diagram_context)
visualizer.load()
def hook_fn(kernel, samples, stage, i):
# Set MBP context to
for node, body_ids in node_to_free_body_ids_map.items():
for body_id in body_ids:
mbp.SetFreeBodyPose(mbp_context, mbp.get_body(body_id),
torch_tf_to_drake_tf(node.tf))
diagram.Publish(diagram_context)
draw_scene_tree_structure_meshcat(scene_tree,
zmq_url=zmq_url,
prefix=prefix + "/structure",
delete=False,
**structure_vis_kwargs)
time.sleep(0.1)
else:
hook_fn = None
mcmc = MCMC(kernel,
num_samples=num_samples,
warmup_steps=min(int(num_samples / 2), 50),
num_chains=1,
disable_progbar=(verbose == -1),
hook_fn=hook_fn)
mcmc.run()
if verbose > 1:
mcmc.summary(prob=0.5)
samples = mcmc.get_samples()
sampled_trees = []
for k in range(0, num_samples, subsample_step):
condition = {key: value[k, ...] for key, value in samples.items()}
with pyro.condition(data=condition):
model()
sampled_trees.append(deepcopy(scene_tree))
return sampled_trees