5_optimal_control_humanoid.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to `robotoc` \n",
    "# 5: Optimal Control of a Humanoid"
   ]
  },
  {
   "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 humanoid robot iCub."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we construct `robotoc.Robot` of a humanoid robot iCub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import robotoc\n",
    "\n",
    "model_info = robotoc.RobotModelInfo()\n",
    "model_info.urdf_path = 'urdf/icub_description/urdf/icub.urdf'\n",
    "model_info.base_joint_type = robotoc.BaseJointType.FloatingBase\n",
    "baumgarte_time_step = 0.05\n",
    "model_info.surface_contacts = [robotoc.ContactModelInfo('l_sole', baumgarte_time_step),\n",
    "                               robotoc.ContactModelInfo('r_sole', baumgarte_time_step)]\n",
    "icub = robotoc.Robot(model_info)"
   ]
  },
  {
   "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)**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Jumping Pattern\n",
    "Before constructing the cost function, we define parameters for the jumping pattern.   \n",
    "Here, we consider a dynamic jumping of the robot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "jump_length = np.array([0.5, 0, 0])\n",
    "flying_time = 0.25\n",
    "ground_time = 0.7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration Space Cost\n",
    "First, we define a configuration-space cost component.\n",
    "We set the reference values and weight parameters of the configuration space cost component. \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\n",
    "- `q_weight_impact`, `v_weight_impact`: weights on the configuration and generalized velocity at the impact instant (just before the impact, i.e., touch-down of the feet onto the ground)\n",
    "- `dv_weight_impact`: weight on the impulse change in the generalized velocity at the impact stage due to the impact event"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_standing = np.array([0, 0, 0.592, 0, 0, 1, 0,\n",
    "                       0.20944, 0.08727, 0, -0.1745, -0.0279, -0.08726, # left leg\n",
    "                       0.20944, 0.08727, 0, -0.1745, -0.0279, -0.08726, # right leg\n",
    "                       0, 0, 0, # torso\n",
    "                       0, 0.35, 0.5, 0.5, 0, 0, 0, # left arm \n",
    "                       0, 0.35, 0.5, 0.5, 0, 0, 0]) # right arm \n",
    "q_ref = q_standing.copy()\n",
    "q_ref[0:3] += jump_length\n",
    "q_weight = np.array([0, 1, 1, 100, 100, 100, \n",
    "                     0.001, 0.001, 0.001, 0.001, 0.001, 0.001, \n",
    "                     0.001, 0.001, 0.001, 0.001, 0.001, 0.001, \n",
    "                     0.001, 1, 1,\n",
    "                     0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, \n",
    "                     0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001])\n",
    "q_weight_terminal = q_weight\n",
    "v_weight = np.full(icub.dimv(), 1.0e-03)\n",
    "a_weight = np.full(icub.dimv(), 1.0e-05)\n",
    "q_weight_impact = 1.0 * q_weight\n",
    "v_weight_impact = 1.0 * v_weight\n",
    "config_cost = robotoc.ConfigurationSpaceCost(icub)\n",
    "config_cost.set_q_ref(q_standing)\n",
    "config_cost.set_q_weight(q_weight)\n",
    "config_cost.set_q_weight_terminal(q_weight)\n",
    "config_cost.set_q_weight_impact(q_weight_impact)\n",
    "config_cost.set_v_weight(v_weight)\n",
    "config_cost.set_v_weight_terminal(v_weight)\n",
    "config_cost.set_v_weight_impact(v_weight_impact)\n",
    "config_cost.set_a_weight(a_weight)"
   ]
  },
  {
   "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**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_position_lower  = robotoc.JointPositionLowerLimit(icub)\n",
    "joint_position_upper  = robotoc.JointPositionUpperLimit(icub)\n",
    "joint_velocity_lower  = robotoc.JointVelocityLowerLimit(icub)\n",
    "joint_velocity_upper  = robotoc.JointVelocityUpperLimit(icub)\n",
    "joint_torques_lower   = robotoc.JointTorquesLowerLimit(icub)\n",
    "joint_torques_upper   = robotoc.JointTorquesUpperLimit(icub)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also use the friction cone constraint component in this example.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "friction_cone = robotoc.FrictionCone(icub)"
   ]
  },
  {
   "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",
    "constraints.add(\"friction_cone\", friction_cone)\n",
    "print(constraints)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contact Sequence\n",
    "The **contact sequence** (`robotoc.ContactSequence`) describes the sequence of the **contact status** (`robotoc.ContactStatus`), which containts contact information such as which contacts are active at each time and contact placements (i.e., positions and rotations of the contacts).   \n",
    "We set the contact sequence to the OCP (optimal control problem) solver by passing `robotoc.ContactSequence` .   \n",
    "Then the OCP solver automatically constructs an OCP in which the robot dynamics follow the input contact sequence."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we construct the contact sequence as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence = robotoc.ContactSequence(icub)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We sequentially build the contact sequence of the bipedal jumings.  \n",
    "First, we make `robotoc.ContactStatus` object.   \n",
    "We set which contacts are active and contact placements in this object.  \n",
    "In addition, we set the friction coefficients that are used in the friction cone constraints.  \n",
    "At the beginning, the robot is standing whose configuration is given as `q_standing` .  \n",
    "The corresponding `robotoc.ContactStatus` is as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "icub.forward_kinematics(q_standing)\n",
    "x3d0_L = icub.frame_placement('l_sole')\n",
    "x3d0_R = icub.frame_placement('r_sole')\n",
    "contact_placements = {'l_sole': x3d0_L, 'r_sole': x3d0_R} \n",
    "mu = 0.6\n",
    "friction_coefficients = {'l_sole': mu, 'r_sole': mu} \n",
    "contact_status_standing = icub.create_contact_status()\n",
    "contact_status_standing.activate_contacts(['l_sole', 'r_sole'])\n",
    "contact_status_standing.set_contact_placements(contact_placements)\n",
    "contact_status_standing.set_friction_coefficients(friction_coefficients)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We initialize the contact sequence with this contact status."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence.init(contact_status_standing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we add the contact sequence in which the contacts in two feet are inactive (that is, the robot is flying).    \n",
    "This is expressed as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_status_flying = icub.create_contact_status()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add (push_back) this contact status into the contact sequence.  \n",
    "In this example, **we optimize the switching times (i.e., contact timings) as well as the trajectory (i.e., the switching time optimization (STO) problems)**.  \n",
    "We specify this in the contact sequence as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence.push_back(contact_status_flying, ground_time, sto=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input switching time is considreded as the initial guess of the switching time.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next phase, the both feet have active contacts.  \n",
    "Now the position of two feet went forward with `jump_length` .   \n",
    "This is described as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_placements['l_sole'].trans = contact_placements['l_sole'].trans + jump_length\n",
    "contact_placements['r_sole'].trans = contact_placements['r_sole'].trans + jump_length \n",
    "contact_status_standing.set_contact_placements(contact_placements)\n",
    "contact_sequence.push_back(contact_status_standing, ground_time+flying_time, sto=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, the robot is again flying, i.e., the contacts in two feet are inactive.   \n",
    "This is written as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence.push_back(contact_status_flying, 2*ground_time+flying_time, sto=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the two feet become having contacts again.   \n",
    "The both feet went forward with jump_length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_placements['l_sole'].trans = contact_placements['l_sole'].trans + jump_length \n",
    "contact_placements['r_sole'].trans = contact_placements['r_sole'].trans + jump_length \n",
    "contact_status_standing.set_contact_placements(contact_placements)\n",
    "contact_sequence.push_back(contact_status_standing, 2*ground_time+2*flying_time, sto=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## STO (Switching Time Optimization) Cost and Constraints \n",
    "Since we consider the **STO problem** in this example, we have to define the STO cost (`robotoc.STOCostFunction`) and STO constraints (`robotoc.STOConstraints`).   \n",
    "Otherwise, the solver does not try to solve the STO problems (just solves the OCP with fixed the switching times)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we define the STO cost, i.e., `robotoc.STOCostFunction` (`std::shared_ptr<robotoc::STOCostFunction>` in C++).   \n",
    "We do not add any cost components in the cost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sto_cost = robotoc.STOCostFunction()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define the STO constraints, i.e., `robotoc.STOConstraints` (`std::shared_ptr<robotoc::STOConstraints>` in C++).   \n",
    "In this object, we set the **minimum dwell-time**, i.e., minimum time-interval of each contact phase.   \n",
    "It is important to set appropriate minimum dwell-times.  \n",
    "Here, we have 4 switches and 5 phases.  \n",
    "We define the STO constraints as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sto_constraints = robotoc.STOConstraints(minimum_dwell_times=[0.6, 0.2, 0.6, 0.2, 0.6])"
   ]
  },
  {
   "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` as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = 0.02\n",
    "T = 2*flying_time + 3*ground_time\n",
    "N = int(np.floor(T/dt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then construct `robotoc.OCP` (`robotoc::OCP` in C++).   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp = robotoc.OCP(robot=icub, cost=cost, constraints=constraints, \n",
    "                  sto_cost=sto_cost, sto_constraints=sto_constraints, \n",
    "                  contact_sequence=contact_sequence, T=T, N=N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we make the OCP solver `robotoc.OCPSolver` (`robotoc::OCPSolver` in C++) for this OCP.  \n",
    "This solver has the following characteristics: \n",
    "- Direct multiple-shooting method and primal-dual interior point method.\n",
    "- Lifted-contact dynamics for robotic systems with contacts.\n",
    "- Constraint-transformation for pure-state equality constraints.\n",
    "- Riccati recursion to compute the Newton steps.\n",
    "- Switching time optimization (STO) via Riccati recursion.\n",
    "\n",
    "We can specify some solver options by `robotoc.SolverOptions`.    \n",
    "In this example, we set\n",
    "- `kkt_tol_mesh`: KKT tolerance for the mesh-refinement\n",
    "- `max_dt_mesh`: Threshold of the maximum discretization size for the mesh-refinement\n",
    "- `initial_sto_reg_iter`: Initial number of iteraions set by this value are regularized for the STO problem.\n",
    "- `max_iter`: Maximum number of the iterations.\n",
    "- `nthreads`: Number of threads used in the parallel computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver_options = robotoc.SolverOptions()\n",
    "solver_options.kkt_tol_mesh = 0.1\n",
    "solver_options.max_dt_mesh = T/N \n",
    "solver_options.max_iter = 350\n",
    "solver_options.initial_sto_reg_iter = 10 \n",
    "solver_options.nthreads = 4\n",
    "ocp_solver = robotoc.OCPSolver(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.\n",
    "q = q_standing.copy()\n",
    "v = np.zeros(icub.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",
    "We set the initial guess of the solution as the above `q` and `v`.   \n",
    "After that, we initializes the variables used in the inequality constraints."
   ]
  },
  {
   "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)\n",
    "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plot_kkt = robotoc.utils.PlotConvergence()\n",
    "plot_kkt.figsize = 12, 9\n",
    "kkt_data = [e.kkt_error for e in ocp_solver.get_solver_statistics().performance_index]\n",
    "ts_data = ocp_solver.get_solver_statistics().ts\n",
    "plot_kkt.plot(kkt_data=kkt_data, ts_data=ts_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualize the solution trajectory.  \n",
    "If we choose `'gepetto'` as the `viewer_type`, we can also see the contact forces and friction cones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = robotoc.utils.TrajectoryViewer(model_info=model_info, viewer_type='gepetto')\n",
    "viewer.set_contact_info(mu=mu)\n",
    "viewer.display(ocp_solver.get_time_discretization(), \n",
    "               ocp_solver.get_solution('q'), \n",
    "               ocp_solver.get_solution('f', 'WORLD'))"
   ]
  }
 ],
 "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
}