2d-reversed-2-link-pendulum-cart-opty.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a527636",
   "metadata": {},
   "source": [
    "A **double pendulum** is attached to the center of cart which rolls in the X direction.\\\n",
    "A force acts on the cart in X direction, the gravity acts in the negative Y direction.\\\n",
    "The goal is to get the pendulum in the upright position by suitably moving the cart to the left and to the right.\\\n",
    "This is basically a copy of the example of the revered pendulum from the docs of opty and of Timo's simulation sent to me, just a two link pendulum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7cc216b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy.physics.mechanics as me\n",
    "from collections import OrderedDict\n",
    "import time\n",
    "import numpy as np\n",
    "import sympy as sm\n",
    "from opty.direct_collocation import Problem\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "matplotlib.rcParams['animation.embed_limit'] = 2**128\n",
    "from matplotlib import patches\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "090b386b",
   "metadata": {},
   "source": [
    "Set up **Kane's equations of motion**\\\n",
    "Very simple, just a two body pendulum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c665b199",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "\n",
    "N, A1, A2      = sm.symbols('N A1, A2', cls = me.ReferenceFrame)\n",
    "t              = me.dynamicsymbols._t\n",
    "O, P1, P2, P3  = sm.symbols('O P1 P2 P3', cls = me.Point)                               # O fixed point in N, P1 center of mass of the cart, P2, P3 centers of mass of the pendulums \n",
    "O.set_vel(N, 0)                                                                         # Velocity of the fixed point is zero   \n",
    "q1, q2, q3, u1, u2, u3, F = me.dynamicsymbols('q1 q2 q3 u1 u2 u3 F')                    # Generalized coordinates of the cart and pendulum; F is the force applied to the cart\n",
    "l, m1, m2, m3, g, iZZ1, iZZ2 = sm.symbols('l, m1, m2, m3 g, iZZ1, iZZ2')                # Length of the pendulum, masses of the pendulum and cart, gravitational constant, and moment of inertia of the pendulums\n",
    "\n",
    "A1.orient_axis(N, q2, N.z)\n",
    "A1.set_ang_vel(N, u2 * N.z)\n",
    "A2.orient_axis(N, q3, N.z)\n",
    "A2.set_ang_vel(N, u3 * N.z)\n",
    "\n",
    "P1.set_pos(O, q1 * N.x)\n",
    "P2.set_pos(P1, l * A1.x)\n",
    "P2.v2pt_theory(P1, N, A1)\n",
    "\n",
    "P3.set_pos(P2, l * A2.x)\n",
    "P3.v2pt_theory(P2, N, A2)\n",
    "\n",
    "P1a = me.Particle('P1a', P1, m1)                                                        # Particle for the cart\n",
    "\n",
    "I = me.inertia(A1, 0, 0, iZZ1)\n",
    "P2a = me.RigidBody('P2a', P2, A1, m2, (I, P2))                                          # Rigid body for the pendulum\n",
    "\n",
    "I = me.inertia(A2, 0, 0, iZZ2)\n",
    "P3a = me.RigidBody('P3a', P3, A2, m3, (I, P3))                                          # Rigid body for the pendulum\n",
    "BODY = [P1a, P2a, P3a]\n",
    "\n",
    "FL = [(P1, F * N.x - m1*g*N.y), (P2, -m2*g*N.y), (P3, -m3*g*N.y)]                       # Force applied to the cart and gravitational force on the pendulum\n",
    "kd = sm.Matrix([q1.diff(t) - u1, q2.diff(t) - u2, q3.diff(t) - u3])                     # Kinematic differential equations\n",
    "\n",
    "q_ind = [q1, q2, q3]\n",
    "u_ind = [u1, u2, u3]\n",
    "\n",
    "KM = me.KanesMethod(N, q_ind=q_ind, u_ind=u_ind, kd_eqs=kd)\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "EOM = kd.col_join(fr + frstar)\n",
    "EOM                                                                                     # Equations of motion       "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0685c32",
   "metadata": {},
   "source": [
    "Set up the machinery for **opty**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc7e7050",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_symbols = tuple((*q_ind, *u_ind))\n",
    "constant_symbols = (l, m1, m2, m3, g, iZZ1, iZZ2)\n",
    "specified_symbols = (F,)\n",
    "\n",
    "\n",
    "target_angle = np.pi /2.0                       # The desired angle for the pendulums to swing up to.\n",
    "duration = 15.0                                 # The duration of the swing up maneuver.\n",
    "num_nodes = 250                                 # The number of nodes to use in the direct collocation problem.\n",
    "save_animation = False\n",
    "\n",
    "interval_value = duration / (num_nodes - 1)\n",
    "\n",
    "# Specify the known system parameters.\n",
    "par_map = OrderedDict()\n",
    "par_map[l]   = 2.0\n",
    "par_map[m1]  = 1.0\n",
    "par_map[m2]  = 1.0\n",
    "par_map[m3]  = 1.0\n",
    "par_map[g]   = 9.81\n",
    "par_map[iZZ1] = 2.0\n",
    "par_map[iZZ2] = 2.0\n",
    "\n",
    "# Specify the objective function and it's gradient.\n",
    "def obj(free): \n",
    "    \"\"\"Minimize the sum of the squares of the control torque.\"\"\"\n",
    "    F = free[6 * num_nodes:]\n",
    "    return interval_value * np.sum(F**2)\n",
    "\n",
    "def obj_grad(free):\n",
    "    grad = np.zeros_like(free)\n",
    "    grad[6 * num_nodes:] = 2.0 * interval_value * free[6 * num_nodes:]\n",
    "    return grad\n",
    "\n",
    "# This portion is copied from Timo's code\n",
    "t0, tf = 0.0, duration\n",
    "\n",
    "# Opty want instance constraints, basically things like `q2(2.0) - 3.141` to specify that q2 should be 3.141 when t=2.0\n",
    "# I personally like to rewrite it to initial state and final state constraints and then translate them into the set of equations opty wants.\n",
    "initial_state_constraints = {\n",
    "                            q1: 0.,  \n",
    "                            q2: -np.pi/2., \n",
    "                            q3: -np.pi/2., \n",
    "                            u1: 0.,\n",
    "                            u2: 0., \n",
    "                            u3: 0.\n",
    "                            }\n",
    "\n",
    "final_state_constraints  = {\n",
    "                           q2: target_angle, \n",
    "                           q3: target_angle, \n",
    "                           u1: 0., \n",
    "                           u2: 0., \n",
    "                           u3: 0.\n",
    "                           }\n",
    "\n",
    "instance_constraints = (\n",
    ") + tuple(\n",
    "    xi.subs({t: t0}) - xi_val for xi, xi_val in initial_state_constraints.items()\n",
    ") + tuple(\n",
    "    xi.subs({t: tf}) - xi_val for xi, xi_val in final_state_constraints.items()\n",
    ")\n",
    "print('instance constraints:', instance_constraints)\n",
    "bounds = {F: (-25., 25.)}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d706a090",
   "metadata": {},
   "source": [
    "Create an **optimization problem** and solve it\\\n",
    "Note:\n",
    "\n",
    "*random.seed(1234)* gives an acceptable solution, *ransom.seed(12346)* gives a nice solution\\\n",
    "I iterate twice, using the solution of the first iteration as initial_guess for the second one. Sometimes this improves the accuracy, sometimes it does not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "079117a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an optimization problem.\n",
    "prob = Problem(obj, obj_grad, EOM, state_symbols, num_nodes, interval_value,\n",
    "               known_parameter_map=par_map,\n",
    "               instance_constraints=instance_constraints,\n",
    "               bounds=bounds)\n",
    "\n",
    "# Use a random positive initial guess.\n",
    "np.random.seed(12342)\n",
    "initial_guess = np.random.randn(prob.num_free)\n",
    "\n",
    "# Find the optimal solution.\n",
    "for i in range(2):\n",
    "    solution, info = prob.solve(initial_guess)\n",
    "    initial_guess = solution\n",
    "    \n",
    "    print(f'{i + 1}st iteration: value of objective function is {obj(solution):.2f}')\n",
    "    prob.plot_objective_value()\n",
    "    print('message from optimizer:', info['status_msg'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fb77b13",
   "metadata": {},
   "source": [
    "Plot some results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6f0398b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(7, 1, figsize=(10, 15), sharex=True)\n",
    "times = np.linspace(0.0, duration, num=num_nodes)\n",
    "\n",
    "for i, j in enumerate(state_symbols + specified_symbols):\n",
    "    ax[i].plot(times, solution[i * num_nodes:(i + 1) * num_nodes])\n",
    "    ax[i].grid()\n",
    "    ax[i].set_ylabel(str(j))\n",
    "ax[-1].set_xlabel('time')\n",
    "ax[0]. set_title('Generalized coordinates and speeds')\n",
    "ax[-1]. set_title('Control force');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c384f66b",
   "metadata": {},
   "source": [
    "**aminate** the pendulum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5546cec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "P1_x = np.empty(num_nodes)\n",
    "P1_y = np.empty(num_nodes)\n",
    "P2_x = np.empty(num_nodes)\n",
    "P2_y = np.empty(num_nodes)\n",
    "P3_x = np.empty(num_nodes)\n",
    "P3_y = np.empty(num_nodes)\n",
    "\n",
    "P1_loc = [me.dot(P1.pos_from(O), uv) for uv in [N.x, N.y]]\n",
    "P2_loc = [me.dot(P2.pos_from(O), uv) for uv in [N.x, N.y]]\n",
    "P3_loc = [me.dot(P3.pos_from(O), uv) for uv in [N.x, N.y]]\n",
    "\n",
    "qL = q_ind + u_ind\n",
    "pL_vals = list(constant_symbols)\n",
    "P1_loc_lam = sm.lambdify(qL + pL_vals, P1_loc, cse=True)\n",
    "P2_loc_lam = sm.lambdify(qL + pL_vals, P2_loc, cse=True)\n",
    "P3_loc_lam = sm.lambdify(qL + pL_vals, P3_loc, cse=True)\n",
    "\n",
    "for i in range(num_nodes):\n",
    "    q_1 = solution[i]\n",
    "    q_2 = solution[i + num_nodes]\n",
    "    q_3 = solution[i + 2 * num_nodes]\n",
    "    u_1 = solution[i + 3 * num_nodes]\n",
    "    u_2 = solution[i + 4 * num_nodes]\n",
    "    u_3 = solution[i + 5 * num_nodes]\n",
    "    P1_x[i], P1_y[i] = P1_loc_lam(q_1, q_2, q_3, u_1, u_2, u_3, *list(par_map.values()))\n",
    "    P2_x[i], P2_y[i] = P2_loc_lam(q_1, q_2, q_3, u_1, u_2, u_3, *list(par_map.values()))\n",
    "    P3_x[i], P3_y[i] = P3_loc_lam(q_1, q_2, q_3, u_1, u_2, u_3, *list(par_map.values()))\n",
    "    \n",
    "\n",
    "# needed to give the picture the right size.\n",
    "xmin = min(np.min(P1_x),np.min(P2_x), np.min(P3_x))\n",
    "xmax = max(np.max(P1_x), np.max(P2_x), np.max(P3_x))\n",
    "ymin = min(np.min(P1_y), np.min(P2_y), np.min(P3_y))\n",
    "ymax = max(np.max(P1_y), np.max(P2_y), np.max(P3_y))\n",
    "\n",
    "width, height = par_map[l]/3., par_map[l]/3.\n",
    "\n",
    "def animate_pendulum(time, P1_x, P1_y, P2_x, P2_y):\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(10, 10), subplot_kw={'aspect': 'equal'})\n",
    "    \n",
    "    ax.axis('on')\n",
    "    ax.set_xlim(xmin - 1., xmax + 1.)\n",
    "    ax.set_ylim(ymin - 1., ymax + 1.)\n",
    "    \n",
    "    ax.set_xlabel('X direction', fontsize=20)\n",
    "    ax.set_ylabel('Y direction', fontsize=20)\n",
    "    ax.axhline(0, color='black', lw=2)\n",
    "\n",
    "\n",
    "\n",
    "    line1, = ax.plot([], [], 'o-', lw=0.5, color='blue')    \n",
    "    line2, = ax.plot([], [], 'o-', lw=0.5, color='green')                    \n",
    "    \n",
    "    recht = patches.Rectangle((P1_x[0] - width/2, P1_y[0] - height/2), width=width, height=height, fill=True, color='red', ec='black')\n",
    "    ax.add_patch(recht)\n",
    "\n",
    "    def animate(i):\n",
    "        message = (f'running time {times[i]:.2f} sec')\n",
    "        ax.set_title(message, fontsize=20)\n",
    "        recht.set_xy((P1_x[i] - width/2., P1_y[i] - height/2.))\n",
    "\n",
    "        wert_x = [P1_x[i], P2_x[i]]\n",
    "        wert_y = [P1_y[i], P2_y[i]]               \n",
    "        line1.set_data(wert_x, wert_y)    \n",
    "\n",
    "        wert_x = [P2_x[i], P3_x[i]]\n",
    "        wert_y = [P2_y[i], P3_y[i]]\n",
    "        line2.set_data(wert_x, wert_y)              \n",
    "        return line1, line2,\n",
    "\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=num_nodes,\n",
    "                                   interval=2000*np.max(times) / num_nodes,\n",
    "                                   blit=False)\n",
    "    plt.close(fig)\n",
    "    return anim\n",
    "\n",
    "anim = animate_pendulum(times, P1_x, P1_y, P2_x, P2_y)\n",
    "display(HTML(anim.to_jshtml()))    \n",
    "print(f'it took {time.time() - start:.3f} seconds to run the code.')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}