From eeb404038c8e93983a605c9169540fed86b20ced Mon Sep 17 00:00:00 2001 From: Bel <644574+tehbelinda@users.noreply.github.com> Date: Thu, 7 Nov 2019 13:49:11 -0800 Subject: [PATCH] Port planar_scenegraph_visualizer from underactuated (#12218) Mesh support is removed for now See https://github.com/RobotLocomotion/drake/issues/11104 Co-authored-by: Russ Tedrake Co-authored-by: Eric Cousineau Co-authored-by: Greg Izatt --- .../pydrake/examples/multibody/BUILD.bazel | 15 + .../run_planar_scenegraph_visualizer.py | 121 +++++++ bindings/pydrake/systems/BUILD.bazel | 53 +++ bindings/pydrake/systems/all.py | 2 + .../pydrake/systems/meshcat_visualizer.py | 2 + .../systems/planar_scenegraph_visualizer.py | 310 ++++++++++++++++++ bindings/pydrake/systems/pyplot_visualizer.py | 107 ++++++ .../test/planar_scenegraph_visualizer_test.py | 136 ++++++++ .../systems/test/pyplot_visualizer_test.py | 97 ++++++ 9 files changed, 843 insertions(+) create mode 100644 bindings/pydrake/examples/multibody/run_planar_scenegraph_visualizer.py create mode 100644 bindings/pydrake/systems/planar_scenegraph_visualizer.py create mode 100644 bindings/pydrake/systems/pyplot_visualizer.py create mode 100644 bindings/pydrake/systems/test/planar_scenegraph_visualizer_test.py create mode 100644 bindings/pydrake/systems/test/pyplot_visualizer_test.py diff --git a/bindings/pydrake/examples/multibody/BUILD.bazel b/bindings/pydrake/examples/multibody/BUILD.bazel index c3eb92316644..90ee2858d7bb 100644 --- a/bindings/pydrake/examples/multibody/BUILD.bazel +++ b/bindings/pydrake/examples/multibody/BUILD.bazel @@ -67,9 +67,24 @@ drake_py_binary( ], ) +drake_py_binary( + name = "run_planar_scenegraph_visualizer", + srcs = ["run_planar_scenegraph_visualizer.py"], + add_test_rule = 1, + data = ["//examples/pendulum:models"], + visibility = ["//visibility:public"], + deps = [ + ":module_py", + "//bindings/pydrake/examples:manipulation_station_py", + "//bindings/pydrake/multibody", + "//bindings/pydrake/systems", + ], +) + PY_LIBRARIES = [ ":cart_pole_passive_simulation", ":pendulum_lqr_monte_carlo_analysis", + ":run_planar_scenegraph_visualizer", ] # Package roll-up (for Bazel dependencies). diff --git a/bindings/pydrake/examples/multibody/run_planar_scenegraph_visualizer.py b/bindings/pydrake/examples/multibody/run_planar_scenegraph_visualizer.py new file mode 100644 index 000000000000..5a9e444c6b36 --- /dev/null +++ b/bindings/pydrake/examples/multibody/run_planar_scenegraph_visualizer.py @@ -0,0 +1,121 @@ +""" +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 FindResourceOrThrow +from pydrake.examples.manipulation_station 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 ( + PlanarSceneGraphVisualizer) + + +def run_pendulum_example(args): + builder = DiagramBuilder() + plant, scene_graph = AddMultibodyPlantSceneGraph(builder) + parser = Parser(plant) + parser.AddModelFromFile(FindResourceOrThrow( + "drake/examples/pendulum/Pendulum.urdf")) + plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("base")) + plant.Finalize() + + pose_bundle_output_port = scene_graph.get_pose_bundle_output_port() + Tview = np.array([[1., 0., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]], + dtype=np.float64) + visualizer = builder.AddSystem(PlanarSceneGraphVisualizer( + scene_graph, T_VW=Tview, xlim=[-1.2, 1.2], ylim=[-1.2, 1.2])) + builder.Connect(pose_bundle_output_port, visualizer.get_input_port(0)) + + 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_context.FixInputPort( + plant.get_actuation_input_port().get_index(), + np.zeros(plant.num_actuators())) + plant_context.SetContinuousState([0.5, 0.1]) + simulator.AdvanceTo(args.duration) + + +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() + pose_bundle_output_port = station.GetOutputPort("pose_bundle") + + # Side-on view of the station. + Tview = np.array([[1., 0., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]], + dtype=np.float64) + visualizer = builder.AddSystem(PlanarSceneGraphVisualizer( + scene_graph, T_VW=Tview, xlim=[-0.5, 1.0], ylim=[-1.2, 1.2], + draw_period=0.1)) + builder.Connect(pose_bundle_output_port, visualizer.get_input_port(0)) + + 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) + + +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"]) + 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() diff --git a/bindings/pydrake/systems/BUILD.bazel b/bindings/pydrake/systems/BUILD.bazel index 9f2d9fe2d44c..108c10d45fdd 100644 --- a/bindings/pydrake/systems/BUILD.bazel +++ b/bindings/pydrake/systems/BUILD.bazel @@ -269,6 +269,31 @@ drake_py_library( ], ) +drake_py_library( + name = "planar_scenegraph_visualizer_py", + srcs = ["planar_scenegraph_visualizer.py"], + imports = PACKAGE_INFO.py_imports, + deps = [ + ":framework_py", + ":module_py", + ":pyplot_visualizer_py", + ":rendering_py", + "//bindings/pydrake:lcm_py", + ], +) + +drake_py_library( + name = "pyplot_visualizer_py", + srcs = ["pyplot_visualizer.py"], + imports = PACKAGE_INFO.py_imports, + deps = [ + ":framework_py", + ":module_py", + ":primitives_py", + "//bindings/pydrake:trajectories_py", + ], +) + PY_LIBRARIES_WITH_INSTALL = [ ":analysis_py", ":controllers_py", @@ -286,6 +311,8 @@ PY_LIBRARIES = [ ":scalar_conversion_py", ":module_py", ":perception_py", + ":planar_scenegraph_visualizer_py", + ":pyplot_visualizer_py", ] # Symbol roll-up (for user ease). @@ -450,6 +477,32 @@ drake_py_unittest( ], ) +drake_py_unittest( + name = "planar_scenegraph_visualizer_test", + data = [ + "//examples/multibody/cart_pole:models", + "//manipulation/models/iiwa_description:models", + ], + deps = [ + ":analysis_py", + ":framework_py", + ":planar_scenegraph_visualizer_py", + "//bindings/pydrake:geometry_py", + "//bindings/pydrake/multibody", + ], +) + +drake_py_unittest( + name = "pyplot_visualizer_test", + deps = [ + ":analysis_py", + ":framework_py", + ":primitives_py", + ":pyplot_visualizer_py", + "//bindings/pydrake:trajectories_py", + ], +) + drake_py_unittest( name = "rendering_test", data = ["//multibody/benchmarks/acrobot:models"], diff --git a/bindings/pydrake/systems/all.py b/bindings/pydrake/systems/all.py index f20b5422826a..3b5b4bd05989 100644 --- a/bindings/pydrake/systems/all.py +++ b/bindings/pydrake/systems/all.py @@ -4,7 +4,9 @@ from .lcm import * from .meshcat_visualizer import * from .perception import * +from .planar_scenegraph_visualizer import * from .primitives import * +from .pyplot_visualizer import * from .rendering import * from .scalar_conversion import * from .sensors import * diff --git a/bindings/pydrake/systems/meshcat_visualizer.py b/bindings/pydrake/systems/meshcat_visualizer.py index 3e6680990dfa..1da2af1cd9bc 100644 --- a/bindings/pydrake/systems/meshcat_visualizer.py +++ b/bindings/pydrake/systems/meshcat_visualizer.py @@ -243,6 +243,8 @@ def __init__(self, self.draw_period = draw_period # Pose bundle (from SceneGraph) input port. + # TODO(tehbelinda): Rename the `lcm_visualization` port to match + # SceneGraph once its output port has been updated. See #12214. self.DeclareAbstractInputPort("lcm_visualization", AbstractValue.Make(PoseBundle(0))) diff --git a/bindings/pydrake/systems/planar_scenegraph_visualizer.py b/bindings/pydrake/systems/planar_scenegraph_visualizer.py new file mode 100644 index 000000000000..d29fac68f927 --- /dev/null +++ b/bindings/pydrake/systems/planar_scenegraph_visualizer.py @@ -0,0 +1,310 @@ +import math +import warnings + +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +with warnings.catch_warnings(): # noqa + # N.B. We must suppress this to appease `all_test`. + # TODO(eric.cousineau): Remove this once all supported platform ships + # `scipy>=1.0.0` by default. + warnings.simplefilter("ignore", ImportWarning) + import scipy as sp + from scipy import spatial + +from drake import lcmt_viewer_load_robot +from pydrake.common.eigen_geometry import Quaternion +from pydrake.geometry import DispatchLoadMessage +from pydrake.lcm import DrakeMockLcm, Subscriber +from pydrake.math import RigidTransform, RotationMatrix +from pydrake.systems.framework import AbstractValue +from pydrake.systems.pyplot_visualizer import PyPlotVisualizer +from pydrake.systems.rendering import PoseBundle + + +class PlanarSceneGraphVisualizer(PyPlotVisualizer): + """ + Given a SceneGraph and a view plane, provides a view of the robot by + projecting all geometry onto the view plane. + + This is intended to be used for robots that operate in the plane, but + should render any robot approximately correctly. It has the following + caveats: + - z-ordering of objects is done based on the object centroid, which + is not perfect for non-planar scenes. + - Object geometry is projected onto the view plane, then a chull is taken, + and finally that chull is drawn as a patch. Nonconvex geometry will thus be + drawn incorrectly, and geometry with many vertices will slow down the + visualizer. + + Specifics on view setup: + + T_VW specifies the 3x4 view projection matrix. For planar orthographic + projection, use: + [ x_axis_shift + y_axis_shift + 0, 0, 0, 1] % homogenizer + + e.g. + + [ 1 0 0 0.5 + 0 1 0 0 + 0 0 0 1]. + + would give a top-down view (i.e squashing the z axis), and would shift + things in the x axis positively by 0.5. + + T_VW can be any valid view projection matrix. If the bottom row is + [0, 0, 0, 1], the view projection will be an orthographic projection. + + xlim and ylim don't technically provide extra functionality, but it's + easier to keep handling scaling with xlim, ylim, and view plane selection + and *maybe* offsetting with the projection matrix. + """ + + def __init__(self, + scene_graph, + draw_period=1./30, + T_VW=np.array([[1., 0., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]), + xlim=[-1., 1], + ylim=[-1, 1], + facecolor=[1, 1, 1], + use_random_colors=False, + ax=None): + """ + Args: + scene_graph: A SceneGraph object. + draw_period: The rate at which this class publishes to the + visualizer. + T_VW: The view projection matrix from world to view coordinates. + xlim: View limit into the scene. + ylim: View limit into the scene. + facecolor: Passed through to figure() and sets background color. + Both color name strings and RGB triplets are allowed. Defaults + to white. + use_random_colors: If set to True, will render each body with a + different color. (Multiple visual elements on the same body + will be the same color.) + ax: If supplied, the visualizer will draw onto those axes instead + of creating a new set of axes. The visualizer will still change + the view range and figure size of those axes. + """ + default_size = matplotlib.rcParams['figure.figsize'] + scalefactor = (ylim[1]-ylim[0]) / (xlim[1]-xlim[0]) + figsize = (default_size[0], default_size[0]*scalefactor) + + PyPlotVisualizer.__init__(self, facecolor=facecolor, figsize=figsize, + ax=ax, draw_period=draw_period) + self.set_name('planar_multibody_visualizer') + + self._scene_graph = scene_graph + self._T_VW = T_VW + + # Pose bundle (from SceneGraph) input port. + # TODO(tehbelinda): Rename the `lcm_visualization` port to match + # SceneGraph once its output port has been updated. See #12214. + self._pose_bundle_port = self.DeclareAbstractInputPort( + "lcm_visualization", AbstractValue.Make(PoseBundle(0))) + + self.ax.axis('equal') + self.ax.axis('off') + + # Achieve the desired view limits. + self.ax.set_xlim(xlim) + self.ax.set_ylim(ylim) + default_size = self.fig.get_size_inches() + self.fig.set_size_inches(figsize[0], figsize[1]) + + # Populate body patches. + self._build_body_patches(use_random_colors) + + # Populate the body fill list -- which requires doing most of a draw + # pass, but with an ax.fill() command rather than an in-place + # replacement of vertex positions to initialize the draw patches. The + # body fill list stores the ax patch objects in the order they were + # spawned (i.e. by body, and then by order of view_patches). Drawing + # the tree should update them by iterating over bodies and patches in + # the same order. + self._body_fill_dict = {} + X_WB_initial = RigidTransform.Identity() + for full_name in self._patch_Blist.keys(): + patch_Wlist, view_colors = self._get_view_patches(full_name, + X_WB_initial) + self._body_fill_dict[full_name] = [] + for patch_W, color in zip(patch_Wlist, view_colors): + # Project the full patch the first time, to initialize a vertex + # list with enough space for any possible convex hull of this + # vertex set. + patch_V = self._T_VW @ patch_W + self._body_fill_dict[full_name] += self.ax.fill( + patch_V[0, :], patch_V[1, :], zorder=0, + edgecolor='k', facecolor=color, closed=True) + + def _build_body_patches(self, use_random_colors): + """ + Generates body patches. self._patch_Blist stores a list of patches for + each body (starting at body id 1). A body patch is a list of all 3D + vertices of a piece of visual geometry. + """ + self._patch_Blist = {} + self._patch_Blist_colors = {} + + mock_lcm = DrakeMockLcm() + mock_lcm_subscriber = Subscriber(lcm=mock_lcm, + channel="DRAKE_VIEWER_LOAD_ROBOT", + lcm_type=lcmt_viewer_load_robot) + # TODO(SeanCurtis-TRI): Use SceneGraph inspection instead of mocking + # LCM and inspecting the generated message. + DispatchLoadMessage(self._scene_graph, mock_lcm) + mock_lcm.HandleSubscriptions(0) + assert mock_lcm_subscriber.count > 0 + load_robot_msg = mock_lcm_subscriber.message + + # Spawn a random color generator, in case we need to pick random colors + # for some bodies. Each body will be given a unique color when using + # this random generator, with each visual element of the body colored + # the same. + color = iter(plt.cm.rainbow(np.linspace(0, 1, + load_robot_msg.num_links))) + + for i in range(load_robot_msg.num_links): + link = load_robot_msg.link[i] + + this_body_patches = [] + this_body_colors = [] + this_color = next(color) + + for j in range(link.num_geom): + geom = link.geom[j] + # MultibodyPlant currently sets alpha=0 to make collision + # geometry "invisible". Ignore those geometries here. + if geom.color[3] == 0: + continue + + X_BG = RigidTransform( + RotationMatrix(Quaternion(geom.quaternion)), + geom.position) + + if geom.type == geom.BOX: + assert geom.num_float_data == 3 + + # Draw a bounding box. + patch_G = np.vstack(( + geom.float_data[0]/2.*np.array( + [-1, -1, 1, 1, -1, -1, 1, 1]), + geom.float_data[1]/2.*np.array( + [-1, 1, -1, 1, -1, 1, -1, 1]), + geom.float_data[2]/2.*np.array( + [-1, -1, -1, -1, 1, 1, 1, 1]))) + + elif geom.type == geom.SPHERE: + assert geom.num_float_data == 1 + radius = geom.float_data[0] + lati, longi = np.meshgrid(np.arange(0., 2.*math.pi, 0.5), + np.arange(0., 2.*math.pi, 0.5)) + lati = lati.ravel() + longi = longi.ravel() + patch_G = np.vstack([ + np.sin(lati)*np.cos(longi), + np.sin(lati)*np.sin(longi), + np.cos(lati)]) + patch_G *= radius + + elif geom.type == geom.CYLINDER: + assert geom.num_float_data == 2 + radius = geom.float_data[0] + length = geom.float_data[1] + + # In the lcm geometry, cylinders are along +z + # https://github.com/RobotLocomotion/drake/blob/last_sha_with_original_matlab/drake/matlab/systems/plants/RigidBodyCylinder.m + # Two circles: one at bottom, one at top. + sample_pts = np.arange(0., 2.*math.pi, 0.25) + patch_G = np.hstack( + [np.array([ + [radius*math.cos(pt), + radius*math.sin(pt), + -length/2.], + [radius*math.cos(pt), + radius*math.sin(pt), + length/2.]]).T + for pt in sample_pts]) + + else: + print("UNSUPPORTED GEOMETRY TYPE {} IGNORED".format( + geom.type)) + continue + + # Compute pose in body. + patch_B = X_BG @ patch_G + + # Close path if not closed. + if (patch_B[:, -1] != patch_B[:, 0]).any(): + patch_B = np.hstack((patch_B, patch_B[:, 0][np.newaxis].T)) + + this_body_patches.append(patch_B) + if use_random_colors: + this_body_colors.append(this_color) + else: + this_body_colors.append(geom.color) + + self._patch_Blist[link.name] = this_body_patches + self._patch_Blist_colors[link.name] = this_body_colors + + def _get_view_patches(self, full_name, X_WB): + """ + Pulls out the view patch verts for the given body index after applying + the appropriate transform, X_WB. X_WB needs to be a RigidTransform. + """ + patch_Wlist = [] + for patch_B in self._patch_Blist[full_name]: + patch_W = X_WB @ patch_B + # Add homogeneous row. + patch_W = np.vstack((patch_W, np.ones((1, patch_W.shape[1])))) + patch_Wlist.append(patch_W) + + colors = self._patch_Blist_colors[full_name] + return (patch_Wlist, colors) + + def draw(self, context): + """Overrides base with the implementation.""" + + pose_bundle = self._pose_bundle_port.Eval(context) + view_dir = np.cross(self._T_VW[0, :3], self._T_VW[1, :3]) + for frame_i in range(pose_bundle.get_num_poses()): + # SceneGraph currently sets the name in PoseBundle as + # "get_source_name::frame_name". + full_name = pose_bundle.get_name(frame_i) + model_id = pose_bundle.get_model_instance_id(frame_i) + + X_WB = RigidTransform(pose_bundle.get_pose(frame_i)) + patch_Wlist, _ = self._get_view_patches(full_name, X_WB) + for i, patch_W in enumerate(patch_Wlist): + # Project the object vertices from 3d in world frame W to 2d in + # view frame V (keeps homogeneous portion, removing it later). + patch_V = self._T_VW @ patch_W + # Applies normalization in the perspective transformation + # to make each projected point have z = 1. If the bottom row + # of T_VW is [0, 0, 0, 1], this will result in an + # orthographic projection. + patch_V[0, :] /= patch_V[2, :] + patch_V[1, :] /= patch_V[2, :] + # Cut patch_V down to 2xN. + patch_V = patch_V[:2, :] + # Take a convex hull to get an accurate shape for drawing, + # with verts coming out in ccw order. + if patch_V.shape[1] > 3: + hull = spatial.ConvexHull(patch_V.T) + patch_V = np.vstack( + [patch_V[:, v] for v in hull.vertices]).T + body_fill = self._body_fill_dict[full_name][i] + n_verts = body_fill.get_path().vertices.shape[0] + # Update the verts, padding out to the appropriate full # of + # verts by replicating the final vertex. + patch_V = np.pad( + patch_V, ((0, 0), (0, n_verts - patch_V.shape[1])), + mode="edge") + body_fill.get_path().vertices[:, :] = patch_V.T + body_fill.zorder = X_WB.translation() @ view_dir + self.ax.set_title('t = {:.1f}'.format(context.get_time())) diff --git a/bindings/pydrake/systems/pyplot_visualizer.py b/bindings/pydrake/systems/pyplot_visualizer.py new file mode 100644 index 000000000000..38322debccdf --- /dev/null +++ b/bindings/pydrake/systems/pyplot_visualizer.py @@ -0,0 +1,107 @@ +import matplotlib +import matplotlib.animation as animation +import matplotlib.pyplot as plt +from matplotlib.widgets import Slider +import numpy as np + +from pydrake.systems.framework import LeafSystem, PublishEvent, TriggerType +from pydrake.systems.primitives import SignalLogger +from pydrake.trajectories import Trajectory + + +class PyPlotVisualizer(LeafSystem): + """ + Base class from planar visualization that relies on pyplot. + + In the configuration set up here, this visualizer provides one + visualization window (self.fig) with axes (self.ax). The axes can + optionally be supplied externally to allow other visualizers to overlay + additional information. + + Subclasses must: + - During initialization, set up the figure bounds and register input port + with the appropriate message type. + - Override the draw method to parse the input and draw the robot in the + appropriate state. + """ + + def __init__(self, draw_period=1./30, facecolor=[1, 1, 1], + figsize=None, ax=None): + LeafSystem.__init__(self) + + self.set_name('pyplot_visualization') + self.timestep = draw_period + self.DeclarePeriodicPublish(draw_period, 0.0) + + if ax is None: + self.fig = plt.figure(facecolor=facecolor, figsize=figsize) + self.ax = plt.axes() + self.fig.add_axes(self.ax) + else: + self.ax = ax + self.fig = ax.get_figure() + + self.ax.axis('equal') + self.ax.axis('off') + + self._show = (matplotlib.get_backend().lower() != 'template') + + def on_initialize(context, event): + if self._show: + self.fig.show() + + self.DeclareInitializationEvent( + event=PublishEvent( + trigger_type=TriggerType.kInitialization, + callback=on_initialize)) + + def DoPublish(self, context, event): + LeafSystem.DoPublish(self, context, event) + if self._show: + self.draw(context) + self.fig.canvas.draw() + plt.pause(1e-10) + + def draw(self, context): + """Draws a single frame. + `context` can either be a Context object, or a raw vector (for ease of + interpolation). + """ + raise NotImplementedError + + def animate(self, log, resample=True, repeat=False): + """ + Args: + log: A reference to a pydrake.systems.primitives.SignalLogger that + contains the plant state after running a simulation. + resample: Whether we should do a resampling operation to make the + samples more consistent in time. This can be disabled if you + know the draw_period passed into the constructor exactly + matches the sample timestep of the log. + repeat: Whether the resulting animation should repeat. + """ + if isinstance(log, SignalLogger): + t = log.sample_times() + x = log.data() + + if resample: + import scipy.interpolate + + t_resample = np.arange(0, t[-1], self.timestep) + x = scipy.interpolate.interp1d(t, x, kind='linear', axis=1)(t_resample) # noqa + t = t_resample + + elif isinstance(log, Trajectory): + t = np.arange(log.start_time(), log.end_time(), self.timestep) + x = np.hstack([log.value(time) for time in t]) + + def animate_update(i): + self.draw(x[:, i]) + + ani = animation.FuncAnimation(self.fig, + animate_update, + t.shape[0], + # Convert from s to ms. + interval=1000*self.timestep, + repeat=repeat) + return ani diff --git a/bindings/pydrake/systems/test/planar_scenegraph_visualizer_test.py b/bindings/pydrake/systems/test/planar_scenegraph_visualizer_test.py new file mode 100644 index 000000000000..c8a3ffd3c34c --- /dev/null +++ b/bindings/pydrake/systems/test/planar_scenegraph_visualizer_test.py @@ -0,0 +1,136 @@ +import unittest + +import numpy as np + +from pydrake.common import FindResourceOrThrow +from pydrake.geometry import Box +from pydrake.math import RigidTransform +from pydrake.multibody.parsing import Parser +from pydrake.multibody.plant import ( + AddMultibodyPlantSceneGraph, CoulombFriction) +from pydrake.multibody.tree import SpatialInertia, UnitInertia +from pydrake.systems.analysis import Simulator +from pydrake.systems.framework import DiagramBuilder +from pydrake.systems.planar_scenegraph_visualizer import ( + PlanarSceneGraphVisualizer) + + +class TestPlanarSceneGraphVisualizer(unittest.TestCase): + def test_cart_pole(self): + """Cart-Pole with simple geometry.""" + file_name = FindResourceOrThrow( + "drake/examples/multibody/cart_pole/cart_pole.sdf") + builder = DiagramBuilder() + cart_pole, scene_graph = AddMultibodyPlantSceneGraph(builder) + Parser(plant=cart_pole).AddModelFromFile(file_name) + cart_pole.Finalize() + self.assertTrue(cart_pole.geometry_source_is_registered()) + + visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph)) + builder.Connect(scene_graph.get_pose_bundle_output_port(), + visualizer.get_input_port(0)) + + diagram = builder.Build() + + diagram_context = diagram.CreateDefaultContext() + cart_pole_context = diagram.GetMutableSubsystemContext( + cart_pole, diagram_context) + vis_context = diagram.GetMutableSubsystemContext( + visualizer, diagram_context) + + cart_pole.get_actuation_input_port().FixValue(cart_pole_context, 0) + + cart_slider = cart_pole.GetJointByName("CartSlider") + pole_pin = cart_pole.GetJointByName("PolePin") + cart_slider.set_translation(context=cart_pole_context, translation=0.) + pole_pin.set_angle(context=cart_pole_context, angle=2.) + + simulator = Simulator(diagram, diagram_context) + simulator.set_publish_every_time_step(False) + simulator.AdvanceTo(.1) + + visualizer.draw(vis_context) + self.assertEqual(visualizer.ax.get_title(), "t = 0.1",) + + def test_kuka(self): + """Kuka IIWA with mesh geometry.""" + file_name = FindResourceOrThrow( + "drake/manipulation/models/iiwa_description/sdf/" + "iiwa14_no_collision.sdf") + builder = DiagramBuilder() + kuka, scene_graph = AddMultibodyPlantSceneGraph(builder) + Parser(plant=kuka).AddModelFromFile(file_name) + kuka.Finalize() + + # Make sure that the frames to visualize exist. + iiwa = kuka.GetModelInstanceByName("iiwa14") + kuka.GetFrameByName("iiwa_link_7", iiwa) + kuka.GetFrameByName("iiwa_link_6", iiwa) + + frames_to_draw = {"iiwa14": {"iiwa_link_7", "iiwa_link_6"}} + visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph)) + builder.Connect(scene_graph.get_pose_bundle_output_port(), + visualizer.get_input_port(0)) + + diagram = builder.Build() + + diagram_context = diagram.CreateDefaultContext() + kuka_context = diagram.GetMutableSubsystemContext( + kuka, diagram_context) + vis_context = diagram.GetMutableSubsystemContext( + visualizer, diagram_context) + + kuka_actuation_port = kuka.get_actuation_input_port() + kuka_actuation_port.FixValue(kuka_context, + np.zeros(kuka_actuation_port.size())) + + simulator = Simulator(diagram, diagram_context) + simulator.set_publish_every_time_step(False) + simulator.AdvanceTo(.1) + + visualizer.draw(vis_context) + self.assertEqual(visualizer.ax.get_title(), "t = 0.1",) + + def test_procedural_geometry(self): + """ + This test ensures we can draw procedurally added primitive + geometry that is added to the world model instance (which has + a slightly different naming scheme than geometry with a + non-default / non-world model instance). + """ + builder = DiagramBuilder() + mbp, scene_graph = AddMultibodyPlantSceneGraph(builder) + world_body = mbp.world_body() + box_shape = Box(1., 2., 3.) + # This rigid body will be added to the world model instance since + # the model instance is not specified. + box_body = mbp.AddRigidBody("box", SpatialInertia( + mass=1.0, p_PScm_E=np.array([0., 0., 0.]), + G_SP_E=UnitInertia(1.0, 1.0, 1.0))) + mbp.WeldFrames(world_body.body_frame(), box_body.body_frame(), + RigidTransform()) + mbp.RegisterVisualGeometry( + box_body, RigidTransform.Identity(), box_shape, "ground_vis", + np.array([0.5, 0.5, 0.5, 1.])) + mbp.RegisterCollisionGeometry( + box_body, RigidTransform.Identity(), box_shape, "ground_col", + CoulombFriction(0.9, 0.8)) + mbp.Finalize() + + frames_to_draw = {"world": {"box"}} + visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph)) + builder.Connect(scene_graph.get_pose_bundle_output_port(), + visualizer.get_input_port(0)) + + diagram = builder.Build() + + diagram_context = diagram.CreateDefaultContext() + vis_context = diagram.GetMutableSubsystemContext( + visualizer, diagram_context) + + simulator = Simulator(diagram, diagram_context) + simulator.set_publish_every_time_step(False) + simulator.AdvanceTo(.1) + + visualizer.draw(vis_context) + self.assertEqual(visualizer.ax.get_title(), "t = 0.1",) diff --git a/bindings/pydrake/systems/test/pyplot_visualizer_test.py b/bindings/pydrake/systems/test/pyplot_visualizer_test.py new file mode 100644 index 000000000000..253f80c08704 --- /dev/null +++ b/bindings/pydrake/systems/test/pyplot_visualizer_test.py @@ -0,0 +1,97 @@ +import unittest + +import matplotlib.animation as animation +import matplotlib.pyplot as plt +import numpy as np + +from pydrake.systems.analysis import Simulator +from pydrake.systems.framework import ( + Context, DiagramBuilder, PortDataType, VectorSystem) +from pydrake.systems.primitives import SignalLogger +from pydrake.systems.pyplot_visualizer import PyPlotVisualizer +from pydrake.trajectories import PiecewisePolynomial + + +# TODO(tehbelinda): Augment this test with a Jupyter notebook to make this +# easier to visualize. +class TestVisualizer(PyPlotVisualizer): + # Set limits of view port. + XLIM = (-20., 20.) + YLIM = (-6., 6.) + TICK_DIMS = (0.2, 0.8) + PATCH_WIDTH = 5. + PATCH_HEIGHT = 1. + + def __init__(self, size): + PyPlotVisualizer.__init__(self) + self.DeclareInputPort(PortDataType.kVectorValued, size) + + self.ax.set_xlim(*self.XLIM) + self.ax.set_ylim(*self.YLIM) + self.ax.set_aspect('auto') + + self._make_background() + + self.patch = plt.Rectangle((0.0, 0.0), + self.PATCH_WIDTH, self.PATCH_HEIGHT, + fc='#A31F34', ec='k') + self.patch.set_x(-self.PATCH_WIDTH / 2) # Center at x. + + def _make_background(self): + # X-axis. + plt.plot(self.XLIM, np.zeros_like(self.XLIM), 'k') + # Tick mark centered at the origin. + tick_pos = -0.5 * np.asarray(self.TICK_DIMS) + self.ax.add_patch(plt.Rectangle(tick_pos, *self.TICK_DIMS, fc='k')) + + def draw(self, context): + try: + x = self.EvalVectorInput(context, 0).get_value()[0] + except TypeError: + x = context[0] + self.patch.set_x(x - self.PATCH_WIDTH / 2) + + +class SimpleContinuousTimeSystem(VectorSystem): + def __init__(self): + VectorSystem.__init__(self, + 0, # Zero inputs. + 1) # One output. + self.DeclareContinuousState(1) # One state variable. + + # xdot(t) = -x(t) + x^3(t) + def DoCalcVectorTimeDerivatives(self, context, u, x, xdot): + xdot[:] = -x + x**3 + + # y(t) = x(t) + def DoCalcVectorOutput(self, context, u, x, y): + y[:] = x + + +class TestPyplotVisualizer(unittest.TestCase): + def test_simple_visualizer(self): + builder = DiagramBuilder() + system = builder.AddSystem(SimpleContinuousTimeSystem()) + logger = builder.AddSystem(SignalLogger(1)) + builder.Connect(system.get_output_port(0), logger.get_input_port(0)) + visualizer = builder.AddSystem(TestVisualizer(1)) + builder.Connect(system.get_output_port(0), + visualizer.get_input_port(0)) + diagram = builder.Build() + + context = diagram.CreateDefaultContext() + context.SetContinuousState([0.9]) + + simulator = Simulator(diagram, context) + simulator.AdvanceTo(.1) + + ani = visualizer.animate(logger, repeat=True) + self.assertIsInstance(ani, animation.FuncAnimation) + + def test_trajectory(self): + builder = DiagramBuilder() + visualizer = builder.AddSystem(TestVisualizer(1)) + ppt = PiecewisePolynomial.FirstOrderHold( + [0., 1.], [[2., 3.], [2., 1.]]) + ani = visualizer.animate(ppt) + self.assertIsInstance(ani, animation.FuncAnimation)