{
  "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Joint-Space Inverse Dynamics\n",
    "\n",
    "This notebook shows how to use the Python bindings of the C++ library TSID to control a manipulator in joint space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "import numpy as np\n",
    "from numpy.linalg import norm\n",
    "import matplotlib.pyplot as plt\n",
    "import plot_utils as plut\n",
    "import time\n",
    "import pinocchio as pin\n",
    "import tsid\n",
    "\n",
    "# import gepetto.corbaserver\n",
    "\n",
    "import talos_arm_conf as conf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a `RobotWrapper` specifying the URDF file describing the robot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "robot = tsid.RobotWrapper(conf.urdf, [conf.path], False)\n",
    "model = robot.model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an `InverseDynamicsFormulationAccForce` object to collect all tasks together and transform them into a QP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "formulation = tsid.InverseDynamicsFormulationAccForce(\"tsid\", robot, False)\n",
    "q0 = conf.q0\n",
    "v0 = np.zeros(robot.nv)\n",
    "formulation.computeProblemData(0.0, q0, v0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a `TaskJointPosture` to control the joint configuration of the robot. Set the proportional and derivative gains $k_p$ and $k_d$. Add the task to the formulation object with a user-specified weight to the priority level 1 (1 corresponds to the cost function, 0 corresponds to the constraints)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "postureTask = tsid.TaskJointPosture(\"task-posture\", robot)\n",
    "postureTask.setKp(conf.kp_posture * np.ones(robot.nv))\n",
    "postureTask.setKd(2.0 * np.sqrt(conf.kp_posture) * np.ones(robot.nv))\n",
    "priorityLevel = 1\n",
    "transitionTime = 0.0\n",
    "formulation.addMotionTask(postureTask, conf.w_posture, priorityLevel, transitionTime)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a constant trajectory in Euclidian space to use as reference for the joint posture task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trajPosture = tsid.TrajectoryEuclidianConstant(\"traj_joint\", q0)\n",
    "postureTask.setReference(trajPosture.computeNext())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a TaskJointBounds to specify the joint velocity limits and add it to the formulation object with priority level 0 (which means it is added as a hard constraint in the QP)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_max = conf.v_max_scaling * model.velocityLimit\n",
    "v_min = -v_max\n",
    "jointBoundsTask = tsid.TaskJointBounds(\"task-joint-bounds\", robot, conf.dt)\n",
    "jointBoundsTask.setVelocityBounds(v_min, v_max)\n",
    "priorityLevel = 0\n",
    "formulation.addMotionTask(\n",
    "    jointBoundsTask, conf.w_joint_bounds, priorityLevel, transitionTime\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a QP solver for solving the TSID QP problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = tsid.SolverHQuadProgFast(\"qp solver\")\n",
    "solver.resize(formulation.nVar, formulation.nEq, formulation.nIn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the viewer if it's not already running."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "robot_display = pin.RobotWrapper.BuildFromURDF(\n",
    "    conf.urdf, [str(Path(conf.path) / \"../..\")]\n",
    ")\n",
    "# Viewer = pin.visualize.GepettoVisualizer\n",
    "Viewer = pin.visualize.MeshcatVisualizer\n",
    "viz = Viewer(\n",
    "    robot_display.model, robot_display.collision_model, robot_display.visual_model\n",
    ")\n",
    "viz.initViewer(loadModel=True)\n",
    "viz.display(q0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hasattr(viz.viewer, \"jupyter_cell\") and viz.viewer.jupyter_cell()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create empty arrays to store the simulation trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = conf.N_SIMULATION\n",
    "tau = np.full((robot.na, N), np.nan)\n",
    "q = np.full((robot.nq, N + 1), np.nan)\n",
    "v = np.full((robot.nv, N + 1), np.nan)\n",
    "dv = np.full((robot.nv, N + 1), np.nan)\n",
    "q_ref = np.full((robot.nq, N), np.nan)\n",
    "v_ref = np.full((robot.nv, N), np.nan)\n",
    "dv_ref = np.full((robot.nv, N), np.nan)\n",
    "dv_des = np.full((robot.nv, N), np.nan)\n",
    "samplePosture = trajPosture.computeNext()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specify amplitude, phase and frequency of the sinusoidal joint trajectory to track."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "amp = np.array([0.2, 0.3, 0.4, 0.0, 0.0, 0.0, 0.0])  # amplitude\n",
    "phi = np.array([0.0, 0.5 * np.pi, 0.0, 0.0, 0.0, 0.0, 0.0])  # phase\n",
    "two_pi_f = (\n",
    "    2 * np.pi * np.array([1.0, 0.5, 0.3, 0.0, 0.0, 0.0, 0.0])\n",
    ")  # frequency (time 2 PI)\n",
    "two_pi_f_amp = np.multiply(two_pi_f, amp)\n",
    "two_pi_f_squared_amp = np.multiply(two_pi_f, two_pi_f_amp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start simulating the system. At each loop:\n",
    "* compute the reference joint sinusoidal trajectory and set it to the posture task\n",
    "* compute the QP problem data using formulation\n",
    "* solve the QP and get the joint accelerations\n",
    "* integrate the accelerations to update the robot state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 0.0\n",
    "dt = conf.dt\n",
    "q[:, 0], v[:, 0] = q0, v0\n",
    "\n",
    "for i in range(N):\n",
    "    time_start = time.time()\n",
    "\n",
    "    # set reference trajectory\n",
    "    q_ref[:, i] = q0 + amp * np.sin(two_pi_f * t + phi)\n",
    "    v_ref[:, i] = two_pi_f_amp * np.cos(two_pi_f * t + phi)\n",
    "    dv_ref[:, i] = -two_pi_f_squared_amp * np.sin(two_pi_f * t + phi)\n",
    "    samplePosture.pos(q_ref[:, i])\n",
    "    samplePosture.vel(v_ref[:, i])\n",
    "    samplePosture.acc(dv_ref[:, i])\n",
    "    postureTask.setReference(samplePosture)\n",
    "\n",
    "    HQPData = formulation.computeProblemData(t, q[:, i], v[:, i])\n",
    "    sol = solver.solve(HQPData)\n",
    "    if sol.status != 0:\n",
    "        print(f\"Time {t:.3f} QP problem could not be solved! Error code:\", sol.status)\n",
    "        break\n",
    "\n",
    "    tau[:, i] = formulation.getActuatorForces(sol)\n",
    "    dv[:, i] = formulation.getAccelerations(sol)\n",
    "    dv_des[:, i] = postureTask.getDesiredAcceleration\n",
    "\n",
    "    if i % conf.PRINT_N == 0:\n",
    "        print(f\"Time {t:.3f}\")\n",
    "        print(\n",
    "            \"\\ttracking err {}: {:.3f}\".format(\n",
    "                postureTask.name.ljust(20, \".\"), norm(postureTask.position_error, 2)\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # numerical integration\n",
    "    v_mean = v[:, i] + 0.5 * dt * dv[:, i]\n",
    "    v[:, i + 1] = v[:, i] + dt * dv[:, i]\n",
    "    q[:, i + 1] = pin.integrate(model, q[:, i], dt * v_mean)\n",
    "    t += conf.dt\n",
    "\n",
    "    if i % conf.DISPLAY_N == 0:\n",
    "        viz.display(q[:, i])\n",
    "\n",
    "    time_spent = time.time() - time_start\n",
    "    if time_spent < conf.dt:\n",
    "        time.sleep(conf.dt - time_spent)\n",
    "\n",
    "# PLOT STUFF\n",
    "time = np.arange(0.0, N * conf.dt, conf.dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(f, ax) = plut.create_empty_figure(robot.nv, 1, figsize=(10, 20))\n",
    "for i in range(robot.nv):\n",
    "    ax[i].plot(time, q[i, :-1], label=f\"q {i}\")\n",
    "    ax[i].plot(time, q_ref[i, :], \"--\", label=f\"q ref {i}\")\n",
    "    ax[i].set_xlabel(\"Time [s]\")\n",
    "    ax[i].set_ylabel(\"q [rad]\")\n",
    "    leg = ax[i].legend()\n",
    "    leg.get_frame().set_alpha(0.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(f, ax) = plut.create_empty_figure(robot.nv, 1, figsize=(10, 20))\n",
    "for i in range(robot.nv):\n",
    "    ax[i].plot(time, v[i, :-1], label=f\"v {i} \")\n",
    "    ax[i].plot(time, v_ref[i, :], \"--\", label=f\"v ref {i}\")\n",
    "    ax[i].plot([time[0], time[-1]], 2 * [v_min[i]], \":\")\n",
    "    ax[i].plot([time[0], time[-1]], 2 * [v_max[i]], \":\")\n",
    "    ax[i].set_xlabel(\"Time [s]\")\n",
    "    ax[i].set_ylabel(\"v [rad/s]\")\n",
    "    leg = ax[i].legend()\n",
    "    leg.get_frame().set_alpha(0.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(f, ax) = plut.create_empty_figure(robot.nv, 1, figsize=(10, 20))\n",
    "for i in range(robot.nv):\n",
    "    ax[i].plot(time, dv[i, :-1], label=\"dv \" + str(i))\n",
    "    ax[i].plot(time, dv_ref[i, :], \"--\", label=f\"dv ref {i}\")\n",
    "    ax[i].plot(time, dv_des[i, :], \":\", label=f\"dv des {i}\")\n",
    "    ax[i].set_xlabel(\"Time [s]\")\n",
    "    ax[i].set_ylabel(\"dv [rad/s^2]\")\n",
    "    leg = ax[i].legend()\n",
    "    leg.get_frame().set_alpha(0.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(f, ax) = plut.create_empty_figure(robot.nv, 1, figsize=(10, 20))\n",
    "for i in range(robot.nv):\n",
    "    ax[i].plot(time, tau[i, :], label=f\"tau {i}\")\n",
    "    ax[i].set_xlabel(\"Time [s]\")\n",
    "    ax[i].set_ylabel(\"tau [Nm]\")\n",
    "    leg = ax[i].legend()\n",
    "    leg.get_frame().set_alpha(0.5)\n",
    "plt.show()"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
