book/robot/rubiks_cube.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/robot.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    ContactVisualizer,\n",
    "    DiagramBuilder,\n",
    "    ExternallyAppliedSpatialForce,\n",
    "    LeafSystem,\n",
    "    List,\n",
    "    MeshcatVisualizer,\n",
    "    ModelVisualizer,\n",
    "    Parser,\n",
    "    Simulator,\n",
    "    SpatialForce,\n",
    "    StartMeshcat,\n",
    "    Value,\n",
    ")\n",
    "\n",
    "from manipulation import ConfigureParser, FindResource, running_as_notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wJWL-ldv5REK"
   },
   "source": [
    "# Rubik's Cube (2x2)\n",
    "\n",
    "TODO(russt): Use the quaternion ball joint when it's available:\n",
    "https://github.com/RobotLocomotion/drake/issues/12404\n",
    "(rotating around y can hit a singularity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ILYLouFTjv6e"
   },
   "outputs": [],
   "source": [
    "cube_file = \"package://manipulation/rubiks_cube_2_by_2.sdf\"\n",
    "\n",
    "meshcat.Delete()\n",
    "\n",
    "visualizer = ModelVisualizer(meshcat=meshcat)\n",
    "ConfigureParser(visualizer.parser())\n",
    "visualizer.parser().AddModelsFromUrl(cube_file)\n",
    "\n",
    "visualizer.Run(loop_once=not running_as_notebook)\n",
    "\n",
    "meshcat.DeleteAddedControls()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5SjOClhTltPk"
   },
   "outputs": [],
   "source": [
    "rotate_about = \"x\"\n",
    "# rotate_about = 'y'\n",
    "# rotate_about = 'z'\n",
    "\n",
    "meshcat.Delete()\n",
    "builder = DiagramBuilder()\n",
    "\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, 0.001)\n",
    "parser = Parser(plant)\n",
    "ConfigureParser(parser)\n",
    "parser.AddModelsFromUrl(cube_file)\n",
    "\n",
    "plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"center\"))\n",
    "plant.Finalize()\n",
    "\n",
    "visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "ContactVisualizer.AddToBuilder(builder, plant, meshcat)\n",
    "# logger = LogVectorOutput(plant.get_state_output_port(), builder)\n",
    "\n",
    "\n",
    "class CubePusher(LeafSystem):\n",
    "    def __init__(self):\n",
    "        LeafSystem.__init__(self)\n",
    "        forces_cls = Value[List[ExternallyAppliedSpatialForce]]\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"applied_force\", lambda: forces_cls(), self.CalcOutput\n",
    "        )\n",
    "\n",
    "    def CalcOutput(self, context, output):\n",
    "        forces = []\n",
    "        for x, y, z in product([0, 1], repeat=3):\n",
    "            force = ExternallyAppliedSpatialForce()\n",
    "            force.body_index = plant.GetBodyByName(f\"box_{x}_{y}_{z}\").index()\n",
    "            # shift from [0, 1] to [-1, 1]\n",
    "            x = 2 * x - 1\n",
    "            y = 2 * y - 1\n",
    "            z = 2 * z - 1\n",
    "            force.p_BoBq_B = -0.0125 * np.array([x, y, z])  # world 0, 0, 0\n",
    "            if rotate_about == \"x\":\n",
    "                force.F_Bq_W = SpatialForce(\n",
    "                    tau=-0.2 * np.array([1 if x < 0 else -1, 0, 0]),\n",
    "                    f=[0, 0, 0],\n",
    "                )\n",
    "            elif rotate_about == \"y\":\n",
    "                force.F_Bq_W = SpatialForce(\n",
    "                    tau=0.2 * np.array([0, 1 if y > 0 else -1, 0]), f=[0, 0, 0]\n",
    "                )\n",
    "            else:\n",
    "                force.F_Bq_W = SpatialForce(\n",
    "                    tau=0.2 * np.array([0, 0, 1 if z > 0 else -1]), f=[0, 0, 0]\n",
    "                )\n",
    "        forces.append(force)\n",
    "        output.set_value(forces)\n",
    "\n",
    "\n",
    "pusher = builder.AddSystem(CubePusher())\n",
    "builder.Connect(pusher.get_output_port(), plant.get_applied_spatial_force_input_port())\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "simulator = Simulator(diagram)\n",
    "context = simulator.get_context()\n",
    "\n",
    "if running_as_notebook:\n",
    "    simulator.set_target_realtime_rate(1.0)\n",
    "    simulator.AdvanceTo(10)\n",
    "else:\n",
    "    simulator.AdvanceTo(0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rubik's Cube (3x3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_3_by_3():\n",
    "    filename = FindResource(\"models/rubiks_cube.sdf\")\n",
    "    box_size = 0.025\n",
    "    box_mass = 0.03\n",
    "    inertia = box_mass * (box_size**2) / 6\n",
    "    inertia = inertia * 100  # scale inertia to help numerics\n",
    "\n",
    "    with open(filename, \"w\") as f:\n",
    "        f.write(\n",
    "            \"\"\"\n",
    "<?xml version=\"1.0\"?>\n",
    "<!-- Autogenerated by rubiks_cube.ipynb -->\n",
    "<sdf version=\"1.7\">\n",
    "  <model name=\"rubiks_cube\">\n",
    "    <link name=\"center\"/>\n",
    "\"\"\"\n",
    "        )\n",
    "\n",
    "        # TODO(russt): I could reduce the dofs for the center cubes, since\n",
    "        # box_m1_0_0 and box_1_0_0, must have their relative poses\n",
    "        # constrained up to a single-axis rotation, etc.\n",
    "\n",
    "        # TODO(russt): Currently I'm adding some faces on the inside of the\n",
    "        # internal blocks (e.g. the x face when x==0). They are inelegant and\n",
    "        # mildly inefficiently, but won't be visible.\n",
    "        for x, y, z in product([-1, 0, 1], repeat=3):\n",
    "            if x == 0 and y == 0 and z == 0:\n",
    "                # No cube in the very center\n",
    "                continue\n",
    "            sx = x if x >= 0 else \"m1\"\n",
    "            sy = y if y >= 0 else \"m1\"\n",
    "            sz = z if z >= 0 else \"m1\"\n",
    "            suffix = f\"_{sx}_{sy}_{sz}\"\n",
    "            x_color = [0, 0, 1] if x < 0 else [0, 1, 0]\n",
    "            y_color = [1, 0.5, 0] if y < 0 else [1, 0, 0]\n",
    "            z_color = [1, 1, 0] if z < 0 else [1, 1, 1]\n",
    "            f.write(\n",
    "                f\"\"\"\n",
    "    <link name=\"box{suffix}\">\n",
    "      <pose>{box_size*x} {box_size*y} {box_size*z} 0 0 0</pose>\n",
    "      <inertial>\n",
    "        <mass>0.03</mass>\n",
    "        <inertia>\n",
    "          <ixx>0.0003125</ixx>\n",
    "          <ixy>0.0</ixy>\n",
    "          <ixz>0.0</ixz>\n",
    "          <iyy>0.0003125</iyy>\n",
    "          <iyz>0.0</iyz>\n",
    "          <izz>0.0003125</izz>\n",
    "        </inertia>\n",
    "      </inertial>\n",
    "      <collision name=\"collision\">\n",
    "        <geometry>\n",
    "          <box>\n",
    "            <size>{box_size} {box_size} {box_size}</size>\n",
    "          </box>\n",
    "        </geometry>\n",
    "      </collision>\n",
    "      <visual name=\"black\">\n",
    "        <geometry>\n",
    "          <box>\n",
    "            <size>{box_size} {box_size} {box_size}</size>\n",
    "          </box>\n",
    "        </geometry>\n",
    "        <material>\n",
    "          <diffuse>0 0 0 1</diffuse>\n",
    "        </material>\n",
    "      </visual>\n",
    "      <visual name=\"x\">\n",
    "        <pose>{x*box_size/2} 0 0 0 0 0</pose>\n",
    "        <geometry>\n",
    "          <box>\n",
    "            <size>0.0002 {box_size*4/5} {box_size*4/5}</size>\n",
    "          </box>\n",
    "        </geometry>\n",
    "        <material>\n",
    "          <diffuse>{x_color[0]} {x_color[1]} {x_color[2]} 1</diffuse>\n",
    "        </material>\n",
    "      </visual>\n",
    "      <visual name=\"y\">\n",
    "        <pose>0 {y*box_size/2} 0 0 0 0</pose>\n",
    "        <geometry>\n",
    "          <box>\n",
    "            <size>{box_size*4/5} 0.0002 {box_size*4/5}</size>\n",
    "          </box>\n",
    "        </geometry>\n",
    "        <material>\n",
    "          <diffuse>{y_color[0]} {y_color[1]} {y_color[2]} 1</diffuse>\n",
    "        </material>\n",
    "      </visual>\n",
    "      <visual name=\"z\">\n",
    "        <pose>0 0 {z*box_size/2} 0 0 0</pose>\n",
    "        <geometry>\n",
    "          <box>\n",
    "            <size>{box_size*4/5} {box_size*4/5} 0.0002</size>\n",
    "          </box>\n",
    "        </geometry>\n",
    "        <material>\n",
    "          <diffuse>{z_color[0]} {z_color[1]} {z_color[2]} 1</diffuse>\n",
    "        </material>\n",
    "      </visual>\n",
    "    </link>      \n",
    "    <joint name=\"ball{suffix}\" type=\"ball\">\n",
    "        <pose>{-x*box_size} {-y*box_size} {-z*box_size} 0 0 0</pose> <!-- in child frame -->\n",
    "        <parent>center</parent>\n",
    "        <child>box{suffix}</child>\n",
    "        <axis>\n",
    "            <dynamics><damping>0.1</damping></dynamics>\n",
    "            <limit><effort>0</effort></limit>\n",
    "        </axis>\n",
    "    </joint>\n",
    "\"\"\"\n",
    "            )\n",
    "\n",
    "        f.write(\n",
    "            \"\"\"\n",
    "  </model>\n",
    "</sdf>\n",
    "\"\"\"\n",
    "        )\n",
    "\n",
    "\n",
    "if running_as_notebook:\n",
    "    generate_3_by_3()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cube_file = \"package://manipulation/rubiks_cube.sdf\"\n",
    "\n",
    "meshcat.Delete()\n",
    "\n",
    "visualizer = ModelVisualizer(meshcat=meshcat)\n",
    "ConfigureParser(visualizer.parser())\n",
    "visualizer.parser().AddModelsFromUrl(cube_file)\n",
    "\n",
    "visualizer.Run(loop_once=not running_as_notebook)\n",
    "\n",
    "meshcat.DeleteAddedControls()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Let's get you a robot.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "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.11.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}