constrained_lqr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import cvxpy as cp\n",
    "import numpy as np\n",
    "from cvxpylayers.torch import CvxpyLayer\n",
    "from scipy.linalg import sqrtm\n",
    "from scipy.linalg import solve_discrete_are\n",
    "\n",
    "torch.set_default_dtype(torch.double)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "\n",
    "# generate problem data\n",
    "n, m = 8, 2\n",
    "noise = np.sqrt(.25)\n",
    "u_max = .1\n",
    "Q0 = np.eye(n)\n",
    "R0 = np.eye(m)\n",
    "A = np.random.randn(n, n)\n",
    "A /= np.max(np.abs(np.linalg.eig(A)[0]))\n",
    "B = np.random.randn(n, m)\n",
    "W = noise**2 * np.eye(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.037499214903595\n"
     ]
    }
   ],
   "source": [
    "# compute lqr solution\n",
    "P = cp.Variable((n, n), PSD=True)\n",
    "R0cvxpy = cp.Parameter((m, m), PSD=True)\n",
    "\n",
    "objective = cp.trace(P@W)\n",
    "constraints = [cp.bmat([\n",
    "    [R0cvxpy + B.T@P@B, B.T@P@A],\n",
    "    [A.T@P@B, Q0+A.T@P@A-P]\n",
    "]) >> 0, P >> 0]\n",
    "R0cvxpy.value = R0\n",
    "result = cp.Problem(cp.Maximize(objective), constraints).solve()\n",
    "P_lqr = P.value\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "P_are = solve_discrete_are(A, B, Q0, R0)\n",
    "np.testing.assert_allclose(P_are, P_lqr, atol=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.794702792539413\n"
     ]
    }
   ],
   "source": [
    "# compute lower bound\n",
    "P = cp.Variable((n, n), PSD=True)\n",
    "R = cp.Variable((m, m), PSD=True)\n",
    "lam = cp.Variable(m, nonneg=True)\n",
    "\n",
    "objective = cp.trace(P@W) - (u_max**2)*cp.sum(lam)\n",
    "constraints = [R - R0 << cp.diag(lam), P >> 0, R >> 0, lam >= 0]\n",
    "constraints += [cp.bmat([\n",
    "    [R + B.T@P@B, B.T@P@A],\n",
    "    [A.T@P@B, Q0 + A.T@P@A-P]\n",
    "]) >> 0]\n",
    "result = cp.Problem(cp.Maximize(objective), constraints).solve()\n",
    "P_lb = P.value\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_lqr = -np.linalg.solve(R0 + B.T @ P_lqr @ B, B.T @ P_lqr @ A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up policy\n",
    "x = cp.Parameter((n, 1))\n",
    "P_sqrt = cp.Parameter((n, n))\n",
    "q = cp.Parameter(n)\n",
    "\n",
    "u = cp.Variable((m, 1))\n",
    "xnext = cp.Variable((n, 1))\n",
    "\n",
    "objective = cp.quad_form(u, R0) + cp.sum_squares(P_sqrt @ xnext) + q @ xnext\n",
    "constraints = [xnext == A @ x + B @ u, cp.norm(u, \"inf\") <= u_max]\n",
    "prob = cp.Problem(cp.Minimize(objective), constraints)\n",
    "policy = CvxpyLayer(prob, [x, P_sqrt, q], [u])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "Qt, Rt, At, Bt = map(torch.from_numpy, [Q0, R0, A, B])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss(time_horizon, batch_size, K, seed=None):\n",
    "    if seed is not None:\n",
    "        torch.manual_seed(seed)\n",
    "    x_batch = noise * torch.randn(batch_size, n, 1).double()\n",
    "    K_batch = K.repeat(batch_size, 1, 1)\n",
    "    Qt_batch = Qt.repeat(batch_size, 1, 1)\n",
    "    Rt_batch = Rt.repeat(batch_size, 1, 1)\n",
    "    At_batch = At.repeat(batch_size, 1, 1)\n",
    "    Bt_batch = Bt.repeat(batch_size, 1, 1)\n",
    "    loss = 0.0\n",
    "    for _ in range(time_horizon):\n",
    "        u_batch = torch.clamp(K @ x_batch, min=-u_max, max=u_max)\n",
    "        state_cost = torch.bmm(torch.bmm(Qt_batch, x_batch).transpose(2, 1), x_batch)\n",
    "        control_cost = torch.bmm(torch.bmm(Rt_batch, u_batch).transpose(2, 1), u_batch)\n",
    "        cost_batch = (state_cost.squeeze() + control_cost.squeeze())\n",
    "        loss += cost_batch.sum() / (time_horizon * batch_size)\n",
    "        x_batch = torch.bmm(At_batch, x_batch) + \\\n",
    "            torch.bmm(Bt_batch, u_batch) + \\\n",
    "            noise * torch.randn(batch_size, n, 1).double()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_clipped = loss(100, 6, torch.from_numpy(K_lqr), seed=0).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss(time_horizon, batch_size, P_sqrt, q, seed=None):\n",
    "    if seed is not None:\n",
    "        torch.manual_seed(seed)\n",
    "    x_batch = noise * torch.randn(batch_size, n, 1).double()\n",
    "    P_sqrt_batch = P_sqrt.repeat(batch_size, 1, 1)\n",
    "    q_batch = q.repeat(batch_size, 1)\n",
    "    Qt_batch = Qt.repeat(batch_size, 1, 1)\n",
    "    Rt_batch = Rt.repeat(batch_size, 1, 1)\n",
    "    At_batch = At.repeat(batch_size, 1, 1)\n",
    "    Bt_batch = Bt.repeat(batch_size, 1, 1)\n",
    "    loss = 0.0\n",
    "    for _ in range(time_horizon):\n",
    "        u_batch, = policy(x_batch, P_sqrt_batch, q_batch, solver_args={\"acceleration_lookback\": 0, \"eps\":1e-8, \"max_iters\":10000})\n",
    "        state_cost = torch.bmm(torch.bmm(Qt_batch, x_batch).transpose(2, 1), x_batch)\n",
    "        control_cost = torch.bmm(torch.bmm(Rt_batch, u_batch).transpose(2, 1), u_batch)\n",
    "        cost_batch = (state_cost.squeeze() + control_cost.squeeze())\n",
    "        loss += cost_batch.sum() / (time_horizon * batch_size)\n",
    "        x_batch = torch.bmm(At_batch, x_batch) + \\\n",
    "            torch.bmm(Bt_batch, u_batch) + \\\n",
    "            noise * torch.randn(batch_size, n, 1).double()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shane/miniconda3/envs/cvxpylayers/lib/python3.7/site-packages/diffcp/cone_program.py:259: UserWarning: Solved/Inaccurate.\n",
      "  warnings.warn(\"Solved/Inaccurate.\")\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    clf_lqr = loss(100, 6, torch.from_numpy(sqrtm(P_lqr)), torch.zeros(n, dtype=torch.double), seed=0).item()\n",
    "    clf_lb = loss(100, 6, torch.from_numpy(sqrtm(P_lb)), torch.zeros(n, dtype=torch.double), seed=0).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(16.940040755077668, 12.975429366507662)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clf_lqr, clf_lb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize parameters\n",
    "torch.manual_seed(0)\n",
    "P_sqrt = torch.from_numpy(sqrtm(P_lqr)); P_sqrt.requires_grad_(True);\n",
    "q = torch.zeros(n).double(); q.requires_grad_(True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "it: 001, loss: 16.940, dist: 20.232\n",
      "it: 002, loss: 14.258, dist: 18.828\n",
      "it: 003, loss: 14.035, dist: 18.613\n",
      "it: 004, loss: 13.795, dist: 18.323\n",
      "it: 005, loss: 13.677, dist: 17.990\n",
      "it: 006, loss: 13.436, dist: 17.716\n",
      "it: 007, loss: 13.311, dist: 17.600\n",
      "it: 008, loss: 13.265, dist: 17.450\n",
      "it: 009, loss: 13.331, dist: 17.374\n",
      "it: 010, loss: 13.231, dist: 17.224\n",
      "it: 011, loss: 13.333, dist: 17.065\n",
      "it: 012, loss: 13.397, dist: 17.053\n",
      "it: 013, loss: 13.474, dist: 17.010\n",
      "it: 014, loss: 13.248, dist: 16.918\n",
      "it: 015, loss: 13.383, dist: 17.065\n",
      "it: 016, loss: 13.296, dist: 17.016\n",
      "it: 017, loss: 13.347, dist: 16.944\n",
      "it: 018, loss: 13.295, dist: 16.676\n",
      "it: 019, loss: 13.265, dist: 16.674\n",
      "it: 020, loss: 13.175, dist: 16.662\n",
      "it: 021, loss: 13.159, dist: 16.625\n",
      "it: 022, loss: 13.211, dist: 16.636\n",
      "it: 023, loss: 13.231, dist: 16.481\n",
      "it: 024, loss: 13.234, dist: 16.443\n",
      "it: 025, loss: 13.150, dist: 16.425\n",
      "it: 026, loss: 13.065, dist: 16.363\n",
      "it: 027, loss: 13.053, dist: 16.355\n",
      "it: 028, loss: 13.062, dist: 16.311\n",
      "it: 029, loss: 13.102, dist: 16.175\n",
      "it: 030, loss: 13.066, dist: 16.092\n",
      "it: 031, loss: 13.236, dist: 16.033\n",
      "it: 032, loss: 13.273, dist: 15.978\n",
      "it: 033, loss: 13.232, dist: 15.981\n",
      "it: 034, loss: 13.398, dist: 16.056\n",
      "it: 035, loss: 13.273, dist: 16.043\n",
      "it: 036, loss: 13.270, dist: 16.134\n",
      "it: 037, loss: 13.166, dist: 16.077\n",
      "it: 038, loss: 13.134, dist: 16.085\n",
      "it: 039, loss: 13.091, dist: 16.217\n",
      "it: 040, loss: 13.106, dist: 16.156\n",
      "it: 041, loss: 13.066, dist: 16.273\n",
      "it: 042, loss: 13.164, dist: 16.026\n",
      "it: 043, loss: 13.083, dist: 16.015\n",
      "it: 044, loss: 13.021, dist: 15.920\n",
      "it: 045, loss: 13.044, dist: 15.888\n",
      "it: 046, loss: 13.039, dist: 15.816\n",
      "it: 047, loss: 13.058, dist: 15.771\n",
      "it: 048, loss: 13.089, dist: 15.847\n",
      "it: 049, loss: 13.081, dist: 15.821\n",
      "it: 050, loss: 13.149, dist: 15.707\n",
      "it: 051, loss: 13.145, dist: 15.753\n",
      "it: 052, loss: 13.073, dist: 15.599\n",
      "it: 053, loss: 13.065, dist: 15.589\n",
      "it: 054, loss: 13.065, dist: 15.591\n",
      "it: 055, loss: 13.065, dist: 15.589\n",
      "it: 056, loss: 13.058, dist: 15.584\n",
      "it: 057, loss: 13.058, dist: 15.581\n",
      "it: 058, loss: 13.062, dist: 15.586\n",
      "it: 059, loss: 13.064, dist: 15.587\n",
      "it: 060, loss: 13.065, dist: 15.586\n",
      "it: 061, loss: 13.068, dist: 15.591\n",
      "it: 062, loss: 13.061, dist: 15.590\n",
      "it: 063, loss: 13.061, dist: 15.590\n",
      "it: 064, loss: 13.060, dist: 15.595\n",
      "it: 065, loss: 13.058, dist: 15.596\n",
      "it: 066, loss: 13.060, dist: 15.598\n",
      "it: 067, loss: 13.056, dist: 15.600\n",
      "it: 068, loss: 13.056, dist: 15.610\n",
      "it: 069, loss: 13.047, dist: 15.602\n",
      "it: 070, loss: 13.044, dist: 15.596\n",
      "it: 071, loss: 13.041, dist: 15.601\n",
      "it: 072, loss: 13.040, dist: 15.604\n",
      "it: 073, loss: 13.041, dist: 15.615\n",
      "it: 074, loss: 13.040, dist: 15.605\n",
      "it: 075, loss: 13.040, dist: 15.603\n",
      "it: 076, loss: 13.034, dist: 15.592\n",
      "it: 077, loss: 13.033, dist: 15.584\n",
      "it: 078, loss: 13.030, dist: 15.588\n",
      "it: 079, loss: 13.024, dist: 15.590\n",
      "it: 080, loss: 13.019, dist: 15.591\n",
      "it: 081, loss: 13.024, dist: 15.585\n",
      "it: 082, loss: 13.020, dist: 15.578\n",
      "it: 083, loss: 13.022, dist: 15.582\n",
      "it: 084, loss: 13.029, dist: 15.596\n",
      "it: 085, loss: 13.023, dist: 15.585\n",
      "it: 086, loss: 13.019, dist: 15.582\n",
      "it: 087, loss: 13.024, dist: 15.581\n",
      "it: 088, loss: 13.014, dist: 15.575\n",
      "it: 089, loss: 13.016, dist: 15.571\n",
      "it: 090, loss: 13.012, dist: 15.557\n",
      "it: 091, loss: 13.013, dist: 15.552\n",
      "it: 092, loss: 13.012, dist: 15.555\n",
      "it: 093, loss: 13.012, dist: 15.555\n",
      "it: 094, loss: 13.017, dist: 15.561\n",
      "it: 095, loss: 13.019, dist: 15.565\n",
      "it: 096, loss: 13.018, dist: 15.551\n",
      "it: 097, loss: 13.017, dist: 15.544\n",
      "it: 098, loss: 13.013, dist: 15.541\n",
      "it: 099, loss: 13.015, dist: 15.539\n",
      "it: 100, loss: 13.014, dist: 15.537\n"
     ]
    }
   ],
   "source": [
    "opt = torch.optim.SGD([P_sqrt, q], lr=.1)\n",
    "losses = []\n",
    "for k in range(100):\n",
    "    with torch.no_grad():\n",
    "        test_loss = loss(100, 6, P_sqrt.detach(), q.detach(), seed=0).item()\n",
    "        losses.append(test_loss)\n",
    "        P_np = (P_sqrt.t() @ P_sqrt).detach().numpy()\n",
    "        dist = np.linalg.norm(P_np - P_lb)\n",
    "        print(\"it: %03d, loss: %3.3f, dist: %3.3f\" % (k+1, test_loss, dist))\n",
    "    opt.zero_grad()\n",
    "    l = loss(100, 6, P_sqrt, q, seed=k+1)\n",
    "    l.backward()\n",
    "    opt.step()\n",
    "    if k == 50:\n",
    "        opt = torch.optim.SGD([P_sqrt, q], lr=.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASYAAAC6CAYAAAADF+BWAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8QZhcZAAAem0lEQVR4nO3dfWxT1/kH8O/JK0kIMSal6wtlcWjQpAqBY7rBxpq1DtuqRkyQhD/YJtoUZ+vWCVZ+ZAGqTistTUpZWatSu5saqUVasJk0iOjWeFOZWpiaxD/YCwhY/Kt4kSoUghkQ0iT28/vD916uHTuxE78cO89Hsojv68G5fnLOvec8RxARGGNMJjnpLgBjjIXjwMQYkw4HJsaYdDgwMcakw4GJMSYdDkyMMenkpeOkQggzAAsROXTLrAAMAHwAQETudJSNMZZ+Ka8xKQGoOWyZAYCBiFxKQDKnulyMMXmkPDApgac7bJkPQKsQwqQsmpfqcjHG5JGWplwULQD6hBBuImqItIEQog5AXWlp6aaqqqrUlo4xlnB9fX0DRHRX+HKRjiEpQoh6AMawe0zbAHgBtAIYJKLaaPtbLBbq7e1NfkEZY0klhOgjIkv4cimeyimByqPcY6oG4FXuRTHGZiApAhMAI4K1JZUz7D1jbAZJ+T0mpSa0HoBBCOElIjcROYQQNiEEAAwC8BERBybGZqiUByblqdy4Pkr6+02MsZlNpqdyjMUkEAjg0qVLuHXrVrqLwiaRn5+P+fPnY86cOXHtx4GJZZyBgQEIIbB48WLk5Mhym5SFIyLcvn0bly9fBoC4ghP/VlnG8fl8uPvuuzkoSU4IgeLiYtx33324cuVKXPvyb5ZlHL/fj/z8/HQXg8WoqKgIo6Ojce2TlsAkhDALIWy69wYhxDUhRJ/utS0dZWOZQXmCyzLAVH5X6eou0ACgT7fYBKBCGTMHIUQ9EblSXTbGmBzS0l1AySZg1C3zqD8rKVE8kfZljM0MMj6VsxJReyIOdPToUdx///1YsmRJIg7HWEzcbjc8Hg/M5mD2HqPRiM7OTrS1tWnrfT4fTCYTBgcH4fP5UF9fH9Mxmpub0dLSAq/Xi9bWVgBAT08PamtrYbVm0SguIkr5C0A9AFuE5VYA9RPsVwfAsWjRIopFWVkZ/exnP4tpW5Y5Tp8+ne4iRGW322nbtm0hy/r7+8lsNmvr7XZ7yHqn0xmyz2THcDqdZLPZQtYbDAa6du1awv4fiRbtdwaglyJ812WrMTUAsEdbSURHAByxWCybYjlYWVkZ/vvf/yaqbIxNqrm5Wf0jqjGZTFi/fj0AoKWlBdeuXQtZX19fj5aWFrS2tsJgMEx6jEiMRiO8Xq9Ww8p0sgUmC4J5mRJizpw5HJhmgM2bN+PkyZNJPcfSpUvx+uuvT7iN2+2GyWSKuG7btm0TrjeZTHC73TAYDBMeIxKvNzisNFuCEiDJIF79elKezCXCnDlzcP369UQdjrFp8/kiX96Dg4NxHcfr9Wr3qgYHB9Hf35+I4klDmkG8yrrqRJ6rrKwMAwMDiTwkk9BkNZlUsVqtWu0lnNvthsViiRqA1GaYyWSa8BjqDW6TyZRdN7vDZHXPb27KsVRzOp1obg6ZawM+nw9GoxEGgwFtbW1obw996OxwOGCz2bQm3ETHmClku8eUUNyUY6lWX18Pg8GA9vZ2mEwmGAwGANBqNzabDS6XCy6XCwaDQWvaqV0JJjuG1+tFZ2cnvF5vSHeCbJP1gYlrTCzVrFbrhM2s8D5L8RzDZDLB6XROq3yZIKubcmVlZRgaGsLY2Fi6i8IYi0NWByY1/wvXmhjLLFJkF1CWGZS839bwdVPFgYmxzCTFFOEKJxE5lO4EUeeUi0dZWRkADkyMZRopsgsoGQV8um0izsQbL64xMZaZZHkqZwEwqNSmDABACcjHpAYm7jLAWGaRJTAZEJwy3A0AQohuIYSHwuaWE0LUAahbtGhRTAflphxjmUmWp3JehM6860Mwq2UIIjpCRDY14EyGm3KMZSZZApMboYHIhARMEc5NOcYykxTZBYjIJ4To1HUT2B3ejJuK4uJi5Obmco2JsQyT8hqTEogaiKhWn/KEiFxKdwFHIm58A8HZGXhYCptJ3G43KisrU35er9eL2tpaeDyJSdcvS1MuaXggL5tJrFarNug3lUwmU0IHFM+IwMQ1JjaTpCs9yrx58xJ2rKwPTJz3e2aoqakZ93rrrbcAAENDQxHXd3R0AAAGBgYiru/s7AQAXLx4MaYy6JtRatPG7XZr6+bOnavNfqLOdDLZuvb2drjdbjgcDi1rZWVlJdxuNxoaGqJmxHS73XC73Whvbw9JPOdwOODxeOB2u+FyuSYtt8fj0c7ndrvR3Nwcck6Hw6GVu6enJ6bPKRay9GNKmjlz5sQ9bzpjU2G1WrVkbyaTCbW1tePWqalMDAYDGhoa0NfXF3WdmixOXd7Q0ACn0wmTyQSj0Thh+hN1H6vViurqavT19WmBSG1ytbS0aOeNVm41q6b68ng86O3thdVq1VL76rNqJsqMCEznz59PdzFYkn300UdR1xUXF0+4vry8fML1CxYsmHrBdPRNrPAUupHW9fX1obq6WtsuPGDESq3hdHd3o6HhzmivyspK9Pb2xnSsSM3D7u5uLF++POZyxIObcowlkD4rZX9/f9Smljrh5UTr1ECk1lYaGxunXCYAIUEOAPr6+mCxWGIqd6Qb6suXLx+XnzzeSRWikSntSZuS+sSQqLQnAN/8Zqm1fv16HDx4UHts3tnZGfIld7lc8Hg8cDgc45pi4evUTJculwtut1tLp9vb2wuHwxG1DGazWbvvoz+PzWaDz+fT7i/V1tZqtaVo5fZ4PPB6vXA4HPD5fOju7obT6QyZPVg9l8fjgd1ujxqM4xJpFsxkvhCcbdeOsJl4AXQD6AfgBGCY6BjV1dUxzwD64osvEgAaGRmJeR8mN5ln4p2I1Wqd0rpsEO9MvGnpYKkEoXB2IqqkYOfLhM0txwN5mQzUmo56AzrWdTOVTDe/TYlOewKEDuRNZD8LxuJhNpvHTQ0ey7qZSprARETaZFtCiD4hhDu85hRv2hOAMwwwlomkeConhKgXQugnZh9EAtKeAHeacjwshbHMIUuNKTwfk5GIEjIakGtMjGUeWdKeeJRakwnBmlJLos7HgYmxzJOWyQgQTAwXvjwpjyQ4WRxjmUeKe0zJxN0FGMs8WR+YZs2ahby8PA5MjGWQrA9MahZLbsoxljmyPjABPJCXpYaauygd0pFSN9HpdPVmRGDigbwsFdTcRemQjpS6iU6nq8eBibEskY6Uuska5iVLB8ukuv/++/GnP/0JAwMDKC8vT3dxWBLU1NRMus0TTzyBrVu3attv3LgRGzduxMDAgJbCI5qJEslF43A4YLFYMDg4qKUJcblc2LRpE/7yl7/A7Xajs7NT+9lut8PpdMLhcMBsNsPr9cJqtcLr9aK5uRl2ux12ux3vvPNO1NqRPiVufX09TCZTxHKoaXL7+/u147e0tMBqtcLj8aChoQF2ux0A4HQ60dbWpp3T4XBoWTR7enq0DJaJJE0+Jt26eqUTZsJs374dN27cQEtLwvptMjYhfRpbq9WKnp4eLViozS6bLfgVMBgMMJlMcDqdOHjwIIBg08xms2nBQp9Od6Imm9VqhdVqxbZt29DQ0BC1HLGm07VarVqmSwAh6XST2XRNV8/vBgB9EdYZEOwVbk/kOR966CFs2bIFr776Kp566il8/etfT+ThmQTirdHot58ste5UTJTGVq391NbWorm5WavlmM3mhKXTBaAldsuUdLp6MuVjAgALgMRNtaDzwgsv4IEHHsCPfvQjjI2NJeMUjGkmSmOrJvIHgMbGRq3JBCBh6XSBYE0sk9Lp6klz81sIYSXdzLyJVlJSgl//+tf417/+hd/97nfJOg2bwfRpaCdKYwvcmTXFYDDAaDRq92mmmk4XiJxSN6PS6eqIYHbL1BJC1COYQcChvDco771K+hNPpCCly8e0aSoznxARHnnkEZw9exbnz5/XxtGxzHLmzBl85StfSXcxWByi/c6EEH1EZAlfHleNSQixVPfzY0KIR6dUyvEsuJPBcjmAWiXTQIip5GPSE0Jg7969uHLlCtra2qZXYsZY0sTblNOCBRH9BRGSuU2FkvrErdSSvAC6icg72X5TYbFYsGHDBuzduxf//Oc/k3EKxtg0xRSYhBCbhBC9ANqFED3K689TOaEuH1NDeLcAIYQZwVlUmiPVmBJl9+7dmDt3LlatWoVjx44l6zSMsSmKqbsAEb0jhDgIwERE/zudE0bLx6Ss8wCons7xY7FgwQKcOHEC3/3ud7F69Wp0dXWFPJZljKVXzE05IrquD0pCiC8no0CpsnDhQnz88ceoqKjAT3/6U+5CkGHS8dCGTU0gEIh7n3hvfu8WQiwTQuwH8CMhxNNxn1EiRqMR7e3tOHfuHHchyCCzZs3C1atXOThJjogwMjKCy5cvo6SkJK594+ouIISoIKL/E0L0EpFFCPGYchM8pSwWC6ld5KeLiLBq1Sr09/fjP//5T9wfoHqMU6dOYcmSJcjJkaZrWNYaHR3FpUuXMDw8nO6isEnk5eWhrKwM5eXlEb8b0boLxDu99zoAawH8j/J+bTz7J+oVzxThsfjkk08IAO3cuTPufU+dOkWrVq0iALR79+6ElouxbIcETRHuBbCciF4VQmxCgroLpNvKlSuxYcMG7Nq1C/v27YtpH6/Xi6amJpjNZpw+fRrLli3DSy+9hM8//xwA8Pbbb2PHjh0YGhrS9vniiy+SUv7pOHr0KEwmE/785yk9ZGUsOSJFq4leAJYBeBrA0nj3TdQr0TUmIqIvvviC1q5dSwDo+eefp8HBQQoEAnT48GH69re/TQcOHCAiotHRUdq2bRvl5uZSYWEhPfvsszQwMEDnzp2j/Px8evrpp+mNN94gAASAHnzwQXrjjTeopqaGANAjjzxCv//97+nmzZsJ/z/Eq7+/nwwGAwkhqKCggLq6utJdJDbDIEqNaSpNud0ANgF4G1nSlFONjo7Shg0bCAAVFBRQVVUVAaDi4mICQD/5yU/o0UcfJQDU1NREly9fDtl/y5YtJIQgALRmzRr68MMPaeHChQSAKioq6Nlnn6WKigoCQLm5uVRdXU179+6lsbGxqGUaHh6mK1eu0OXLl2loaGjK/7fLly/T/v376Xvf+x7V19fTe++9R8uWLSODwUC9vb1UXV1N+fn59LWvfY2++c1v0p49eygQCEz5fIzFIlGB6bGw9+vi2T9Rr2QFJiKiQCBAvb29tHnzZvrGN75Bv/3tb2loaIi2bNlCAKiwsJDefffdiPsODg7SvffeS6tXr6bh4WEiIrpx4wZ5PB7y+/1EROT3+6m7u5t27NhBK1asIABUU1NDn332GRERjY2NUVdXFzU2NtKDDz5IOTk5Wu3LaDTS+++/H3fAOH36NM2ePZsA0MKFC+nee+/VjnnkyBEiIrp27Ro1NTXR6tWryWw2EwBqbGyUombHsleiAtPSsPePxrO/bj8zAFvYMquyvD58XfgrmYFpIt3d3fSPf/xjwm1u3boVc+AIBAL07rvvUklJiRZ45s+fTwBo/vz5VF9fT88//zy9+eabtH//fi2Q1dXVjautRTM0NEQPPfQQlZeX08mTJykQCJDf76cTJ05Qd3d31HK1t7eTEIIWLFhATU1NdODAAbp+/XrE7f1+P73yyiv08MMPk9vtjqlckYyNjVFHRwc1NTXRM888Qy+88AJ9/vnnUz4ek1+0wBRvd4FNyl9aL4I3vg1EtCfmAyA0URyFZhdwElGt8p6ISEQ7RiK7C8jA6/XC5XLhs88+g8/nw9q1a7FmzRrk5+eHbOf3+7Fv3z7s2LEDs2bNwr59+7BmzRotS8LQ0BD+/ve/49ChQzh37hy+9a1v4cyZMzhw4AA++OADfOc734mrXN3d3XjzzTfxt7/9DT6fD4WFhbBarfD7/bh06RJMJhPWrFmDP/7xjzh8+LCW3+epp57Crl27cM8990Q99vXr1/Hxxx/j+PHjEEKgpKQEBw4cwL///W8tj/S1a9dQUlKC1tZWrFq1CgaDAYWFhRBCaK/c3FyUlpaitLQUOTk5ICKMjo5ieHgYIyMjGBsbw9jYGPx+P8bGxrRlRUVFKCsrw6xZs0LKlZOTg7y8POTl5SEQCMDv9yMvLw8FBQUQIuolyaYoWneBuNOeCCHWAahFMNL9doqFCUl7ErbODKCZiJqj7Z9tgSleZ8+exZNPPokTJ04AAIqKiuD3+zEyMgIAKC4uRmVlpTZIuaWlBa+88sqUz+f3+/Hpp5+is7MTH3zwAUpLS3Hffffh5MmTuHDhAvLy8rB37148/fTT+NWvfoVXX30VeXl5aGpqgtFoxKlTp3D9+nUUFxdjdHQU58+fx8WLF0FEyMvLAxHB7/ejqqoKL730EtatWwchBM6dO4etW7fiyJEj0//QEiA3NxdCCOTn58NoNGLu3LkoLCxEfn4+CgoKkJ+fj/z8fOTm5iInJweBQACBQCAkkAJAYWEhSktLUVxcjJycHOTm5qKgoAAFBQUoKSnB7NmzUVJSgqKiIgghcPv2bYyMjKCgoACFhYUIBAIYHR3FrVu3cOPGDRCRts+sWbOQn5+P4eFh3L59G7Nnz0Z5eTkKCwvh9/sBBDuoFhYWAgi2mIQQyMnJ0cqnHvv27dsYGxvTem6r2+Xm5iIvL0+bTFYN/Or6iooKLFiwIKbPNCGBSQixDMFm1o+FEGUAKojoZMwHuHOciIFJCUrriWjC5NylpaVUXR06pK6xsRHPPPMMhoaG8Pjjj4/bZ7LE8z/+8Y+xfv16XLx4ET/4wQ/GrX/uuedQV1eHs2fPorl5fMzcuXMnrFYrTp48ic2bN49b//LLL2PlypU4fvw4tm/fPm7966+/jqVLl8LtdmPXrl3j1tvtdixevBhHjhzBa6+9BiLC1atXtYu2rq4OCxcuxJUrV9DT04OcnByMjIzg5s2bmDt3Lg4dOoTy8nJ0dHSgo6Nj3PGPHj2K4uJivPXWW1reaT019eyePXvQ1dWlLb958yZKSkq0wdAvvvgiurq6cOHCBa3rRGlpKcxmM4aGhnDp0iX4/X4UFxejrKwMc+bMwQMPPIC3334bxcXF+PnPf46TJ0Mvqbvuugs2mw3Xrl3D/v37teMCwS/W3XffjdraWhARXC4Xbty4oX3RcnJysGjRIqxfvx65ubn4zW9+g1u3biEQCGhfuqqqKqxevRpEBLvdjpGREe0LK4RAVVUVHn74YYyNjeH999/X9lX3nz9/Pr70pS9haGgIZ86cudMcUb7o5eXlMBqNGBsbg9frDTk3AK1mPDw8jHgrCjKqqKjAAw88oL13uVxRr71jx45FDExTyfn9CyA4dk4IkdD5YojII4TwCiH6iWjc7H1qorjw6vdMJIQImfGltbUVCxYsQGdnJ/r6gunUCwoKkj6lz+zZs1FUVBSyrKioCIsXL0ZlZSVycnJw11134dChQ1o51Zpe+HGimTt3rjbI2u12j/vyLl26FDt37gQQrE1eunQpZP2SJUvwwx/+EADwhz/8AVevXg1Zv2LFCu2PyYcffojbt2+HrH/ssce02VU++eSTceVL1B/FCxcu4Pvf/z78fr/WjAQAm82Gxx9/HGfOnEFra2tIMzY3Nxe/+MUv8NWvfhWffvopfvnLXyIQCGi1Nr/fjyeffBKLFy/GmTNn8M4772g1OdWmTZuwcOFCnDp1CgcPHtSOrQb3nTt34stf/jKOHTuG9957DwC0YxARWltbcc899+Cvf/0rDh8+PK55PBXx1pgeJaK/6t6vJaI/xH3S8Rkszcp7t/K+H8HmXMQsBDO9KcdYtojWlIu3xlQphKhFcMKA5QCuTrJ9rCwAwjOaJyVRHGNMfnEFJgrmZXoMwZvfPUR0KN4T6hLFGYQQXgpmrnQIIWxKTcoEoIWSlMGSMSa/uO8xUTCbwJQzClCURHGRntAxxmYmztHBGJMOBybGmHQ4MDHGpMOBiTEmHQ5MjDHpcGBijEmHAxNjTDocmBhj0uHAxBiTzlSyC0ybMmjXou/trQxHAYJj8HqIyJWOsjHG0i/lgUmfwVK3zAzApwxXcQkhrgkh3ETkS3X5GGPpl/KmnBJ8usMWmxAcGKwaRJbMWccYi19amnLhlGabC9Dyf4OIPGktFGMsbWS8+d2G0NqTRghRJ4RwXL9+PcVFYoylklSBSbkB3hYtFxMRHSEiW1lZWYpLxhhLJWkCk3JT3ENEXiGEQQjB95gYm6HS9VQuJIOl8lTODsCnzCxhIqK5qS4bY0wOKQ9MkTJYKje6x82KwhibmaRpyjHGmIoDE2NMOhyYGGPS4cDEGJMOBybGmHQ4MDHGpMOBiTEmHQ5MjDHpcGBijEmHAxNjTDocmBhj0uHAxBiTTloCkxDCLISwxbqcMTazpDwwKWlPmmNdzhibedKS9kTJ622MZTljbObJqHtMnPObsZkhowIT5/xmbGbIqMDEGJsZODAxxqQjxWQEEy2P5OzZs6ipqZnwPE888QS2bt0KAKipqcHGjRuxceNGDAwMoL6+ftJyhm//3HPPoa6uDmfPnkVz8+QPD8O3f/nll7Fy5UocP34c27dvn3T/8O3tdjsWL16MI0eO4LXXXpt0//DtXS4XysvL0dHRgY6Ojkn3D9/+o48+AgDs2bMHXV1dk+6v3/7EiRM4dOgQAKC1tRUnTpyYcN958+aFbH/16lU4HA4AgM1mw7lz5ybcv6qqKmT7efPmYffu3QCAdevW4erVqxPuv2LFipDtV6xYEXItTYavvelde4AkkxFMtJwxNvMIIkp3GeJmsViot7c33cVgjE2TEKKPiCzhy/keE2NMOhyYGGPS4cDEGJMOBybGmHQ4MDHGpMOBiTEmHQ5MjDHpcGBijEmHAxNjTDocmBhj0pEiMAkhDEIIqxBiW7rLwhhLPykCExH5APDgN8YYgCRnFxBCmAFYiMihW1YPwAfAAMBLRJ5kloExlnmSVmOKNOuJMtlALRG5icgFoDVZ52eMZa6kBSYlv1J32OJGBGtLGqVWxRhjmlTfYzIA6Ne9HwRgUn5uBFArhDCN24sxNqOkPINlNMp9KEe09coMveosvTeFEGdjPHQ5gIFpFi+VMq28AJc5FTKtvEBsZV4YaWGqA5N601tlBOCNZcfJAlc0QojeSBnyZJVp5QW4zKmQaeUFplfmVDflDgKo1L038FM5xli4pNWYIs16QkQ+IYRTWQcAbck6P2MscyUtME0yG0oqxd38S7NMKy/AZU6FTCsvMI0yZ+QsKYyx7CbFkBTGGNOTprsAyyzq0CLdTMo81GiGiWfIWbzXR9YGpkz5oijlBIDlAHqUoTpSl18ZWrQegF33vpaImpX3TgAN6SthKKV8jQh2TTGpXyRZP2Pl4ZAByigJGYO/UsYGAH26ZRGvg6lcH1nZlMuUMXnKXxwfEbmIqAXAO0oKGNnLbwHQo3sv+1AjJxE5lC94LSDvNaKUy6BcE24AZt1yacob55CzuK+PrAxMkP+LojJB+aIo1CE60pZfCGGN8GR1oqFGaaUGf/U9Eal/qaX8jJUUQK26oVnzlH+lLG+YaNdB3NdHtgYmab8oerqakvoXEUr1XMryK2WMqae+RCwABpVEhPW6prOUn7GiBUCfEMKpXh+Qu7wJl7X3mDJQG0JrTzKyAIDy13w5gHlCCC+mMdQoBQwAjLr7NN1CCCnuJU3ADGATgjWnbiKS/bpQRbsOjFGWR5WtNSaZvyjjKH/F24hILaOU5Vfub7iVL7kXQLdSZpmHGnkR+tn5EKxpSPkZK9eCR6lNVwPwKjeapSxvmGjXQdzXR7YGJpm/KCGUi85DRF7lxrcJkpdfubdhBdAshDAp90WcSnPJCrmGGrkR2uQxIfiFlvUzDg84TkhYXt2QswZ1iFm062Aq10fW9vzWjccDkJahMJNSvuBO3LmpaSKiuco66cufKZRaiFF5O6jrkiHlZ6yk+AGC95H0fcWkLG8yZG1gYoxlrmxtyjHGMhgHJsaYdDgwMcakw4GJMSYdDkxs2oQQ/ZNvlfnnZKnDgYlNGxFp/Wt0j7oTJtIx9edk2YcDE5sWIYRJHX+mjKVrnmSXeI8/7pj6c7LsxIGJTdcg7vTkNSE4+US9fuS7EGKb0uvXpgQVqxCiX/nXqQ5gVt6bhRBtutH1kY6pPyeU45rVgbrKMrPuHFYhhF09D5MfD+Jl06LMfONVfvYIIXxqz2rgTjNM13vZSUQNyj6DujQkANCsrAOCtaSWSMfUn1MNRLpMiW3KrDweZRuvMtzHjOAg5KztLZ1NuMbEkq0agE+pKZmgSy4WYaxXizLswoLQAasTqUXo2LJ+ZX/VYPxFZunGgYkl2iAQMq6rGwCIyKvLRDCOsn2LftovXXMu/Jh6fQgdpFsNoFd9owwgZRmGAxObFqWJZNI9OXPqn6Lpc5grgcWkNqvCnrYN6o5nQHDQrSnSMfXnVPJ3G3T3l7qVZpy2jZqWFnfyTzPJ8SBexph0uMbEGJMOBybGmHQ4MDHGpMOBiTEmHQ5MjDHpcGBijEmHAxNjTDocmBhj0uHAxBiTzv8DvFC4aN9TwbMAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 324x200.243 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from utils import latexify\n",
    "\n",
    "latexify(fig_width=4.5)\n",
    "plt.semilogy(losses, c='k', label='COCP')\n",
    "plt.gca().yaxis.set_minor_formatter(matplotlib.ticker.ScalarFormatter())\n",
    "plt.axhline(clf_lb, linestyle='--', c='k', label='upper bound')\n",
    "plt.axhline(result, linestyle='-.', c='k', label='lower bound')\n",
    "plt.xlabel('iteration')\n",
    "plt.ylabel('cost')\n",
    "plt.subplots_adjust(left=.15, bottom=.2)\n",
    "plt.ylim(10, 18)\n",
    "plt.legend(loc='upper right')\n",
    "plt.savefig(\"lqr_constrained.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from torch.nn import functional\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, n, K):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(n, 64)\n",
    "        self.fc2 = nn.Linear(64, 64)\n",
    "        self.fc3 = nn.Linear(64, m)\n",
    "        self.fc_K = nn.Linear(n, m, bias=False)\n",
    "        self.fc_K.weight.data = K.data\n",
    "#         for fc in [self.fc1, self.fc2, self.fc3]:\n",
    "#             fc.weight.data.zero_()\n",
    "#             fc.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x):\n",
    "        u = self.fc_K(x) \n",
    "        x = functional.relu(self.fc1(x))\n",
    "        x = functional.relu(self.fc2(x))\n",
    "        x = torch.clamp(self.fc3(x), -u_max, u_max)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = Net(n, torch.from_numpy(K_lqr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss(time_horizon, batch_size, seed=None):\n",
    "    if seed is not None:\n",
    "        torch.manual_seed(seed)\n",
    "    x_batch = noise * torch.randn(batch_size, n, 1).double()\n",
    "    Qt_batch = Qt.repeat(batch_size, 1, 1)\n",
    "    Rt_batch = Rt.repeat(batch_size, 1, 1)\n",
    "    At_batch = At.repeat(batch_size, 1, 1)\n",
    "    Bt_batch = Bt.repeat(batch_size, 1, 1)\n",
    "    loss = 0.0\n",
    "    for _ in range(time_horizon):\n",
    "        u_batch = policy(x_batch.squeeze(-1))\n",
    "        u_batch = u_batch.unsqueeze(-1)\n",
    "        state_cost = torch.bmm(torch.bmm(Qt_batch, x_batch).transpose(2, 1), x_batch)\n",
    "        control_cost = torch.bmm(torch.bmm(Rt_batch, u_batch).transpose(2, 1), u_batch)\n",
    "        cost_batch = (state_cost.squeeze() + control_cost.squeeze())\n",
    "        loss += cost_batch.sum() / (time_horizon * batch_size)\n",
    "        x_batch = torch.bmm(At_batch, x_batch) + \\\n",
    "            torch.bmm(Bt_batch, u_batch) + \\\n",
    "            noise * torch.randn(batch_size, n, 1).double()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "it: 001, loss: 13.605\n",
      "it: 002, loss: 13.605\n",
      "it: 003, loss: 13.604\n",
      "it: 004, loss: 13.596\n",
      "it: 005, loss: 13.575\n",
      "it: 006, loss: 13.551\n",
      "it: 007, loss: 13.527\n",
      "it: 008, loss: 13.505\n",
      "it: 009, loss: 13.501\n",
      "it: 010, loss: 13.491\n",
      "it: 011, loss: 13.489\n",
      "it: 012, loss: 13.482\n",
      "it: 013, loss: 13.471\n",
      "it: 014, loss: 13.464\n",
      "it: 015, loss: 13.456\n",
      "it: 016, loss: 13.448\n",
      "it: 017, loss: 13.440\n",
      "it: 018, loss: 13.431\n",
      "it: 019, loss: 13.415\n",
      "it: 020, loss: 13.397\n",
      "it: 021, loss: 13.379\n",
      "it: 022, loss: 13.360\n",
      "it: 023, loss: 13.343\n",
      "it: 024, loss: 13.330\n",
      "it: 025, loss: 13.321\n",
      "it: 026, loss: 13.313\n",
      "it: 027, loss: 13.309\n",
      "it: 028, loss: 13.304\n",
      "it: 029, loss: 13.297\n",
      "it: 030, loss: 13.293\n",
      "it: 031, loss: 13.289\n",
      "it: 032, loss: 13.283\n",
      "it: 033, loss: 13.279\n",
      "it: 034, loss: 13.275\n",
      "it: 035, loss: 13.273\n",
      "it: 036, loss: 13.275\n",
      "it: 037, loss: 13.276\n",
      "it: 038, loss: 13.279\n",
      "it: 039, loss: 13.280\n",
      "it: 040, loss: 13.280\n",
      "it: 041, loss: 13.278\n",
      "it: 042, loss: 13.278\n",
      "it: 043, loss: 13.277\n",
      "it: 044, loss: 13.276\n",
      "it: 045, loss: 13.272\n",
      "it: 046, loss: 13.267\n",
      "it: 047, loss: 13.260\n",
      "it: 048, loss: 13.254\n",
      "it: 049, loss: 13.250\n",
      "it: 050, loss: 13.251\n",
      "it: 051, loss: 13.249\n",
      "it: 052, loss: 13.247\n",
      "it: 053, loss: 13.245\n",
      "it: 054, loss: 13.246\n",
      "it: 055, loss: 13.247\n",
      "it: 056, loss: 13.254\n",
      "it: 057, loss: 13.259\n",
      "it: 058, loss: 13.264\n",
      "it: 059, loss: 13.269\n",
      "it: 060, loss: 13.275\n",
      "it: 061, loss: 13.279\n",
      "it: 062, loss: 13.282\n",
      "it: 063, loss: 13.284\n",
      "it: 064, loss: 13.285\n",
      "it: 065, loss: 13.286\n",
      "it: 066, loss: 13.287\n",
      "it: 067, loss: 13.285\n",
      "it: 068, loss: 13.280\n",
      "it: 069, loss: 13.274\n",
      "it: 070, loss: 13.267\n",
      "it: 071, loss: 13.257\n",
      "it: 072, loss: 13.244\n",
      "it: 073, loss: 13.236\n",
      "it: 074, loss: 13.232\n",
      "it: 075, loss: 13.227\n",
      "it: 076, loss: 13.218\n",
      "it: 077, loss: 13.214\n",
      "it: 078, loss: 13.210\n",
      "it: 079, loss: 13.202\n",
      "it: 080, loss: 13.197\n",
      "it: 081, loss: 13.189\n",
      "it: 082, loss: 13.186\n",
      "it: 083, loss: 13.179\n",
      "it: 084, loss: 13.173\n",
      "it: 085, loss: 13.169\n",
      "it: 086, loss: 13.166\n",
      "it: 087, loss: 13.169\n",
      "it: 088, loss: 13.171\n",
      "it: 089, loss: 13.172\n",
      "it: 090, loss: 13.168\n",
      "it: 091, loss: 13.165\n",
      "it: 092, loss: 13.165\n",
      "it: 093, loss: 13.169\n",
      "it: 094, loss: 13.174\n",
      "it: 095, loss: 13.178\n",
      "it: 096, loss: 13.182\n",
      "it: 097, loss: 13.186\n",
      "it: 098, loss: 13.187\n",
      "it: 099, loss: 13.188\n",
      "it: 100, loss: 13.193\n",
      "it: 101, loss: 13.197\n",
      "it: 102, loss: 13.197\n",
      "it: 103, loss: 13.199\n",
      "it: 104, loss: 13.201\n",
      "it: 105, loss: 13.196\n",
      "it: 106, loss: 13.190\n",
      "it: 107, loss: 13.183\n",
      "it: 108, loss: 13.180\n",
      "it: 109, loss: 13.185\n",
      "it: 110, loss: 13.189\n",
      "it: 111, loss: 13.195\n",
      "it: 112, loss: 13.210\n",
      "it: 113, loss: 13.230\n",
      "it: 114, loss: 13.245\n",
      "it: 115, loss: 13.252\n",
      "it: 116, loss: 13.250\n",
      "it: 117, loss: 13.248\n",
      "it: 118, loss: 13.244\n",
      "it: 119, loss: 13.241\n",
      "it: 120, loss: 13.237\n",
      "it: 121, loss: 13.230\n",
      "it: 122, loss: 13.226\n",
      "it: 123, loss: 13.223\n",
      "it: 124, loss: 13.220\n",
      "it: 125, loss: 13.231\n",
      "it: 126, loss: 13.242\n",
      "it: 127, loss: 13.251\n",
      "it: 128, loss: 13.259\n",
      "it: 129, loss: 13.262\n",
      "it: 130, loss: 13.260\n",
      "it: 131, loss: 13.252\n",
      "it: 132, loss: 13.239\n",
      "it: 133, loss: 13.227\n",
      "it: 134, loss: 13.214\n",
      "it: 135, loss: 13.206\n",
      "it: 136, loss: 13.199\n",
      "it: 137, loss: 13.197\n",
      "it: 138, loss: 13.195\n",
      "it: 139, loss: 13.197\n",
      "it: 140, loss: 13.198\n",
      "it: 141, loss: 13.199\n",
      "it: 142, loss: 13.198\n",
      "it: 143, loss: 13.197\n",
      "it: 144, loss: 13.199\n",
      "it: 145, loss: 13.205\n",
      "it: 146, loss: 13.206\n",
      "it: 147, loss: 13.204\n",
      "it: 148, loss: 13.203\n",
      "it: 149, loss: 13.197\n",
      "it: 150, loss: 13.192\n",
      "it: 151, loss: 13.182\n",
      "it: 152, loss: 13.174\n",
      "it: 153, loss: 13.166\n",
      "it: 154, loss: 13.167\n",
      "it: 155, loss: 13.167\n",
      "it: 156, loss: 13.167\n",
      "it: 157, loss: 13.169\n",
      "it: 158, loss: 13.173\n",
      "it: 159, loss: 13.175\n",
      "it: 160, loss: 13.176\n",
      "it: 161, loss: 13.185\n",
      "it: 162, loss: 13.192\n",
      "it: 163, loss: 13.196\n",
      "it: 164, loss: 13.202\n",
      "it: 165, loss: 13.206\n",
      "it: 166, loss: 13.209\n",
      "it: 167, loss: 13.209\n",
      "it: 168, loss: 13.207\n",
      "it: 169, loss: 13.200\n",
      "it: 170, loss: 13.188\n",
      "it: 171, loss: 13.174\n",
      "it: 172, loss: 13.161\n",
      "it: 173, loss: 13.154\n",
      "it: 174, loss: 13.151\n",
      "it: 175, loss: 13.146\n",
      "it: 176, loss: 13.143\n",
      "it: 177, loss: 13.141\n",
      "it: 178, loss: 13.135\n",
      "it: 179, loss: 13.127\n",
      "it: 180, loss: 13.123\n",
      "it: 181, loss: 13.122\n",
      "it: 182, loss: 13.123\n",
      "it: 183, loss: 13.123\n",
      "it: 184, loss: 13.126\n",
      "it: 185, loss: 13.131\n",
      "it: 186, loss: 13.136\n",
      "it: 187, loss: 13.141\n",
      "it: 188, loss: 13.151\n",
      "it: 189, loss: 13.158\n",
      "it: 190, loss: 13.163\n",
      "it: 191, loss: 13.166\n",
      "it: 192, loss: 13.165\n",
      "it: 193, loss: 13.164\n",
      "it: 194, loss: 13.164\n",
      "it: 195, loss: 13.165\n",
      "it: 196, loss: 13.166\n",
      "it: 197, loss: 13.165\n",
      "it: 198, loss: 13.161\n",
      "it: 199, loss: 13.156\n",
      "it: 200, loss: 13.151\n",
      "it: 201, loss: 13.148\n",
      "it: 202, loss: 13.145\n",
      "it: 203, loss: 13.143\n",
      "it: 204, loss: 13.142\n",
      "it: 205, loss: 13.141\n",
      "it: 206, loss: 13.141\n",
      "it: 207, loss: 13.142\n",
      "it: 208, loss: 13.144\n",
      "it: 209, loss: 13.145\n",
      "it: 210, loss: 13.148\n",
      "it: 211, loss: 13.150\n",
      "it: 212, loss: 13.153\n",
      "it: 213, loss: 13.153\n",
      "it: 214, loss: 13.154\n",
      "it: 215, loss: 13.154\n",
      "it: 216, loss: 13.156\n",
      "it: 217, loss: 13.158\n",
      "it: 218, loss: 13.163\n",
      "it: 219, loss: 13.167\n",
      "it: 220, loss: 13.173\n",
      "it: 221, loss: 13.178\n",
      "it: 222, loss: 13.178\n",
      "it: 223, loss: 13.178\n",
      "it: 224, loss: 13.177\n",
      "it: 225, loss: 13.181\n",
      "it: 226, loss: 13.180\n",
      "it: 227, loss: 13.179\n",
      "it: 228, loss: 13.176\n",
      "it: 229, loss: 13.174\n",
      "it: 230, loss: 13.171\n",
      "it: 231, loss: 13.172\n",
      "it: 232, loss: 13.172\n",
      "it: 233, loss: 13.176\n",
      "it: 234, loss: 13.179\n",
      "it: 235, loss: 13.179\n",
      "it: 236, loss: 13.183\n",
      "it: 237, loss: 13.185\n",
      "it: 238, loss: 13.188\n",
      "it: 239, loss: 13.191\n",
      "it: 240, loss: 13.195\n",
      "it: 241, loss: 13.196\n",
      "it: 242, loss: 13.192\n",
      "it: 243, loss: 13.191\n",
      "it: 244, loss: 13.191\n",
      "it: 245, loss: 13.197\n",
      "it: 246, loss: 13.199\n",
      "it: 247, loss: 13.197\n",
      "it: 248, loss: 13.193\n",
      "it: 249, loss: 13.188\n",
      "it: 250, loss: 13.180\n",
      "it: 251, loss: 13.173\n",
      "it: 252, loss: 13.167\n",
      "it: 253, loss: 13.165\n",
      "it: 254, loss: 13.163\n",
      "it: 255, loss: 13.162\n",
      "it: 256, loss: 13.161\n",
      "it: 257, loss: 13.159\n",
      "it: 258, loss: 13.159\n",
      "it: 259, loss: 13.160\n",
      "it: 260, loss: 13.164\n",
      "it: 261, loss: 13.168\n",
      "it: 262, loss: 13.176\n",
      "it: 263, loss: 13.188\n",
      "it: 264, loss: 13.200\n",
      "it: 265, loss: 13.214\n",
      "it: 266, loss: 13.226\n",
      "it: 267, loss: 13.235\n",
      "it: 268, loss: 13.241\n",
      "it: 269, loss: 13.242\n",
      "it: 270, loss: 13.244\n",
      "it: 271, loss: 13.247\n",
      "it: 272, loss: 13.254\n",
      "it: 273, loss: 13.260\n",
      "it: 274, loss: 13.263\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-128-6a9e4fb645ff>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[0ml\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m6\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mseed\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m     \u001b[0ml\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m     \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;31m#     if k == 200:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/cvxpylayers/lib/python3.7/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m    148\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m         \"\"\"\n\u001b[0;32m--> 150\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    151\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    152\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/cvxpylayers/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     97\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     98\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 99\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    100\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "opt = torch.optim.Adam(policy.parameters(), lr=3e-4)\n",
    "losses = []\n",
    "for k in range(10000):\n",
    "    with torch.no_grad():\n",
    "        test_loss = loss(100, 6, seed=0).item()\n",
    "        losses.append(test_loss)\n",
    "        print(\"it: %03d, loss: %3.3f\" % (k+1, test_loss))\n",
    "    opt.zero_grad()\n",
    "    l = loss(100, 6, seed=k+1)\n",
    "    l.backward()\n",
    "    opt.step()\n",
    "#     if k == 200:\n",
    "#         opt = torch.optim.SGD(policy.parameters(), lr=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f3832723410>]"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plt.plot(losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Piecewise quadratic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up policy\n",
    "x = cp.Parameter(n)\n",
    "P_sqrt_cp = [cp.Parameter((n, n), PSD=True) for _ in range(K)]\n",
    "q_cp = [cp.Parameter(n) for _ in range(K)]\n",
    "r_cp = [cp.Parameter(1) for _ in range(K)]\n",
    "\n",
    "u = cp.Variable(m)\n",
    "xnext = cp.Variable(n)\n",
    "t = cp.Variable(1)\n",
    "\n",
    "objective = cp.quad_form(u, R0) + t\n",
    "constraints = [xnext == A @ x + B @ u, cp.norm(u, \"inf\") <= u_max]\n",
    "constraints += [cp.sum_squares(P_sqrt_cp[i] @ xnext) + q_cp[i] @ xnext + r_cp[i] <= t for i in range(K)]\n",
    "prob = cp.Problem(cp.Minimize(objective), constraints)\n",
    "policy = CvxpyLayer(prob, [x] + P_sqrt_cp + q_cp + r_cp, [u])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "P_sqrt = [torch.from_numpy(sqrtm(P_lqr) + 1e-3*np.random.randn(n, n)) for _ in range(K)]\n",
    "for i in range(K):\n",
    "    P_sqrt[i].requires_grad_(True)\n",
    "q = [torch.randn(n, requires_grad=True) for _ in range(K)]\n",
    "r = [torch.zeros(1, requires_grad=True) for _ in range(K)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss(time_horizon, batch_size, seed=None):\n",
    "    if seed is not None:\n",
    "        torch.manual_seed(seed)\n",
    "    x_batch = noise * torch.randn(batch_size, n).double()\n",
    "    Qt_batch = Qt.repeat(batch_size, 1, 1)\n",
    "    Rt_batch = Rt.repeat(batch_size, 1, 1)\n",
    "    At_batch = At.repeat(batch_size, 1, 1)\n",
    "    Bt_batch = Bt.repeat(batch_size, 1, 1)\n",
    "    loss = 0.0\n",
    "    for _ in range(time_horizon):\n",
    "        inputs = [x_batch] + P_sqrt + q + r\n",
    "        u_batch, = policy(*inputs, solver_args={\"acceleration_lookback\": 0, \"eps\":1e-5, \"max_iters\":10000})\n",
    "        state_cost = torch.bmm(torch.bmm(Qt_batch, x_batch.unsqueeze(-1)).transpose(2, 1), x_batch.unsqueeze(-1))\n",
    "        control_cost = torch.bmm(torch.bmm(Rt_batch, u_batch.unsqueeze(-1)).transpose(2, 1), u_batch.unsqueeze(-1))\n",
    "        cost_batch = (state_cost.squeeze() + control_cost.squeeze())\n",
    "        loss += cost_batch.sum() / (time_horizon * batch_size)\n",
    "        x_batch = torch.bmm(At_batch, x_batch.unsqueeze(-1)) + \\\n",
    "            torch.bmm(Bt_batch, u_batch.unsqueeze(-1)) + \\\n",
    "            noise * torch.randn(batch_size, n, 1).double()\n",
    "        x_batch = x_batch.squeeze(-1)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "def norm(list_of_pytorch_tensors):\n",
    "    with torch.no_grad():\n",
    "        return torch.cat([t.view(-1) for t in list_of_pytorch_tensors]).norm().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "it: 001, loss: 13.368, grad_norm: 0.000\n",
      "it: 002, loss: 13.218, grad_norm: 2.971\n",
      "it: 003, loss: 13.293, grad_norm: 2.776\n",
      "it: 004, loss: 13.455, grad_norm: 5.518\n",
      "it: 005, loss: 13.715, grad_norm: 3.672\n",
      "it: 006, loss: 13.730, grad_norm: 3.191\n",
      "it: 007, loss: 13.817, grad_norm: 3.788\n",
      "it: 008, loss: 13.735, grad_norm: 4.652\n",
      "it: 009, loss: 13.792, grad_norm: 4.192\n",
      "it: 010, loss: 13.480, grad_norm: 4.183\n",
      "it: 011, loss: 13.299, grad_norm: 6.466\n",
      "it: 012, loss: 13.519, grad_norm: 3.917\n",
      "it: 013, loss: 13.386, grad_norm: 4.412\n",
      "it: 014, loss: 13.521, grad_norm: 2.561\n",
      "it: 015, loss: 13.845, grad_norm: 3.820\n",
      "it: 016, loss: 13.671, grad_norm: 1.890\n",
      "it: 017, loss: 13.580, grad_norm: 2.818\n",
      "it: 018, loss: 13.508, grad_norm: 6.822\n",
      "it: 019, loss: 13.659, grad_norm: 8.734\n",
      "it: 020, loss: 13.598, grad_norm: 3.174\n",
      "it: 021, loss: 13.446, grad_norm: 4.614\n",
      "it: 022, loss: 13.518, grad_norm: 4.279\n",
      "it: 023, loss: 13.473, grad_norm: 5.674\n",
      "it: 024, loss: 13.466, grad_norm: 3.205\n",
      "it: 025, loss: 13.397, grad_norm: 4.298\n",
      "it: 026, loss: 13.527, grad_norm: 3.969\n",
      "it: 027, loss: 13.528, grad_norm: 4.837\n",
      "it: 028, loss: 13.424, grad_norm: 2.147\n",
      "it: 029, loss: 13.414, grad_norm: 3.381\n",
      "it: 030, loss: 13.374, grad_norm: 4.972\n",
      "it: 031, loss: 13.466, grad_norm: 4.453\n",
      "it: 032, loss: 13.887, grad_norm: 4.018\n",
      "it: 033, loss: 13.674, grad_norm: 6.750\n",
      "it: 034, loss: 13.524, grad_norm: 4.174\n",
      "it: 035, loss: 13.533, grad_norm: 3.244\n",
      "it: 036, loss: 13.760, grad_norm: 7.816\n",
      "it: 037, loss: 13.812, grad_norm: 2.345\n",
      "it: 038, loss: 13.576, grad_norm: 6.047\n",
      "it: 039, loss: 13.528, grad_norm: 4.764\n",
      "it: 040, loss: 13.522, grad_norm: 2.829\n",
      "it: 041, loss: 13.400, grad_norm: 2.327\n",
      "it: 042, loss: 13.427, grad_norm: 3.172\n",
      "it: 043, loss: 13.596, grad_norm: 1.978\n",
      "it: 044, loss: 13.360, grad_norm: 3.592\n",
      "it: 045, loss: 13.376, grad_norm: 1.863\n",
      "it: 046, loss: 13.648, grad_norm: 10.000\n",
      "it: 047, loss: 13.634, grad_norm: 2.687\n",
      "it: 048, loss: 13.604, grad_norm: 2.535\n",
      "it: 049, loss: 13.847, grad_norm: 6.168\n",
      "it: 050, loss: 13.795, grad_norm: 3.298\n",
      "it: 051, loss: 13.851, grad_norm: 2.573\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'param' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-63-353835ae8f59>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mk\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m50\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m         \u001b[0mopt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSGD\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.005\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'param' is not defined"
     ]
    }
   ],
   "source": [
    "params = P_sqrt + q + r\n",
    "opt = torch.optim.SGD(params, lr=0.1)\n",
    "losses = []\n",
    "for k in range(100):\n",
    "    with torch.no_grad():\n",
    "        test_loss = loss(100, 6, seed=0).item()\n",
    "        losses.append(test_loss)\n",
    "        if params[0].grad is None:\n",
    "            grad_norm = np.nan\n",
    "        else:\n",
    "            grad_norm = norm([p.grad.data for p in params])\n",
    "        print(\"it: %03d, loss: %3.3f, grad_norm: %3.3f\" % (k+1, test_loss, grad_norm))\n",
    "    opt.zero_grad()\n",
    "    l = loss(100, 6, seed=k+1)\n",
    "    l.backward()\n",
    "    torch.nn.utils.clip_grad_norm_(params, 10)\n",
    "    opt.step()\n",
    "    if k == 50:\n",
    "        opt = torch.optim.SGD(params, lr=0.005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}