policy_search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](https://underactuated.csail.mit.edu/policy_search.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import mpld3\n",
    "import numpy as np\n",
    "from IPython.display import clear_output, display\n",
    "from pydrake.all import Integrator  # wrap_to,\n",
    "from pydrake.all import (AutoDiffXd, BasicVector_, DiagramBuilder,\n",
    "                         ExtractGradient, ExtractValue, InitializeAutoDiff,\n",
    "                         LeafSystem_, LogVectorOutput, MeshcatVisualizer,\n",
    "                         Rgba, SceneGraph, Simulator, Simulator_, StartMeshcat,\n",
    "                         TemplateSystem)\n",
    "from pydrake.examples import PendulumGeometry, PendulumPlant\n",
    "\n",
    "from underactuated import running_as_notebook\n",
    "from underactuated.utils import running_as_test\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer (run this cell only once, each instance consumes a port)\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Policy Optimization for LQR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trajectory-based (stochastic) gradient descent\n",
    "\n",
    "Note: This code is super slow right now, but will be optimized soon!\n",
    "\n",
    "## Pendulum swing-up\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@TemplateSystem.define(\"ParameterizedController_\")\n",
    "def ParameterizedController_(T):\n",
    "\n",
    "    class Impl(LeafSystem_[T]):\n",
    "\n",
    "        def _construct(self, converter=None):\n",
    "            LeafSystem_[T].__init__(self, converter)\n",
    "            self.DeclareVectorInputPort(\"state\", BasicVector_[T](2))\n",
    "            self.DeclareNumericParameter(BasicVector_[T](np.zeros((9,1))))\n",
    "            self.DeclareVectorOutputPort(\"command\", BasicVector_[T](1),self.CommandOutput)\n",
    "\n",
    "        def _construct_copy(self, other, converter=None):\n",
    "            Impl._construct(self, converter=converter)\n",
    "\n",
    "        def CommandOutput(self, context, output):\n",
    "            x = self.get_input_port().Eval(context)\n",
    "            q = x[0]\n",
    "            qd = x[1]\n",
    "            basis = np.array([1, np.sin(q), np.cos(q), qd, np.sin(q)**2, np.sin(q)*np.cos(q), qd*np.sin(q), qd*np.cos(q), qd**2])\n",
    "            alpha = context.get_numeric_parameter(0).CopyToVector()\n",
    "            output[0] = alpha.dot(basis)\n",
    "\n",
    "    return Impl\n",
    "\n",
    "ParameterizedController = ParameterizedController_[None]  # Default instantiation\n",
    "\n",
    "def wrap_to(value, low, high):\n",
    "    range = high - low\n",
    "    return value - range * np.floor((value - low) / range)\n",
    "\n",
    "\n",
    "@TemplateSystem.define(\"RunningCost_\")\n",
    "def RunningCost_(T):\n",
    "\n",
    "    class Impl(LeafSystem_[T]):\n",
    "\n",
    "        def _construct(self, converter=None):\n",
    "            LeafSystem_[T].__init__(self, converter)\n",
    "            self.DeclareVectorInputPort(\"state\", BasicVector_[T](2))\n",
    "            self.DeclareVectorInputPort(\"command\", BasicVector_[T](1))\n",
    "            self.DeclareVectorOutputPort(\"cost\", BasicVector_[T](1),self.CostOutput)\n",
    "\n",
    "        def _construct_copy(self, other, converter=None):\n",
    "            Impl._construct(self, converter=converter)\n",
    "\n",
    "        def CostOutput(self, context, output):\n",
    "            x = self.get_input_port(0).Eval(context)\n",
    "            x[0] = wrap_to(x[0] - np.pi, -np.pi, np.pi)\n",
    "            u = self.get_input_port(1).Eval(context)[0]\n",
    "            Q = np.diag([10, 1])\n",
    "            R = 1\n",
    "            output[0] = x.dot(Q.dot(x)) + R*u**2\n",
    "\n",
    "    return Impl\n",
    "\n",
    "RunningCost = RunningCost_[None]  # Default instantiation\n",
    "\n",
    "\n",
    "def simulate(alpha, N=10, set_planar=True):\n",
    "    builder = DiagramBuilder()\n",
    "    plant = builder.AddSystem(PendulumPlant())\n",
    "    pi = builder.AddSystem(ParameterizedController())\n",
    "    builder.Connect(plant.get_state_output_port(), pi.get_input_port())\n",
    "    builder.Connect(pi.get_output_port(), plant.get_input_port())\n",
    "    scene_graph = builder.AddSystem(SceneGraph())\n",
    "    PendulumGeometry.AddToBuilder(builder, plant.get_state_output_port(),\n",
    "                                  scene_graph)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    meshcat.Delete()\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    rng = np.random.default_rng(123)\n",
    "\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(0 if running_as_test else 1.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "    pi_context = pi.GetMyContextFromRoot(context)\n",
    "    pi_context.get_mutable_numeric_parameter(0).SetFromVector(alpha)\n",
    "\n",
    "    for n in range(1 if running_as_test else N):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(rng.standard_normal((2,1)))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(0.1 if running_as_test else 2.0)\n",
    "\n",
    "\n",
    "def plot_rollouts(alpha, N=10):\n",
    "    builder = DiagramBuilder()\n",
    "    plant = builder.AddSystem(PendulumPlant())\n",
    "    pi = builder.AddSystem(ParameterizedController())\n",
    "    builder.Connect(plant.get_state_output_port(), pi.get_input_port())\n",
    "    builder.Connect(pi.get_output_port(), plant.get_input_port())\n",
    "    logger = LogVectorOutput(plant.get_state_output_port(), builder)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    rng = np.random.default_rng(123)\n",
    "\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_mutable_context()\n",
    "    pi_context = pi.GetMyContextFromRoot(context)\n",
    "    pi_context.get_mutable_numeric_parameter(0).SetFromVector(alpha)\n",
    "\n",
    "    meshcat.Delete()\n",
    "    tf = 0.1 if running_as_test else 2.5\n",
    "    vertices = np.vstack([[0, tf], [0, 0], np.array([1, 1])*np.pi])\n",
    "    meshcat.SetLine(\"goal\", vertices, rgba=Rgba(0,.5,0))\n",
    "    vertices = np.vstack([[0, tf], [0, 0], -np.array([1, 1])*np.pi])\n",
    "    meshcat.SetLine(\"goal_wrapped\", vertices, rgba=Rgba(0,.5, 0))\n",
    "    for n in range(1 if running_as_test else N):\n",
    "        context.SetTime(0.)\n",
    "        context.SetContinuousState(rng.standard_normal((2,1)))\n",
    "        simulator.Initialize()\n",
    "        simulator.AdvanceTo(tf)\n",
    "        log = logger.FindLog(context)\n",
    "        times = log.sample_times()\n",
    "        vertices = np.vstack([times, 0*times, log.data()[0,:]])\n",
    "        meshcat.SetLine(f\"traj{n}\", vertices, rgba=Rgba(0, 0, .5))\n",
    "\n",
    "\n",
    "def trajectory_gradient_descent(alpha0):\n",
    "    builder = DiagramBuilder()\n",
    "    plant = builder.AddSystem(PendulumPlant())\n",
    "    plant.set_name(\"plant\")\n",
    "    pi = builder.AddSystem(ParameterizedController())\n",
    "    pi.set_name(\"pi\")\n",
    "    builder.Connect(plant.get_state_output_port(), pi.get_input_port())\n",
    "    builder.Connect(pi.get_output_port(), plant.get_input_port())\n",
    "\n",
    "    running_cost = builder.AddSystem(RunningCost())\n",
    "    running_cost.set_name(\"running_cost\")\n",
    "    builder.Connect(plant.get_state_output_port(),\n",
    "                    running_cost.get_input_port(0))\n",
    "    builder.Connect(pi.get_output_port(), running_cost.get_input_port(1))\n",
    "    cost = builder.AddSystem(Integrator(1))\n",
    "    cost.set_name(\"integrator\")\n",
    "    builder.Connect(running_cost.get_output_port(), cost.get_input_port())\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    diagram.set_name(\"trajectory gradient descent\")\n",
    "    # Useful if you want to poke around:\n",
    "    #display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString())[0].create_svg()))\n",
    "\n",
    "    diagram_ad = diagram.ToAutoDiffXd()\n",
    "    plant_ad = diagram_ad.GetSubsystemByName(\"plant\")\n",
    "    pi_ad = diagram_ad.GetSubsystemByName(\"pi\")\n",
    "    cost_ad = diagram_ad.GetSubsystemByName(\"integrator\")\n",
    "\n",
    "    simulator = Simulator_[AutoDiffXd](diagram_ad)\n",
    "    integrator = simulator.get_mutable_integrator()\n",
    "    integrator.set_fixed_step_mode(True)\n",
    "    integrator.set_maximum_step_size(0.05)\n",
    "\n",
    "    rng = np.random.default_rng(123)\n",
    "\n",
    "    context = simulator.get_mutable_context()\n",
    "    plant_context = plant_ad.GetMyContextFromRoot(context)\n",
    "    pi_context = pi_ad.GetMyContextFromRoot(context)\n",
    "    cost_context = cost_ad.GetMyContextFromRoot(context)\n",
    "\n",
    "    alpha = alpha0\n",
    "    N = 3 if running_as_test else 50\n",
    "    x0 = rng.standard_normal((2, N))\n",
    "    learning_rate = 2e-5\n",
    "    for epoch in range(2 if running_as_test else 2000):\n",
    "        J = 0\n",
    "        pi_context.get_mutable_numeric_parameter(0).SetFromVector(\n",
    "            InitializeAutoDiff(alpha))\n",
    "        for n in range(N):  # TODO(russt): Do multi-threading here.\n",
    "            context.SetTime(0.)\n",
    "            plant_context.SetContinuousState(x0[:,n])\n",
    "            cost_context.SetContinuousState([0])\n",
    "            simulator.Initialize()\n",
    "            simulator.AdvanceTo(2.5)\n",
    "            J += cost_context.get_continuous_state_vector().CopyToVector()/N\n",
    "\n",
    "        dJdalpha = ExtractGradient(J)\n",
    "        # TODO(russt): Use Adam here instead of SGD\n",
    "        alpha -= learning_rate*dJdalpha.T\n",
    "\n",
    "        clear_output(wait=True)\n",
    "        print(f\"epoch {epoch}: J = {ExtractValue(J)}\")\n",
    "        if epoch % 100 == 99:\n",
    "            simulate(alpha, N=5, set_planar=False)\n",
    "            print(alpha.T)\n",
    "        if epoch % 10 == 9:\n",
    "            plot_rollouts(alpha, N=5)\n",
    "\n",
    "    return alpha\n",
    "\n",
    "meshcat.Set2dRenderMode()\n",
    "#alpha=0.01*np.ones((9,1))\n",
    "# From a previous run:\n",
    "alpha = np.array([0.71027108, 2.61142863, 0.6379523, -0.23501137, 0.64446781, -1.45158088, 0.37477416, 1.05945338, -0.04337835]).reshape((9,1))\n",
    "alpha = trajectory_gradient_descent(alpha)\n",
    "print(alpha.T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_rollouts(alpha, N=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulate(alpha, N=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Static Output Feedback Counter-example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plt.rcParams['text.usetex'] = True\n",
    "def static_output_feedback():\n",
    "    A = np.array([[0, 0, 2], [1, 0, 0], [0, 1, 0]])\n",
    "    B = np.array([[1],[0],[0]])\n",
    "    C = np.array([[1, 1, 3]])\n",
    "\n",
    "    k = np.linspace(0, 3, 101)\n",
    "    lambda_max_real = 0*k\n",
    "    for i in range(101):\n",
    "        A_cl = (A - B @ (k[i]*C))\n",
    "        lambda_max_real[i] = np.amax(np.real(np.linalg.eigvals(A_cl)))\n",
    "        \n",
    "    fig, ax = plt.subplots(figsize=(6,3))\n",
    "    ax.plot(k, lambda_max_real)\n",
    "    ax.plot(k, 0*k, 'k')\n",
    "    ax.set_xlabel('$k$')\n",
    "    ax.set_ylabel('$\\Re(\\lambda_{max})$')\n",
    "    display(mpld3.display())\n",
    "\n",
    "static_output_feedback()"
   ]
  },
  {
   "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.10.9 (main, Dec 15 2022, 10:44:50) [Clang 14.0.0 (clang-1400.0.29.202)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}