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": [
    "## 本章の目的\n",
    "ここでは，４脚ロボット ANYmal の最適制御問題（OCP）を定式化し，解くことが目的です．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "まず，ANYmal の `robotoc.Robot` モデルを作成します．"
   ]
  },
  {
   "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",
    "次に，評価関数 (cost function) を作成します．   \n",
    "`robotoc`では，様々な評価関数のコンポーネント（**cost components**）を組み合わせて評価関数, `robotoc.CostFunction` (`std::shared_ptr<robotoc::CostFunction>` in C++), を作ります．  \n",
    "各 **cost component** はある特定の評価関数の項を表しており，`robotoc.CostFunctionComponentBase` (`robotoc::CostFunctionComponentBase` in C++) を継承しています.   \n",
    "**cost component** を評価関数 `robotoc.CostFunction` に集め， 評価関数 `robotoc.CostFunction` をソルバーに渡します．   \n",
    "このようにすることで，複数の**cost components** を簡単に組み合わせて評価関数が設計できるようになっています."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基本的な **cost components** は `robotoc` のパッケージに既に多数実装されています.  \n",
    "この例では， configuration つまり関節角，角速度，トルクなどについての **cost component**, task-space つまり end-effector (振り脚の軌道)についての **cost component**，そして重心 (CoM) 位置についての **cost component** を考えます．   \n",
    "以下がその **cost component** です．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaite Pattern (歩行パターン)\n",
    "評価関数を作る前に，まずは歩行パターンを設定します．　　   \n",
    "以下のようなシンプルなトロットを考えます． "
   ]
  },
  {
   "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",
    "\n",
    "レファレンスと重みパラメータを設定していきます．   \n",
    "以下のような命名規則があります．   \n",
    "\n",
    "- `q_ref`: configuration のレファレンス (目標値）\n",
    "- `q_weight`: configuration の重みパラメータ\n",
    "- `v_weight`: 一般化速度 (関節角速度) の重みパラメータ\n",
    "- `a_weight`: 一般化加速度 (関節角加速度) の重みパラメータ\n",
    "- `u_weight`: 関節トルクの重みパラメータ\n",
    "- `q_weight_terminal`, `v_weight_terminal`: ホライゾン終端での configuration と一般化速度の重みパラメータ\n",
    "- `q_weight_impact`, `v_weight_impact`: 衝突の瞬間での configuration と一般化速度の重みパラメータ\n",
    "- `dv_weight_impact`: 衝突の瞬間での 一般化速度変化（状態ジャンプ）の重みパラメータ"
   ]
  },
  {
   "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",
    "次に脚先の位置追従を表す，タスク空間の cost component を，各脚に合計4つ定義していきます．   \n",
    "このような cost component は， `robotoc.TaskSpace3DCost` (`std::shared_ptr<robotoc::TaskSpace3DCost>` in C++) を用いて定義できます.  \n",
    "各 `robotoc.TaskSpace3DCost`では, `robotoc.TaskSpace3DRefBase` (`robotoc::TaskSpace3DRefBase` in C++) を継承したレファレンスクラスを指定する必要があります．   \n",
    "シンプルな歩行に対して， `robotoc.PeriodicSwingFootRef` (`robotoc::PeriodicSwingFootRef` in C++) が既に実装されています．   \n",
    "これには， 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": [
    "まず各脚の初期位置を得るために，直立している姿勢（デフォルトの姿勢）での各脚の位置を forward kinematics から得ます．"
   ]
  },
  {
   "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": [
    "これと歩行パターンに基づいて時変の脚先位置のレファレンスをそれぞれの脚について作ります．"
   ]
  },
  {
   "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": [
    "このレファレンスを用いて脚先位置追従の cost component を作ります．"
   ]
  },
  {
   "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",
    "脚先と同様に，重心位置 (CoM) の cost component を作ります．   \n",
    "まずは，CoMの位置についての時変のレファレンスを作ります．"
   ]
  },
  {
   "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": [
    "このレファレンスから，CoM位置の cost component を作ります．"
   ]
  },
  {
   "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": [
    "以上の **cost components** を`robotoc.CostFunction`にまとめます．"
   ]
  },
  {
   "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",
    "制約についても評価関数と同様に作成します．   \n",
    "`robotoc.Constraints` (`std::shared_ptr<robotoc::Constraints>` in C++) は **constraint components**のコレクションになっています.\n",
    "各 **constraints component** はある特定の制約を表しており， `robotoc.ConstraintComponentBase` (`robotoc::ConstraintComponentBase` in C++) を継承しています．     \n",
    "これにより，様々な **constraints components** を簡単に組み合わせて制約を設計できます."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基本的な **constraint components** のは既に `robotoc` に多数実装されています.  \n",
    "この例では，関節角度，角速度，トルクの制約を考えます．     "
   ]
  },
  {
   "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": [
    "さらに，摩擦錐 (friction cone) 制約をつくります．    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "friction_cone = robotoc.FrictionCone(anymal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "制約を作り，上記のコンポーネントを加えます．"
   ]
  },
  {
   "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",
    "**Contact sequence** (`robotoc.ContactSequence`) は **contact status**  (`robotoc.ContactStatus`) の系列です．   \n",
    "**Contact status** は接触についての情報，どの接触がアクティブなのか？ や接触の位置など，を記述したものです.  \n",
    "これをOCPソルバーに渡すことで，OCPソルバーは自動的に問題を contact sequence に合わせて定式化してくれます．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "この例では，次の回数だけのトロットを考えます．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cycle = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以下のように contact sequence を作ります．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence = robotoc.ContactSequence(anymal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ここから，逐次的に contact sequence を作り上げていきます．   \n",
    "まずは，初期状態での `robotoc.ContactStatus`をつくります．   \n",
    "この`robotoc.ContactStatus`に接触状況（接触がアクティブか否か，接触の位置）を記述します．   \n",
    "初期状態ではロボットは `q_standing`の姿勢をとって直立しています．　  \n",
    "つまり，全ての脚で接触はアクティブです．      \n",
    "従って，以下のように記述します．"
   ]
  },
  {
   "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": [
    "これを用いて contact sequence を初期化します．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contact_sequence.init(contact_status_standing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "次に， `LF_FOOT` と `RH_FOOT` の接触がアクティブ（地面についている）で， `LH_FOOT`　と `RF_FOOT` が振り脚になっている状況を考えます． "
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "こちらを contact sequence に加えます．(push_backします．)"
   ]
  },
  {
   "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": [
    "次に，再び４つの脚が地面につきます．   \n",
    "振り脚であった `LH_FOOT` と　`RF_FOOT` は 0.5 x `step_length` だけ進んでいます.   \n",
    "こちらは以下のように書くことができます．"
   ]
  },
  {
   "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": [
    "次は， `LF_FOOT` と `RH_FOOT` が地面についていて（接触がアクティブ）， `LF_FOOT` と `RH_FOOT` が振り脚（接触がアクティブでない）状況を考えます．   "
   ]
  },
  {
   "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_sequence.push_back(contact_status_lfrh_swing, \n",
    "                           swing_start_time+swing_time+double_support_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "その次は，再び4脚が地面につきます．  \n",
    "振り脚だった脚は `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": [
    "これを指定した cycle の数だけ繰返します．"
   ]
  },
  {
   "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",
    "ここでは最適制御問題 (OCP) とそのソルバーを作ります．  \n",
    "まずはホライゾン長さ `T` と離散化のグリッドの数 `N` を設定しましょう． "
   ]
  },
  {
   "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": [
    "最適制御問題 (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": [
    "ではこのOCPのソルバー`robotoc.OCPSolver` (`robotoc::OCPSolver` in C++)　を作ります．   \n",
    "このソルバーは以下の特徴があります: \n",
    "- Direct multiple-shooting 法　と primal-dual interior point 法 （主双対内点法）\n",
    "- Inverse dynamics を用いた **unconstrained** OCP　 (i.e., no contacts or floating bases) の効率的なアルゴリズム.\n",
    "- 状態のみの等式制約のための制約変換法\n",
    "- Riccati recursion によるニュートンステップの計算\n",
    "- スイッチ時刻最適化可能な Riccati recursion アルゴリズム（この例では用いていません）\n",
    "\n",
    "種々のソルバーのオプションを `robotoc.SolverOptions` により設定できます．(ここではデフォルトのパラメータを用います).  \n",
    "例えば，ソルバー作成時に並列計算で用いるスレッドの数を`nthreads`で指定します."
   ]
  },
  {
   "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": [
    "初期状態 (`q`, `v`) と初期時刻 (`t`) が以下で与えられるOCP（最適制御問題）を考えます．"
   ]
  },
  {
   "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": [
    "まずはソルバーを初期化します．   \n",
    "以下は初期推定解を上の`q`, `v`に設定します．    \n",
    "さらに，接触力`f`の推定解をロボットの weight から設定します．   \n",
    "これは非常にシンプルですが高速な収束に割と重要です．   \n",
    "(例えば状態方程式を用いるといったような，より良い初期化の方法もあると思います．しかし適当な初期化でも direct multiple-shooting の収束性能は良く実装は面倒なため実装していません．)"
   ]
  },
  {
   "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": [
    "また，不等式制約に関してソルバーを初期化します．    \n",
    "(具体的には，主双対内点法のスラック変数とラグランジュ乗数を初期化しますが，特に内部について意識する必要はありません)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.init_constraints()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ソルバーの現在の推定解がどの程度最適解に近いか，ということを判断するために，本ソルバーでは **KKT error** という指標を用います．   \n",
    "これは， Karush–Kuhn–Tucker (KKT) 条件（1次の最適性の必要条件）のl2ノルムです．(最適解でこれは0になります)      \n",
    "まず，初期化後の KKT error を見てみます． "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.KKT_error(t, q, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "次にOCPを解きます．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.solve(t, q, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "KKT error がとても小さくなっていることが分かります．    \n",
    "(具体的には，`robotoc.SolverOption`で設定した閾値より小さくなっています)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocp_solver.KKT_error(t, q, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "また，各反復での KKT error やステップサイズといった情報も以下で見れます．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = ocp_solver.get_solver_statistics()\n",
    "print(stats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最適軌道のconfiguration, 一般化速度，関節トルクをプロットしてみます．   "
   ]
  },
  {
   "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` は接触力を friction cone 制約とともにプロットするモジュールも提供しています．   \n",
    "グレーのハッチが fx と fy の infeasible な領域を表しています."
   ]
  },
  {
   "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": [
    "また収束をプロットするモジュールもあります．"
   ]
  },
  {
   "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": [
    "また，以下で最適化された軌道を可視化できます．   \n",
    "`'gepetto'` を `viewer_type` に選べば， 接触力および friction cone 制約も可視化できます．"
   ]
  },
  {
   "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
}