3_optimal_control_robot_manipulator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to `robotoc` \n",
    "# 3: Optimal Control of a Robot Manipulator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The goal of this chapter\n",
    "Here, we learn how to construct and solve an optimal control problem (OCP) for a robot manipulator KUKA iiwa14."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we construct `robotoc.Robot` of a robot manipulator KUKA iiwa14."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import robotoc\n",
    "model_info = robotoc.RobotModelInfo()\n",
    "model_info.urdf_path = \"urdf/iiwa_description/urdf/iiwa14.urdf\"\n",
    "iiwa14 = robotoc.Robot(model_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we change the joint torque and velocity limits of the robot manipulator.  \n",
    "Like the following, you can change such limits depending on your goal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "iiwa14.set_joint_effort_limit(np.full(iiwa14.dimu(), 50))\n",
    "iiwa14.set_joint_velocity_limit(np.full(iiwa14.dimv(), 0.5*np.pi))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can check that these limits are set in the robot model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(iiwa14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cost Function\n",
    "Next, we create the cost function of the optimal control problem (OCP).   \n",
    "In `robotoc`, we create the cost function `robotoc.CostFunction` (`std::shared_ptr<robotoc::CostFunction>` in C++) by combining various **cost components**.  \n",
    "That is, `robotoc.CostFunction` (`std::shared_ptr<robotoc::CostFunction>` in C++) is a collection of the **cost components** .   \n",
    "Each **cost component** inherits `robotoc.CostFunctionComponentBase` (`robotoc::CostFunctionComponentBase` in C++).    \n",
    "By doing so, we can easily combine many kinds of **cost components**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basic and useful **cost components** are already implemented in robotoc.  \n",
    "In this example, we utilize the **cost component on the configuration (i.e., cost on the joint-space)** as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_cost = robotoc.ConfigurationSpaceCost(iiwa14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set the reference values and weight parameters of the configuration space cost compoonent.  \n",
    "Naming conventions are as follows:\n",
    "- `q_ref`: reference configuration (joint positions)\n",
    "- `q_weight`: weight on the configuration \n",
    "- `v_weight`: weight on the generalized velocity (joint velocities)\n",
    "- `a_weight`: weight on the generalized acceleration (joint accelerations)\n",
    "- `u_weight`: weight on the joint torques\n",
    "- `q_weight_terminal`, `v_weight_terminal`: weights on the configuration and generalized velocity at the terminal time of the horizon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_ref = np.array([0, 0.5*np.pi, 0, 0.5*np.pi, 0, 0.5*np.pi, 0]) \n",
    "config_cost.set_q_ref(q_ref)\n",
    "config_cost.set_q_weight(np.full(iiwa14.dimv(), 10))\n",
    "config_cost.set_v_weight(np.full(iiwa14.dimv(), 0.01))\n",
    "config_cost.set_a_weight(np.full(iiwa14.dimv(), 0.01))\n",
    "config_cost.set_q_weight_terminal(np.full(iiwa14.dimv(), 10))\n",
    "config_cost.set_v_weight_terminal(np.full(iiwa14.dimv(), 0.01))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then make the **cost function** `robotoc.CostFunction` and add the **configuration cost component** to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost = robotoc.CostFunction()\n",
    "cost.add(\"config_cost\", config_cost)\n",
    "print(cost)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constraints \n",
    "We can construct constraints in the same way as the cost function.  \n",
    "In `robotoc`, we create the constraints `robotoc.Constraints` (`std::shared_ptr<robotoc::Constraints>` in C++) by combining various **constraint components**.  \n",
    "That is, the `robotoc.Constraints` (`std::shared_ptr<robotoc::Constraints>` in C++) is a collection of the **constraint components** .  \n",
    "Each **constraints component** inherits `robotoc.ConstraintComponentBase` (`robotoc::ConstraintComponentBase` in C++).   \n",
    "By doing so, we can easily combine many kinds of **constraint components**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basic and useful **constraint components** are already implemented in robotoc.  \n",
    "In this example, we utilize **constraint components on the joint position, velocity, and torque limits** as   \n",
    "(limits set by `iiwa14.set_joint_effort_limit()` `iiwa14.set_joint_velocity_limit()` are used in these constraint components)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_position_lower  = robotoc.JointPositionLowerLimit(iiwa14)\n",
    "joint_position_upper  = robotoc.JointPositionUpperLimit(iiwa14)\n",
    "joint_velocity_lower  = robotoc.JointVelocityLowerLimit(iiwa14)\n",
    "joint_velocity_upper  = robotoc.JointVelocityUpperLimit(iiwa14)\n",
    "joint_torques_lower   = robotoc.JointTorquesLowerLimit(iiwa14)\n",
    "joint_torques_upper   = robotoc.JointTorquesUpperLimit(iiwa14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then make the **constraints** `robotoc.Constraints` and add these **constraint components**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "constraints = robotoc.Constraints()\n",
    "constraints.add(\"joint_position_lower\", joint_position_lower)\n",
    "constraints.add(\"joint_position_upper\", joint_position_upper)\n",
    "constraints.add(\"joint_velocity_lower\", joint_velocity_lower)\n",
    "constraints.add(\"joint_velocity_upper\", joint_velocity_upper)\n",
    "constraints.add(\"joint_torques_lower\", joint_torques_lower)\n",
    "constraints.add(\"joint_torques_upper\", joint_torques_upper)\n",
    "print(constraints)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimal Control Problem (OCP) and Solver\n",
    "\n",
    "Now we can construct the optimal control problem (OCP) and solvers.  \n",
    "First, set the length of the horizon `T` and the number of discretization grids `N` and construct the optimal control problem (OCP) `robotoc.OCP` (`robotoc::OCP` in C++) as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 3.0\n",
    "N = 100\n",
    "ocp = robotoc.OCP(robot=iiwa14, cost=cost, constraints=constraints, T=T, N=N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this OCP, first we try `robotoc.UnconstrOCPSolver` (`robotoc::UnconstrOCPSolver` in C++).  \n",
    "This solver has the following characteristics: \n",
    "- Direct multiple-shooting method and primal-dual interior point method\n",
    "- Inverse dynamics-based solution method for **unconstrained OCPs (i.e., no contacts or floating bases)**.\n",
    "- Riccati recursion to compute the Newton-steps\n",
    "\n",
    "We can specify some solver options by `robotoc.SolverOptions` (we use the default settings in the following) such as `nthreads`, the number of threads used in parallel computing with the direct multiple shooting method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver_options = robotoc.SolverOptions()\n",
    "solver_options.nthreads = 4\n",
    "ocp_solver = robotoc.UnconstrOCPSolver(ocp=ocp, solver_options=solver_options)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solve the OCP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We solve the OCP in which the initial state (`q`, `v`) and initial time (`t`) are given as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 0.0\n",
    "q = np.array([0.5*np.pi, 0, 0.5*np.pi, 0, 0.5*np.pi, 0, 0.5*np.pi]) \n",
    "v = np.zeros(iiwa14.dimv())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We initialize the OCP solver.  \n",
    "The following appropriate initialization is very simple but important to solve the OCP efficiently.   \n",
    "First, we set the initial guess of the solution as the above `q` and `v`.   \n",
    "Before set the initial guesses, we discretize the OCP by `discretize(t)` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.discretize(t)\n",
    "ocp_solver.set_solution(\"q\", q)\n",
    "ocp_solver.set_solution(\"v\", v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also initialize the solver to treat the inequality constraints.   \n",
    "The following methods do it. (Specifically, it initializes slack and dual variables of the primal-dual interior point method)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.init_constraints()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see how the current iterate is close to the stationary points, we see the **KKT error**, which is an l2-norm of the Karush–Kuhn–Tucker (KKT) conditions (first-order necessary conditions of the optimality).   \n",
    "After the above initialization, the KKT error is "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.KKT_error(t, q, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we solve the OCP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.solve(t, q, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the KKT error is very small. (Specifically, smaller than the threshold set in `robotoc.SolverOption`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.KKT_error(t, q, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also see the stats of the solver as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = ocp_solver.get_solver_statistics()\n",
    "print(stats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we plot the solution trajectory. We can get he solution from the solver by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_opt = ocp_solver.get_solution('q')\n",
    "v_opt = ocp_solver.get_solution('v')\n",
    "u_opt = ocp_solver.get_solution('u')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the plot, we can see that the solution strictly satisfies the constraints, which are plotted as dashed lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn\n",
    "\n",
    "plt.rcParams['figure.figsize'] = 10, 9\n",
    "fig = plt.figure()\n",
    "ax_q = fig.add_subplot(3, 1, 1)\n",
    "ax_v = fig.add_subplot(3, 1, 2)\n",
    "ax_u = fig.add_subplot(3, 1, 3)\n",
    "\n",
    "# Plot results\n",
    "ax_q.plot(np.linspace(t, t+T, N+1), q_opt)\n",
    "ax_v.plot(np.linspace(t, t+T, N+1), v_opt)\n",
    "ax_u.plot(np.linspace(t, t+T-(T/N), N), u_opt)\n",
    "\n",
    "# Plot constraints boundaries\n",
    "ax_v.plot(np.linspace(t, t+T, N+1), np.full(N+1, 0.5*np.pi), linestyle='--', color='gray')\n",
    "ax_v.plot(np.linspace(t, t+T, N+1), np.full(N+1, -0.5*np.pi), linestyle='--', color='gray')\n",
    "ax_u.plot(np.linspace(t, t+T-(T/N), N), np.full(N, 50.0), linestyle='--', color='gray')\n",
    "ax_u.plot(np.linspace(t, t+T-(T/N), N), np.full(N, -50.0), linestyle='--', color='gray')\n",
    "\n",
    "ax_q.set_xlabel('')\n",
    "ax_v.set_xlabel('')\n",
    "ax_u.set_xlabel('t [s]')\n",
    "\n",
    "ax_q.set_ylabel('q [rad]')\n",
    "ax_v.set_ylabel('v [rad/s]')\n",
    "ax_u.set_ylabel('u [Nm]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualize the solution trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = robotoc.utils.TrajectoryViewer(model_info=model_info,\n",
    "                                        viewer_type='meshcat')\n",
    "viewer.display(ocp_solver.get_time_discretization(), \n",
    "               ocp_solver.get_solution('q'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try Another Solver\n",
    "\n",
    "Next, we try another solver `robotoc.UnconstrParNMPCSolver` (`robotoc::UnconstrParNMPCSolver` in C++).  \n",
    "This solver has the following characteristics: \n",
    "- Direct multiple-shooting method and primal-dual interior point method\n",
    "- Inverse dynamics-based solution method for **unconstrained** OCPs (i.e., no contacts or floating bases).\n",
    "- ParNMPC algorithm (highly parallelizable Newton-type method) to compute the Newton-steps\n",
    "- Each Newton-step computation is very fast if ta large number of processors are available.\n",
    "- Convergence speed (i.e., required number of iterations until convergence) is worse than `robotoc.UnconstrOCPSolver`. It also means that the solver sometimes lacks robustness.\n",
    "\n",
    "We can specify some solver options by `robotoc.SolverOptions` (we use the default settings in the following) including `nthreads`, the number of threads used in parallel computing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver_options = robotoc.SolverOptions()\n",
    "solver_options.nthreads = 8\n",
    "parnmpc_solver = robotoc.UnconstrParNMPCSolver(ocp=ocp, solver_options=solver_options)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parnmpc_solver.discretize(t)\n",
    "parnmpc_solver.set_solution(\"q\", q)\n",
    "parnmpc_solver.set_solution(\"v\", v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parnmpc_solver.init_constraints()\n",
    "parnmpc_solver.init_backward_correction()\n",
    "parnmpc_solver.solve(t, q, v)\n",
    "stats = parnmpc_solver.get_solver_statistics()\n",
    "print(stats)"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}