manipulation_station.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eeMrMI0-1Dhu"
      },
      "outputs": [],
      "source": [
        "from IPython.display import display, SVG\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import os\n",
        "import pydot\n",
        "import sys\n",
        "\n",
        "from pydrake.all import (Adder, AddMultibodyPlantSceneGraph, Demultiplexer,\n",
        "                         DiagramBuilder, InverseDynamicsController,\n",
        "                         MakeMultibodyStateToWsgStateSystem,\n",
        "                         MeshcatVisualizerCpp, MultibodyPlant, Parser,\n",
        "                         PassThrough, RigidTransform, RollPitchYaw,\n",
        "                         SchunkWsgPositionController,\n",
        "                         StateInterpolatorWithDiscreteDerivative)\n",
        "from manipulation.meshcat_cpp_utils import StartMeshcat\n",
        "from manipulation.scenarios import AddCameraBox, AddIiwa, AddWsg, AddRgbdSensors\n",
        "from manipulation.utils import FindResource\n",
        "from manipulation import running_as_notebook\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Start the visualizer.\n",
        "meshcat = StartMeshcat()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Set up a basic ManipulationStation diagram\n",
        "\n",
        "Completely in pydrake.  Feel free to modify it as you see fit.  You might also look at the [C++ version](https://github.com/RobotLocomotion/drake/blob/master/examples/manipulation_station/manipulation_station.cc#L193) if you want inspiration for tables / cupboards / bins, etc that you might add.  Here is [a link to the scenarios file](https://github.com/RussTedrake/manipulation/blob/master/manipulation/scenarios.py), in case you need to modify `AddIiwa` and friends."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def MakeManipulationStation(time_step=0.002):\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    # Add (only) the iiwa, WSG, and cameras to the scene.\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(\n",
        "        builder, time_step=time_step)\n",
        "    iiwa = AddIiwa(plant)\n",
        "    wsg = AddWsg(plant, iiwa)\n",
        "    X_WC = RigidTransform(RollPitchYaw(-5 * np.pi / 8, 0, 0), [0.25, -1.5, 1])\n",
        "    AddCameraBox(plant, X_WC, \"camera0\")\n",
        "    plant.Finalize()\n",
        "\n",
        "    num_iiwa_positions = plant.num_positions(iiwa)\n",
        "\n",
        "    # I need a PassThrough system so that I can export the input port.\n",
        "    iiwa_position = builder.AddSystem(PassThrough(num_iiwa_positions))\n",
        "    builder.ExportInput(iiwa_position.get_input_port(), \"iiwa_position\")\n",
        "    builder.ExportOutput(iiwa_position.get_output_port(), \"iiwa_position_command\")\n",
        "\n",
        "    # Export the iiwa \"state\" outputs.\n",
        "    demux = builder.AddSystem(Demultiplexer(\n",
        "        2 * num_iiwa_positions, num_iiwa_positions))\n",
        "    builder.Connect(plant.get_state_output_port(iiwa), demux.get_input_port())\n",
        "    builder.ExportOutput(demux.get_output_port(0), \"iiwa_position_measured\")\n",
        "    builder.ExportOutput(demux.get_output_port(1), \"iiwa_velocity_estimated\")\n",
        "    builder.ExportOutput(plant.get_state_output_port(iiwa), \"iiwa_state_estimated\")\n",
        "\n",
        "    # Make the plant for the iiwa controller to use.\n",
        "    controller_plant = MultibodyPlant(time_step=time_step)\n",
        "    controller_iiwa = AddIiwa(controller_plant)\n",
        "    AddWsg(controller_plant, controller_iiwa, welded=True)\n",
        "    controller_plant.Finalize()\n",
        "\n",
        "    # Add the iiwa controller\n",
        "    iiwa_controller = builder.AddSystem(\n",
        "        InverseDynamicsController(\n",
        "            controller_plant,\n",
        "            kp=[100]*num_iiwa_positions,\n",
        "            ki=[1]*num_iiwa_positions,\n",
        "            kd=[20]*num_iiwa_positions,\n",
        "            has_reference_acceleration=False))\n",
        "    iiwa_controller.set_name(\"iiwa_controller\")\n",
        "    builder.Connect(\n",
        "        plant.get_state_output_port(iiwa), iiwa_controller.get_input_port_estimated_state())\n",
        "\n",
        "    # Add in the feed-forward torque\n",
        "    adder = builder.AddSystem(Adder(2, num_iiwa_positions))\n",
        "    builder.Connect(iiwa_controller.get_output_port_control(),\n",
        "                    adder.get_input_port(0))\n",
        "    # Use a PassThrough to make the port optional (it will provide zero values if not connected).\n",
        "    torque_passthrough = builder.AddSystem(PassThrough([0]*num_iiwa_positions))\n",
        "    builder.Connect(torque_passthrough.get_output_port(),\n",
        "                    adder.get_input_port(1))\n",
        "    builder.ExportInput(torque_passthrough.get_input_port(),\n",
        "                        \"iiwa_feedforward_torque\")\n",
        "    builder.Connect(adder.get_output_port(),\n",
        "                    plant.get_actuation_input_port(iiwa))\n",
        "\n",
        "    # Add discrete derivative to command velocities.\n",
        "    desired_state_from_position = builder.AddSystem(\n",
        "        StateInterpolatorWithDiscreteDerivative(\n",
        "            num_iiwa_positions, time_step, suppress_initial_transient=True))\n",
        "    desired_state_from_position.set_name(\"desired_state_from_position\")\n",
        "    builder.Connect(desired_state_from_position.get_output_port(),\n",
        "                    iiwa_controller.get_input_port_desired_state())\n",
        "    builder.Connect(iiwa_position.get_output_port(),\n",
        "                    desired_state_from_position.get_input_port())\n",
        "\n",
        "    # Export commanded torques.\n",
        "    builder.ExportOutput(adder.get_output_port(), \"iiwa_torque_commanded\")\n",
        "    builder.ExportOutput(adder.get_output_port(), \"iiwa_torque_measured\")\n",
        "\n",
        "    builder.ExportOutput(plant.get_generalized_contact_forces_output_port(iiwa),\n",
        "                         \"iiwa_torque_external\")\n",
        "\n",
        "    # Wsg controller.\n",
        "    wsg_controller = builder.AddSystem(SchunkWsgPositionController())\n",
        "    wsg_controller.set_name(\"wsg_controller\")\n",
        "    builder.Connect(\n",
        "        wsg_controller.get_generalized_force_output_port(),\n",
        "        plant.get_actuation_input_port(wsg))\n",
        "    builder.Connect(plant.get_state_output_port(wsg),\n",
        "                    wsg_controller.get_state_input_port())\n",
        "    builder.ExportInput(wsg_controller.get_desired_position_input_port(),\n",
        "                        \"wsg_position\")\n",
        "    builder.ExportInput(wsg_controller.get_force_limit_input_port(),\n",
        "                        \"wsg_force_limit\")\n",
        "    wsg_mbp_state_to_wsg_state = builder.AddSystem(\n",
        "        MakeMultibodyStateToWsgStateSystem())\n",
        "    builder.Connect(plant.get_state_output_port(wsg),\n",
        "                    wsg_mbp_state_to_wsg_state.get_input_port())\n",
        "    builder.ExportOutput(wsg_mbp_state_to_wsg_state.get_output_port(),\n",
        "                         \"wsg_state_measured\")\n",
        "    builder.ExportOutput(wsg_controller.get_grip_force_output_port(),\n",
        "                         \"wsg_force_measured\")\n",
        "\n",
        "    # Cameras.\n",
        "    AddRgbdSensors(builder, plant, scene_graph)\n",
        "\n",
        "    # Export \"cheat\" ports.\n",
        "    builder.ExportOutput(scene_graph.get_query_output_port(), \"geometry_query\")\n",
        "    builder.ExportOutput(plant.get_contact_results_output_port(),\n",
        "                         \"contact_results\")\n",
        "    builder.ExportOutput(plant.get_state_output_port(),\n",
        "                         \"plant_continuous_state\")\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    return diagram\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def DrawBlockDiagram():\n",
        "    diagram = MakeManipulationStation()\n",
        "    display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString())[0].create_svg()))\n",
        "\n",
        "DrawBlockDiagram()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def TestWithMeshcat():\n",
        "    builder = DiagramBuilder()\n",
        "    station = builder.AddSystem(MakeManipulationStation())\n",
        "\n",
        "    MeshcatVisualizerCpp.AddToBuilder(\n",
        "        builder, station.GetOutputPort(\"geometry_query\"), meshcat)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "    diagram.Publish(context)\n",
        "\n",
        "TestWithMeshcat()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def TestCamera():\n",
        "    station = MakeManipulationStation()\n",
        "    context = station.CreateDefaultContext()\n",
        "    image = station.GetOutputPort(\"camera0_rgb_image\").Eval(\n",
        "        context).data[:, :, :3]\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.imshow(image)\n",
        "\n",
        "TestCamera()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Robotic Manipulation - Force Control.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}