{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mathematical Program MultibodyPlant Tutorial\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\n"
   ]
  },
  {
   "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 Meshcat.\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.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 import MathematicalProgram, Solve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant_f = MultibodyPlant(0.0)\n",
    "iiwa_url = (\n",
    "   \"package://drake_models/iiwa_description/sdf/iiwa14_no_collision.sdf\"\n",
    ")\n",
    "(iiwa,) = Parser(plant_f).AddModels(url=iiwa_url)\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 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
}