zono_nn.ipynb

{
 "metadata": {
  "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": 2,
  "kernelspec": {
   "name": "python3810jvsc74a57bd00e4d859fb8a0ad1c8f202d01ca00ae23b252102f36a988fa3f0899d5572ebd11",
   "display_name": "Python 3.8.10 64-bit ('reach-net': conda)"
  },
  "metadata": {
   "interpreter": {
    "hash": "1482548d84cd7364095636d674edf388250952b5e53e00df7916bd6a74ee53b5"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import cvxpy as cp\n",
    "from cvxpylayers.torch import CvxpyLayer\n",
    "\n",
    "from util.collision_loss import center_param_collision_check, center_param_collision_loss\n",
    "from util.zonotope import Zonotope\n",
    "from util.constrained_zonotope import ConstrainedZonotope"
   ]
  },
  {
   "source": [
    "# Collision check"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_out_tch = 1.5*torch.ones(2, 1, requires_grad=True)\n",
    "G_out = np.array([[0.5, 0.5],[-0.5, 0.5]])\n",
    "obstacle = Zonotope(np.zeros((2,1)), np.eye(2))\n",
    "v = center_param_collision_check(c_out_tch, G_out, obstacle)\n",
    "print(\"collision check: \", v)\n",
    "fig, ax = plt.subplots()\n",
    "obstacle.plot(ax=ax, color='r')\n",
    "Z_out = Zonotope(c_out_tch.detach().numpy(), G_out)\n",
    "Z_out.plot(ax=ax, alpha=0.2)"
   ]
  },
  {
   "source": [
    "# Gradient updates to push out of collision"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_out = torch.ones(2, 1, requires_grad=True)\n",
    "G_out = np.array([[0.5, 0.5],[-0.5, 0.5]])\n",
    "obstacle = Zonotope(np.zeros((2,1)), np.eye(2))\n",
    "\n",
    "num_iters = 10\n",
    "learning_rate = 0.1\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10,10))\n",
    "obstacle.plot(ax=ax, color='r')\n",
    "\n",
    "# gradient updates:\n",
    "for i in range(num_iters):\n",
    "    #c_out.grad.zero_()   # zero the gradient buffers\n",
    "\n",
    "    # plot the current zonotope\n",
    "    Z_out = Zonotope(c_out.detach().numpy(), G_out)\n",
    "    Z_out.plot(ax=ax, alpha=0.2)\n",
    "\n",
    "    # compute loss and gradient\n",
    "    loss = center_param_collision_loss(c_out, G_out, obstacle)\n",
    "    loss.backward()\n",
    "\n",
    "    # apply gradient update\n",
    "    c_out.data.sub_(c_out.grad.data * learning_rate)\n",
    "    "
   ]
  },
  {
   "source": [
    "# Feedforward network"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3 Layer ReLU Feedforward Network\n",
    "# Output: zonotope of dimension 2, order 2 (6 total parameters)\n",
    "\n",
    "class Net(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        # an affine operation: y = Wx + b\n",
    "        self.fc1 = nn.Linear(4, 16)  # 5*5 from image dimension\n",
    "        self.fc2 = nn.Linear(16, 10)\n",
    "        self.fc3 = nn.Linear(10, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "net = Net()\n",
    "print(net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# arbitary input\n",
    "inpt = torch.tensor([1,2,3,4], dtype=torch.float)\n",
    "\n",
    "G_out = np.random.rand(2,4)\n",
    "obstacle = Zonotope(np.zeros((2,1)), np.random.rand(2,3))\n",
    "\n",
    "# \"train\" the network\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.1)\n",
    "\n",
    "num_iters = 10\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10,10))\n",
    "obstacle.plot(ax=ax, color='r')\n",
    "\n",
    "# training loop:\n",
    "for i in range(num_iters):\n",
    "    optimizer.zero_grad()   # zero the gradient buffers\n",
    "\n",
    "    # generate a new zonotope (center) from the network\n",
    "    out = net(inpt)\n",
    "    out = torch.reshape(out, (2,1))\n",
    "\n",
    "    # plot the resulting zonotope\n",
    "    Z_out = Zonotope(out.detach().numpy(), G_out)\n",
    "    Z_out.plot(ax=ax, alpha=0.2)\n",
    "\n",
    "    loss = center_param_collision_loss(out, G_out, obstacle)\n",
    "    loss.backward()\n",
    "    optimizer.step() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ]
}