check_05_controller.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": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "from rnn.arch import RNN\n",
    "from vae.arch import VAE\n",
    "import matplotlib.pyplot as plt\n",
    "from gym.utils import seeding\n",
    "from IPython import display\n",
    "import time\n",
    "from model import make_model\n",
    "\n",
    "import config\n",
    "\n",
    "\n",
    "np.set_printoptions(precision=4, suppress = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mixture_coef(z_pred):\n",
    "\n",
    "    log_pi, mu, log_sigma = np.split(z_pred, 3, 1)\n",
    "    log_pi = log_pi - np.log(np.sum(np.exp(log_pi), axis = 1, keepdims = True))\n",
    "\n",
    "    return log_pi, mu, log_sigma\n",
    "\n",
    "def get_pi_idx(x, pdf):\n",
    "    # samples from a categorial distribution\n",
    "    N = pdf.size\n",
    "    accumulate = 0\n",
    "    for i in range(0, N):\n",
    "        accumulate += pdf[i]\n",
    "        if (accumulate >= x):\n",
    "            return i\n",
    "    random_value = np.random.randint(N)\n",
    "    #print('error with sampling ensemble, returning random', random_value)\n",
    "    return random_value\n",
    "\n",
    "def sample_z(mu, log_sigma):\n",
    "    z =  mu + (np.exp(log_sigma)) * np.random.randn(*log_sigma.shape) \n",
    "    return z\n",
    "\n",
    "\n",
    "def get_z_from_rnn_output(y_pred):\n",
    "    HIDDEN_UNITS = 256\n",
    "    GAUSSIAN_MIXTURES = 5\n",
    "    Z_DIM = 32\n",
    "    d = GAUSSIAN_MIXTURES * Z_DIM\n",
    "\n",
    "    z_pred = y_pred[:(3*d)]\n",
    "    rew_pred = y_pred[-1]\n",
    "\n",
    "    z_pred = np.reshape(z_pred, [-1, GAUSSIAN_MIXTURES * 3])\n",
    "\n",
    "    log_pi, mu, log_sigma = get_mixture_coef(z_pred)\n",
    "\n",
    "    chosen_log_pi = np.zeros(Z_DIM)\n",
    "    chosen_mu = np.zeros(Z_DIM)\n",
    "    chosen_log_sigma = np.zeros(Z_DIM)\n",
    "\n",
    "    # adjust temperatures\n",
    "    logmix2 = np.copy(log_pi)\n",
    "    logmix2 -= logmix2.max()\n",
    "    logmix2 = np.exp(logmix2)\n",
    "    logmix2 /= logmix2.sum(axis=1).reshape(Z_DIM, 1)\n",
    "\n",
    "\n",
    "    for j in range(Z_DIM):\n",
    "        idx = get_pi_idx(np.random.rand(), logmix2[j])\n",
    "        chosen_log_pi[j] = idx\n",
    "        chosen_mu[j] = mu[j, idx]\n",
    "        chosen_log_sigma[j] = log_sigma[j,idx]\n",
    "\n",
    "    next_z = sample_z(chosen_mu, chosen_log_sigma)\n",
    "\n",
    "    # print(next_z)\n",
    "    # print(rew_pred)\n",
    "    if rew_pred > 0:\n",
    "        next_reward = 1\n",
    "    else:\n",
    "        next_reward = 0\n",
    "        \n",
    "    return next_z, next_reward, chosen_mu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = make_model()\n",
    "model.make_env('car_racing')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_model('./controller/car_racing.cma.1.4.best.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_weight = model.weight[0][:32,2]\n",
    "h_weight = model.weight[0][32:,2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rollout_files = np.load('./data/rollout/35106062.npz') \n",
    "obs_file = rollout_files['obs']\n",
    "action_file = rollout_files['action']\n",
    "reward_file = rollout_files['reward']\n",
    "done_file = rollout_files['done']\n",
    "\n",
    "series_files = np.load('./data/series/35106062.npz') \n",
    "mu_file = series_files['mu']\n",
    "log_var_file = series_files['log_var']\n",
    "action_2_file = series_files['action']\n",
    "reward_2_file = series_files['reward']\n",
    "done_2_file = series_files['done']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = obs_file[4]\n",
    "action = [0,0,0]\n",
    "reward = 0\n",
    "model.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_encoded_obs = model.update(obs, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_encoded_obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recon = model.vae.decoder.predict(np.array([vae_encoded_obs]))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(recon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obs = obs_file[0]\n",
    "action = [0,1,0]\n",
    "reward = 0\n",
    "model.reset()\n",
    "total_reward = 0\n",
    "total_pseudo_reward = 0\n",
    "t = 0\n",
    "\n",
    "obs = model.env.reset()\n",
    "\n",
    "model.env.render('rgb_array')\n",
    "\n",
    "actions0 = []\n",
    "actions1 = []\n",
    "actions2 = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "while (1):\n",
    "   \n",
    "    ######\n",
    "\n",
    "    obs = config.adjust_obs(obs)\n",
    "    reward = config.adjust_reward(reward)\n",
    "    \n",
    "    total_pseudo_reward+= reward\n",
    "\n",
    "    vae_encoded_obs = model.update(obs, 0)\n",
    "\n",
    "    recon = model.vae.decoder.predict(np.array([vae_encoded_obs]))[0]\n",
    "    \n",
    "#     input_to_rnn = [np.array([[np.concatenate([vae_encoded_obs, action, [reward]])]]),np.zeros(shape=(1,256)),np.zeros(shape=(1,256))]\n",
    "#     input_to_rnn = [np.array([[np.concatenate([np.zeros(32), action, [reward]])]]),np.array([model.hidden]),np.array([model.cell_values])]\n",
    "\n",
    "\n",
    "    input_to_rnn = [np.array([[np.concatenate([vae_encoded_obs, action, [reward]])]]),np.array([model.hidden]),np.array([model.cell_values])]\n",
    "\n",
    "    #     print(np.array([[np.concatenate([vae_encoded_obs, action, [reward]])]]).shape)\n",
    "    #     print(np.array([model.hidden]).shape)\n",
    "    #     print(np.array([model.cell_values]).shape)\n",
    "\n",
    "    out = model.rnn.forward.predict(input_to_rnn)\n",
    "\n",
    "    y_pred = out[0][0][0]\n",
    "    h = out[1][0]\n",
    "    c = out[2][0]\n",
    "\n",
    "    model.hidden = h\n",
    "    model.cell_values = c\n",
    "\n",
    "    next_z, next_reward, chosen_mu = get_z_from_rnn_output(y_pred)\n",
    "\n",
    "    recon_next_z = model.vae.decoder.predict(np.array([next_z]))[0]\n",
    "\n",
    "    controller_obs = np.concatenate([vae_encoded_obs,model.hidden])\n",
    "    action = model.get_action(controller_obs, t=0, add_noise=0)\n",
    "    #     actions0.append(action[0])\n",
    "    #     actions1.append(action[1])\n",
    "    #     actions2.append(action[2])\n",
    "\n",
    "    #     action = model.activations(action)\n",
    "#     action = [1,1,0]\n",
    "    obs, reward, done, _ = model.env.step(action)\n",
    "    \n",
    "    total_reward+= reward\n",
    "\n",
    "\n",
    "    plt.gca().cla() \n",
    "    plt.imshow( obs)\n",
    "\n",
    "    display.clear_output(wait=True)\n",
    "    display.display(plt.gcf()) \n",
    "    \n",
    "    print(total_reward)\n",
    "    print(total_pseudo_reward)\n",
    "    print(t)\n",
    "    print(action)\n",
    "    t += 1\n",
    "\n",
    "\n",
    "    #     print(action)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "worldmodels",
   "language": "python",
   "name": "worldmodels"
  },
  "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
}