disc_quaternion_test.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",
    "from scipy.linalg import expm, null_space\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, ZeroOrderHold\n",
    "from pydrake.multibody.inverse_kinematics import PositionConstraint, AddUnitQuaternionConstraintOnPlant\n",
    "\n",
    "from pydrake.multibody.plant import ContactResults, CoulombFriction\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Meshcat listening for connections at http://localhost:7001\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 = 6\n",
    "\n",
    "x0 = np.zeros(num_disc_states)\n",
    "x0[0], x0[6] = 1, 0.4\n",
    "\n",
    "plant.SetPositionsAndVelocities(context, x0)\n",
    "plant.get_actuation_input_port().FixValue(context, np.zeros(num_act))\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",
    "    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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8573310072313118\n"
     ]
    }
   ],
   "source": [
    "np.random.seed(5)\n",
    "\n",
    "# Change between 0 input and random input\n",
    "generic_input = builder.AddSystem(ConstantVectorSource(np.random.randn(num_act)))\n",
    "\n",
    "print(np.max(np.random.randn(num_act)))\n",
    "\n",
    "# generic_input = builder.AddSystem(ConstantVectorSource(np.zeros(num_act)))\n",
    "builder.Connect(generic_input.get_output_port(), plant.get_actuation_input_port())\n",
    "logger_states = LogVectorOutput(plant.get_state_output_port(), builder, publish_period=1/1000)\n",
    "diagram = builder.Build()\n",
    "\n",
    "simulator = Simulator(diagram)\n",
    "sim_context = simulator.get_mutable_context()\n",
    "\n",
    "x0[6] += 0.3\n",
    "\n",
    "sim_context.SetTime(0.)\n",
    "sim_context.SetDiscreteState(x0)\n",
    "vis.StartRecording()\n",
    "simulator.AdvanceTo(2)\n",
    "vis.StopRecording()\n",
    "vis.PublishRecording()\n",
    "\n",
    "meshcat.DeleteAddedControls()\n",
    "\n",
    "state_log = logger_states.FindLog(sim_context).data()\n",
    "t = np.arange(0, len(state_log[0, :]))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(t, state_log[0, :], label=\"q0\")\n",
    "ax.plot(t, state_log[1, :], label=\"q1\")\n",
    "ax.plot(t, state_log[2, :], label=\"q2\")\n",
    "ax.plot(t, state_log[3, :], label=\"q3\")\n",
    "ax.legend()\n",
    "\n",
    "plt.show()\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}