Tut 3 - Contacts.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 3 - Contacts\n",
    "**Aim:** To finally make a friggin' leg! Aw yisss. But also to properly confront the problem of hard contacts, and what we do about them.\n",
    "\n",
    "**Further reading:** There are two important papers I'll be referring to, but I'd suggest rather reading them afterwards:\n",
    "1. <a href=\"https://jeb.biologists.org/content/jexbio/202/23/3325.full.pdf\"> Templates and Anchors: Neuromechanical Hypotheses of Legged Locomotion on Land - RJ Full and DE Koditschek (1999) </a>\n",
    "2. <a href=https://journals.sagepub.com/doi/pdf/10.1177/0278364913506757> A direct method for trajectory optimization of rigid bodies through contact - M Posa, C Cantu and R Tedrake (2014) </a>\n",
    "\n",
    "If you have a bit more time and want to immerse yourself in one of the religious texts of legged robotics, check out *Legged Robots that Balance* by Marc Raibert. (The dude founded Boston Dynamics, so I reckon he knows what he's talking about.)  \n",
    "\n",
    "## Templates and Anchors\n",
    "The model we'll be constructing consists of one leg (two links connected by a prismatic joint) and a body:\n",
    "\n",
    "<img src = \"hopper.png\" width = \"400\">\n",
    "\n",
    "A monopedal hopper probably isn't the most useful configuration in real life, but it's still an important model that helps biologists and roboticists alike to understand legged locomotion.\n",
    "\n",
    "Considering that most of the macroscopic, multicellular life forms on Earth are insects, the average number of legs that an animal has is... in my estimation... more than one. So you may be wondering how a one-legged model can tell us anything useful about any real critter. But imagine somebody told you to come up with the shortest possible \"executive summary\" of the running motion of, say, a cheetah. (It's always a cheetah... ;P) The first thing you'd do is remove all the non-essential bits: so head, flexy torso and tail become just one rigid body. Or maybe even a point mass. And then you notice that only one leg is ever really in contact with the ground at a time, and hey, even if more than one was in contact, there is only one centre of pressure, so why do we even need more than one 'leg?\\*'... And soon enough, you're left with a hopper. The outcome would be the same if you tried it for a human, or a cockroach, or pretty much any other legged animal you could think of (except maybe a millipede... I suspect locomotion fundamentals change when the number of legs hits double figures...) \n",
    "\n",
    "\\* this single \"leg\" (well, limb betweent the body and centre of pressure) used to represent the action of more than one leg is referred to as a *virtual leg.*\n",
    "\n",
    "In biomechanics, this \"executive summary\" model for a motion - i.e. the lowest degree-of-freedom model that still captures its essential characteristics - is called a *template*. More complex models are called *anchors*, since they anchor the template behaviour into the more realistic, detailed behaviour of the system. In addition to the obvious simplification benefit, templates are beneficial because they allow us to compare the motion of systems with widely differing morphologies.  \n",
    "\n",
    "...So that's why we care about monopeds: from Raibert's bouncing robots to the spring-loaded inverted pendulum (SLIP) models frequently applied to represent running, it's clear they are a good template for legged locomotion. \n",
    "\n",
    "**The takeaway:** Rather than just giving you some cool biomechanical jargon to throw around, the thing I want you to get out of this template idea is that dynamic modelling should be a modular process. Trajectory optimization is extremely computationally demanding, so we have to be as economical as possible with our degrees of freedom. Being efficient with your coordinates can make the difference between finding a solution in 15 minutes, or three hours.\n",
    "\n",
    "Of course, that doesn't mean we should never use high-order models: anchors are absolutely essential to understanding motion. It just means you should get the simplest model working first, and then build complexity slowly and iteratively from there."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hopper Equations of Motion\n",
    "Let's quickly get the dynamics of the hopper out of the way. We'll use 5 generalized coordinates: the position and angle of the body, the hip angle, and the length of the prismatic joint at the knee. We'll add actuator forces at both joints, and ground reaction forces at the foot. The activation of these forces will be controlled with complementarity constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DERIVE THE EOMs SYMBOLICALLY ----------------------------------------------------------------------------------------------\n",
    "\n",
    "# import libraries\n",
    "import sympy as sym\n",
    "import numpy as np\n",
    "\n",
    "sym.init_printing()\n",
    "from IPython.display import display #for pretty printing\n",
    "\n",
    "# create symbolic variables\n",
    "\n",
    "# system parameters\n",
    "g = sym.symbols('g')\n",
    "mb,ml1,ml2 = sym.symbols(['m_{body}','m_{leg1}','m_{leg2}']) # mass\n",
    "lb,ll1,ll2 = sym.symbols(['l_{body}','l_{leg1}','l_{leg2}']) # length\n",
    "Inb,Inl1,Inl2 = sym.symbols(['I_{body}','I_{leg1}','I_{leg2}']) # moment of intertia\n",
    "\n",
    "# generalized coordinates\n",
    "x,y,thb,thl,r = sym.symbols(['x','y','\\\\theta_{body}','\\\\theta_{leg}','r']) \n",
    "dx,dy,dthb,dthl,dr = sym.symbols(['\\dot{x}','\\dot{y}','\\dot{\\\\theta}_{body}','\\dot{\\\\theta}_{leg}','\\dot{r}']) \n",
    "ddx,ddy,ddthb,ddthl,ddr = sym.symbols(['\\ddot{x}','\\ddot{y}','\\ddot{\\\\theta}_{body}','\\ddot{\\\\theta}_{leg}','\\ddot{r}']) \n",
    "\n",
    "q = sym.Matrix([[x],[y],[thb],[thl],[r]])\n",
    "dq = sym.Matrix([[dx],[dy],[dthb],[dthl],[dr]])\n",
    "ddq = sym.Matrix([[ddx],[ddy],[ddthb],[ddthl],[ddr]])\n",
    "\n",
    "# forces\n",
    "# total joint action = actuator + rebound, but that will be dealt with elsewhere\n",
    "F,tau,GRFx,GRFy = sym.symbols(['F','\\\\tau','G_x','G_y']) \n",
    "\n",
    "# STEP 1: position vectors ri = [x,y,theta] (world frame)\n",
    "rb = sym.Matrix([[x],\n",
    "                [y],\n",
    "                [thb]])\n",
    "\n",
    "rl1 = sym.Matrix([[x + 0.5*ll1*sym.sin(thb + thl)],\n",
    "                [y - 0.5*ll1*sym.cos(thb + thl)],\n",
    "                [thb + thl]])\n",
    "\n",
    "rl2 = sym.Matrix([[x + (0.5*ll1+r)*sym.sin(thb + thl)],\n",
    "                [y - (0.5*ll1+r)*sym.cos(thb + thl)],\n",
    "                [thb + thl]])\n",
    "\n",
    "# the Jacobians\n",
    "Jb = rb.jacobian(q)\n",
    "Jl1 = rl1.jacobian(q)\n",
    "Jl2 = rl2.jacobian(q)\n",
    "\n",
    "# STEP 2: generate expressions for the system space velocities from the jacobians\n",
    "vb = Jb*dq\n",
    "vl1 = Jl1*dq\n",
    "vl2 = Jl2*dq\n",
    "\n",
    "# STEP 3: generate expressions for the kinetic and potential energy\n",
    "# mass vectors\n",
    "Mb = sym.Matrix([[mb,mb,Inb]])\n",
    "Ml1 = sym.Matrix([[ml1,ml1,Inl1]])\n",
    "Ml2 = sym.Matrix([[ml2,ml2,Inl2]])\n",
    "\n",
    "T = 0.5*Mb*sym.matrix_multiply_elementwise(vb,vb) + 0.5*Ml1*sym.matrix_multiply_elementwise(vl1,vl1) + 0.5*Ml2*sym.matrix_multiply_elementwise(vl2,vl2)\n",
    "T = T[0]\n",
    "V = mb*g*rb[1] + ml1*g*rl1[1] + ml2*g*rl2[1]\n",
    "\n",
    "\n",
    "# STEP 4: calculate each term of the Lagrange equation\n",
    "# term 1\n",
    "Lg1 = sym.zeros(len(q),1)\n",
    "for i in range(len(q)):\n",
    "    dT_ddq = sym.Matrix([sym.diff(T,dq[i])]) # get partial of T in dq_i\n",
    "    Lg1[i] = dT_ddq.jacobian(q)*dq + dT_ddq.jacobian(dq)*ddq #...then get time derivative of that partial\n",
    "\n",
    "# term 3\n",
    "Lg3 = sym.Matrix([T]).jacobian(q).transpose() # partial of T in q\n",
    "\n",
    "# term 4\n",
    "Lg4 = sym.Matrix([V]).jacobian(q).transpose() # partial of U in q\n",
    "\n",
    "# STEP 5: generalized forces\n",
    "\n",
    "# force vectors for each link\n",
    "tau_b = sym.Matrix([[0],[0],[-tau]])\n",
    "tau_l1 = sym.Matrix([[0],[0],[tau]])\n",
    "\n",
    "F_l1 = sym.Matrix([[-F*sym.sin(thb+thl)],[F*sym.cos(thb+thl)],[0]])\n",
    "F_l2 = sym.Matrix([[F*sym.sin(thb+thl)],[-F*sym.cos(thb+thl)],[0]])\n",
    "\n",
    "GRF_l2 = sym.Matrix([[GRFx],[GRFy],[0.5*ll2*GRFx*sym.cos(thb+thl)+0.5*ll2*GRFy*sym.sin(thb+thl)]])\n",
    "\n",
    "Q = sym.zeros(len(q),1)\n",
    "for j in range(len(q)):\n",
    "    Q[j] = tau_b.transpose()*Jb[:,j]+(tau_l1+F_l1).transpose()*Jl1[:,j]+(F_l2+GRF_l2).transpose()*Jl2[:,j]\n",
    "\n",
    "# AND combine!\n",
    "EOM = Lg1 - Lg3 + Lg4 - Q\n",
    "\n",
    "EOMs = sym.zeros(len(q),1)\n",
    "for j in range(len(q)):\n",
    "    EOMs[j] = EOM[j].simplify()\n",
    "    \n",
    "# the simplification step is a little time-consuming so try to avoid re-running this cell if possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lambdify\n",
    "from pyomo.environ import*\n",
    "from pyomo.opt import SolverFactory\n",
    "from pyomo.opt import SolverStatus, TerminationCondition\n",
    "\n",
    "func_map = {'sin':sin, 'cos':cos} \n",
    "\n",
    "sym_list = [g,mb,ml1,ml2,lb,ll1,ll2,Inb,Inl1,Inl2,\n",
    "            x,y,thb,thl,r,\n",
    "            dx,dy,dthb,dthl,dr,\n",
    "            ddx,ddy,ddthb,ddthl,ddr,\n",
    "            F,tau,GRFx,GRFy]\n",
    "\n",
    "lambEOM = {}\n",
    "\n",
    "DOFs = ['x','y','thb','thl','r']\n",
    "\n",
    "for dof_i, dof in enumerate(DOFs):\n",
    "    lambEOM[dof] = sym.lambdify(sym_list,EOMs[dof_i],modules = [func_map])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Protip: lambdify your aux variable definitions\n",
    "As you add links to the leg, the expressions for the positions, and especially the velocities of the contacts get very unweildly very quickly. Luckily, the LORD hath given unto us sympy, that we need fear no kerfuffle. I lambdify the expressions for the positions and velocities of the contact points to make my life easier when I construct the relevant constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TDOFs = ['x','y'] # translational DOFs - i.e. the ones that the foot has\n",
    "\n",
    "# foot position\n",
    "pfoot = sym.Matrix([[x + (0.5*ll1+r+0.5*ll2)*sym.sin(thb + thl)],\n",
    "                    [y - (0.5*ll1+r+0.5*ll2)*sym.cos(thb + thl)]])\n",
    "# foot velocity\n",
    "vfoot = pfoot.jacobian(q)*dq\n",
    "\n",
    "lamb_pfoot = {}\n",
    "lamb_vfoot = {}\n",
    "\n",
    "for dof_i, dof in enumerate(TDOFs):\n",
    "    lamb_pfoot[dof] = sym.lambdify(sym_list,pfoot[dof_i],modules = [func_map])\n",
    "    lamb_vfoot[dof] = sym.lambdify(sym_list,vfoot[dof_i],modules = [func_map])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Pyomo Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "# Pyomo stuff\n",
    "from pyomo.environ import*\n",
    "from pyomo.opt import SolverFactory\n",
    "from pyomo.opt import SolverStatus, TerminationCondition\n",
    "import numpy as np\n",
    "\n",
    "from IPython.display import display #for pretty printing\n",
    "\n",
    "# create the model\n",
    "m = ConcreteModel()\n",
    "\n",
    "# SETS-----------------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# for sprint task\n",
    "N = 100 \n",
    "\n",
    "# for sanity checks\n",
    "# N = 20\n",
    "m.N = RangeSet(N) \n",
    "\n",
    "# sets can have multidimensional indices\n",
    "# (probably a little gratuitous for such a simple model, but thought I'd show you anyway)\n",
    "links = [('body',1),('leg',1),('leg',2)]\n",
    "m.L = Set(dimen=2, initialize = links)\n",
    "\n",
    "m.DOF = Set(initialize = DOFs) # the coordinates for each link\n",
    "m.TDOF = Set(initialize = TDOFs) # world-frame coordinates for contact variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***Annoying sidenote on the subject of set creation**\n",
    "\n",
    "When you create a set, always do it like this:\n",
    "\n",
    "```python\n",
    "SpiceGirls = ['ginger','baby','posh','scary','sporty']\n",
    "m.sg = Set(initialize = SpiceGirls)\n",
    "```\n",
    "or perhaps like this:\n",
    "\n",
    "```python\n",
    "m.sg = Set(initialize = ['ginger','baby','posh','scary','sporty'])\n",
    "```\n",
    "(though the first one is nicer, because having a list of the names it makes it easier to iterate over members of that set later if you have to.)\n",
    "\n",
    "But **never** like this:\n",
    "\n",
    "```python\n",
    "m.sg = Set(['ginger','baby','posh','scary','sporty'])\n",
    "```\n",
    "Daft as it is, that last one actually creates five empty *sets*, instead of what you want: a single set with five values.\n",
    "\n",
    "You can probably get away with it in many circumstances, but attempting to index using a set obviously makes absolutely no sense, so you're going to run into an error sooner or later. (If you ever come across the error **\"cannot index a component with an indexed set\"**, this may be why."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PARAMETERS-----------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "m.g = Param(initialize = 9.81)\n",
    "\n",
    "# whenever a multidimensional set is passed to a function, the index is expanded (infuriatingly enough...)\n",
    "# so every function indexed on L needs to expect two inputs from that set \n",
    "# I've called them l = [lb, ln] for 'branch' and 'number'\n",
    "\n",
    "def get_m(n, lb, ln):\n",
    "    if lb == 'body':\n",
    "        return 0.5\n",
    "    else: return 0.25\n",
    "# note that the masses add up to 1\n",
    "m.m = Param(m.L, initialize = get_m) # mass of links\n",
    "\n",
    "def get_len(n, lb, ln):\n",
    "    if lb == 'body':\n",
    "        return 1.0\n",
    "    else: return 0.5\n",
    "m.len = Param(m.L, initialize = get_len) # length of links\n",
    "\n",
    "def calculate_In(m, lb, ln): \n",
    "    l = (lb,ln)\n",
    "    # yes, that does mean you have to rebuild the tuple inside the function.\n",
    "    return m.m[l]*m.len[l]**2/12 \n",
    "m.In = Param(m.L, initialize = calculate_In) # moment of inertia\n",
    "\n",
    "# VARIABLES -----------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# system coordinates\n",
    "m.q = Var(m.N, m.DOF) # position\n",
    "m.dq = Var(m.N, m.DOF) # velocity\n",
    "m.ddq = Var(m.N, m.DOF) # acceleration\n",
    "\n",
    "# joint forces\n",
    "joints = ['hip','knee']\n",
    "m.J = Set(initialize = joints)\n",
    "m.fj = Var(m.N, m.J) # net force at each joint\n",
    "\n",
    "# ground reaction forces\n",
    "m.GRF = Var(m.N, m.TDOF)\n",
    "\n",
    "# bound variables\n",
    "for n in range(1,N+1):\n",
    "    for l in links:\n",
    "        m.q[n,'y'].setlb(0.0)\n",
    "        m.q[n,'thb'].setlb(-0.5*np.pi)\n",
    "        m.q[n,'thb'].setub(0.5*np.pi)\n",
    "        m.q[n,'thl'].setlb(-0.5*np.pi)\n",
    "        m.q[n,'thl'].setub(0.5*np.pi)\n",
    "        m.q[n,'r'].setlb(0)\n",
    "        m.q[n,'r'].setub(0.5)\n",
    "        \n",
    "    m.fj[n,'hip'].setlb(-2)\n",
    "    m.fj[n,'hip'].setub(2)\n",
    "    m.fj[n,'knee'].setlb(-10)\n",
    "    m.fj[n,'knee'].setub(10)\n",
    "        \n",
    "    m.GRF[n,'y'].setlb(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contacts\n",
    "The problems at the fixed point in the last tutorial started to reveal some of the difficulties around formulating hard boundaries in an optimizaion problem. This tutorial is going to elaborate on those issues and explain the workarounds we use to overcome them. Well, maybe \"overcome\" is a strong word... often the best we can do is just making the situation less dire so it can actually solve within a human lifetime and only break for half the random seeds we throw at it instead of most of them.\n",
    "\n",
    "I'll get more specific as we go, but as a general overview, contacts are difficult to deal with for two primary reasons:\n",
    "\n",
    "**They're discontinuous.** Imagine trying to draw a freebody diagram representing the forces acting on a four-legged animal: you'd have the weight, the torques acting at each of its joints, the normal force on each of its paws... But only if it actually happened to be standing on all four paws. It could have just one foot on the ground, or two, or three, or *the other three*, or none at all or... you get the point. Changing the contact state essentially changes the dynamic model you're working with by switching different forces 'off' or 'on'. \n",
    "\n",
    "There is an approach called *hybrid modelling* that makes this explicit: each contact state gets its own dedicated dynamic model with switches between models happening at specified events, such as foot touch-downs and lift-offs. But this isn't suitable for our needs because:\n",
    "1. The number of discrete models is prone to explode. For a biped, you need four: right foot down, left foot down, neither, both. For a quadruped, you need sixteen. And that's just assuming the foot sticks in place and doesn't slide. \n",
    "2. You need to know the contact order up front. Not a problem if you're modelling a well-established motion pattern like a gallop, but our research questions are often around less-defined manoeuvres like rapid stops or turns, or involve finding out the best movement to achieve something, and therefore requires that the model isn't limited to a specific contact order.\n",
    "\n",
    "...So we need a single model that can apply the right combination of forces under the right conditions. If I had to ask you to describe such a thing using pseudocode or a flow chart, it would probably involve a lot of IF statements:\n",
    "\n",
    "```python\n",
    "if right_foot_contact:\n",
    "    right_GRF_active = True\n",
    "    ```\n",
    "...that sort of thing. Problem is, NLP solvers are always (at least, to my knowledge) gradient-based to some extent, and like nice, smooth derivatives they can slide down to an answer. Conditionals like that create sharp cliffs that they absolutely can't deal with. To get around this, we cheat using something called *complementarity*... but we'll get to that just now.\n",
    "\n",
    "**They're fast.** Which brings me to... \n",
    "\n",
    "### Contacts and timing\n",
    "Hard boundaries themselves aren't really the problem - it's actually the dynamics of collisions that's tricky. Hard boundaries merely create a site where collisions must occur.\n",
    "\n",
    "Collisions are hard because they're *impulsive events:* in continuous time world, their dynamics involve very large forces that act over infinitesimally short time intervals. But in trajectory optimization world, there's a finite minimum timestep over which forces must act. \n",
    "\n",
    "Another way to think about it, if you'll allow me a moment to put my Signals and Systems hat on, is that we're attempting to deal with high frequency signals within the confines of a slow sample rate. An IF statement is just a step function, and think of all the high-frequency harmonics you need to make a step look truly *crispy*... Now imagine what it would look like aliased to the 50 - 100 Hz sample rate we typically limit ourselves to. Nyquist is shaking.\n",
    "\n",
    "So... why not just sample faster? To quote an ancient meme proverb: \"Ain't nobody got time for that.\" Shorter timestep = more nodes to cover the same duration = more variables and constraints = huge problem size = centuries of solving time.\n",
    "\n",
    "Ideally, what you want is a slow sample rate for the continuous parts and a fast sample rate for when collisions happen. And that's exactly what we do! Within limitations. \n",
    "\n",
    "But let's start with the good news: a variable timestep allows the flexible sampling rate we desire, and also allows collisions to occur more naturally. See, the model can only change states at the boundaries of a timestep, meaning that if you're using a fixed step of say 0.01 seconds, the foot can only touch down or the joint hit its limit or whatever at $t=$ some precise multiple of 0.01. Putting some bend in the timestep theoretically allows it to stay in a state for an arbitrary duration.\n",
    "\n",
    "Unfortunately, there is only so much bend we can allow. Introducing a variable timestep makes the problem even more nonlinear, and we have been warned by Prof Larry Biegler (the First of his Name, High Archmage of trajectory optimization, Coder of IPOPT) that allowing it too much flexibility can cancel out the benefits and make the problem harder to solve. Based on his rule of thumb, we usually only allow it to vary within 20% of a maximum timestep.\n",
    "_______\n",
    "## The scaling problem\n",
    "If you do any reading on the topic of optimization, you'll probably hear about the *conditioning* of problems. A *well-conditioned* optimization problem is something an NLP solver is well-equipped to deal with: we've already mentioned that solvers like smooth, convex functions. Another thing they really like is a *well-scaled* problem: that is, one where all the variables and constraint derivatives are a similar order of magnitude.\n",
    "\n",
    "Alas, we find ourselves in an especially difficult predicament here. If you had to describe the worst possible system to do trajectory optimization on... The 9th Circle of Badly-Conditioned Hell... whatever you come up with would probably be pretty close to the systems we work with.\n",
    "\n",
    "Consider the velocity integration constraint. It includes:\n",
    "* velocities with magnitudes in 1's or 10's\n",
    "* accelerations that might have magnitudes similar to the velocities, but can also have magitudes in the 100's, 1000's or even 10 000's when collisions occur\n",
    "* our now-variable timestep in the order 0.01\n",
    "\n",
    "In combination with the horribly infeasible starting points our random initialization method creates, it's not surprising that these problems take ages to solve, frequently come back infeasible, or crash the algorithm entirely. Even once your model has passed a bunch of sanity checks and been debugged to the extent that you're confident all the equations, etc. are fundamentally correct, you might still hit the dreaded *restoration failed* error. I don't want to get too in-depth into the technical details of IPOPT, but when you see that, what's happened is the optimization gnome tasked with solving your problem has had a breakdown, given up, flipped his desk and stormed out, viciously cursing you, your robot, your ancestors, and the entire science of mathematics all the while.\n",
    "\n",
    "### Scaling - an overview\n",
    "One of the things we can do to make our poor gnome's life easier is attempt to scale our problem better. The inherent properties of our system mean its probably never going to be well-scaled, but a bit of tinkering can make it slightly less than a complete trash fire.\n",
    "\n",
    "One of my few complaints about Pyomo is that it doesn't have an inbuilt scaling feature, meaning we have to introduce scale factors manually. This can be a good thing, since it's much more flexible, and you know exactly how the scaling is implemented, but it does introduce the potential for human error so you need to keep your wits about you.\n",
    "\n",
    "Scaling is yet another one of those *art* things (I'm sure you're getting really tired of those by now) so I don't really have a precise, step-by-step theory of scaling you can follow. The best I can do is tell you the general guidelines and explain a few specific cases as I go along. \n",
    "\n",
    "Here's what we're trying to achieve by scaling:\n",
    "1. the values of all the variables in the solution should be around 1 (\"around 1\" here meaning order of magnitude 1 ideally, but somewhere between 0.01 and 100 pragmatically.)\n",
    "2. The values of all constraint derivatives should be around 1. For linear constraints, this is as simple as making sure all the coefficients are within the established \"around 1\" range. For nonlinear constraints, it's much trickier, since the \"coefficients\" (when you think of the partial derivative in a particular variable) become functions of the variables. \n",
    "3. The value of each term in a constraint equation should be around 1.\n",
    "\n",
    "With our problems, it's almost impossible to satisfy all of these at once. You might have to make trade-offs. For variables with wide ranges, work with the mean values you expect them to take on and accept that the nodes where the extreme values occur will be nasty. Most of your variables shouldn't need scaling: start with the most egregious cases (e.g. the timestep or the contact forces, which tend to have magnitudes furthest outside these bounds) and go from there.  \n",
    "\n",
    "This is all complicated further by the inescapable tangle of interdependency: how well your model works is tightly tied up with the initial guess, the timestep and the duration of the simulation. Finding the right balance of all these things is the process of *conditioning* your model a.k.a. all the kerfuffle that comes after writing and basic sanity checks, but before the model is ready to be rolled out for your experiment. Conditioning is the worst, most rage-inducing part of trajectory optimization, and often takes far longer than actually writing the model. My approach to it is to try to separate each of the aspects into a separate step:\n",
    "1. **scaling:** I choose whichever scale factors seem logical based on the stated guidelines, adding them iteratively and keeping those that don't appear to actively worsen the performance. Normalization is your friend: you're much more likely to come out with a well-scaled model if you work in units of body mass (i.e. defining the mass of the limb links by the total mass such that the combined mass of the model comes out as 1) than trying to work with measurements in kilograms. Testing the scaling on your target problem can be tricky, since it might not solve at all until you've tweaked the timing, so use a sanity check you know it can pass already. (I provide a few examples at the end of this notebook). The `m.pprint()` readout is super helpful for pinpointing the variables and equations that need the most work.\n",
    "\n",
    "2. **timing:** Now you're ready to try your target problem. You probably have some idea of how long the motion should take, so start with that. If it doesn't solve, add or remove nodes based on whether $h$ is tending to hit the upper or lower bound. (Technically, you could change the duration by changing the timestep, but I find that causes scaling issues, so I prefer to keep $h_m$ somewhere between 0.01 and 0.02 seconds and adjust the duration using the number of nodes instead.)\n",
    "\n",
    "3. **initialization:** This is where the model gets most temperamental, so I wouldn't recommend even touching this step until you have a solution converged with the default initial values. Before wasting time here, you need to know that the motion is achievable and how long it takes. Remember: unless you're *warm-starting* from a previous solution, whatever you add here is likely to be less feasible than the default input - if it can't solve from the default point, it definitely can't solve from whatever weird point you inflict on it. \n",
    "\n",
    "You could really keep fiddling forever, so it's important to know when to stop: you're not going for *perfect*, or even *good* - you're going for *usable*.\n",
    "\n",
    "### Scaling example 1: velocity integration\n",
    "We want a master timestep $h_m = 0.02$, meaning the variable timestep has a range $0.016 \\leq h[n] \\leq 0.02$. We could just create an $h$ variable with those as the bounds, but it fits much better into the Around 1 Zone if we think of the timestep as $h_m h[n]$, since the variable part then has the range $0.8 \\leq h[n] \\leq 1.0$.\n",
    "\n",
    "This means that $h_m$ becomes our scaling factor for $h$: wherever we use it in an equation, it will be as the product of the scalar part and the variable part. The integration constraint therefore becomes: $$\\dot{x}[n] = \\dot{x}[n-1] + 0.02h[n]*\\ddot{x}[n]$$\n",
    "\n",
    "The typical velocity and acceleration values should be around 10, maybe around 100 for the latter, so the partial derivatives and values of individual terms should come out Around 1 ($10 = 10 + 0.02 \\times 1 \\times 10$). Even if we allow for the more extreme accelerations, the value of that tricky last term should still be on the high side of Around 1 ($0.02 \\times 1 \\times 10000$). That's probably good enough."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# variable timestep\n",
    "hm = 0.02 # master timestep\n",
    "m.h = Var(m.N, bounds = (0.8,1.0))\n",
    "\n",
    "# Integration constraints \n",
    "def BwEuler_p(m,n,dof): # for positions\n",
    "    if n > 1:\n",
    "        return m.q[n,dof] == m.q[n-1,dof] + hm*m.h[n]*m.dq[n,dof]\n",
    "    else:\n",
    "        return Constraint.Skip #use this to leave out members of a set that the constraint doesn't apply to\n",
    "m.integrate_p = Constraint(m.N, m.DOF, rule = BwEuler_p)\n",
    "\n",
    "def BwEuler_v(m,n,dof): # for velocities\n",
    "    if n > 1:\n",
    "        return m.dq[n,dof] == m.dq[n-1,dof] + hm*m.h[n]*m.ddq[n-1,dof]\n",
    "    else:\n",
    "        return Constraint.Skip \n",
    "m.integrate_v = Constraint(m.N, m.DOF, rule = BwEuler_v)\n",
    "\n",
    "# Equations of motion -----------------------------------------------------\n",
    "\n",
    "# if you have a bunch of lambdify functions, it's helpful to create a function to deliver the list of variables each time\n",
    "def get_var_list(m,n):\n",
    "    var_list = [m.g,m.m[('body',1)],m.m[('leg',1)],m.m[('leg',2)],\n",
    "                m.len[('body',1)],m.len[('leg',1)],m.len[('leg',2)],\n",
    "                m.In[('body',1)],m.In[('leg',1)],m.In[('leg',2)],\n",
    "                m.q[n,'x'],m.q[n,'y'],m.q[n,'thb'],m.q[n,'thl'],m.q[n,'r'],\n",
    "                m.dq[n,'x'],m.dq[n,'y'],m.dq[n,'thb'],m.dq[n,'thl'],m.dq[n,'r'],\n",
    "                m.ddq[n,'x'],m.ddq[n,'y'],m.ddq[n,'thb'],m.ddq[n,'thl'],m.ddq[n,'r'],\n",
    "                m.fj[n,'knee'],m.fj[n,'hip'],m.GRF[n,'x'],m.GRF[n,'y']]\n",
    "    return var_list\n",
    "\n",
    "def dynamics(m,n,dof):\n",
    "    var_list = get_var_list(m,n)\n",
    "    return lambEOM[dof](*var_list) == 0\n",
    "m.dynamics = Constraint(m.N, m.DOF, rule = dynamics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complementarity\n",
    "Complementarity is how we finesse conditional behaviour out of functions that are still somewhat compatible with an NLP solver. *Complementing* two quantities just means specifying that their product is zero. This effectively creates a NAND relationship between them: the constraint is satisfied if one or both of the variables is zero, but FALSE if they're both nonzero at the same time. For our purposes, the complemented quantities **must be positive.**\n",
    "\n",
    "In an ideal world, you'd simply be able to write $AB = 0$ and be done with it, but as you've seen already, trajectory optimization world is far from ideal. Constraints of this form are actually extremely difficult for most NLP solvers to deal with. IPOPT in particular is an interior point solver, which means it looks for the feasible interior region between the multidimensional surfaces formed by the constraints. Complementarity constraints don't have an interior, though: the solution lies along the axis of one complemented variable or the other, since the value of at least one of them must be squashed down to zero. To improve IPOPT's chances, we create a false interior by setting $AB = \\varepsilon$, where $\\varepsilon$ is a penalty variable we subsequently minimize as part of our cost function. \n",
    "\n",
    "Throughout the rest of this tut, you'll see how we construct complementarity constraints to deal with assorted hard contact problems.\n",
    "\n",
    "## Ground interactions\n",
    "\n",
    "The ground reaction force has two components:\n",
    "1. the vertical normal force $G_y$. This can only act upwards, so $G_y \\geq 0$.\n",
    "2. the horizontal friction force $G_x$. This can act in either direction, but we represent it using two positive variables $G_x^+ > 0$ and $G_x^- > 0$ such that $G_x = G_x^+ - G_x^-$. Breaking a variable into positive and negative components is something you'll see happening all over the place to facilitate complementarity. Most of the time, it's not necessary to complement these variables with each other, but occassionally you might need to do so to ensure that they can't make weird things happen by being nonzero at the same time. A few back-of-the-envelope calculations never go amiss.\n",
    "\n",
    "We use four complementarity constraints to bring about this behaviour:\n",
    "\n",
    "### Contact\n",
    "**Purpose:** Makes sure $G_y$ only acts when the foot is on the ground.\n",
    "\n",
    "**Constraint:** $G_y[n]y_{foot}[n+1] \\leq \\varepsilon[n]$\n",
    "\n",
    "We like complementarity equations to be as simple as possible, so we use an auxiliary variable $y_{foot}$ to represent the foot height. Note that the normal force at *this* timestep is complemented with the foot height at the *next* timestep. Sure this means that the force technically starts acting one node before the foot is directly in contact with the ground, but it is the only way that the nonpenetration condition is actually solveable using an implicit Euler integration.\n",
    "\n",
    "### Friction\n",
    "**Purpose:** Only allows the foot to have a horizontal velocity (that is, to slide) when it maxes out the static friction limit.\n",
    "\n",
    "**Constraint:** $(\\dot{x}_{foot}^+[n] + \\dot{x}_{foot}^-[n])f[n] \\leq \\varepsilon[n]$\n",
    "\n",
    "Notice that we've split the foot velocity into positive and negative components, so their sum indicates whether the magnitude is nonzero. \n",
    "\n",
    "The aux variable $f$ represents what we call the *friction cone*, though in 2D, it's more like the friction triangle: \n",
    "\n",
    "<img src = \"friction_cone.png\" width = \"300\">\n",
    "\n",
    "if the magnitude of the velocity is inside the cone, the foot sticks. If it hits the edge of the cone, the foot slides. Another way of putting it: if the foot is moving while in contact with the ground, the friction force must be at its maximum value $\\mu G_y$. The difference between the static friction limit and the friction's actual magnitude is what we define as $f$: $f = \\mu G_y - (G_x^+[n] + G_x^-[n])$.\n",
    "\n",
    "### Sliding\n",
    "\n",
    "**Purpose:** Makes the friction force act in the opposite direction to the foot's velocity.\n",
    "\n",
    "**Constraint:** This one has two parts - one for each direction: $\\dot{x}_{foot}^+[n]G_x^+[n] = p[n]$ and $\\dot{x}_{foot}^-[n]G_x^-[n] \\leq \\varepsilon[n]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GROUND INTERACTIONS -------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# paramters\n",
    "m.mu = Param(initialize = 0.5) # friction coefficient\n",
    "\n",
    "# sign set for positive and negative components\n",
    "signs = ['ps','ng'] \n",
    "m.sgn = Set(initialize = signs)\n",
    "\n",
    "# variables\n",
    "# (I had that set error while trying to set up these default bounds. That's why I pointed it out.)\n",
    "\n",
    "m.footy = Var(m.N, bounds = (0.0,None)) # foot position\n",
    "m.footdx = Var(m.N, m.sgn, bounds = (0.0,None)) # foot velocity\n",
    "\n",
    "m.friction_cone = Var(m.N, bounds = (0.0,None))\n",
    "\n",
    "m.Gx = Var(m.N, m.sgn, bounds = (0.0,None)) # signed components of friction force\n",
    "\n",
    "ground_constraints = ['contact','friction','slip_ps','slip_ng'] \n",
    "m.ground_constraints = Set(initialize = ground_constraints) # set for indexing ground-related penalties\n",
    "m.ground_penalty = Var(m.N, m.ground_constraints, bounds = (0.0,None))\n",
    "\n",
    "# constraints: aux variables\n",
    "\n",
    "def def_footy(m,n):\n",
    "    var_list = get_var_list(m,n)\n",
    "    return m.footy[n] == lamb_pfoot['y'](*var_list)\n",
    "m.def_footy = Constraint(m.N, rule = def_footy)\n",
    "\n",
    "def def_footdx(m,n):\n",
    "    var_list = get_var_list(m,n)\n",
    "    return m.footdx[n,'ps'] - m.footdx[n,'ng'] == lamb_vfoot['x'](*var_list)\n",
    "m.def_footdx = Constraint(m.N, rule = def_footdx)\n",
    "\n",
    "def get_Gx(m,n):\n",
    "    return m.GRF[n,'x'] == m.Gx[n,'ps'] - m.Gx[n,'ng']\n",
    "m.get_Gx = Constraint(m.N, rule = get_Gx)\n",
    "\n",
    "def def_friction_cone(m,n):\n",
    "    return m.friction_cone[n] == m.mu*m.GRF[n,'y'] - (m.Gx[n,'ps'] + m.Gx[n,'ng'])\n",
    "m.def_friction_cone = Constraint(m.N, rule = def_friction_cone)\n",
    "\n",
    "# constraints: complementarity\n",
    "\n",
    "# contact\n",
    "def ground_contact(m,n):\n",
    "    if n < N:\n",
    "        return m.footy[n+1]*m.GRF[n,'y']  <= m.ground_penalty[n,'contact']\n",
    "        # notice that the GRF is complemented with the foot height at the NEXT node\n",
    "    else:\n",
    "        return Constraint.Skip\n",
    "m.ground_contact = Constraint(m.N, rule = ground_contact)\n",
    "\n",
    "# friction\n",
    "def ground_friction(m,n):\n",
    "    return (m.footdx[n,'ps']+m.footdx[n,'ng'])*m.friction_cone[n] <= m.ground_penalty[n,'friction'] \n",
    "m.ground_friction = Constraint(m.N, rule = ground_friction)\n",
    "\n",
    "# slipping\n",
    "def ground_slip_ps(m,n):\n",
    "    return m.footdx[n,'ps']*m.Gx[n,'ps'] <= m.ground_penalty[n,'slip_ps']\n",
    "m.ground_slip_ps = Constraint(m.N, rule = ground_slip_ps)\n",
    "\n",
    "def ground_slip_ng(m,n):\n",
    "    return m.footdx[n,'ng']*m.Gx[n,'ng'] <= m.ground_penalty[n,'slip_ng']\n",
    "m.ground_slip_ng = Constraint(m.N, rule = ground_slip_ng)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Joints with hard end stops\n",
    "\n",
    "You can also use complementarity constraints to create hard boundaries at the end ranges of motion (ROM) for the joints. Constraints on a system's motion are only enforceable if they are literally _enforced_ - as in, if forces or torques are made available to reverse the velocity before they are exceeded. Without hard joint stops, the only forces that can do this will be the actuator forces, which are typically bounded to some approximation of the capabilities of a real-world system. This is not necessarily a problem, but it does remove the possibility of solutions where the joint actuates at full tilt until it hits its limit, as the force would need to reverse direction to keep the joint within the ROM constraints. This can potentially cause oscillatory \"bang bang\" force profiles to emerge, where the force alternates between its maximum values in opposite directions.\n",
    "\n",
    "In theory, enforcing joint ROMs as collisions should make a wider range of solutions possible, but personally, I find that the added overhead of extra complementarity constraints negates any positive effect this could have on solver performance. For this reason, I'm not putting them in the Pyomo model, but I'll still show you how to do it just because it might be important to realistically model some system you care about:\n",
    "\n",
    "### Hard stops in a rotary joint\n",
    "Hard joint stops work exactly like the contact constraint in the ground block, only now, instead of the foot height, it's the distance to the endpoint, and instead of the ground reaction force, it's a rebound action opposing motion beyond the limit.\n",
    "\n",
    "Only two complementarity constraints are needed - one for each boundary:\n",
    "\n",
    "On the upper side: $\\tau_r^-[n](\\theta_{max} - \\theta[n+1]) \\leq \\varepsilon[n]$\n",
    "\n",
    "And on the lower side: $\\tau_r^+[n](\\theta[n+1] - \\theta_{min}) \\leq \\varepsilon[n]$\n",
    "\n",
    "You could also do something like this to set hard limits on the prismatic joint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Minimizing the penalty\n",
    "There are three ways to approach minimizing the complementarity penalties:\n",
    "1. Simultaneous: add them as a term to the cost function, typically magnified by some large scale factor so flattening them is prioritized. I normally go with 1000, or whatever amounts to around three orders of magnitude greater than whatever I'm trying to minimize.\n",
    "2. Iterative: here, the complementarity penalties aren't added to the cost function. Rather, you give them an upper bound and then solve the problem repeatedly, decreasing this bound each time.\n",
    "3. Two stage: first, only the sum of the penalties is minimized to get a feasible solution. The penalties are then bounded to below an acceptable value (say, $1e-4$) and the problem is then solved again with the intended objective being minimized.\n",
    "\n",
    "Personally, I find 3 to be the most reliable, so that's what I use in my research. I'll be using 1 here, because it's the simplest and the only option that completes in a single solve stage, but be warned that scaling issues can arise due to the large scale factor on the penalties. It often takes a bit of tweaking to find a suitable weight that results in sufficiently small penalties but doesn't break the whole problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COST FUNCTION -------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# minimum time\n",
    "\n",
    "def CostFun(m):\n",
    "    T = sum(m.h[n] for n in range(1,N+1))\n",
    "    penalty_sum = sum([m.ground_penalty[n,gc] for n in range(1,N+1) for gc in ground_constraints])\n",
    "    # ain't single-line for loops grand?\n",
    "#     return T+1000*penalty_sum\n",
    "\n",
    "    return penalty_sum # for sanity checks\n",
    "m.Cost = Objective(rule = CostFun)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity checks\n",
    "One of the tricky and often frustrating things about trajectory optimization is that you have to code everything before you can test anything. Due to the aforementioned interdependency issue, it can be very difficult to pinpoint the exact cause of strange behaviour or isolate any one part to check if it's working correctly.\n",
    "\n",
    "To help with this, there are a few simple tests I like to run right after finishing a model to help find and kill some of the more obvious bugs. Because the cost function can potentially lead the solver into a difficult place, I'd recommend running these with a *feasibility* objective: that is, no other goal besides satisfying the penalty. \n",
    "\n",
    "### The high drop test\n",
    "**How you do it:** Set the initial condition to rest in a known pose (say, all angles = 0) high enough above the ground that it can fall without hitting it, fix all actuator forces/ torques to zero and then just let it drop.\n",
    "\n",
    "**What's the point?** This test will reveal a few things:\n",
    "1. Big, stupid problems: if it comes back infeasible, your model is broken somewhere. Time to make yourself some coffee, dust off the ol' `m.pprint()` statement and try to find that one equation you copied and forgot to change a variable name in.\n",
    "2. Ground interaction problems: if a ground reaction force is anything other than zero when it's three metres in the air, something's wrong.\n",
    "3. Equations of motion problems: check `m.ddq.pprint()`. You want to see zero acceleration everywhere, and -9.81 in the $y$ direction. If you see anything else, something's wrong. This is especially important when you're working in system space coordinates, where you have all those tricky constraint forces to balance; you have to make sure they're not somehow conspiring to create motion where none should happen (this is why we set the actuator actions to zero).\n",
    "\n",
    "### The low drop test\n",
    "**How you do it:** Same as the high drop, but now you're starting with the foot only a short way above the ground (e.g. 1 cm) and fixing the foot to be grounded at some time, so it's guaranteed to land.\n",
    "\n",
    "**What's the point?** To check if your ground contacts work to activate the ground reaction forces and successfully stop the feet.\n",
    "\n",
    "### The standing test\n",
    "**How you do it:** Deactivate the actuators as in the other tests, and start the model at rest in a known pose with its feet (well, foot in this case) on the ground. It might also be worth trying a few postures that push the limits of the joints e.g. have the prismatic knee joint compressed all the way in.\n",
    "\n",
    "**What's the point?** While the low drop is meant to check if your contacts behave as expected in impulsive collisions, the standing test checks whether they can handle sustained contact.\n",
    "\n",
    "### The hop test\n",
    "**How you do it:** Now we get to add some power :) Start the model at rest in contact with the ground, and set the final condition to some $x$ or $y$ that will either force it to move forward or jump in the air. \n",
    "\n",
    "**What's the point?** To see if the actuator force limits you've set allow enough *oomph* to get the dude off the ground. If you're brave and want to test the limits in the other direction, you can tell it to maximize the final $x$ or $y$ position (by making $-x[N]$ or $-y[N]$ your objective) to make sure it can't shoot itself into orbit, either.\n",
    "\n",
    "**Note:** To save on time, use as few nodes as possible to do these tests. Especially for something like the high drop, you should be able to see what you're trying to see with just $N = 10$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # HIGH DROP -----------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# # initial condition\n",
    "# for dof in DOFs:\n",
    "#     m.dq[1,dof].fix(0) # rest\n",
    "#     if dof != 'y':\n",
    "#         m.q[1,dof].fix(0) # neutral posture\n",
    "# m.footy[1].fix(10)\n",
    "\n",
    "# for n in range(1,N+1):\n",
    "#     for j in joints:\n",
    "#         m.fj[n,j].fix(0) # no forces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # LOW DROP -----------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# # initial condition\n",
    "# for dof in DOFs:\n",
    "#     m.dq[1,dof].fix(0) # rest\n",
    "#     if dof not in ['y','r']:\n",
    "#         m.q[1,dof].fix(0) # neutral posture\n",
    "# m.footy[1].setlb(0.1) # start near the ground\n",
    "\n",
    "# m.footy[15].fix(0) # force it to land\n",
    "\n",
    "# for n in range(1,N+1):\n",
    "#     for j in joints:\n",
    "#         m.fj[n,j].fix(0) # no forces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # # STANDING ------------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# # initial condition\n",
    "# for dof in DOFs:\n",
    "#     m.dq[1,dof].fix(0) # rest\n",
    "#     if dof not in ['y','r']:\n",
    "#         m.q[1,dof].fix(0) # neutral posture\n",
    "\n",
    "# for n in range(1,N+1):\n",
    "#     m.footy[n].fix(0) # must stay grounded for the full time\n",
    "\n",
    "# for n in range(1,N+1):\n",
    "#     for j in joints:\n",
    "#         if j != 'knee': # needs knee force to maintain ROM limits of prismatic joint\n",
    "#             m.fj[n,j].fix(0) # no forces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # # HOP -----------------------------------------------------------------------------------------------------------------------\n",
    "# # initial condition\n",
    "# for dof in DOFs:\n",
    "#     m.dq[1,dof].fix(0) # rest\n",
    "#     if dof not in ['y','r']:\n",
    "#         m.q[1,dof].fix(0) # neutral posture\n",
    "# m.footy[1].fix(0)\n",
    "\n",
    "# # midpoint\n",
    "# m.footy[10].setlb(0.2)\n",
    "\n",
    "# # final condition\n",
    "# m.footy[N].fix(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SPRINT --------------------------------------------------------------------------------------------------------------------\n",
    "#sprint 5m from rest\n",
    "# initial condition\n",
    "for dof in DOFs:\n",
    "    m.dq[1,dof].fix(0) # rest\n",
    "    if dof not in ['y','r']:\n",
    "        m.q[1,dof].fix(0) # neutral posture\n",
    "m.footy[1].fix(0)\n",
    "\n",
    "# final condition\n",
    "m.q[N,'x'].fix(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# solving\n",
    "#opt = SolverFactory('ipopt') # standard issue, garden variety ipopt\n",
    "\n",
    "opt = SolverFactory('ipopt',executable = 'C:/cygwin64/home/Stacey/CoinIpopt/build/bin/ipopt.exe')\n",
    "opt.options[\"linear_solver\"] = 'ma86'\n",
    "\n",
    "# solver options\n",
    "opt.options[\"expect_infeasible_problem\"] = 'yes'\n",
    "#pt.options[\"linear_system_scaling\"] = 'none'\n",
    "#opt.options[\"mu_strategy\"] = \"adaptive\"\n",
    "opt.options[\"print_level\"] = 5 # prints a log with each iteration (you want to this - it's the only way to see progress.)\n",
    "opt.options[\"max_iter\"] = 30000 # maximum number of iterations\n",
    "opt.options[\"max_cpu_time\"] = 600 # maximum cpu time in seconds\n",
    "opt.options[\"Tol\"] = 1e-6 # the tolerance for feasibility. Considers constraints satisfied when they're within this margin.\n",
    "    \n",
    "results = opt.solve(m, tee = True) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if the problem is infeasible, this is how you can see which constraints weren't satisfied\n",
    "from pyomo.util.infeasible import log_infeasible_constraints\n",
    "log_infeasible_constraints(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(results.solver.status) # tells you if the solver had any errors/ warnings\n",
    "print(results.solver.termination_condition) # tells you if the solution was (locally) optimal, feasible, or neither.\n",
    "\n",
    "penalty_sum = 0\n",
    "for n in range(1,N+1):\n",
    "    for gc in ground_constraints:\n",
    "        penalty_sum += m.ground_penalty[n,gc].value\n",
    "print(penalty_sum)\n",
    "\n",
    "#m.pprint() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#animate it\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as ani\n",
    "from IPython.display import HTML\n",
    "%matplotlib inline\n",
    "\n",
    "fig1, ax1 = plt.subplots(1,1) #create axes\n",
    "ax1.set_aspect('equal')\n",
    "\n",
    "xmax = np.max([m.q[n,'x'].value for n in range(1,N+1)])\n",
    "ymax = np.max([m.q[n,'y'].value for n in range(1,N+1)])\n",
    "\n",
    "def plot_robot(i,m,ax): #update function for animation\n",
    "    ax.clear()\n",
    "    ax.set_xlim([-1,xmax+1]) # adjust limits to solution\n",
    "    ax.set_ylim([0,ymax+0.5])\n",
    "    \n",
    "    #plot body\n",
    "    body_xb = m.q[i,'x'].value - 0.5*m.len[('body',1)]*cos(m.q[i,'thb'].value)\n",
    "    body_yb = m.q[i,'y'].value - 0.5*m.len[('body',1)]*sin(m.q[i,'thb'].value)\n",
    "    body_xf = m.q[i,'x'].value + 0.5*m.len[('body',1)]*cos(m.q[i,'thb'].value)\n",
    "    body_yf = m.q[i,'y'].value + 0.5*m.len[('body',1)]*sin(m.q[i,'thb'].value)  \n",
    "    ax.plot([body_xb,body_xf],[body_yb,body_yf],color='xkcd:black')\n",
    "      \n",
    "    #plot leg 1\n",
    "    thA = m.q[i,'thb'].value+m.q[i,'thl'].value\n",
    "    leg1_xt = m.q[i,'x'].value\n",
    "    leg1_yt = m.q[i,'y'].value\n",
    "    leg1_xb = m.q[i,'x'].value + m.len[('leg',1)]*sin(thA)\n",
    "    leg1_yb = m.q[i,'y'].value - m.len[('leg',1)]*cos(thA)\n",
    "    ax.plot([leg1_xt,leg1_xb],[leg1_yt,leg1_yb],color='xkcd:black')\n",
    "    \n",
    "    #plot leg 2\n",
    "    Lt = 0.5*m.len[('leg',1)] + m.q[i,'r'].value - 0.5*m.len[('leg',2)]\n",
    "    Lb = 0.5*m.len[('leg',1)] + m.q[i,'r'].value + 0.5*m.len[('leg',2)]\n",
    "    leg2_xt = m.q[i,'x'].value + Lt*sin(thA)\n",
    "    leg2_yt = m.q[i,'y'].value - Lt*cos(thA)\n",
    "    leg2_xb = m.q[i,'x'].value + Lb*sin(thA)\n",
    "    leg2_yb = m.q[i,'y'].value - Lb*cos(thA)\n",
    "    ax.plot([leg2_xt,leg2_xb],[leg2_yt,leg2_yb],color='xkcd:black')\n",
    "    \n",
    "update = lambda i: plot_robot(i,m,ax1) #lambdify update function\n",
    "\n",
    "animate = ani.FuncAnimation(fig1,update,range(1,N+1),interval = 50,repeat=True)\n",
    "\n",
    "HTML(animate.to_html5_video()) #you need to convert the animation to HTML5 to embed it in the notebook\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}