forked from pz4kybsvg/Conception
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
137 lines
4.7 KiB
137 lines
4.7 KiB
"""
|
|
Run examples of PlanarSceneGraphVisualizer, e.g. to visualize a pendulum.
|
|
Usage demo: load a URDF, rig it up with a constant torque input, and draw it.
|
|
"""
|
|
|
|
import argparse
|
|
|
|
import numpy as np
|
|
|
|
from pydrake.common import temp_directory
|
|
from pydrake.examples import ManipulationStation
|
|
from pydrake.multibody.parsing import Parser
|
|
from pydrake.multibody.plant import AddMultibodyPlantSceneGraph
|
|
from pydrake.systems.analysis import Simulator
|
|
from pydrake.systems.framework import DiagramBuilder
|
|
from pydrake.systems.planar_scenegraph_visualizer import (
|
|
ConnectPlanarSceneGraphVisualizer)
|
|
|
|
|
|
def run_pendulum_example(args):
|
|
builder = DiagramBuilder()
|
|
plant, scene_graph = AddMultibodyPlantSceneGraph(builder, 0.0)
|
|
parser = Parser(plant)
|
|
parser.AddModelsFromUrl(
|
|
url="package://drake/examples/pendulum/Pendulum.urdf")
|
|
plant.Finalize()
|
|
|
|
T_VW = np.array([[1., 0., 0., 0.],
|
|
[0., 0., 1., 0.],
|
|
[0., 0., 0., 1.]])
|
|
visualizer = ConnectPlanarSceneGraphVisualizer(
|
|
builder, scene_graph, T_VW=T_VW, xlim=[-1.2, 1.2], ylim=[-1.2, 1.2])
|
|
|
|
if args.playback:
|
|
visualizer.start_recording()
|
|
|
|
diagram = builder.Build()
|
|
simulator = Simulator(diagram)
|
|
simulator.Initialize()
|
|
simulator.set_target_realtime_rate(1.0)
|
|
|
|
# Fix the input port to zero.
|
|
plant_context = diagram.GetMutableSubsystemContext(
|
|
plant, simulator.get_mutable_context())
|
|
plant.get_actuation_input_port().FixValue(
|
|
plant_context, np.zeros(plant.num_actuators()))
|
|
plant_context.SetContinuousState([0.5, 0.1])
|
|
simulator.AdvanceTo(args.duration)
|
|
|
|
if args.playback:
|
|
visualizer.stop_recording()
|
|
ani = visualizer.get_recording_as_animation()
|
|
# Playback the recording and save the output.
|
|
ani.save("{}/pend_playback.mp4".format(temp_directory()), fps=30)
|
|
|
|
|
|
def run_manipulation_example(args):
|
|
builder = DiagramBuilder()
|
|
station = builder.AddSystem(ManipulationStation(time_step=0.005))
|
|
station.SetupManipulationClassStation()
|
|
station.Finalize()
|
|
|
|
plant = station.get_multibody_plant()
|
|
scene_graph = station.get_scene_graph()
|
|
query_object_output_port = station.GetOutputPort("geometry_query")
|
|
|
|
# Side-on view of the station.
|
|
T_VW = np.array([[1., 0., 0., 0.],
|
|
[0., 0., 1., 0.],
|
|
[0., 0., 0., 1.]])
|
|
visualizer = ConnectPlanarSceneGraphVisualizer(
|
|
builder, scene_graph, query_object_output_port, T_VW=T_VW,
|
|
xlim=[-0.5, 1.0], ylim=[-1.2, 1.2], draw_period=0.1)
|
|
|
|
if args.playback:
|
|
visualizer.start_recording()
|
|
|
|
diagram = builder.Build()
|
|
simulator = Simulator(diagram)
|
|
simulator.Initialize()
|
|
simulator.set_target_realtime_rate(1.0)
|
|
|
|
# Fix the control inputs to zero.
|
|
station_context = diagram.GetMutableSubsystemContext(
|
|
station, simulator.get_mutable_context())
|
|
station.GetInputPort("iiwa_position").FixValue(
|
|
station_context, station.GetIiwaPosition(station_context))
|
|
station.GetInputPort("iiwa_feedforward_torque").FixValue(
|
|
station_context, np.zeros(7))
|
|
station.GetInputPort("wsg_position").FixValue(
|
|
station_context, np.zeros(1))
|
|
station.GetInputPort("wsg_force_limit").FixValue(
|
|
station_context, [40.0])
|
|
simulator.AdvanceTo(args.duration)
|
|
|
|
if args.playback:
|
|
visualizer.stop_recording()
|
|
ani = visualizer.get_recording_as_animation()
|
|
# Playback the recording and save the output.
|
|
ani.save("{}/manip_playback.mp4".format(temp_directory()), fps=30)
|
|
|
|
|
|
def main():
|
|
np.set_printoptions(precision=5, suppress=True)
|
|
parser = argparse.ArgumentParser(
|
|
description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter
|
|
)
|
|
parser.add_argument("-T", "--duration",
|
|
type=float,
|
|
help="Duration to run sim in seconds.",
|
|
default=1.0)
|
|
parser.add_argument("-m", "--models",
|
|
type=str,
|
|
nargs="*",
|
|
help="Models to run, at least one of [pend, manip]",
|
|
default=["pend"])
|
|
parser.add_argument("-p", "--playback",
|
|
type=bool,
|
|
help="Whether to record and playback the animation",
|
|
default=False)
|
|
args = parser.parse_args()
|
|
|
|
for model in args.models:
|
|
if model == "pend":
|
|
run_pendulum_example(args)
|
|
elif model == "manip":
|
|
run_manipulation_example(args)
|
|
else:
|
|
print("Unrecognized model %s." % model)
|
|
parser.print_usage()
|
|
exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|