Double_Leg_Contact.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using matplotlib backend: TkAgg\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "%matplotlib\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 pydrake.math import RollPitchYaw_, RotationMatrix_, RigidTransform_\n",
    "from underactuated.meshcat_cpp_utils import MeshcatSliders, MeshcatJointSliders\n",
    "from Custom_LeafSystems import *\n",
    "\n",
    "from tqdm import tqdm as tqdm\n",
    "\n",
    "from Linearize import getGradients\n",
    "\n",
    "from pydrake.solvers.mathematicalprogram import MathematicalProgram, Solve\n",
    "\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": [
    "# Start the visualizer (run this cell only once, each instance consumes a port)\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_step = 1/1000.\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=time_step)\n",
    "body = Parser(plant).AddModelFromFile(\"Double_Leg/Double_Leg_Collisions.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",
    "fb_num_pos = 7\n",
    "fb_num_vel = 6\n",
    "n = 3\n",
    "\n",
    "x0 = np.zeros(num_disc_states)\n",
    "x0[0], x0[6] = 1, 0.4\n",
    "\n",
    "plant.SetPositionsAndVelocities(context, x0)\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",
    "foot_L = [\"collision_dummy_L_BR\", \"collision_dummy_L_BL\", \"collision_dummy_L_FR\", \"collision_dummy_L_FL\"]\n",
    "foot_R = [\"collision_dummy_R_BR\", \"collision_dummy_R_BL\", \"collision_dummy_R_FR\", \"collision_dummy_R_FL\"]\n",
    "feet = [\"foot_L\", \"foot_R\"]\n",
    "\n",
    "num_contact = len(foot_L) + len(foot_R)\n",
    "\n",
    "def sethome(plant, context, ad=1.):\n",
    "\n",
    "    shaft_roll = 0.3\n",
    "    hip_pitch = 0.3\n",
    "    knee_pitch = 0.6\n",
    "    foot_pitch = 0.3\n",
    "\n",
    "    plant.GetJointByName(\"shaft_roll\").set_angle(context, shaft_roll)\n",
    "    plant.GetJointByName(\"Hip_Pitch_R\").set_angle(context, hip_pitch)\n",
    "    plant.GetJointByName(\"Hip_Pitch_L\").set_angle(context, hip_pitch)\n",
    "    plant.GetJointByName(\"knee_R\").set_angle(context, -knee_pitch)\n",
    "    plant.GetJointByName(\"knee_L\").set_angle(context, -knee_pitch)\n",
    "    plant.GetJointByName(\"foot_leg_pitch_R\").set_angle(context, foot_pitch)\n",
    "    plant.GetJointByName(\"foot_leg_pitch_L\").set_angle(context, -foot_pitch)\n",
    "\n",
    "    foot_pos = plant.GetBodyByName(\"collision_dummy_L_C\").EvalPoseInWorld(context).GetAsMatrix34()[:, -1].reshape(3)\n",
    "    body_pos = plant.GetBodyByName(\"base_link\").EvalPoseInWorld(context).GetAsMatrix34()[:, -1].reshape(3)\n",
    "\n",
    "    # print(dir(body_pos), body_pos, foot_pos)\n",
    "    # print(body_pos, body_pos.GetAsMatrix34())\n",
    "\n",
    "    body_pos[2] = body_pos[2] - foot_pos[2] + 0.001 # Add radius of collision sphere\n",
    "    body_pos[1] = -foot_pos[1]\n",
    "\n",
    "    if np.array([ad]).dtype == float:\n",
    "        transform = RigidTransform_[float](body_pos)\n",
    "    else:\n",
    "        transform = RigidTransform_[AutoDiffXd](body_pos)\n",
    "\n",
    "    plant.SetFreeBodyPose(context, plant.GetBodyByName(\"base_link\"), transform)\n",
    "\n",
    "\n",
    "sethome(plant, context)\n",
    "\n",
    "x0 = plant.GetPositionsAndVelocities(context)\n",
    "\n",
    "# print(plant.GetBodyByName(foot_L).EvalPoseInWorld(context))\n",
    "# display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString(max_depth=2))[0].create_svg()))\n",
    "\n",
    "# diagram = builder.Build()\n",
    "\n",
    "# def PoseOptimization():\n",
    "#     context_diagram = diagram.CreateDefaultContext()\n",
    "#     plant_context_from_diagram = diagram.GetSubsystemContext(plant, context_diagram)\n",
    "#     plant_context_from_diagram.get_mutable_discrete_state_vector().SetFromVector(x0)\n",
    "\n",
    "#     sliders = MeshcatJointSliders(meshcat, plant, context_diagram)\n",
    "#     sliders.Run(vis, context_diagram)\n",
    "#     meshcat.DeleteAddedControls()\n",
    "#     meshcat.Delete()\n",
    "\n",
    "# PoseOptimization()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 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",
    "plant_f.SetPositionsAndVelocities(context_f, x0)\n",
    "plant_ad.SetPositionsAndVelocities(context_ad, InitializeAutoDiff(x0))\n",
    "\n",
    "B = plant.MakeActuationMatrix()\n",
    "\n",
    "def makePosConstraint(plant, context):\n",
    "    # Position constraint feet\n",
    "    \n",
    "    p_AQ_lower = np.array([-np.inf, -np.inf, 0.001])\n",
    "    p_AQ_upper = np.array([np.inf, np.inf, 0.001])\n",
    "    p_BQ = np.zeros(3)\n",
    "\n",
    "    frame_world = plant.world_frame()\n",
    "    constraints = []\n",
    "\n",
    "    for point in foot_L:\n",
    "        frame_foot = plant.GetFrameByName(point)\n",
    "\n",
    "        pos_constraint = PositionConstraint(plant=plant,\n",
    "                            frameA=frame_world,\n",
    "                            p_AQ_lower=p_AQ_lower,\n",
    "                            p_AQ_upper=p_AQ_upper,\n",
    "                            frameB=frame_foot,\n",
    "                            p_BQ=p_BQ,\n",
    "                            plant_context=context\n",
    "                            )\n",
    "\n",
    "        constraints.append(pos_constraint)\n",
    "\n",
    "    for point in foot_R:\n",
    "        frame_foot = plant.GetFrameByName(point)\n",
    "\n",
    "        pos_constraint = PositionConstraint(plant=plant,\n",
    "                            frameA=frame_world,\n",
    "                            p_AQ_lower=p_AQ_lower,\n",
    "                            p_AQ_upper=p_AQ_upper,\n",
    "                            frameB=frame_foot,\n",
    "                            p_BQ=p_BQ,\n",
    "                            plant_context=context\n",
    "                            )\n",
    "\n",
    "        constraints.append(pos_constraint)\n",
    "\n",
    "    return constraints\n",
    "\n",
    "pos_constraint = makePosConstraint(plant_ad, context_ad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_foot_Jacobian(var):\n",
    "\n",
    "    if var.dtype == float:\n",
    "        plant = plant_f\n",
    "        context = context_f\n",
    "        p_BoBi_B = np.zeros(3)\n",
    "    else:\n",
    "        plant = plant_ad\n",
    "        context = context_ad \n",
    "        p_BoBi_B = InitializeAutoDiff(np.zeros(3))\n",
    "\n",
    "    frame_world = plant.world_frame()\n",
    "\n",
    "    for index, point in enumerate(foot_L):\n",
    "        \n",
    "        frame_foot = plant.GetFrameByName(point)\n",
    "\n",
    "        Jc_trans = plant.CalcJacobianTranslationalVelocity(\n",
    "                context,\n",
    "                JacobianWrtVariable(1),\n",
    "                frame_foot,\n",
    "                p_BoBi_B,\n",
    "                frame_world,\n",
    "                frame_world\n",
    "                )\n",
    "\n",
    "        if index == 0:\n",
    "            Jc = Jc_trans\n",
    "        else:\n",
    "            Jc = np.vstack((Jc, Jc_trans))\n",
    "\n",
    "    for point in foot_R:\n",
    "        \n",
    "        frame_foot = plant.GetFrameByName(point)\n",
    "\n",
    "        Jc_trans = plant.CalcJacobianTranslationalVelocity(\n",
    "                context,\n",
    "                JacobianWrtVariable(1),\n",
    "                frame_foot,\n",
    "                p_BoBi_B,\n",
    "                frame_world,\n",
    "                frame_world\n",
    "                )\n",
    "\n",
    "        Jc = np.vstack((Jc, Jc_trans))\n",
    "\n",
    "    return Jc\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[:, :6].T @ 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[:, 6:].T @ 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution found? False.\n",
      "['Dynamics constraint[0]: -0.000010 <= -0.231693 <= 0.000010', 'Torque constraint[1]: -0.000100 <= 0.218763 <= 0.000100', 'Torque constraint[3]: -0.000100 <= -0.218763 <= 0.000100', 'Torque constraint[4]: -0.000100 <= 0.201165 <= 0.000100', 'Torque constraint[5]: -0.000100 <= 0.201165 <= 0.000100', 'Torque constraint[6]: -0.000100 <= -0.334000 <= 0.000100', 'Torque constraint[8]: -0.000100 <= -0.334000 <= 0.000100', 'Torque constraint[11]: -0.000100 <= -0.436308 <= 0.000100']\n"
     ]
    }
   ],
   "source": [
    "# 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*n, \"lambda\")\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-5\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",
    "# Unit quaternion constraint\n",
    "unit_quat = AddUnitQuaternionConstraintOnPlant(plant_ad, pos, prog)\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",
    "# Static coefficient of friction\n",
    "mu = 1\n",
    "X = 0.03\n",
    "Y = 0.02\n",
    "\n",
    "# Origin points of the foot\n",
    "O_L = plant.GetBodyByName(\"collision_dummy_L_C\").EvalPoseInWorld(context).translation()\n",
    "O_R = plant.GetBodyByName(\"collision_dummy_R_C\").EvalPoseInWorld(context).translation()\n",
    "\n",
    "\n",
    "# CWC constraint\n",
    "def CWC(foot, Origin, side):\n",
    "    \n",
    "    if side == \"L\":\n",
    "        offset = 0\n",
    "    else:\n",
    "        offset = n*len(foot_L)\n",
    "\n",
    "    for index, point in enumerate(foot):\n",
    "    \n",
    "        # OC_i\n",
    "        C_i = plant.GetBodyByName(point).EvalPoseInWorld(context).translation()\n",
    "        OC_i = C_i - Origin\n",
    "\n",
    "        # Contact wrench cone f = sum(f_i)\n",
    "        # Contact wrench cone tau = sum(OC_i x f_i)\n",
    "        if index == 0:\n",
    "            f = np.array([lambdas[offset + n*index], lambdas[offset + 1 + n*index], lambdas[offset + 2 + n*index]])\n",
    "            tau = np.cross(OC_i, np.array([lambdas[offset + n*index], lambdas[offset + 1 + n*index], lambdas[offset + 2 + n*index]]))\n",
    "        else:\n",
    "            f += np.array([lambdas[offset + n*index], lambdas[offset + 1 + n*index], lambdas[offset + 2 + n*index]])\n",
    "            tau += np.cross(OC_i, np.array([lambdas[offset + n*index], lambdas[offset + 1 + n*index], lambdas[offset + 2 + n*index]]))\n",
    "\n",
    "        # Positive normal force \n",
    "        prog.AddConstraint(lambdas[offset + 2 + n*index] >= 1e-5)\n",
    "\n",
    "        # tau = OCxf\n",
    "        # OCxf = np.cross(OC_i, np.array([lambdas[offset + n*index], lambdas[offset + 1 + n*index], lambdas[offset + 2 + n*index]]))\n",
    "        # prog.AddConstraint(lambdas[offset + 3 + n*index] == OCxf[0]).evaluator().set_description(f\"OCxf 1 {point}\")\n",
    "        # prog.AddConstraint(lambdas[offset + 4 + n*index] == OCxf[1]).evaluator().set_description(f\"OCxf 2 {point}\")\n",
    "        # prog.AddConstraint(lambdas[offset + 5 + n*index] == OCxf[2]).evaluator().set_description(f\"OCxf 3 {point}\")\n",
    "\n",
    "    return f, tau\n",
    "    \n",
    "\n",
    "f_x_L, tau_L = CWC(foot_L, O_L, \"L\")\n",
    "f_x_R, tau_R = CWC(foot_R, O_R, \"R\")\n",
    "\n",
    "for i in range(3):\n",
    "    prog.AddConstraint(plant.gravity_field().gravity_vector()[i]*plant.CalcTotalMass(context) + f_x_L[i] + f_x_R[i] == 0).evaluator().set_description(\"static force eq.\")\n",
    "\n",
    "def CWC_constraint(force, torque, foot):\n",
    "\n",
    "    # https://scaron.info/publications/icra-2015.html\n",
    "    # Linear friction cone\n",
    "    prog.AddConstraint(-force[0] - mu*force[2] <= 0).evaluator().set_description(f\"friction cone x pos {foot}\")\n",
    "    prog.AddConstraint(force[0] - mu*force[2] <= 0).evaluator().set_description(f\"friction cone x neg {foot}\")\n",
    "    prog.AddConstraint(-force[1] - mu*force[2] <= 0).evaluator().set_description(f\"friction cone y pos {foot}\")\n",
    "    prog.AddConstraint(force[1] - mu*force[2] <= 0).evaluator().set_description(f\"friction cone y neg {foot}\")\n",
    "\n",
    "    # Center of pressure constraint\n",
    "    prog.AddConstraint(-Y*force[0] - torque[0] <= 0).evaluator().set_description(f\"CoP x pos {foot}\")\n",
    "    prog.AddConstraint(-Y*force[0] + torque[0] <= 0).evaluator().set_description(f\"CoP x neg {foot}\")\n",
    "    prog.AddConstraint(-X*force[0] - torque[1] <= 0).evaluator().set_description(f\"CoP y pos {foot}\")\n",
    "    prog.AddConstraint(-X*force[0] - torque[1] <= 0).evaluator().set_description(f\"CoP y neg {foot}\")\n",
    "\n",
    "    # Now yaw slippage\n",
    "    prog.AddConstraint(-Y*force[0] - X*force[1] -(X+Y)*mu*force[2] + mu*torque[0] + mu*torque[1] - torque[2] <= 0).evaluator().set_description(f\"yaw 1 {foot}\")\n",
    "    prog.AddConstraint(-Y*force[0] + X*force[1] -(X+Y)*mu*force[2] + mu*torque[0] - mu*torque[1] - torque[2] <= 0).evaluator().set_description(f\"yaw 2 {foot}\")\n",
    "    prog.AddConstraint(Y*force[0] - X*force[1] -(X+Y)*mu*force[2] - mu*torque[0] + mu*torque[1] - torque[2] <= 0).evaluator().set_description(f\"yaw 3 {foot}\")\n",
    "    prog.AddConstraint(Y*force[0] + X*force[1] -(X+Y)*mu*force[2] - mu*torque[0] - mu*torque[1] - torque[2] <= 0).evaluator().set_description(f\"yaw 4 {foot}\")\n",
    "    prog.AddConstraint(Y*force[0] + X*force[1] -(X+Y)*mu*force[2] + mu*torque[0] + mu*torque[1] + torque[2] <= 0).evaluator().set_description(f\"yaw 5 {foot}\")\n",
    "    prog.AddConstraint(Y*force[0] - X*force[1] -(X+Y)*mu*force[2] + mu*torque[0] - mu*torque[1] + torque[2] <= 0).evaluator().set_description(f\"yaw 6 {foot}\")\n",
    "    prog.AddConstraint(-Y*force[0] + X*force[1] -(X+Y)*mu*force[2] - mu*torque[0] + mu*torque[1] + torque[2] <= 0).evaluator().set_description(f\"yaw 7 {foot}\")\n",
    "    prog.AddConstraint(-Y*force[0] - X*force[1] -(X+Y)*mu*force[2] - mu*torque[0] - mu*torque[1] + torque[2] <= 0).evaluator().set_description(f\"yaw 8 {foot}\")\n",
    "\n",
    "\n",
    "CWC_constraint(f_x_L, tau_L, \"foot_L\")\n",
    "CWC_constraint(f_x_R, tau_R, \"foot_R\")\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",
    "# Foot position constraints\n",
    "for index, constraints in enumerate(pos_constraint):\n",
    "    prog.AddConstraint(constraints, pos).evaluator().set_description(f\"foot position {index}\")\n",
    "\n",
    "w_CoM = 0.1\n",
    "prog.AddCost(CostTorque, vars=torques)\n",
    "prog.AddCost(CostCoM, vars=pos)\n",
    "\n",
    "fc_guess = np.zeros(n)\n",
    "fc_guess = -plant.gravity_field().gravity_vector() * plant.CalcTotalMass(context) / num_contact\n",
    "fc_guess = np.tile(fc_guess, num_contact)\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/Double_Leg/{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 to still enable simulation\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.SetPositionsAndVelocities(context, x0)\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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"2123pt\" height=\"495pt\" viewBox=\"0.00 0.00 2123.00 495.00\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 491)\">\n<title>_62997856</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-491 2119,-491 2119,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster62997856diagram</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"8,-8 8,-479 2107,-479 2107,-8 8,-8\"/>\n<text text-anchor=\"middle\" x=\"1057.5\" y=\"-463.8\" font-family=\"Times,serif\" font-size=\"14.00\">62997856</text>\n</g>\n<g id=\"clust4\" class=\"cluster\">\n<title>cluster62997856subsystems</title>\n<polygon fill=\"none\" stroke=\"white\" points=\"16,-16 16,-448 2099,-448 2099,-16 16,-16\"/>\n</g>\n<!-- 64309840 -->\n<g id=\"node1\" class=\"node\">\n<title>64309840</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"502,-25 502,-439 1124,-439 1124,-25 502,-25\"/>\n<text text-anchor=\"middle\" x=\"813\" y=\"-423.8\" font-family=\"Times,serif\" font-size=\"14.00\">plant</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-416 1124,-416 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-384.8\" font-family=\"Times,serif\" font-size=\"14.00\">geometry_query</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-361 750,-361 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-329.3\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-305 750,-305 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-273.3\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-249 750,-249 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-217.3\" font-family=\"Times,serif\" font-size=\"14.00\">Double_Leg_Collision_actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-193 750,-193 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-161.3\" font-family=\"Times,serif\" font-size=\"14.00\">actuation</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-137 750,-137 \"/>\n<text text-anchor=\"middle\" x=\"626\" y=\"-105.3\" font-family=\"Times,serif\" font-size=\"14.00\">applied_generalized_force</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"502,-81 750,-81 \"/>\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,-25 750,-416 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-400.8\" font-family=\"Times,serif\" font-size=\"14.00\">geometry_pose</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-393 1124,-393 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-377.8\" font-family=\"Times,serif\" font-size=\"14.00\">continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-370 1124,-370 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-354.8\" font-family=\"Times,serif\" font-size=\"14.00\">body_poses</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-347 1124,-347 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-331.8\" font-family=\"Times,serif\" font-size=\"14.00\">spatial_velocities</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-324 1124,-324 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-308.8\" font-family=\"Times,serif\" font-size=\"14.00\">spatial_accelerations</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-301 1124,-301 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-285.8\" font-family=\"Times,serif\" font-size=\"14.00\">generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-278 1124,-278 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-262.8\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-255 1124,-255 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-239.8\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-232 1124,-232 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-216.8\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-209 1124,-209 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-193.8\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-186 1124,-186 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-170.8\" font-family=\"Times,serif\" font-size=\"14.00\">Double_Leg_Collision_continuous_state</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-163 1124,-163 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-147.8\" font-family=\"Times,serif\" font-size=\"14.00\">Double_Leg_Collision_generalized_acceleration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-140 1124,-140 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-124.8\" font-family=\"Times,serif\" font-size=\"14.00\">WorldModelInstance_generalized_contact_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-117 1124,-117 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-101.8\" font-family=\"Times,serif\" font-size=\"14.00\">DefaultModelInstance_generalized_contact_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-94 1124,-94 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-78.8\" font-family=\"Times,serif\" font-size=\"14.00\">Double_Leg_Collision_generalized_contact_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-71 1124,-71 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">reaction_forces</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"750,-48 1124,-48 \"/>\n<text text-anchor=\"middle\" x=\"937\" y=\"-32.8\" font-family=\"Times,serif\" font-size=\"14.00\">contact_results</text>\n</g>\n<!-- 64296912 -->\n<g id=\"node2\" class=\"node\">\n<title>64296912</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"1290,-368.5 1290,-437.5 1503,-437.5 1503,-368.5 1290,-368.5\"/>\n<text text-anchor=\"middle\" x=\"1396.5\" y=\"-422.3\" font-family=\"Times,serif\" font-size=\"14.00\">scene_graph</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1290,-414.5 1503,-414.5 \"/>\n<text text-anchor=\"middle\" x=\"1367.5\" y=\"-399.3\" font-family=\"Times,serif\" font-size=\"14.00\">plant_pose</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1290,-391.5 1445,-391.5 \"/>\n<text text-anchor=\"middle\" x=\"1367.5\" y=\"-376.3\" font-family=\"Times,serif\" font-size=\"14.00\">plant_configuration</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1445,-368.5 1445,-414.5 \"/>\n<text text-anchor=\"middle\" x=\"1474\" y=\"-387.8\" font-family=\"Times,serif\" font-size=\"14.00\">query</text>\n</g>\n<!-- 64309840&#45;&gt;64296912 -->\n<g id=\"edge3\" class=\"edge\">\n<title>64309840:y0-&gt;64296912:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1124,-405C1193.55,-405 1213.47,-403.19 1278.17,-403.01\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"1278.5,-406.51 1288.5,-403 1278.5,-399.51 1278.5,-406.51\"/>\n</g>\n<!-- 62965728 -->\n<g id=\"node5\" class=\"node\">\n<title>62965728</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"1160,-24.5 1160,-93.5 1633,-93.5 1633,-24.5 1160,-24.5\"/>\n<text text-anchor=\"middle\" x=\"1396.5\" y=\"-78.3\" font-family=\"Times,serif\" font-size=\"14.00\">Custom_LeafSystems.ReadContactResults@0000000003c0c7e0</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1160,-70.5 1633,-70.5 \"/>\n<text text-anchor=\"middle\" x=\"1292.5\" y=\"-43.8\" font-family=\"Times,serif\" font-size=\"14.00\">contact_results_input</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1425,-24.5 1425,-70.5 \"/>\n<text text-anchor=\"middle\" x=\"1529\" y=\"-55.3\" font-family=\"Times,serif\" font-size=\"14.00\">contact_force</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1425,-47.5 1633,-47.5 \"/>\n<text text-anchor=\"middle\" x=\"1529\" y=\"-32.3\" font-family=\"Times,serif\" font-size=\"14.00\">is_contact</text>\n</g>\n<!-- 64309840&#45;&gt;62965728 -->\n<g id=\"edge2\" class=\"edge\">\n<title>64309840:y16-&gt;62965728:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1124,-36C1136.68,-36 1141.21,-42.32 1149.94,-45.38\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"1149.57,-48.86 1160,-47 1150.69,-41.95 1149.57,-48.86\"/>\n</g>\n<!-- 64296912&#45;&gt;64309840 -->\n<g id=\"edge5\" class=\"edge\">\n<title>64296912:y0-&gt;64309840:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1474.5,-367C1474.5,-285 1370.27,-351.22 1290,-368 1211.37,-384.44 1202.68,-423.28 1124,-439.5 1056.31,-453.46 551.35,-487.88 502,-439.5 488.98,-426.73 481.64,-401.37 491.38,-392.31\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"492.69,-395.57 501,-389 490.41,-388.95 492.69,-395.57\"/>\n</g>\n<!-- 62434208 -->\n<g id=\"node3\" class=\"node\">\n<title>62434208</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"1669,-380 1669,-426 2091,-426 2091,-380 1669,-380\"/>\n<text text-anchor=\"middle\" x=\"1880\" y=\"-410.8\" font-family=\"Times,serif\" font-size=\"14.00\">drake/geometry/MeshcatVisualizer@0000000003b8aba0</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1669,-403 2091,-403 \"/>\n<text text-anchor=\"middle\" x=\"1796\" y=\"-387.8\" font-family=\"Times,serif\" font-size=\"14.00\">query_object</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1923,-380 1923,-403 \"/>\n<text text-anchor=\"middle\" x=\"2007\" y=\"-387.8\" font-family=\"Times,serif\" font-size=\"14.00\"> </text>\n</g>\n<!-- 64296912&#45;&gt;62434208 -->\n<g id=\"edge1\" class=\"edge\">\n<title>64296912:y0-&gt;62434208:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1504.5,-391C1574.04,-391 1593.97,-391 1658.67,-391\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"1659,-394.5 1669,-391 1659,-387.5 1659,-394.5\"/>\n</g>\n<!-- 34857520 -->\n<g id=\"node4\" class=\"node\">\n<title>34857520</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"24,-154 24,-200 466,-200 466,-154 24,-154\"/>\n<text text-anchor=\"middle\" x=\"245\" y=\"-184.8\" font-family=\"Times,serif\" font-size=\"14.00\">drake/systems/ConstantVectorSource@000000000213e230</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"24,-177 466,-177 \"/>\n<text text-anchor=\"middle\" x=\"131\" y=\"-161.8\" font-family=\"Times,serif\" font-size=\"14.00\"> </text>\n<polyline fill=\"none\" stroke=\"black\" points=\"238,-154 238,-177 \"/>\n<text text-anchor=\"middle\" x=\"352\" y=\"-161.8\" font-family=\"Times,serif\" font-size=\"14.00\">y0</text>\n</g>\n<!-- 34857520&#45;&gt;64309840 -->\n<g id=\"edge6\" class=\"edge\">\n<title>34857520:y0-&gt;64309840:u4</title>\n<path fill=\"none\" stroke=\"black\" d=\"M466,-165C478,-165 483.25,-165 491.88,-165\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"492,-168.5 502,-165 492,-161.5 492,-168.5\"/>\n</g>\n<!-- 64304960 -->\n<g id=\"node6\" class=\"node\">\n<title>64304960</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"1687,-25 1687,-71 2073,-71 2073,-25 1687,-25\"/>\n<text text-anchor=\"middle\" x=\"1880\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">drake/systems/VectorLogSink@0000000003d53740</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1687,-48 2073,-48 \"/>\n<text text-anchor=\"middle\" x=\"1790\" y=\"-32.8\" font-family=\"Times,serif\" font-size=\"14.00\">data</text>\n<polyline fill=\"none\" stroke=\"black\" points=\"1893,-25 1893,-48 \"/>\n<text text-anchor=\"middle\" x=\"1983\" y=\"-32.8\" font-family=\"Times,serif\" font-size=\"14.00\"> </text>\n</g>\n<!-- 62965728&#45;&gt;64304960 -->\n<g id=\"edge4\" class=\"edge\">\n<title>62965728:y1-&gt;64304960:u0</title>\n<path fill=\"none\" stroke=\"black\" d=\"M1633,-36C1652.69,-36 1660.04,-36 1675.68,-36\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"1676,-39.5 1686,-36 1676,-32.5 1676,-39.5\"/>\n</g>\n</g>\n</svg>",
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 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()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "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(10)\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",
    "# TODO: Read sources from Shivesh\n",
    "# TODO: 6 DOF wrench?\n",
    "# TODO: Apply Zero Moment Point principle to optimization\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
}