examples/example_robot_systems_design_v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eac7f0b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load modules directly from the folder -> Easy to develop\n",
    "import sys\n",
    "sys.path.append(\"../src/DoD\")\n",
    "from DoD import *\n",
    "from Node import *\n",
    "from pydrake.all import *\n",
    "from IPython.display import display, SVG, clear_output, HTML\n",
    "import time\n",
    "from manipulation.scenarios import AddMultibodyTriad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9b4bed0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot settings\n",
    "import matplotlib.pyplot as plt\n",
    "# set latex like plotting\n",
    "plt.rcParams.update({\n",
    "    \"text.usetex\": True,\n",
    "    \"font.family\": \"serif\",\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "94171f4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Meshcat listening for connections at http://localhost:7000\n"
     ]
    }
   ],
   "source": [
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4358faca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modifying the function type instantiation to A* problem class\n",
    "class stowingProblem:\n",
    "    def __init__(self, pathPoints, collisionsInfo):\n",
    "        self.req = {\n",
    "            'tcyc': 8,\n",
    "            'path': pathPoints,\n",
    "            'N':1000,\n",
    "            \"W\": 1e-3,\n",
    "        }\n",
    "        # Send in both the collision urdf and the spawn pose\n",
    "        self.collision_info = collisionsInfo\n",
    "        self.urdf_path = \"./node_urdfs/\"\n",
    "        self.marker = './urdfs/helper/sphere.urdf'\n",
    "        self.marker_link = 'base_link'\n",
    "        self.gh_weighting = [1e-3,10]\n",
    "        self.robot_base_loc = RigidTransform([0,0,0])\n",
    "        \n",
    "        \n",
    "    # visualise the scene with all the collisions\n",
    "    def visualise_scene(self, meshcat, urdf_path=False, node_instance=False, visReq=False):\n",
    "        builder = DiagramBuilder()\n",
    "        plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=1e-4)\n",
    "        parser = Parser(plant, scene_graph)\n",
    "        parser.package_map().Add(\"mods\", \"../src/DoD/\")\n",
    "        \n",
    "        for collision in self.collision_info:\n",
    "            # collision info has name, urdf_path and its pose\n",
    "            coll = parser.AddModelFromFile(collision[1],collision[0])\n",
    "            plant.WeldFrames(plant.world_frame(), plant.get_body(plant.GetBodyIndices(coll)[0]).body_frame(), collision[2])\n",
    "        \n",
    "        # add visualisation for the points to be traversed\n",
    "        for pose in self.req['path']:\n",
    "            # collision info has name, urdf_path and its pose\n",
    "            sp = parser.AddModelFromFile(self.marker,'sp'+str(np.random.randint(9999)))\n",
    "            plant.WeldFrames(plant.world_frame(), plant.get_body(plant.GetBodyIndices(sp)[0]).body_frame(), pose)\n",
    "            AddMultibodyTriad(plant.GetFrameByName(self.marker_link, sp), scene_graph)          \n",
    "        \n",
    "        if visReq == True:\n",
    "            assert (urdf_path == False) or (node_instance == False), \"Robot information missing\"\n",
    "            model_name = 'visRobot'\n",
    "            if not urdf_path:\n",
    "                # load the robot directly from the instance\n",
    "                temp_node_instance = copy.deepcopy(node_instance.current_parent_node)\n",
    "                temp_node_instance.add_module(node_instance.parent_node.eef1l)\n",
    "                temp_node_instance.add_eef()\n",
    "                model = parser.AddModelsFromString(str(temp_node_instance.robot), 'urdf')     \n",
    "                model = model[0]\n",
    "\n",
    "            else:\n",
    "                model = parser.AddModelFromFile(urdf_path, model_name)           \n",
    "            \n",
    "            plant.WeldFrames(plant.world_frame(), \\\n",
    "                             plant.get_body(plant.GetBodyIndices(model)[0]).body_frame(),self.robot_base_loc)\n",
    "            AddMultibodyTriad(plant.GetFrameByName('eef', model), scene_graph)          \n",
    "            gripper_frame = plant.GetBodyByName(\"eef\").body_frame()\n",
    "        \n",
    "        plant.Finalize()\n",
    "        meshcat.Delete()\n",
    "        meshcat.DeleteAddedControls()\n",
    "\n",
    "        visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "        diagram = builder.Build()\n",
    "        context = diagram.CreateDefaultContext()\n",
    "        plant_context = plant.GetMyMutableContextFromRoot(context)\n",
    "        q0 = plant.GetPositions(plant_context)\n",
    "        \n",
    "        if visReq == True:\n",
    "            for pose in self.req['path']:\n",
    "                ik = InverseKinematics(plant, plant_context)\n",
    "                ik.AddPositionConstraint(\n",
    "                            gripper_frame, [0, 0, 0], plant.world_frame(),\n",
    "                            pose.translation(), pose.translation())\n",
    "#                 ik.AddOrientationConstraint(\n",
    "#                         gripper_frame, RotationMatrix(), plant.world_frame(),\n",
    "#                         pose.rotation(), 0)\n",
    "#                 ik.AddMinimumDistanceConstraint(0.001, 0.1)\n",
    "                prog = ik.get_mutable_prog()\n",
    "                q = ik.q()\n",
    "                prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)\n",
    "                prog.SetInitialGuess(q, q0)\n",
    "                result = Solve(ik.prog())\n",
    "                qr = result.GetSolution(ik.q())\n",
    "                qr = (np.arctan2(np.sin(qr), np.cos(qr)))\n",
    "                plant.SetPositions(plant.GetMyContextFromRoot(context),model,qr)\n",
    "                diagram.ForcedPublish(context)\n",
    "                time.sleep(1)            \n",
    "    \n",
    "        diagram.ForcedPublish(context)\n",
    "        if visReq == True:\n",
    "            return urdf_path\n",
    "        pass\n",
    "\n",
    "    def simulate_robot(self, urdf_path, sim_time=1.3):        \n",
    "        builder = DiagramBuilder()\n",
    "        plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=1e-4)\n",
    "        for collision in self.collision_info:\n",
    "            # collision info has name, urdf_path and its pose\n",
    "            coll = Parser(plant, scene_graph).AddModelFromFile(collision[1],collision[0])\n",
    "            plant.WeldFrames(plant.world_frame(), plant.get_body(plant.GetBodyIndices(coll)[0]).body_frame(), collision[2])\n",
    "        \n",
    "        # adds visualisation for the points to be traversed\n",
    "        for pose in self.req['path']:\n",
    "            # collision info has name, urdf_path and its pose\n",
    "            sp = Parser(plant, scene_graph).AddModelFromFile(self.marker,'sp'+str(np.random.randint(9999)))\n",
    "            plant.WeldFrames(plant.world_frame(), plant.get_body(plant.GetBodyIndices(sp)[0]).body_frame(), pose)\n",
    "\n",
    "#         urdf_path = robot_urdf\n",
    "        model_name = 'visRobot'\n",
    "        model = Parser(plant, scene_graph).AddModelFromFile(urdf_path, model_name)\n",
    "        plant.WeldFrames(plant.world_frame(), plant.get_body(plant.GetBodyIndices(model)[0]).body_frame())\n",
    "#         AddMultibodyTriad(plant.GetFrameByName('eef', model), scene_graph)          \n",
    "        gripper_frame = plant.GetBodyByName(\"eef\").body_frame()\n",
    "        \n",
    "        plant.Finalize()\n",
    "        meshcat.Delete()\n",
    "        meshcat.DeleteAddedControls()\n",
    "\n",
    "        visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "        na = plant.num_actuators()\n",
    "        # Initialise controller params, for now limited to Kp and Kd search\n",
    "        Kp = np.full(na, 50)\n",
    "        Ki = np.full(na, 0)\n",
    "        Kd = np.full(na, 14)\n",
    "        iiwa_controller = builder.AddSystem(InverseDynamicsController(plant, Kp, Ki, Kd, False))\n",
    "        iiwa_controller.set_name(\"collision_controller\");\n",
    "        # Complete connections\n",
    "        builder.Connect(plant.get_state_output_port(model),\n",
    "                    iiwa_controller.get_input_port_estimated_state())\n",
    "        builder.Connect(iiwa_controller.get_output_port_control(), \n",
    "                        plant.get_actuation_input_port())\n",
    "        # Connecting a data logger\n",
    "        # Commanded torque\n",
    "        logger1 = LogVectorOutput(iiwa_controller.get_output_port_control(), builder)\n",
    "        # Joint state\n",
    "        logger2 = LogVectorOutput(plant.get_state_output_port(model), builder)\n",
    "\n",
    "        \n",
    "        diagram = builder.Build()\n",
    "        context = diagram.CreateDefaultContext()\n",
    "        plant_context = plant.GetMyMutableContextFromRoot(context)\n",
    "        q0 = plant.GetPositions(plant_context)\n",
    "\n",
    "        res = True\n",
    "        move_poses = []\n",
    "        for pose in self.req['path']:\n",
    "            ik = InverseKinematics(plant, plant_context)\n",
    "            ik.AddPositionConstraint(\n",
    "                        gripper_frame, [0, 0, 0], plant.world_frame(),\n",
    "                        pose.translation(), pose.translation())\n",
    "#             ik.AddMinimumDistanceConstraint(0.001, 0.1)\n",
    "            prog = ik.get_mutable_prog()\n",
    "            q = ik.q()\n",
    "            prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)\n",
    "            prog.SetInitialGuess(q, q0)\n",
    "            result = Solve(ik.prog())\n",
    "            res = result and res\n",
    "            qr = result.GetSolution(ik.q())\n",
    "            qr = (np.arctan2(np.sin(qr), np.cos(qr)))\n",
    "            move_poses.append(qr)\n",
    "\n",
    "        print(res)\n",
    "    \n",
    "        xd = np.hstack((move_poses[1], 0*move_poses[1]))\n",
    "        plant.SetPositions(plant_context, move_poses[0])\n",
    "        iiwa_controller.GetInputPort('desired_state').FixValue(\n",
    "        iiwa_controller.GetMyMutableContextFromRoot(context), xd);\n",
    "        simulator = Simulator(diagram, context)\n",
    "        simulator.AdvanceTo(sim_time);\n",
    "        simulator.set_target_realtime_rate(0.05)\n",
    "    \n",
    "        log1 = logger1.FindLog(context)\n",
    "        log2 = logger2.FindLog(context)\n",
    "\n",
    "        return log2.data()\n",
    "\n",
    "        # Routine to compute the heuristic cost\n",
    "    def compute_heuristic_cost(self, node_instance=None, node_path=None, model_name=None):\n",
    "        builder = DiagramBuilder()\n",
    "        plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=1e-4)\n",
    "        parser = Parser(plant, scene_graph)\n",
    "        # retain load by path for troubleshooting\n",
    "        if node_path != None:\n",
    "            assert(model_name!=None), \"Please enter a model_name\"\n",
    "            model = parser.AddModelFromFile(node_path, model_name)\n",
    "        # in case the robot is loaded by direct node instance\n",
    "        if node_instance !=None:\n",
    "            parser.package_map().Add(\"mods\", \"../src/DoD/\")\n",
    "            # the end effector is added by default\n",
    "            temp_instance = copy.deepcopy(node_instance)\n",
    "            temp_instance.add_module(temp_instance.eef1l)\n",
    "            temp_instance.add_eef()\n",
    "            # now load the temp_insatnce           \n",
    "            model = parser.AddModelsFromString(str(temp_instance.robot), 'urdf')     \n",
    "            # select one model from the list of models loaded\n",
    "            model = model[0]\n",
    "        \n",
    "        plant.WeldFrames(plant.world_frame(), \\\n",
    "                         plant.get_body(plant.GetBodyIndices(model)[0]).body_frame(), \\\n",
    "                         self.robot_base_loc)\n",
    "        \n",
    "        for collision in self.collision_info:\n",
    "            # collision info has name, urdf_path and its pose\n",
    "            coll = Parser(plant, scene_graph).AddModelFromFile(collision[1],collision[0])\n",
    "            plant.WeldFrames(plant.world_frame(), plant.get_body(plant.GetBodyIndices(coll)[0]).body_frame(), collision[2])\n",
    "                        \n",
    "        plant.Finalize()\n",
    "        # Select the last frame as the eef_frame\n",
    "        gripper_frame = plant.GetBodyByName(\"eef\").body_frame()\n",
    "        diagram = builder.Build()\n",
    "        context = diagram.CreateDefaultContext()\n",
    "        plant_context = plant.GetMyMutableContextFromRoot(context)\n",
    "        q0 = plant.GetPositions(plant_context)\n",
    "\n",
    "        total_result = True\n",
    "        total_cost = 0\n",
    "        # Check design        \n",
    "        for pose in self.req['path']:\n",
    "            ik = InverseKinematics(plant, plant_context)\n",
    "            ik.AddPositionConstraint(\n",
    "                        gripper_frame, [0, 0, 0], plant.world_frame(),\n",
    "                        pose.translation(), pose.translation())\n",
    "#             ik.AddOrientationConstraint(\n",
    "#                         gripper_frame, RotationMatrix(), plant.world_frame(),\n",
    "#                         pose.rotation(), 0)\n",
    "#             ik.AddMinimumDistanceConstraint(0.001, 0.1)\n",
    "            prog = ik.get_mutable_prog()\n",
    "            q = ik.q()\n",
    "            prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)\n",
    "            prog.SetInitialGuess(q, q0)\n",
    "            result = Solve(ik.prog())\n",
    "            qr = result.GetSolution(ik.q())\n",
    "            qr = (np.arctan2(np.sin(qr), np.cos(qr)))\n",
    "            # Visualise the best solution\n",
    "            plant.SetPositions(plant.GetMyContextFromRoot(context),model,qr)\n",
    "            diagram.ForcedPublish(context)\n",
    "            total_result = total_result and result.is_success()\n",
    "            eef_pos_realised = plant.GetBodyByName('eef').EvalPoseInWorld(plant_context).translation()\n",
    "            total_cost += np.linalg.norm(eef_pos_realised-pose.translation())\n",
    "    \n",
    "        return total_cost, total_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "482a5cb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Provide example points to reach\n",
    "pathPoints = [\n",
    "    RigidTransform(np.array([0.15, 0.1, 0.45])),\n",
    "    RigidTransform(np.array([-0.05, 0.35, 0.3]))\n",
    "]\n",
    "\n",
    "# Provide example collisions\n",
    "# However, passing collisions is expensive\n",
    "# TODO: Primitive collision geometries needs to be implemented for the modules\n",
    "collisionInfo = [\n",
    "            ['ball1', './urdfs/collisionScene/ball1.urdf', RigidTransform(np.array([0,0,0.55]))],\n",
    "            ]\n",
    "\n",
    "# instantiate problem class \n",
    "sp = stowingProblem(pathPoints, [])\n",
    "\n",
    "# visualise initial scene\n",
    "sp.visualise_scene(meshcat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3ec4108c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# intiate helper classes\n",
    "roboDoD = robotDoD('example-problem')\n",
    "roboDoD.add_root()\n",
    "roboDoD.add_module(roboDoD.base360)\n",
    "# define modules to be used\n",
    "roboDoD.declared_modules = [              \n",
    "                            roboDoD.j90,\n",
    "                            roboDoD.c90,\n",
    "                            roboDoD.l50,\n",
    "                            roboDoD.l100,\n",
    "                            roboDoD.l150, \n",
    "                           ]\n",
    "node = robotNode(roboDoD)\n",
    "prob = sp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e1ba158a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "def search_design_space(roboDoD, node, prob, meshcat):\n",
    "    end_game = False\n",
    "    end_flag = False\n",
    "    iters = 0\n",
    "    search_cost_dict = {}\n",
    "    while not end_game:\n",
    "        node.expand_node()\n",
    "        for i in range(len(node.child_nodes)):\n",
    "            cost, end_flag = node.get_node_cost(prob.compute_heuristic_cost, \\\n",
    "                                                node.child_nodes[i], \\\n",
    "                                                prob.gh_weighting)\n",
    "            if end_flag:\n",
    "                node.current_parent_node = node.child_nodes[i]\n",
    "                urdf_path = prob.visualise_scene(meshcat, node_instance=node, visReq=True)\n",
    "                print(\"Search completed with: \", cost,\" \", urdf_path)\n",
    "                end_game = True\n",
    "                break\n",
    "            else:\n",
    "                search_cost_dict[cost] = node.child_nodes[i]\n",
    "\n",
    "        if not end_game:\n",
    "            search_cost_dict_array = np.array(list(search_cost_dict.keys()))\n",
    "            least_cost_node_val = np.min(search_cost_dict_array)\n",
    "            best_child = search_cost_dict[least_cost_node_val]\n",
    "            node.current_parent_node = best_child\n",
    "            search_cost_dict.pop(least_cost_node_val)\n",
    "\n",
    "            if np.mod(iters, 10)==0:\n",
    "                urdf_path = prob.visualise_scene(meshcat, node_instance=node, visReq=True)\n",
    "        # display the iterations\n",
    "        clear_output(wait=True)\n",
    "        display(iters)\n",
    "        if iters>500:\n",
    "            end_game=True\n",
    "            print('Max iterations reached...')\n",
    "        iters+=1\n",
    "        print(end_game)\n",
    "\n",
    "# search through the design space\n",
    "search_design_space(roboDoD, node, prob, meshcat)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}