training_evaluation.ipynb

{
 "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",
    "from timeit import default_timer as timer\n",
    "\n",
    "import cvxpy as cp\n",
    "from cvxpylayers.torch import CvxpyLayer\n",
    "\n",
    "from util.collision_loss import torch_collision_check, NN_constraint_step\n",
    "from util.zonotope import Zonotope, TorchZonotope\n",
    "from util.constrained_zonotope import TorchConstrainedZonotope\n",
    "from util.NN_con_zono import forward_pass_NN_torch, forward_pass_NN_con_zono_torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "torch.cuda.set_device(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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(2, 10)  \n",
    "        self.fc2 = nn.Linear(10, 2)  \n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "net = Net()\n",
    "net.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load saved parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.load_state_dict(torch.load(\"log/constrained_1_1000its.pt\"))\n",
    "#net.load_state_dict(torch.load(\"log/constrained_2_400its.pt\"))\n",
    "net.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Constrained training\n",
    "Train function approximator under some obstacle constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input zonotope\n",
    "Z_in = TorchZonotope(torch.zeros(2,1).to(device),torch.eye(2).to(device))\n",
    "\n",
    "# output constraint zonotope (\"obstacle\")\n",
    "c_obs = torch.tensor([[1.5],[1.5]]).to(device)\n",
    "G_obs = torch.diag(torch.tensor([0.5,0.5]).to(device))\n",
    "Z_obs = TorchConstrainedZonotope(c_obs, G_obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# recompute reachable set\n",
    "Z_out = forward_pass_NN_torch(Z_in, net)\n",
    "\n",
    "# plot zonotopes\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20,10))\n",
    "for z in Z_out:\n",
    "    z.plot(ax1)\n",
    "Z_obs.plot(ax1, 'r')\n",
    "# plot samples\n",
    "N_samples = 10000\n",
    "X_in = np.random.uniform(-1, 1, (N_samples,2))\n",
    "X_in = torch.as_tensor(X_in, dtype=torch.float).to(device)\n",
    "Y_out = net(X_in)\n",
    "color_vec = torch.sum(X_in,dim=1).cpu()\n",
    "ax2.scatter(Y_out[:,0].cpu().detach().numpy(), Y_out[:,1].cpu().detach().numpy(),c=color_vec,cmap=\"gist_rainbow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "values = []\n",
    "losses = []\n",
    "for z in Z_out:\n",
    "    v = torch_collision_check(z, Z_obs)\n",
    "    if v <= 1: # in collision\n",
    "        print(\"in collision\")\n",
    "        values.append(v.item())\n",
    "        loss = torch.square(1 - v)\n",
    "        losses.append(loss.item())\n",
    "\n",
    "print(\"collision values: \",values)\n",
    "print(\"losses: \",losses)\n",
    "print(\"total loss: \",sum(losses))"
   ]
  }
 ],
 "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.8.5"
  },
  "metadata": {
   "interpreter": {
    "hash": "1482548d84cd7364095636d674edf388250952b5e53e00df7916bd6a74ee53b5"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}