4_optimal_control_quadruped.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to `robotoc` \n",
    "# 4: Optimal Control of a Quadruped"
   ]
  },
  {
   "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 quadruped ANYmal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, construct `robotoc.Robot` of a quadruped robot ANYmal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import robotoc\n",
    "\n",
    "model_info = robotoc.RobotModelInfo()\n",
    "model_info.urdf_path = 'urdf/anymal_b_simple_description/urdf/anymal.urdf'\n",
    "model_info.base_joint_type = robotoc.BaseJointType.FloatingBase\n",
    "baumgarte_time_step = 0.05\n",
    "model_info.point_contacts = [robotoc.ContactModelInfo('LF_FOOT', baumgarte_time_step),\n",
    "                             robotoc.ContactModelInfo('LH_FOOT', baumgarte_time_step),\n",
    "                             robotoc.ContactModelInfo('RF_FOOT', baumgarte_time_step),\n",
    "                             robotoc.ContactModelInfo('RH_FOOT', baumgarte_time_step)]\n",
    "anymal = 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)**, **cost components on the end-effector position (i.e., cost on the task-space position)**, and **cost component on the center of mass (CoM) position**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gait Pattern\n",
    "Before constructing the cost function, we define parameters for the gait pattern.   \n",
    "Here, we consider a simple trot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "step_length = np.array([0.15, 0, 0])\n",
    "step_height = 0.1\n",
    "swing_time = 0.5\n",
    "double_support_time = 0.04\n",
    "swing_start_time = 0.04"
   ]
  },
  {
   "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.4792, 0, 0, 0, 1, \n",
    "                       -0.1,  0.7, -1.0, \n",
    "                       -0.1, -0.7,  1.0, \n",
    "                        0.1,  0.7, -1.0, \n",
    "                        0.1, -0.7,  1.0])\n",
    "q_weight = np.array([0, 0, 0, 250000, 250000, 250000, \n",
    "                     0.0001, 0.0001, 0.0001, \n",
    "                     0.0001, 0.0001, 0.0001,\n",
    "                     0.0001, 0.0001, 0.0001,\n",
    "                     0.0001, 0.0001, 0.0001])\n",
    "v_weight = np.array([100, 100, 100, 100, 100, 100, \n",
    "                     1, 1, 1, \n",
    "                     1, 1, 1,\n",
    "                     1, 1, 1,\n",
    "                     1, 1, 1])\n",
    "u_weight = np.full(anymal.dimu(), 1.0e-01)\n",
    "q_weight_terminal = q_weight\n",
    "v_weight_terminal = v_weight\n",
    "q_weight_impact = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, \n",
    "                            100, 100, 100, \n",
    "                            100, 100, 100,\n",
    "                            100, 100, 100,\n",
    "                            100, 100, 100])\n",
    "v_weight_impact = np.full(anymal.dimv(), 100)\n",
    "config_cost = robotoc.ConfigurationSpaceCost(anymal)\n",
    "config_cost.set_q_ref(q_standing)\n",
    "config_cost.set_q_weight(q_weight)\n",
    "config_cost.set_v_weight(v_weight)\n",
    "config_cost.set_u_weight(u_weight)\n",
    "config_cost.set_q_weight_terminal(q_weight_terminal)\n",
    "config_cost.set_v_weight_terminal(v_weight_terminal)\n",
    "config_cost.set_q_weight_impact(q_weight_impact)\n",
    "config_cost.set_v_weight_impact(v_weight_impact)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task Space Cost\n",
    "Second, we define task-space cost components.  \n",
    "We set the reference positions of the end-effector (each foot) and weight parameters.  \n",
    "This kind of costs can be constructed as `robotoc.TaskSpace3DCost` (`std::shared_ptr<robotoc::TaskSpace3DCost>` in C++).  \n",
    "In each `robotoc.TaskSpace3DCost`, we specify a periodic position reference that inherits `robotoc.TaskSpace3DRefBase` (`robotoc::TaskSpace3DRefBase` in C++).  \n",
    "For robot walking, we provide `robotoc.PeriodicSwingFootRef` (`robotoc::PeriodicSwingFootRef` in C++), in which we specify the initial position, step length, step height, swing-start time, swing time, support time, and whether the initial step-length is half or not."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To specify the initial positions of all feet, we first obtain the feet' positions at the standing pose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anymal.forward_kinematics(q_standing)\n",
    "x3d0_LF = anymal.frame_position('LF_FOOT')\n",
    "x3d0_LH = anymal.frame_position('LH_FOOT')\n",
    "x3d0_RF = anymal.frame_position('RF_FOOT')\n",
    "x3d0_RH = anymal.frame_position('RH_FOOT')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then construct the feet refs as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LF_t0 = swing_start_time + swing_time + double_support_time\n",
    "LH_t0 = swing_start_time\n",
    "RF_t0 = swing_start_time\n",
    "RH_t0 = swing_start_time + swing_time + double_support_time\n",
    "LF_foot_ref = robotoc.PeriodicSwingFootRef(x3d0_LF, step_length, step_height, \n",
    "                                           LF_t0, swing_time, \n",
    "                                           swing_time+2*double_support_time, False)\n",
    "LH_foot_ref = robotoc.PeriodicSwingFootRef(x3d0_LH, step_length, step_height, \n",
    "                                           LH_t0, swing_time, \n",
    "                                           swing_time+2*double_support_time, True)\n",
    "RF_foot_ref = robotoc.PeriodicSwingFootRef(x3d0_RF, step_length, step_height, \n",
    "                                           RF_t0, swing_time, \n",
    "                                           swing_time+2*double_support_time, True)\n",
    "RH_foot_ref = robotoc.PeriodicSwingFootRef(x3d0_RH, step_length, step_height, \n",
    "                                           RH_t0, swing_time, \n",
    "                                           swing_time+2*double_support_time, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can construct time-varying task space cost components of feet as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LF_cost = robotoc.TaskSpace3DCost(anymal, 'LF_FOOT', LF_foot_ref)\n",
    "LH_cost = robotoc.TaskSpace3DCost(anymal, 'LH_FOOT', LH_foot_ref)\n",
    "RF_cost = robotoc.TaskSpace3DCost(anymal, 'RF_FOOT', RF_foot_ref)\n",
    "RH_cost = robotoc.TaskSpace3DCost(anymal, 'RH_FOOT', RH_foot_ref)\n",
    "swing_foot_weight = np.full(3, 1.0e06)\n",
    "LF_cost.set_weight(swing_foot_weight)\n",
    "LH_cost.set_weight(swing_foot_weight)\n",
    "RF_cost.set_weight(swing_foot_weight)\n",
    "RH_cost.set_weight(swing_foot_weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Center of Mass (CoM) Cost\n",
    "\n",
    "We also define a cost component on the position of center of mass (CoM).   \n",
    "The way to define this cost is similar to the above.   \n",
    "First, we create the time-varying reference of the CoM position as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "com_ref0 = anymal.com()\n",
    "vcom_ref = 0.5 * step_length / swing_time\n",
    "com_ref = robotoc.PeriodicCoMRef(com_ref0, vcom_ref, swing_start_time, swing_time, \n",
    "                                 double_support_time, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this reference, we can create the time-varying CoM cost as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "com_cost = robotoc.CoMCost(anymal, com_ref)\n",
    "com_cost.set_weight(np.full(3, 1.0e06))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then make the **cost function** `robotoc.CostFunction` and add the **configuration cost component**, **task-space cost components**, and **CoM cost components** to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost = robotoc.CostFunction()\n",
    "cost.add(\"config_cost\", config_cost)\n",
    "cost.add(\"LF_cost\", LF_cost)\n",
    "cost.add(\"LH_cost\", LH_cost)\n",
    "cost.add(\"RF_cost\", RF_cost)\n",
    "cost.add(\"RH_cost\", RH_cost)\n",
    "cost.add(\"com_cost\", com_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(anymal)\n",
    "joint_position_upper  = robotoc.JointPositionUpperLimit(anymal)\n",
    "joint_velocity_lower  = robotoc.JointVelocityLowerLimit(anymal)\n",
    "joint_velocity_upper  = robotoc.JointVelocityUpperLimit(anymal)\n",
    "joint_torques_lower   = robotoc.JointTorquesLowerLimit(anymal)\n",
    "joint_torques_upper   = robotoc.JointTorquesUpperLimit(anymal)"
   ]
  },
  {
   "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(anymal)"
   ]
  },
  {
   "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 pass the `robotoc.ContactSequence` to the OCP solver.   \n",
    "Then the OCP solver automatically constructs an OCP in which the robot dynamics follow the input contact sequence."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we first specify the number of cycles of the trotting gait."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cycle = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We construct the contact sequence as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence = robotoc.ContactSequence(anymal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We sequentially build the contact sequence of the biped walking.  \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",
    "This is specified as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anymal.forward_kinematics(q_standing)\n",
    "x3d0_LF = anymal.frame_position('LF_FOOT')\n",
    "x3d0_LH = anymal.frame_position('LH_FOOT')\n",
    "x3d0_RF = anymal.frame_position('RF_FOOT')\n",
    "x3d0_RH = anymal.frame_position('RH_FOOT')\n",
    "contact_positions = {'LF_FOOT': x3d0_LF, 'LH_FOOT': x3d0_LH, 'RF_FOOT': x3d0_RF, 'RH_FOOT': x3d0_RH} \n",
    "mu = 0.7\n",
    "friction_coefficients = {'LF_FOOT': mu, 'LH_FOOT': mu, 'RF_FOOT': mu, 'RH_FOOT': mu} \n",
    "contact_status_standing = anymal.create_contact_status()\n",
    "contact_status_standing.activate_contacts(['LF_FOOT', 'LH_FOOT', 'RF_FOOT', 'RH_FOOT'])\n",
    "contact_status_standing.set_contact_placements(contact_positions)\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 `LF_FOOT` and `RH_FOOT` are active and these in `LH_FOOT` and `RF_FOOT` are inactive (that is, `LH_FOOT` and `RF_FOOT` are swinging).    \n",
    "This is expressed as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_status_lhrf_swing = anymal.create_contact_status()\n",
    "contact_status_lhrf_swing.activate_contacts(['LF_FOOT', 'RH_FOOT'])\n",
    "contact_status_lhrf_swing.set_contact_placements(contact_positions)\n",
    "contact_status_lhrf_swing.set_friction_coefficients(friction_coefficients)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add (push_back) this contact status into the contact sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence.push_back(contact_status_lhrf_swing, swing_start_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next phase, the four feet have active contacts.  \n",
    "Now the positions of `LH_FOOT` and `RF_FOOT` went forward 0.5 times `step_length` (because length of the initial step is half of the other steps).  \n",
    "This is described as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_positions['LH_FOOT'] += 0.5 * step_length\n",
    "contact_positions['RF_FOOT'] += 0.5 * step_length\n",
    "contact_status_standing.set_contact_placements(contact_positions)\n",
    "contact_sequence.push_back(contact_status_standing, swing_start_time+swing_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, `LF_FOOT` and `RH_FOOT` are swinging, i.e., the contacts in `LF_FOOT` and `RH_FOOT` are inactive and these of `LH_FOOT` and `RF_FOOT` are active.  \n",
    "This is written as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_status_lfrh_swing = anymal.create_contact_status()\n",
    "contact_status_lfrh_swing.activate_contacts(['LH_FOOT', 'RF_FOOT'])\n",
    "contact_status_lfrh_swing.set_contact_placements(contact_positions)\n",
    "contact_status_lfrh_swing.set_friction_coefficients(friction_coefficients)\n",
    "contact_sequence.push_back(contact_status_lfrh_swing, \n",
    "                           swing_start_time+swing_time+double_support_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, the four feet become having contacts again.   \n",
    "`LF_FOOT` and `RH_FOOT` went forward with `step_length`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_positions['LF_FOOT'] += step_length\n",
    "contact_positions['RH_FOOT'] += step_length\n",
    "contact_status_standing.set_contact_placements(contact_positions)\n",
    "contact_sequence.push_back(contact_status_standing, \n",
    "                           swing_start_time+2*swing_time+double_support_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is repeated until the specified number of cycles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(cycle-1):\n",
    "    t1 = swing_start_time + (i+1)*(2*swing_time+2*double_support_time)\n",
    "    contact_status_lhrf_swing.set_contact_placements(contact_positions)\n",
    "    contact_sequence.push_back(contact_status_lhrf_swing, t1)\n",
    "\n",
    "    contact_positions['LH_FOOT'] += step_length\n",
    "    contact_positions['RF_FOOT'] += step_length\n",
    "    contact_status_standing.set_contact_placements(contact_positions)\n",
    "    contact_sequence.push_back(contact_status_standing, t1+swing_time)\n",
    "\n",
    "    contact_status_lfrh_swing.set_contact_placements(contact_positions)\n",
    "    contact_sequence.push_back(contact_status_lfrh_swing, \n",
    "                               t1+swing_time+double_support_time)\n",
    "\n",
    "    contact_positions['LF_FOOT'] += step_length\n",
    "    contact_positions['RH_FOOT'] += step_length\n",
    "    contact_status_standing.set_contact_placements(contact_positions)\n",
    "    contact_sequence.push_back(contact_status_standing, \n",
    "                               t1+2*swing_time+double_support_time)"
   ]
  },
  {
   "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": [
    "T = swing_start_time + cycle*(2*double_support_time+2*swing_time)\n",
    "dt = 0.02\n",
    "N = int(np.floor(T/dt))    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then construct the optimal control problem (OCP), `robotoc.OCP` (`robotoc::OCP` in C++).      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp = robotoc.OCP(robot=anymal, cost=cost, constraints=constraints, \n",
    "                  contact_sequence=contact_sequence, T=T, N=N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then makes 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 via Riccati recursion (we do not use this function in this example).\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 in 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.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(anymal.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",
    "Second, we set the initial guess of the solution of contact force `f` by the total weight of the robot."
   ]
  },
  {
   "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",
    "f_init = np.array([0.0, 0.0, 0.25*anymal.total_weight()])\n",
    "ocp_solver.set_solution(\"f\", f_init)"
   ]
  },
  {
   "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": [
    "We plot the joint position, velocities, and torques as "
   ]
  },
  {
   "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",
    "time_discretization = ocp_solver.get_time_discretization()\n",
    "t = [time_discretization[i].t for i in range(len(time_discretization))]\n",
    "ax_q.plot(t, ocp_solver.get_solution('q'))\n",
    "ax_v.plot(t, ocp_solver.get_solution('v'))\n",
    "ax_u.plot(t, ocp_solver.get_solution('u'))\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": [
    "`robotoc` provides a module to plot the contact forces with friction cone constraints.  \n",
    "In the plot, the hatched region shows the infeasible region of fx and fy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_f = robotoc.utils.PlotContactForce(mu=mu)\n",
    "plot_f.figsize = 14, 8\n",
    "plot_f.wspace = 0.2\n",
    "plot_f.hspace = 0.2\n",
    "plot_f.legend_bbox_to_anchor = (-0.2, 2.7)\n",
    "plot_f.plot(f_traj=ocp_solver.get_solution('f', 'WORLD'), \n",
    "            time_discretization=ocp_solver.get_time_discretization())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`robotoc` also provides a module to plot the convergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_kkt = robotoc.utils.PlotConvergence()\n",
    "#plot_kkt.ylim = [0., 1.5]\n",
    "plot_kkt.figsize = 6, 4\n",
    "kkt_data = [e.kkt_error for e in ocp_solver.get_solver_statistics().performance_index]\n",
    "plot_kkt.plot(kkt_data=kkt_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,\n",
    "                                        viewer_type='gepetto')\n",
    "viewer.set_contact_info(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
}