{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mathematical Program MultibodyPlant Tutorial\n",
    "\n",
    "For instructions on how to run these tutorial notebooks, please see\n",
    "the [README](https://github.com/RobotLocomotion/drake/blob/master/tutorials/README.md)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shows examples of:\n",
    "* Creating a `MultibodyPlant` containing an IIWA arm\n",
    "* Solve a simple inverse kinematics problem by writing a custom evaluator\n",
    "for `MathematicalProgram` that can handle both `float` and `AutoDiffXd`\n",
    "inputs\n",
    "* Using the custom evaluator in a constraint\n",
    "* Using the custom evaluator in a cost.\n",
    "\n",
    "***To be added***:\n",
    "* Using `pydrake.multibody.inverse_kinematics`.\n",
    "* Visualizing with Drake Visualizer.\n",
    "\n",
    "### Important Note\n",
    "\n",
    "Please review the\n",
    "[API for `pydrake.multibody.inverse_kinematics`](\n",
    "https://drake.mit.edu/pydrake/pydrake.multibody.inverse_kinematics.html)\n",
    "before you delve too far into writing custom evaluators for use with\n",
    "`MultibodyPlant`. You may find the functionality you want there."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inverse Kinematics Problem\n",
    "\n",
    "In this tutorial, we will be solving a simple inverse kinematics problem to\n",
    "put Link 7's origin at a given distance from a target position. We will use\n",
    "`MathematicalProgram` to solve this problem in two different ways: first\n",
    "using the evaluator as a constraint (with a minimum and maximum distance),\n",
    "and second using the evaluator as a cost (to get as close as possible).\n",
    "\n",
    "For more information about `MathematicalProgram`, please see the\n",
    "[`MathematicalProgram` Tutorial](./mathematical_program.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, we will import the necessary modules and load a `MultibodyPlant`\n",
    "containing an IIWA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from pydrake.common import FindResourceOrThrow\n",
    "from pydrake.math import RigidTransform\n",
    "from pydrake.multibody.parsing import Parser\n",
    "from pydrake.systems.analysis import Simulator\n",
    "from pydrake.all import MultibodyPlant\n",
    "\n",
    "from pydrake.solvers.mathematicalprogram import MathematicalProgram, Solve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant_f = MultibodyPlant(0.0)\n",
    "iiwa_file = FindResourceOrThrow(\n",
    "   \"drake/manipulation/models/iiwa_description/sdf/\"\n",
    "   \"iiwa14_no_collision.sdf\")\n",
    "iiwa = Parser(plant_f).AddModelFromFile(iiwa_file)\n",
    "\n",
    "# Define some short aliases for frames.\n",
    "W = plant_f.world_frame()\n",
    "L0 = plant_f.GetFrameByName(\"iiwa_link_0\", iiwa)\n",
    "L7 = plant_f.GetFrameByName(\"iiwa_link_7\", iiwa)\n",
    "\n",
    "plant_f.WeldFrames(W, L0)\n",
    "plant_f.Finalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing our Custom Evaluator\n",
    "\n",
    "Our evaluator is implemented using the custom evaluator\n",
    "`link_7_distance_to_target`, since the its functionality is not already\n",
    "handled by existing classes in the `inverse_kinematics` submodule.\n",
    "\n",
    "Note that in order to write a custom evaluator in Python, we must explicitly\n",
    "check for `float` and `AutoDiffXd` inputs, as you will see in the implementation\n",
    "of `link_7_distance_to_target`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Allocate float context to be used by evaluators.\n",
    "context_f = plant_f.CreateDefaultContext()\n",
    "# Create AutoDiffXd plant and corresponding context.\n",
    "plant_ad = plant_f.ToAutoDiffXd()\n",
    "context_ad = plant_ad.CreateDefaultContext()\n",
    "\n",
    "def resolve_frame(plant, F):\n",
    "    \"\"\"Gets a frame from a plant whose scalar type may be different.\"\"\"\n",
    "    return plant.GetFrameByName(F.name(), F.model_instance())\n",
    "\n",
    "# Define target position.\n",
    "p_WT = [0.1, 0.1, 0.6]\n",
    "\n",
    "def link_7_distance_to_target(q):\n",
    "    \"\"\"Evaluates squared distance between L7 origin and target T.\"\"\"\n",
    "    # Choose plant and context based on dtype.\n",
    "    if q.dtype == float:\n",
    "        plant = plant_f\n",
    "        context = context_f\n",
    "    else:\n",
    "        # Assume AutoDiff.\n",
    "        plant = plant_ad\n",
    "        context = context_ad\n",
    "    # Do forward kinematics.\n",
    "    plant.SetPositions(context, iiwa, q)\n",
    "    X_WL7 = plant.CalcRelativeTransform(\n",
    "        context, resolve_frame(plant, W), resolve_frame(plant, L7))\n",
    "    p_TL7 = X_WL7.translation() - p_WT\n",
    "    return p_TL7.dot(p_TL7)\n",
    "\n",
    "# WARNING: If you return a scalar for a constraint, or a vector for\n",
    "# a cost, you may get the following cryptic error:\n",
    "# \"Unable to cast Python instance to C++ type\"\n",
    "link_7_distance_to_target_vector = lambda q: [link_7_distance_to_target(q)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Formulating the Optimization Problems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formluation 1: Using the Custom Evaluator in a Constraint\n",
    "\n",
    "We will formulate and solve the problem with a basic cost and our custom\n",
    "evaluator in a constraint.\n",
    "\n",
    "Note that we use the vectorized version of the evaluator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prog = MathematicalProgram()\n",
    "\n",
    "q = prog.NewContinuousVariables(plant_f.num_positions())\n",
    "# Define nominal configuration.\n",
    "q0 = np.zeros(plant_f.num_positions())\n",
    "\n",
    "# Add basic cost. (This will be parsed into a QuadraticCost.)\n",
    "prog.AddCost((q - q0).dot(q - q0))\n",
    "\n",
    "# Add constraint based on custom evaluator.\n",
    "prog.AddConstraint(\n",
    "    link_7_distance_to_target_vector,\n",
    "    lb=[0.1], ub=[0.2], vars=q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = Solve(prog, initial_guess=q0)\n",
    "\n",
    "print(f\"Success? {result.is_success()}\")\n",
    "print(result.get_solution_result())\n",
    "q_sol = result.GetSolution(q)\n",
    "print(q_sol)\n",
    "\n",
    "print(f\"Initial distance: {link_7_distance_to_target(q0):.3f}\")\n",
    "print(f\"Solution distance: {link_7_distance_to_target(q_sol):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formulation 2: Using Custom Evaluator in a Cost\n",
    "\n",
    "We will formulate and solve the problem, but this time we will use our custom\n",
    "evaluator in a cost.\n",
    "\n",
    "Note that we use the scalar version of the evaluator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prog = MathematicalProgram()\n",
    "\n",
    "q = prog.NewContinuousVariables(plant_f.num_positions())\n",
    "# Define nominal configuration.\n",
    "q0 = np.zeros(plant_f.num_positions())\n",
    "\n",
    "# Add custom cost.\n",
    "prog.AddCost(link_7_distance_to_target, vars=q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = Solve(prog, initial_guess=q0)\n",
    "\n",
    "print(f\"Success? {result.is_success()}\")\n",
    "print(result.get_solution_result())\n",
    "q_sol = result.GetSolution(q)\n",
    "print(q_sol)\n",
    "\n",
    "print(f\"Initial distance: {link_7_distance_to_target(q0):.3f}\")\n",
    "print(f\"Solution distance: {link_7_distance_to_target(q_sol):.3f}\")"
   ]
  }
 ],
 "metadata": {
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}