check_02_vae.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from vae.arch import VAE\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from IPython import display\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual, FloatSlider\n",
    "import ipywidgets as widgets\n",
    "from tensorflow.keras.utils import plot_model\n",
    "\n",
    "np.set_printoptions(precision=4, suppress = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = VAE()\n",
    "vae.set_weights('./vae/weights.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIR_NAME = './data/rollout/'\n",
    "file = os.listdir(DIR_NAME)[200]\n",
    "print(file)\n",
    "obs_data = np.load(DIR_NAME + file)['obs']\n",
    "print(obs_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_model(vae.encoder, to_file='./vae/encoder.png', show_shapes=True)\n",
    "plot_model(vae.decoder, to_file='./vae/decoder.png', show_shapes=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### output from one episode\n",
    "\n",
    "for idx in range(0,300):\n",
    "    plt.gca().cla() \n",
    "    \n",
    "    obs = obs_data[idx]\n",
    "    z_decoded = vae.full_model.predict(np.array([obs]))[0]\n",
    "    \n",
    "    plt.subplot(121)\n",
    "    plt.imshow( obs)\n",
    "    plt.subplot(122)\n",
    "    plt.imshow( z_decoded)\n",
    "    \n",
    "#     plt.show()\n",
    "    display.clear_output(wait=True)\n",
    "    display.display(plt.gcf()) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### output from the full_model\n",
    "DIR_NAME = './data/rollout/'\n",
    "file = os.listdir(DIR_NAME)[179]\n",
    "obs_data = np.load(DIR_NAME + file)['obs']\n",
    "\n",
    "obs = obs_data[50]\n",
    "reconstruction = vae.full_model.predict(np.array([obs]))[0]\n",
    "\n",
    "ax1 = plt.subplot(121)\n",
    "plt.imshow( obs)\n",
    "ax1.axis('off')\n",
    "ax1.text(0.5,-0.1, \"INPUT\", size=12, ha=\"center\", \n",
    "         transform=ax1.transAxes)\n",
    "\n",
    "ax2 = plt.subplot(122)\n",
    "plt.imshow( reconstruction)\n",
    "ax2.axis('off')\n",
    "ax2.text(0.5,-0.1, \"OUTPUT\", size=12, ha=\"center\", \n",
    "         transform=ax2.transAxes);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mus, log_vars, _  = vae.encoder.predict(np.array([obs]))\n",
    "mu = mus[0]\n",
    "log_var = log_vars[0]\n",
    "print(\"mu = \" + str(mu))\n",
    "print(\"log_var = \" + str(log_var))\n",
    "\n",
    "encoded_z = vae.encoder.predict(np.array([obs]))[0]\n",
    "print(\"z = \" + str(encoded_z))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(mu);\n",
    "plt.plot(log_var);\n",
    "print('informative dimensions:')\n",
    "hot_zs = np.where(abs(log_var) > 0.5)[0]\n",
    "hot_zs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file1 = os.listdir(DIR_NAME)[100]\n",
    "obs_data1 = np.load(DIR_NAME + file)['obs']\n",
    "file2 = os.listdir(DIR_NAME)[140]\n",
    "obs_data2 = np.load(DIR_NAME + file)['obs']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.encoder.predict(np.array([obs_data1[20]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_left = vae.encoder.predict(np.array([obs_data1[20]]))[2][0]\n",
    "top_right = vae.encoder.predict(np.array([obs_data1[60]]))[2][0]\n",
    "bottom_left = vae.encoder.predict(np.array([obs_data2[80]]))[2][0]\n",
    "sample_z = vae.encoder.predict(np.array([obs_data[60]]))[2][0]\n",
    "\n",
    "fig, axes = plt.subplots(nrows=5, ncols=5, figsize=(32,32), sharex=True, sharey=True)\n",
    "plt.subplots_adjust(hspace = 0.4)\n",
    "\n",
    "for x in range(5):\n",
    "    for y in range(5):\n",
    "        \n",
    "        z = sample_z.copy()\n",
    "\n",
    "#         z = top_left + (top_right - top_left) * x/8 + (bottom_left - top_left) * y/8\n",
    "        z[7] = (x-2)\n",
    "        z[10] = (y-2)\n",
    "        \n",
    "        decoded_obs = vae.decoder.predict(np.array([z]))[0]\n",
    "        \n",
    "        axes[x,y].set_axis_off()\n",
    "        \n",
    "        axes[x,y].imshow(decoded_obs)\n",
    "        \n",
    "        axes[x,y].text(0.5,-0.1, 'z[7] ' + str(np.round(z[7],2)), size=20, ha=\"center\", \n",
    "                       transform=axes[x,y].transAxes)\n",
    "        \n",
    "        axes[x,y].text(0.5,-0.2, 'z[10] = ' + str(np.round(z[10],2)), size=20, ha=\"center\", \n",
    "                       transform=axes[x,y].transAxes)\n",
    "    \n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def play_with_z(z0 = sample_z[0]\n",
    "                ,z2 = sample_z[2]\n",
    "                ,z4 = sample_z[4]\n",
    "                ,z5 = sample_z[5]\n",
    "                ,z6 = sample_z[6]\n",
    "                ,z7 = sample_z[7]\n",
    "                ,z10 = sample_z[10]\n",
    "                ,z15 = sample_z[15]\n",
    "                ,z24= sample_z[24]\n",
    "                ,z25= sample_z[25]\n",
    "                ,z27= sample_z[27]\n",
    "                ):\n",
    "    z = sample_z.copy()\n",
    "    z[0] = z0\n",
    "    z[2] = z2\n",
    "    z[4] = z4\n",
    "    z[5] = z5\n",
    "    z[6] = z6\n",
    "    z[7] = z7\n",
    "    z[10] = z10\n",
    "    z[15] = z15\n",
    "    z[24] = z24\n",
    "    z[25] = z25\n",
    "    z[27] = z27\n",
    "    \n",
    "    decoded_obs = vae.decoder.predict(np.array([z]))[0]\n",
    "    plt.gca().cla() \n",
    "    plt.imshow( decoded_obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interact(play_with_z\n",
    "        , z0=(-3.0,3.0,0.1)\n",
    "        , z2=(-3.0,3.0,0.1)\n",
    "        , z4=(-3.0,3.0,0.1)\n",
    "        , z5=(-3.0,3.0,0.1)\n",
    "        , z6=(-3.0,3.0,0.1)\n",
    "        , z7=(-3.0,3.0,0.1)\n",
    "        , z10=(-3.0,3.0,0.1)\n",
    "        , z15=(-3.0,3.0,0.1)\n",
    "        , z24=(-3.0,3.0,0.1)\n",
    "        , z25=(-3.0,3.0,0.1)\n",
    "        , z27=(-3.0,3.0,0.1)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "worldmodels_3.8",
   "language": "python",
   "name": "worldmodels_3.8"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}