class Visualizer():
def __init__(self, urdf):
self._create_plant(urdf)
def _create_plant(self, urdf):
self.plant = MultibodyPlant(time_step=0.0)
self.scenegraph = SceneGraph()
self.plant.RegisterAsSourceForSceneGraph(self.scenegraph)
self.model_index = Parser(self.plant).AddModelFromFile(FindResource(urdf))
self.builder = DiagramBuilder()
self.builder.AddSystem(self.scenegraph)
def _finalize_plant(self):
self.plant.Finalize()
def _add_trajectory_source(self, traj):
# Create the trajectory source
source = self.builder.AddSystem(TrajectorySource(traj))
pos2pose = self.builder.AddSystem(MultibodyPositionToGeometryPose(self.plant, input_multibody_state=True))
# Wire the source to the scene graph
self.builder.Connect(source.get_output_port(0), pos2pose.get_input_port())
self.builder.Connect(pos2pose.get_output_port(), self.scenegraph.get_source_pose_port(self.plant.get_source_id()))
def _add_renderer(self):
renderer_name = "renderer"
self.scenegraph.AddRenderer(renderer_name, MakeRenderEngineVtk(RenderEngineVtkParams()))
# Add a camera
depth_camera = DepthRenderCamera(
RenderCameraCore(
renderer_name,
CameraInfo(width=640, height=480, fov_y=np.pi/4),
ClippingRange(0.01, 10.0),
RigidTransform()),
DepthRange(0.01,10.0)
)
world_id = self.plant.GetBodyFrameIdOrThrow(self.plant.world_body().index())
X_WB = xyz_rpy_deg([4,0,0],[-90,0,90])
sensor = RgbdSensor(world_id, X_PB=X_WB, depth_camera=depth_camera)
self.builder.AddSystem(sensor)
self.builder.Connect(self.scenegraph.get_query_output_port(), sensor.query_object_input_port())
def _connect_visualizer(self):
DrakeVisualizer.AddToBuilder(self.builder, self.scenegraph)
self.meshcat = ConnectMeshcatVisualizer(self.builder, self.scenegraph, zmq_url="new", open_browser=False)
self.meshcat.vis.jupyter_cell()
def _make_visualization(self, stop_time):
simulator = Simulator(self.builder.Build())
simulator.Initialize()
self.meshcat.vis.render_static()
# Set simulator context
simulator.get_mutable_context().SetTime(0.0)
# Start recording and simulate the diagram
self.meshcat.reset_recording()
self.meshcat.start_recording()
simulator.AdvanceTo(stop_time)
# Publish the recording
self.meshcat.publish_recording()
# Render
self.meshcat.vis.render_static()
input("View visualization. Press <ENTER> to end")
print("Finished")
def visualize_trajectory(self, xtraj=None):
self._finalize_plant()
print("Adding trajectory source")
xtraj = self._check_trajectory(xtraj)
self._add_trajectory_source(xtraj)
print("Adding renderer")
self._add_renderer()
print("Connecting to MeshCat")
self._connect_visualizer()
print("Making visualization")
self._make_visualization(xtraj.end_time())
def _check_trajectory(self, traj):
if traj is None:
plant_context = self.plant.CreateDefaultContext()
pose = self.plant.GetPositions(plant_context)
pose = np.column_stack((pose, pose))
pose = zero_pad_rows(pose, self.plant.num_positions() + self.plant.num_velocities())
return PiecewisePolynomial.FirstOrderHold([0., 1.], pose)
elif type(traj) is np.ndarray:
if traj.ndim == 1:
traj = np.column_stack(traj, traj)
traj = zero_pad_rows(traj, self.plant.num_positions() + self.plant.num_velocities())
return PiecewisePolynomial.FirstOrderHold([0.,1.], traj)
elif type(traj) is PiecewisePolynomial:
breaks = traj.get_segment_times()
values = traj.vector_values(breaks)
values = zero_pad_rows(values, self.plant.num_positions() + self.plant.num_velocities())
return PiecewisePolynomial.FirstOrderHold(breaks, values)
else:
raise ValueError("Trajectory must be a piecewise polynomial, an ndarray, or None")
builder.Connect(controller.get_output_port(0),
plant.get_actuation_input_port())
"""
# Set up a simulator to run this diagram
simulator = Simulator(diagram)
context = simulator.get_mutable_context()
plant_context = plant.GetMyContextFromRoot(context)
# Problem block 2: This is following the linearization of the ballbot
# The error that comes up here is
# RuntimeError: The object named [] of type drake::systems::Diagram<double> does not support ToAutoDiffXd.
# I also get the same error calling diagram.ToAutoDiffXd()
"""
diag_context = diagram.CreateDefaultContext()
diagram.get_input_port().FixValue(diag_context, [0])
result = FirstOrderTaylorApproximation(diagram, context)
"""
plant.get_actuation_input_port(cartpole).FixValue(plant_context, np.array([0]))
plot_system_graphviz(diagram)
plt.show()
# Set the initial conditions
context.SetContinuousState([0, np.pi, 0, 0.00]) # x, theta, xdot, thetadot
context.SetTime(0.0)
visualizer.start_recording()
simulator.AdvanceTo(10.0)
visualizer.publish_recording()
visualizer.vis.render_static()
input()
# WARNING: If you do not cast the `geometry_label` to `int`, this
# comparison will take a long time since NumPy will do
# element-by-element comparison using `RenderLabel.__eq__`.
mask = (label == int(geometry_label))
label_by_model[mask] = int(body.model_instance())
plt.imshow(colorize_labels(label_by_model))
meshcat_vis.vis.render_static()
# Set the context to make the simulation valid and slightly interesting.
simulator.get_mutable_context().SetTime(.0)
plant_context = plant.GetMyContextFromRoot(simulator.get_mutable_context())
plant.get_actuation_input_port(iiwa_1).FixValue(plant_context, np.zeros(
(7, 1)))
plant.get_actuation_input_port(iiwa_2).FixValue(plant_context, np.zeros(
(7, 1)))
plant.SetPositions(plant_context, iiwa_1, [0.2, 0.4, 0, 0, 0, 0, 0])
# Reset the recording (in case you are running this cell more than once).
meshcat_vis.reset_recording()
# Start recording and simulate.
meshcat_vis.start_recording()
simulator.AdvanceTo(1.0)
# Publish the recording to meshcat.
meshcat_vis.publish_recording()
# Render meshcat's current state.
meshcat_vis.vis.render_static()