notebooks/01_building_a_scene.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Drake scene with `airo-drake` 🏗️️\n",
    "\n",
    "Ready to set up your robot scene in Drake? This notebook will show you how `airo-drake` makes the process a breeze! ️🏖️ `airo-drake` is designed to streamline common tasks, but still give you full flexibility to customize your Drake diagrams.\n",
    "\n",
    "For comparison, we've included a more \"native\" way to build a similar Drake scene that directly uses the models from [`airo-models`](https://github.com/airo-ugent/airo-models)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. The `airo-drake` way 🚀\n",
    "\n",
    "Run the cell below and open click on the link to open the visualizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.planning import RobotDiagramBuilder\n",
    "from airo_drake import add_manipulator, add_floor, add_meshcat, finish_build\n",
    "\n",
    "robot_diagram_builder = RobotDiagramBuilder() \n",
    "\n",
    "meshcat = add_meshcat(robot_diagram_builder)\n",
    "arm_index, gripper_index = add_manipulator(robot_diagram_builder, \"ur5e\", \"robotiq_2f_85\")\n",
    "add_floor(robot_diagram_builder)\n",
    "\n",
    "robot_diagram, context = finish_build(robot_diagram_builder, meshcat)\n",
    "del robot_diagram_builder # no longer needed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Understanding Drake's building blocks 🧩\n",
    "\n",
    "Drake doesn't have a single \"Scene\" class. Instead, here are the essential components you'll want to keep around:\n",
    "\n",
    "\n",
    "  - 🤖 `RobotDiagram`: The core of your scene. It houses the Drake systems and is crucial for tasks like collision checking and time parameterization.\n",
    "  - 🐈 `MeshCat`: The browser-based visualizer (e.g., at `localhost:7000`). Our `add_meshcat()` function creates it along with a `MeshCatVisualizer` that syncs the last published context. \n",
    "  - 🗂️ `ModelInstanceIndices`: The model identifiers, allowing you to track and manipulate the models in your scene.\n",
    "\n",
    "### 1.2 Helping you stay organized 🖼️\n",
    "\n",
    "For ease of use, we offer the optional `SingleArmScene` and `DualArmScene` dataclasses to group these components. Feel free to use them. It's not required, all our function operate on the individual objects directly – the choice is yours!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import SingleArmScene\n",
    "\n",
    "scene = SingleArmScene(robot_diagram, arm_index, gripper_index, meshcat)\n",
    "scene"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 💡 **Tip:** Editing Drake diagrams (e.g. adding new obstacles) after finalization is not trivial. It's often easier to create an entirely new scene and treat scenes as disposable 🗑️. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Basic usage 📖\n",
    "\n",
    "More advanced usage and visualizations are shown in the next notebook. Here we'll simply show how you can get some basic information and how you can change the pose of the robot arm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant = scene.robot_diagram.plant()\n",
    "plant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant.num_positions(scene.arm_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant_context = plant.GetMyContextFromRoot(context)\n",
    "plant.GetPositions(plant_context, scene.arm_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joints_new = np.deg2rad([0, -45, -90, -90, 90, 0])\n",
    "\n",
    "plant.SetPositions(plant_context, scene.arm_index, joints_new)\n",
    "scene.robot_diagram.ForcedPublish(context) # updates the meshcat visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant.GetPositions(plant_context, scene.arm_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. The long way ⛰️\n",
    "\n",
    "The code below does almost the same thing as the code above (i.e. without `airo-drake`). \n",
    "\n",
    "It's more hackable, but you lose a lot of convenience and some useful defaults we set with `airo-drake`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import airo_models\n",
    "from pydrake.math import RigidTransform, RollPitchYaw\n",
    "from pydrake.geometry import Meshcat\n",
    "from pydrake.geometry import MeshcatVisualizer\n",
    "\n",
    "robot_diagram_builder = RobotDiagramBuilder()\n",
    "scene_graph = robot_diagram_builder.scene_graph()\n",
    "plant = robot_diagram_builder.plant()\n",
    "builder = robot_diagram_builder.builder()\n",
    "plant = robot_diagram_builder.plant()\n",
    "parser = robot_diagram_builder.parser()\n",
    "parser.SetAutoRenaming(True)\n",
    "\n",
    "meshcat = Meshcat()\n",
    "visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "\n",
    "# Load URDF files\n",
    "ur5e_urdf_path = airo_models.get_urdf_path(\"ur5e\")\n",
    "robotiq_urdf_path = airo_models.get_urdf_path(\"robotiq_2f_85\")\n",
    "\n",
    "floor_thickness = 0.2\n",
    "floor_urdf_path = airo_models.box_urdf_path((3.0, 3.0, floor_thickness), \"floor\")\n",
    "\n",
    "arm_index = parser.AddModels(ur5e_urdf_path)[0]\n",
    "gripper_index = parser.AddModels(robotiq_urdf_path)[0]\n",
    "floor_index = parser.AddModels(floor_urdf_path)[0]\n",
    "\n",
    "# Weld some frames together\n",
    "world_frame = plant.world_frame()\n",
    "floor_frame = plant.GetFrameByName(\"base_link\", floor_index)\n",
    "arm_frame = plant.GetFrameByName(\"base_link\", arm_index)\n",
    "arm_tool_frame = plant.GetFrameByName(\"tool0\", arm_index)\n",
    "gripper_frame = plant.GetFrameByName(\"base_link\", gripper_index)\n",
    "\n",
    "arm_transform = RigidTransform(rpy=RollPitchYaw([0, 0, np.pi]), p=[0, 0, 0])\n",
    "floor_transform = RigidTransform(p=[0, 0, -floor_thickness / 2])\n",
    "\n",
    "X_URTOOL0_ROBOTIQ = RigidTransform(rpy=RollPitchYaw([0, 0, np.pi / 2]), p=[0, 0, 0])\n",
    "\n",
    "plant.WeldFrames(world_frame, arm_frame, arm_transform)\n",
    "plant.WeldFrames(arm_tool_frame, gripper_frame, X_URTOOL0_ROBOTIQ)\n",
    "plant.WeldFrames(world_frame, floor_frame, floor_transform)\n",
    "\n",
    "# Finishing and visualizing\n",
    "diagram = robot_diagram_builder.Build()\n",
    "context = diagram.CreateDefaultContext()\n",
    "plant_context = plant.GetMyContextFromRoot(context)\n",
    "diagram.ForcedPublish(context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cloth-competition",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}