littledog_contact.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "plt.style.use(\"ggplot\")\n",
    "plt.rcParams[\"font.family\"] = \"sans-serif\"\n",
    "sns.set_palette(\"dark\")\n",
    "sns.set_context(\"talk\")\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from IPython.display import HTML, display, SVG\n",
    "\n",
    "import pydot\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph, DiagramBuilder, BasicVector,\n",
    "                        Parser, Saturation, Simulator, PlanarSceneGraphVisualizer, \n",
    "                        LinearQuadraticRegulator, AbstractValue, MeshcatVisualizer, \n",
    "                        StartMeshcat, WrapToSystem, SceneGraph, DiscreteTimeLinearQuadraticRegulator,\n",
    "                        JacobianWrtVariable, eq, SnoptSolver)    \n",
    "\n",
    "from pydrake.systems.primitives import LogVectorOutput, ConstantVectorSource\n",
    "from pydrake.multibody.inverse_kinematics import PositionConstraint, AddUnitQuaternionConstraintOnPlant\n",
    "\n",
    "\n",
    "from pydrake.multibody.plant import ContactResults\n",
    "from pydrake.autodiffutils import AutoDiffXd, InitializeAutoDiff, ExtractGradient, ExtractValue\n",
    "from underactuated.meshcat_cpp_utils import MeshcatSliders, MeshcatJointSliders\n",
    "from Custom_LeafSystems import *\n",
    "\n",
    "from pydrake.math import RollPitchYaw_, RotationMatrix_, RigidTransform_\n",
    "\n",
    "from tqdm import tqdm as tqdm\n",
    "\n",
    "from Linearize import getGradients\n",
    "\n",
    "from pydrake.solvers.mathematicalprogram import MathematicalProgram, Solve\n",
    "import os\n",
    "from datetime import datetime\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Meshcat listening for connections at http://localhost:7000\n"
     ]
    }
   ],
   "source": [
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:93: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n",
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:99: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n",
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:193: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n",
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:199: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n",
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:292: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n",
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:298: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n",
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:391: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n",
      "WARNING:drake:/home/noah/Documents/Systems/littledog/LittleDog.urdf:397: warning: A link has specified the 'group' attribute on the <collision> tag. Drake doesn't make use of that attribute and all instances will be ignored.\n"
     ]
    }
   ],
   "source": [
    "def set_home(plant, context, ad=1.):\n",
    "\n",
    "    hip_roll = 0.1\n",
    "    hip_pitch = 1\n",
    "    knee = 1.55\n",
    "\n",
    "    back_hip_pitch = 0.9\n",
    "    back_knee = 1.\n",
    "\n",
    "    if np.array([ad]).dtype == float:\n",
    "        RPY = RollPitchYaw_[float](np.array([0., -np.pi/2+0.1, 0.]))\n",
    "        transform = RigidTransform_[float]([0, 0, 0.2645])\n",
    "        transform.set_rotation(RPY)\n",
    "    else:\n",
    "        RPY = RollPitchYaw_[AutoDiffXd](np.array([0., -np.pi/2+0.1, 0.]))\n",
    "        transform = RigidTransform_[AutoDiffXd]([0, 0, 0.2645])\n",
    "        transform.set_rotation(RPY)\n",
    "        \n",
    "        # hip_roll = InitializeAutoDiff(hip_roll)\n",
    "        # hip_pitch = InitializeAutoDiff(hip_pitch)\n",
    "        # knee = InitializeAutoDiff(knee)\n",
    "        # back_hip_pitch = InitializeAutoDiff(back_hip_pitch)\n",
    "        # back_knee = InitializeAutoDiff(back_knee)\n",
    "\n",
    "    plant.GetJointByName(\"front_right_hip_roll\").set_angle(context, -hip_roll)\n",
    "    plant.GetJointByName(\"front_right_hip_pitch\").set_angle(context, hip_pitch)\n",
    "    plant.GetJointByName(\"front_right_knee\").set_angle(context, -knee)\n",
    "    plant.GetJointByName(\"front_left_hip_roll\").set_angle(context, hip_roll)\n",
    "    plant.GetJointByName(\"front_left_hip_pitch\").set_angle(context, hip_pitch)\n",
    "    plant.GetJointByName(\"front_left_knee\").set_angle(context, -knee)\n",
    "    plant.GetJointByName(\"back_right_hip_roll\").set_angle(context, -hip_roll)\n",
    "    plant.GetJointByName(\"back_right_hip_pitch\").set_angle(context, back_hip_pitch)\n",
    "    plant.GetJointByName(\"back_right_knee\").set_angle(context, back_knee)\n",
    "    plant.GetJointByName(\"back_left_hip_roll\").set_angle(context, hip_roll)\n",
    "    plant.GetJointByName(\"back_left_hip_pitch\").set_angle(context, back_hip_pitch)\n",
    "    plant.GetJointByName(\"back_left_knee\").set_angle(context, back_knee)\n",
    "\n",
    "    plant.SetFreeBodyPose(context, plant.GetBodyByName(\"body\"), transform)\n",
    "\n",
    "\n",
    "time_step = 1/1000.\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=time_step)\n",
    "body = Parser(plant).AddModelFromFile(\"littledog/LittleDog.urdf\")\n",
    "ground = Parser(plant).AddModelFromFile(\"littledog/ground.urdf\")\n",
    "plant.Finalize()\n",
    "\n",
    "meshcat.Delete()\n",
    "vis = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "context = plant.CreateDefaultContext()\n",
    "\n",
    "num_pos = plant.num_positions()\n",
    "num_vel = plant.num_velocities()\n",
    "num_act = plant.num_actuators()\n",
    "num_disc_states = context.num_total_states() - context.num_continuous_states()\n",
    "num_contact = 2\n",
    "fb_num_pos = 7\n",
    "fb_num_vel = 6\n",
    "\n",
    "set_home(plant, context)\n",
    "\n",
    "CoM_0 = plant.CalcCenterOfMassPositionInWorld(context)\n",
    "\n",
    "plant.get_actuation_input_port().FixValue(context, np.zeros(num_act))\n",
    "\n",
    "# display(SVG(pydot.graph_from_dot_data(plant.GetTopologyGraphvizString())[0].create_svg()))\n",
    "\n",
    "# diagram = builder.Build()\n",
    "# display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString(max_depth=2))[0].create_svg()))\n",
    "\n",
    "x0 = plant.GetPositions(context)\n",
    "\n",
    "foot_L = \"back_left_foot_center\"\n",
    "foot_R = \"back_right_foot_center\"\n",
    "\n",
    "# def PoseOptimization():\n",
    "#     context_diagram = diagram.CreateDefaultContext()\n",
    "#     plant_context_from_diagram = diagram.GetSubsystemContext(plant, context_diagram)\n",
    "#     set_home(plant, plant_context_from_diagram)\n",
    "\n",
    "#     sliders = MeshcatJointSliders(meshcat, plant, context_diagram)\n",
    "#     sliders.Run(vis, context_diagram)\n",
    "#     meshcat.DeleteAddedControls()\n",
    "#     meshcat.Delete()\n",
    "\n",
    "# PoseOptimization()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution found? False.\n",
      "['unit quaternion constraint[0]: 0.999900 <= 1.003663 <= 1.000100', 'Dynamics constraint[0]: -0.000010 <= -0.256868 <= 0.000010', 'Dynamics constraint[1]: -0.000010 <= -0.135643 <= 0.000010', 'Dynamics constraint[2]: -0.000010 <= 0.001247 <= 0.000010', 'Dynamics constraint[5]: -0.000010 <= 6.439132 <= 0.000010', 'Torque constraint[0]: -0.000100 <= 0.003126 <= 0.000100', 'Torque constraint[1]: -0.000100 <= -0.021719 <= 0.000100', 'Torque constraint[2]: -0.000100 <= -0.005395 <= 0.000100', 'Torque constraint[3]: -0.000100 <= -0.003110 <= 0.000100', 'Torque constraint[4]: -0.000100 <= -0.021529 <= 0.000100', 'Torque constraint[5]: -0.000100 <= -0.005398 <= 0.000100', 'Torque constraint[6]: -0.000100 <= -0.029941 <= 0.000100', 'Torque constraint[7]: -0.000100 <= 0.420117 <= 0.000100', 'Torque constraint[8]: -0.000100 <= -0.050337 <= 0.000100', 'Torque constraint[9]: -0.000100 <= 0.035434 <= 0.000100', 'Torque constraint[10]: -0.000100 <= 0.107333 <= 0.000100', 'Torque constraint[11]: -0.000100 <= -0.367886 <= 0.000100', 'left foot position[2]: 0.000000 <= 0.003280 <= 0.000000', 'foot right position[2]: 0.000000 <= 0.005801 <= 0.000000']\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"1864pt\" height=\"563pt\" viewBox=\"0.00 0.00 1864.00 563.00\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 559)\">\n<title>_62810128</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-559 1860,-559 1860,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster62810128diagram</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"8,-8 8,-547 1848,-547 1848,-8 8,-8\"/>\n<text text-anchor=\"middle\" x=\"928\" y=\"-531.8\" font-family=\"Times,serif\" font-size=\"14.00\">62810128</text>\n</g>\n<g id=\"clust4\" class=\"cluster\">\n<title>cluster62810128subsystems</title>\n<polygon fill=\"none\" stroke=\"white\" points=\"16,-16 16,-516 1840,-516 1840,-16 16,-16\"/>\n</g>\n<!-- 63880976 -->\n<g id=\"node1\" class=\"node\">\n<title>63880976</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"502,-24.5 502,-507.5 1124,-507.5 1124,-24.5 502,-24.5\"/>\n<text text-anchor=\"middle\" x=\"813\" y=\"-492.3\" font-family=\"Times,serif\" font-size=\"14.00\">plant</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-484.5 1124,-484.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-452.3\" font-family=\"Times,serif\" font-size=\"14.00\">geometry_query</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-427.5 750,-427.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-394.8\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-369.5 750,-369.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-337.3\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-312.5 750,-312.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-279.8\" font-family=\"Times,serif\" font-size=\"14.00\">LittleDog_actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-254.5 750,-254.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-222.3\" font-family=\"Times,serif\" font-size=\"14.00\">ground_actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-197.5 750,-197.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-164.8\" font-family=\"Times,serif\" font-size=\"14.00\">actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-139.5 750,-139.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-107.3\" font-family=\"Times,serif\" font-size=\"14.00\">applied_generalized_force</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-82.5 750,-82.5 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-49.8\" font-family=\"Times,serif\" font-size=\"14.00\">applied_spatial_force</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-24.5 750,-484.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-469.3\" font-family=\"Times,serif\" font-size=\"14.00\">geometry_pose</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-461.5 1124,-461.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-446.3\" font-family=\"Times,serif\" font-size=\"14.00\">continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-438.5 1124,-438.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-423.3\" font-family=\"Times,serif\" font-size=\"14.00\">body_poses</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-415.5 1124,-415.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-400.3\" font-family=\"Times,serif\" font-size=\"14.00\">spatial_velocities</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-392.5 1124,-392.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-377.3\" font-family=\"Times,serif\" font-size=\"14.00\">spatial_accelerations</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-369.5 1124,-369.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-354.3\" font-family=\"Times,serif\" font-size=\"14.00\">generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-346.5 1124,-346.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-331.3\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-323.5 1124,-323.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-308.3\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-300.5 1124,-300.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-285.3\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-277.5 1124,-277.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-262.3\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-254.5 1124,-254.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-239.3\" font-family=\"Times,serif\" font-size=\"14.00\">LittleDog_continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-231.5 1124,-231.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-216.3\" font-family=\"Times,serif\" font-size=\"14.00\">LittleDog_generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-208.5 1124,-208.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-193.3\" font-family=\"Times,serif\" font-size=\"14.00\">ground_continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-185.5 1124,-185.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-170.3\" font-family=\"Times,serif\" font-size=\"14.00\">ground_generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-162.5 1124,-162.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-147.3\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_generalized_contact_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-139.5 1124,-139.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-124.3\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_generalized_contact_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-116.5 1124,-116.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-101.3\" font-family=\"Times,serif\" font-size=\"14.00\">LittleDog_generalized_contact_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-93.5 1124,-93.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-78.3\" font-family=\"Times,serif\" font-size=\"14.00\">ground_generalized_contact_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-70.5 1124,-70.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-55.3\" font-family=\"Times,serif\" font-size=\"14.00\">reaction_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-47.5 1124,-47.5 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-32.3\" font-family=\"Times,serif\" font-size=\"14.00\">contact_results</text>\n</g>\n<!-- 63909232 -->\n<g id=\"node2\" class=\"node\">\n<title>63909232</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"1160,-435.5 1160,-504.5 1373,-504.5 1373,-435.5 1160,-435.5\"/>\n<text text-anchor=\"middle\" x=\"1266.5\" y=\"-489.3\" font-family=\"Times,serif\" font-size=\"14.00\">scene_graph</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1160,-481.5 1373,-481.5 \"/>\n<text text-anchor=\"middle\" x=\"1237.5\" y=\"-466.3\" font-family=\"Times,serif\" font-size=\"14.00\">plant_pose</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1160,-458.5 1315,-458.5 \"/>\n<text text-anchor=\"middle\" x=\"1237.5\" y=\"-443.3\" font-family=\"Times,serif\" font-size=\"14.00\">plant_configuration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1315,-435.5 1315,-481.5 \"/>\n<text text-anchor=\"middle\" x=\"1344\" y=\"-454.8\" font-family=\"Times,serif\" font-size=\"14.00\">query</text>\n</g>\n<!-- 63880976&#45;&gt;63909232 -->\n<g id=\"edge4\" class=\"edge\">\n<title>63880976:y0-&gt;63909232:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1124,-473C1136.04,-473 1141.24,-471.31 1149.86,-470.47\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"1150.17,-473.96 1160,-470 1149.85,-466.97 1150.17,-473.96\"/>\n</g>\n<!-- 63909232&#45;&gt;63880976 -->\n<g id=\"edge2\" class=\"edge\">\n<title>63909232:y0-&gt;63880976:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1344.5,-434C1344.5,-352 1229.52,-391.51 1160,-435 1129.33,-454.19 1154.72,-488.9 1124,-508 1065.31,-544.5 551.34,-556.4 502,-508 488.46,-494.72 480.91,-468.24 491.39,-459.15\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"492.59,-462.44 501,-456 490.41,-455.79 492.59,-462.44\"/>\n</g>\n<!-- 62103632 -->\n<g id=\"node3\" class=\"node\">\n<title>62103632</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"1409,-447 1409,-493 1832,-493 1832,-447 1409,-447\"/>\n<text text-anchor=\"middle\" x=\"1620.5\" y=\"-477.8\" font-family=\"Times,serif\" font-size=\"14.00\">drake/geometry/MeshcatVisualizer@0000000003b3a050</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1409,-470 1832,-470 \"/>\n<text text-anchor=\"middle\" x=\"1536.5\" y=\"-454.8\" font-family=\"Times,serif\" font-size=\"14.00\">query_object</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1664,-447 1664,-470 \"/>\n<text text-anchor=\"middle\" x=\"1748\" y=\"-454.8\" font-family=\"Times,serif\" font-size=\"14.00\"> </text>\n</g>\n<!-- 63909232&#45;&gt;62103632 -->\n<g id=\"edge1\" class=\"edge\">\n<title>63909232:y0-&gt;62103632:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1373,-458C1385,-458 1390.25,-458 1398.88,-458\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"1399,-461.5 1409,-458 1399,-454.5 1399,-461.5\"/>\n</g>\n<!-- 64434400 -->\n<g id=\"node4\" class=\"node\">\n<title>64434400</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"24,-157 24,-203 466,-203 466,-157 24,-157\"/>\n<text text-anchor=\"middle\" x=\"245\" y=\"-187.8\" font-family=\"Times,serif\" font-size=\"14.00\">drake/systems/ConstantVectorSource@0000000003d730e0</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"24,-180 466,-180 \"/>\n<text text-anchor=\"middle\" x=\"131\" y=\"-164.8\" font-family=\"Times,serif\" font-size=\"14.00\"> </text>\n<polyline fill=\"none\" stroke=\"black\" points=\"238,-157 238,-180 \"/>\n<text text-anchor=\"middle\" x=\"352\" y=\"-164.8\" font-family=\"Times,serif\" font-size=\"14.00\">y0</text>\n</g>\n<!-- 64434400&#45;&gt;63880976 -->\n<g id=\"edge3\" class=\"edge\">\n<title>64434400:y0-&gt;63880976:u5</title>\n<path fill=\"none\" stroke=\"black\" d=\"M466,-168C478,-168 483.25,-168 491.88,-168\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"492,-171.5 502,-168 492,-164.5 492,-171.5\"/>\n</g>\n</g>\n</svg>",
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Pose optimization\n",
    "\n",
    "plant_f = plant\n",
    "context_f = plant_f.CreateDefaultContext()\n",
    "\n",
    "plant_ad = plant.ToAutoDiffXd()\n",
    "context_ad = plant_ad.CreateDefaultContext()\n",
    "\n",
    "set_home(plant_f, context_f)\n",
    "set_home(plant_ad, context_ad, ad=None)\n",
    "\n",
    "B = plant.MakeActuationMatrix()\n",
    "\n",
    "def makePosConstraint(plant, context):\n",
    "    # Position constraint feet\n",
    "    epsilon = 1E-1\n",
    "\n",
    "    p_AQ_lower = np.array([-0.09-epsilon, -epsilon, 0])\n",
    "    p_AQ_upper = np.array([0.09+epsilon, epsilon, 0])\n",
    "    p_BQ = np.zeros(3)\n",
    "\n",
    "    frame_left = plant.GetFrameByName(foot_L)\n",
    "    frame_right = plant.GetFrameByName(foot_R)\n",
    "    frame_world = plant.world_frame()\n",
    "\n",
    "    pos_constraint_left = PositionConstraint(plant=plant,\n",
    "                            frameA=frame_world,\n",
    "                            p_AQ_lower=p_AQ_lower,\n",
    "                            p_AQ_upper=p_AQ_upper,\n",
    "                            frameB=frame_left,\n",
    "                            p_BQ=p_BQ,\n",
    "                            plant_context=context\n",
    "                            )\n",
    "\n",
    "    # Position constraint foot right\n",
    "    pos_constraint_right = PositionConstraint(plant=plant,\n",
    "                            frameA=frame_world,\n",
    "                            p_AQ_lower=p_AQ_lower,\n",
    "                            p_AQ_upper=p_AQ_upper,\n",
    "                            frameB=frame_right,\n",
    "                            p_BQ=p_BQ,\n",
    "                            plant_context=context\n",
    "                            )\n",
    "\n",
    "    return pos_constraint_left, pos_constraint_right\n",
    "\n",
    "pos_constraint_left, pos_constraint_right = makePosConstraint(plant_ad, context_ad)\n",
    "\n",
    "def get_foot_Jacobian(var):\n",
    "\n",
    "    if var.dtype == float:\n",
    "        plant = plant_f\n",
    "        context = context_f\n",
    "        p_BoBi_B_left = np.zeros(3)\n",
    "        p_BoBi_B_right = np.zeros(3)\n",
    "    else:\n",
    "        plant = plant_ad\n",
    "        context = context_ad \n",
    "        p_BoBi_B_left = InitializeAutoDiff(np.zeros(3))\n",
    "        p_BoBi_B_right = InitializeAutoDiff(np.zeros(3))\n",
    "\n",
    "    frame_left = plant.GetFrameByName(foot_L)\n",
    "    frame_right = plant.GetFrameByName(foot_R)\n",
    "    frame_world = plant.world_frame()\n",
    "\n",
    "    Jc_left = plant.CalcJacobianTranslationalVelocity(\n",
    "            context,\n",
    "            JacobianWrtVariable(1),\n",
    "            frame_left,\n",
    "            p_BoBi_B_left,\n",
    "            frame_world,\n",
    "            frame_world\n",
    "            )\n",
    "\n",
    "    Jc_right = plant.CalcJacobianTranslationalVelocity(\n",
    "                context,\n",
    "                JacobianWrtVariable(1),\n",
    "                frame_right,\n",
    "                p_BoBi_B_right,\n",
    "                frame_world,\n",
    "                frame_world\n",
    "                )\n",
    "\n",
    "    return np.vstack((Jc_left, Jc_right))\n",
    "\n",
    "\n",
    "def calcUnactuatedDynamics(vars):\n",
    "\n",
    "    split_at = [num_act, num_act + num_pos]\n",
    "    tau, pos, lambda_c = np.split(vars, split_at)\n",
    "    \n",
    "    if tau.dtype == float:\n",
    "        plant = plant_f\n",
    "        context = context_f\n",
    "    else:\n",
    "        plant = plant_ad\n",
    "        context = context_ad \n",
    "\n",
    "    plant.SetPositions(context, pos)\n",
    "    plant.SetVelocities(context, np.zeros(num_vel))\n",
    "    plant.get_actuation_input_port().FixValue(context, tau)\n",
    "\n",
    "    # M, Cv, taug\n",
    "    taug = plant.CalcGravityGeneralizedForces(context)\n",
    "    Cv = plant.CalcBiasTerm(context)\n",
    "    Jc = get_foot_Jacobian(tau)\n",
    "    \n",
    "    return taug[:6] + Jc.T[:6, :] @ lambda_c - Cv[:6]\n",
    "\n",
    "\n",
    "def torque_constraint(vars):\n",
    "    split_at = [num_act, num_act + num_pos]\n",
    "    tau, pos, lambda_c = np.split(vars, split_at)\n",
    "    \n",
    "    if tau.dtype == float:\n",
    "        plant = plant_f\n",
    "        context = context_f\n",
    "    else:\n",
    "        plant = plant_ad\n",
    "        context = context_ad \n",
    "\n",
    "    plant.SetPositions(context, pos)\n",
    "    plant.SetVelocities(context, np.zeros(num_vel))\n",
    "    plant.get_actuation_input_port().FixValue(context, tau)\n",
    "\n",
    "    # M, Cv, taug\n",
    "    taug = plant.CalcGravityGeneralizedForces(context)\n",
    "    Cv = plant.CalcBiasTerm(context)\n",
    "    Jc = get_foot_Jacobian(tau)\n",
    "\n",
    "    B_act = B[6:, :]\n",
    "    \n",
    "    return tau + np.linalg.inv(B_act) @ (Cv[6:] - taug[6:] - Jc.T[6:, :] @ lambda_c ) \n",
    "\n",
    "\n",
    "def CostTorque(tau):\n",
    "    return tau.dot(tau)\n",
    "\n",
    "\n",
    "def CostCoM(q):\n",
    "\n",
    "    if q.dtype == float:\n",
    "        plant = plant_f\n",
    "        context = context_f\n",
    "    else:\n",
    "        plant = plant_ad\n",
    "        context = context_ad\n",
    "\n",
    "    plant.SetPositions(context, q)\n",
    "    CoM = plant.CalcCenterOfMassPositionInWorld(context)\n",
    "\n",
    "    CoM_xy = CoM[0:2]\n",
    "    CoM_z = CoM[2]\n",
    "\n",
    "    return CoM_xy.dot(CoM_xy) + 1 / (w_CoM*CoM_z)\n",
    "\n",
    "\n",
    "def getCoM(q):\n",
    "\n",
    "    if q.dtype == float:\n",
    "        plant = plant_f\n",
    "        context = context_f\n",
    "    else:\n",
    "        plant = plant_ad\n",
    "        context = context_ad\n",
    "\n",
    "    plant.SetPositions(context, q)\n",
    "    CoM = plant.CalcCenterOfMassPositionInWorld(context)[0]\n",
    "\n",
    "    return CoM\n",
    "\n",
    "\n",
    "# Write Mathematical program\n",
    "prog = MathematicalProgram()\n",
    "\n",
    "# Define decision variables\n",
    "pos = prog.NewContinuousVariables(num_pos, \"positions\")\n",
    "torques = prog.NewContinuousVariables(num_act, \"torques\")\n",
    "lambdas = prog.NewContinuousVariables(num_contact*3, \"lambda\")\n",
    "\n",
    "# Unit quaternion constraint\n",
    "# unit_quat = AddUnitQuaternionConstraintOnPlant(plant_ad, pos, prog)\n",
    "\n",
    "# Bent knee and hip\n",
    "# knee_constraint = prog.AddConstraint(pos[4 + fb_num_pos] <= -0.3).evaluator().set_description(\"Knee bend constraint\")\n",
    "# hip_constraint = prog.AddConstraint(pos[3 + fb_num_pos] >= 0.15).evaluator().set_description(\"Hip bend constraint\")\n",
    "# knee_constraint_L = prog.AddConstraint(pos[9 + fb_num_pos] <= -0.3).evaluator().set_description(\"Knee bend constraint\")\n",
    "# hip_constraint_L = prog.AddConstraint(pos[8 + fb_num_pos] >= 0.15).evaluator().set_description(\"Hip bend constraint\")\n",
    "\n",
    "# Symmetric pose Constraints\n",
    "# prog.AddConstraint(pos[2 + fb_num_pos] == pos[7 + fb_num_pos]).evaluator().set_description(\"Symmetry in hip roll\")\n",
    "# prog.AddConstraint(pos[3 + fb_num_pos] == pos[8 + fb_num_pos]).evaluator().set_description(\"Symmetry in hip pitch\")\n",
    "# prog.AddConstraint(pos[4 + fb_num_pos] == pos[9 + fb_num_pos]).evaluator().set_description(\"Symmetry in knee pitch\")\n",
    "# prog.AddConstraint(pos[5 + fb_num_pos] == pos[10 + fb_num_pos]).evaluator().set_description(\"Symmetry in foot pitch\")\n",
    "# prog.AddConstraint(pos[6 + fb_num_pos] == pos[11 + fb_num_pos]).evaluator().set_description(\"Symmetry in foot roll\")\n",
    "\n",
    "unit_quat_eps = 1e-4\n",
    "\n",
    "# Quaternion unit length constraint\n",
    "(prog.AddConstraint(lambda x: [x @ x], [1 - unit_quat_eps], [1 + unit_quat_eps], pos[0:4]).evaluator().set_description(\"unit quaternion constraint\"))\n",
    "\n",
    "# Limit Inividual parts of quaternion to be between -1 and 1\n",
    "(prog.AddBoundingBoxConstraint(-1.0, 1.0, pos[0]).evaluator().set_description(\"quat0 limit constraint\"))\n",
    "(prog.AddBoundingBoxConstraint(-1.0, 1.0, pos[1]).evaluator().set_description(\"quat1 limit constraint\"))\n",
    "(prog.AddBoundingBoxConstraint(-1.0, 1.0, pos[2]).evaluator().set_description(\"quat2 limit constraint\"))\n",
    "(prog.AddBoundingBoxConstraint(-1.0, 1.0, pos[3]).evaluator().set_description(\"quat3 limit constraint\"))\n",
    "\n",
    "# Friction cone constraint\n",
    "friction_left = prog.AddConstraint(lambdas[0]**2 + lambdas[1]**2 <= lambdas[2]**2).evaluator().set_description(\"left foot friction cone\")\n",
    "friction_right = prog.AddConstraint(lambdas[3]**2 + lambdas[4]**2 <= lambdas[5]**2).evaluator().set_description(\"right foot friction cone\")\n",
    "\n",
    "# Dynamics constraints for unactuated DOFs\n",
    "vars = np.concatenate((torques, pos, lambdas))\n",
    "dyn_constraint = prog.AddConstraint(calcUnactuatedDynamics, lb=[-1e-5]*(6), ub=[1e-5]*(6), vars=vars)\n",
    "dyn_constraint.evaluator().set_description(\"Dynamics constraint\")\n",
    "\n",
    "# Torque constraint for actuated DOFs\n",
    "dyn_constraint = prog.AddConstraint(torque_constraint, lb=[-1E-4]*(num_act), ub=[1E-4]*(num_act), vars=vars) \n",
    "dyn_constraint.evaluator().set_description(\"Torque constraint\")\n",
    "\n",
    "# Postiive contact force\n",
    "prog.AddConstraint(lambdas[2] >= 1E-5), prog.AddConstraint(lambdas[5] >= 1E-5)\n",
    "\n",
    "# Foot position constraints\n",
    "foot_left_pos = prog.AddConstraint(pos_constraint_left, pos).evaluator().set_description(\"left foot position\")\n",
    "foot_right_pos = prog.AddConstraint(pos_constraint_right, pos).evaluator().set_description(\"foot right position\")\n",
    "\n",
    "w_CoM = 0.1\n",
    "prog.AddCost(CostTorque, vars=torques)\n",
    "prog.AddCost(CostCoM, vars=pos)\n",
    "\n",
    "fc_guess = -plant.gravity_field().gravity_vector() * plant.CalcTotalMass(context) / 2\n",
    "fc_guess = np.hstack((fc_guess, fc_guess))\n",
    "\n",
    "tau_guess = np.zeros(num_act)\n",
    "\n",
    "# Setting initial guess for decision variables\n",
    "prog.SetInitialGuess(pos, x0[:num_pos])\n",
    "prog.SetInitialGuess(torques, tau_guess)\n",
    "prog.SetInitialGuess(lambdas, fc_guess)\n",
    "\n",
    "# Write solver status to file\n",
    "path = f\"Logs/littledog/{datetime.now().strftime('%Y-%m-%d_%H-%M_SNOPT.out')}\"\n",
    "snopt = SnoptSolver().solver_id()\n",
    "prog.SetSolverOption(snopt, 'Print file', path)\n",
    "\n",
    "# Solve optimization problem\n",
    "result = Solve(prog)\n",
    "\n",
    "q = result.GetSolution(pos)\n",
    "tau = result.GetSolution(torques)\n",
    "fc = result.GetSolution(lambdas)\n",
    "\n",
    "x0 = np.hstack((q, np.zeros(num_vel)))\n",
    "\n",
    "# In case of infeasible quaternion constraint\n",
    "x0[0:4] = x0[0:4] / (x0[0:4].dot(x0[0:4]))\n",
    "\n",
    "print(f'Solution found? {result.is_success()}.')\n",
    "\n",
    "if not result.is_success():\n",
    "    infeasible = result.GetInfeasibleConstraintNames(prog)\n",
    "    print(infeasible)\n",
    "\n",
    "plant.SetPositions(context, q)\n",
    "\n",
    "CoM = plant.CalcCenterOfMassPositionInWorld(context)\n",
    "\n",
    "with open(path, \"a\") as file:\n",
    "    file.write(\"\\n\")\n",
    "    np.savetxt(file, q, delimiter=\",\", header=\"Array of robot joint angles\")\n",
    "    file.write(\"\\n\")\n",
    "    np.savetxt(file, tau, delimiter=\",\", header=\"Array of torques\")\n",
    "    file.write(\"\\n\")\n",
    "    np.savetxt(file, fc, delimiter=\",\", header=\"Array of contact forces\")\n",
    "    file.write(\"\\n\")\n",
    "    np.savetxt(file, CoM, delimiter=\",\", header=\"CoM location\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Wire contact detection logic\n",
    "generic_input = builder.AddSystem(ConstantVectorSource(np.zeros(num_act)))\n",
    "# contact = builder.AddSystem(ReadContactResults(plant, foot_L))\n",
    "\n",
    "# builder.Connect(plant.get_contact_results_output_port(), contact.get_input_port())\n",
    "builder.Connect(generic_input.get_output_port(), plant.get_actuation_input_port())\n",
    "\n",
    "# logger_is_contact = LogVectorOutput(contact.get_output_port(1), builder, publish_period=1/1000)\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "# display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString(max_depth=2))[0].create_svg()))\n",
    "\n",
    "simulator = Simulator(diagram)\n",
    "sim_context = simulator.get_mutable_context()\n",
    "\n",
    "sim_context.SetTime(0.)\n",
    "sim_context.SetDiscreteState(x0)\n",
    "vis.StartRecording()\n",
    "simulator.AdvanceTo(0.5)\n",
    "vis.StopRecording()\n",
    "vis.PublishRecording()\n",
    "\n",
    "meshcat.DeleteAddedControls()\n",
    "\n",
    "# contact_data = logger_is_contact.FindLog(sim_context).data().flatten()\n",
    "# t = np.arange(0, len(contact_data))\n",
    "\n",
    "# fig, ax = plt.subplots()\n",
    "# ax.plot(t, contact_data)\n",
    "\n",
    "# plt.show()\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}