04_train_PPO_v4-all_curious.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This experiment, higher learning rate. Larger rollouts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:50.211096Z",
     "start_time": "2018-06-18T10:10:49.036092Z"
    }
   },
   "outputs": [],
   "source": [
    "import deep_rl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:50.263915Z",
     "start_time": "2018-06-18T10:10:50.215465Z"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:50.293477Z",
     "start_time": "2018-06-18T10:10:50.266237Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch.autograd import Variable\n",
    "from torch import nn, optim\n",
    "import torch.utils.data\n",
    "\n",
    "# load as dask array\n",
    "import time\n",
    "\n",
    "import logging\n",
    "import sys\n",
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "import datetime\n",
    "from matplotlib import pyplot as plt\n",
    "from tqdm import tqdm_notebook as tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:50.322779Z",
     "start_time": "2018-06-18T10:10:50.295652Z"
    }
   },
   "outputs": [],
   "source": [
    "from deep_rl.utils.logger import get_logger, get_default_log_dir\n",
    "\n",
    "from deep_rl.network.network_heads import CategoricalActorCriticNet, QuantileNet, OptionCriticNet, DeterministicActorCriticNet, GaussianActorCriticNet\n",
    "from deep_rl.network.network_bodies import FCBody\n",
    "from deep_rl.utils.normalizer import RunningStatsNormalizer, RescaleNormalizer\n",
    "from deep_rl.component.task import ParallelizedTask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:50.694723Z",
     "start_time": "2018-06-18T10:10:50.325238Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Importing 0 potential games...\n",
      "Imported 0 games\n"
     ]
    }
   ],
   "source": [
    "from world_models_sonic.models.vae import VAE5, loss_function_vae\n",
    "from world_models_sonic.helpers.summarize import TorchSummarizeDf\n",
    "from world_models_sonic.models.rnn import MDNRNN\n",
    "from world_models_sonic.models.inverse_model import InverseModel, InverseModel2\n",
    "from world_models_sonic.models.world_model import WorldModel\n",
    "from world_models_sonic.custom_envs.env import make_env\n",
    "from world_models_sonic.custom_envs.wrappers import RandomGameReset\n",
    "from world_models_sonic import config\n",
    "from world_models_sonic.helpers.deep_rl import PPOAgent, run_iterations, SonicWorldModelDeepRL, CategoricalWorldActorCriticNet, Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-23T01:54:51.062733Z",
     "start_time": "2018-05-23T01:54:50.908162Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-20T07:15:35.605788Z",
     "start_time": "2018-05-20T07:15:35.588273Z"
    }
   },
   "source": [
    "# Init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:50.740842Z",
     "start_time": "2018-06-18T10:10:50.697477Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./outputs/RNN_v3b_128im_512z_1512_v6l_VAE5_curiosity\n"
     ]
    }
   ],
   "source": [
    "cuda = torch.cuda.is_available()\n",
    "env_name = 'sonic256'\n",
    "z_dim = 512  # latent dimensions\n",
    "channels = 3*4\n",
    "\n",
    "# RNN\n",
    "action_dim = 10\n",
    "image_size = 128\n",
    "\n",
    "verbose = True  # Set this true to render (and make it go slower)\n",
    "\n",
    "NAME = 'RNN_v3b_128im_512z_1512_v6l_VAE5_curiosity'\n",
    "ppo_save_file = './outputs/{NAME}/PPO_512z_all_c.pkl'.format(NAME=NAME)\n",
    "\n",
    "if not os.path.isdir('./outputs/{NAME}'.format(NAME=NAME)):\n",
    "    os.makedirs('./outputs/{NAME}'.format(NAME=NAME))\n",
    "\n",
    "# Log to file and stream\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logger = logging.getLogger(NAME)\n",
    "\n",
    "log_dir = log_dir='./outputs/{NAME}'.format(NAME=NAME)\n",
    "print(log_dir)\n",
    "\n",
    "deep_rl_logger = get_logger(\n",
    "    NAME,\n",
    "    file_name='deep_rl_ppo.log',\n",
    "    level=logging.INFO,\n",
    "    log_dir='./outputs/{NAME}'.format(NAME=NAME), )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-23T02:36:44.288963Z",
     "start_time": "2018-05-23T02:36:40.986Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-20T07:41:17.362954Z",
     "start_time": "2018-05-20T07:41:17.338009Z"
    }
   },
   "source": [
    "# World model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:55.443525Z",
     "start_time": "2018-06-18T10:10:50.743951Z"
    }
   },
   "outputs": [],
   "source": [
    "# Load VAE\n",
    "# TODO swap z and k dim, since it's inconsistent with other models\n",
    "vae = VAE5(image_size=image_size, z_dim=128, conv_dim=64, code_dim=8, k_dim=z_dim, channels=channels)\n",
    "    \n",
    "# Load MDRNN\n",
    "action_dim, hidden_size, n_mixture, temp = action_dim, z_dim*2, 5, 0.0\n",
    "\n",
    "mdnrnn = MDNRNN(z_dim, action_dim, hidden_size, n_mixture, temp)\n",
    "    \n",
    "finv = InverseModel2(z_dim, action_dim, hidden_size=64)\n",
    "    \n",
    "world_model = WorldModel(vae, mdnrnn, finv, logger=deep_rl_logger, lambda_vae_kld=1 / 10., lambda_finv=1/50, lambda_vae=1/8, lambda_loss=1000)\n",
    "world_model = world_model.train()\n",
    "if cuda:\n",
    "    world_model = world_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:55.485326Z",
     "start_time": "2018-06-18T10:10:55.446631Z"
    }
   },
   "outputs": [],
   "source": [
    "# Optimiser for world models\n",
    "import torch.optim.lr_scheduler\n",
    "torch.cuda.empty_cache()\n",
    "optimizer = optim.Adam(world_model.parameters(), lr=5e-6)\n",
    "\n",
    "world_model.optimizer = optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-20T07:18:55.625655Z",
     "start_time": "2018-05-20T07:18:55.606567Z"
    }
   },
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-27T14:37:01.913659Z",
     "start_time": "2018-05-27T14:37:01.828185Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:55.858173Z",
     "start_time": "2018-06-18T10:10:55.487960Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "game: SonicAndKnuckles3-Genesis state: CarnivalNightZone.Act1\n",
      "rollout of  20\n",
      "mini batch 5.0\n",
      "sequence of batch 20\n",
      "loading ppo_save_file ./outputs/RNN_v3b_128im_512z_1512_v6l_VAE5_curiosity/PPO_512z_all_c.pkl modified Mon Jun 18 07:06:17 2018\n"
     ]
    }
   ],
   "source": [
    "z_state_dim=world_model.mdnrnn.z_dim + world_model.mdnrnn.hidden_size\n",
    "\n",
    "\n",
    "def task_fn(log_dir):\n",
    "    return SonicWorldModelDeepRL(\n",
    "        env_fn=lambda: RandomGameReset(make_env(\n",
    "            'sonic', max_episode_steps=1000, to_gray=False, image_size=image_size)),\n",
    "        log_dir=log_dir,\n",
    "        verbose=verbose\n",
    "    )\n",
    "\n",
    "config = Config()\n",
    "\n",
    "verbose = True  # Set this true to render (and make it go slower)\n",
    "config.num_workers = 1 if verbose else 10\n",
    "config.task_fn = lambda: ParallelizedTask(\n",
    "    task_fn, config.num_workers, single_process=config.num_workers == 1)\n",
    "config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 3e-5)\n",
    "config.network_fn = lambda state_dim, action_dim: CategoricalWorldActorCriticNet(\n",
    "    state_dim, action_dim, FCBody(z_state_dim, hidden_units=(64, 64), gate=F.relu), gpu=0 if cuda else -1, world_model_fn=lambda: world_model,\n",
    "    render=(config.num_workers==1 and verbose),\n",
    "    z_shape=(32, 16)\n",
    ")\n",
    "config.discount = 0.99\n",
    "config.logger = deep_rl_logger\n",
    "config.use_gae = True\n",
    "config.gae_tau = 0.95\n",
    "config.entropy_weight = 0.01\n",
    "config.gradient_clip = 50\n",
    "config.rollout_length = 1*20\n",
    "config.optimization_epochs = 3\n",
    "config.num_mini_batches = 4\n",
    "config.ppo_ratio_clip = 0.1\n",
    "config.iteration_log_interval = 20\n",
    "\n",
    "config.train_world_model = True\n",
    "config.world_model_batch_size = 1\n",
    "\n",
    "# I tuned these so the intrinsic reward was 1) within an order of magnitude of the extrinsic. 2) smaller, 3) negative when stuck\n",
    "# TODO use reward normalisers etc to reduce the need for these hyperparameters\n",
    "config.curiosity = True\n",
    "config.curiosity_only = True\n",
    "config.curiosity_weight = 0.000005\n",
    "config.curiosity_boredom = 1 # how many standard deviations above the mean does it's new experience need to be, so it's not bored\n",
    "config.reward_normalizer = RunningStatsNormalizer()\n",
    "config.intrinsic_reward_normalizer = RunningStatsNormalizer()\n",
    "agent = PPOAgent(config)\n",
    "\n",
    "print('rollout of ', config.rollout_length*config.num_workers)\n",
    "print('mini batch', (config.rollout_length*config.num_workers)/config.num_mini_batches)\n",
    "print('sequence of batch', (config.rollout_length))\n",
    "\n",
    "if os.path.isfile(ppo_save_file):\n",
    "    print('loading ppo_save_file', ppo_save_file, 'modified', time.ctime(os.path.getmtime(ppo_save_file)))\n",
    "    agent.load(ppo_save_file)\n",
    "    \n",
    "    # also load normalizers\n",
    "    state_dict = torch.load(ppo_save_file.replace('.pkl', '-intrinsic_reward_normalizer.pkl'))\n",
    "    config.intrinsic_reward_normalizer.load_state_dict(state_dict)\n",
    "\n",
    "    state_dict = torch.load(ppo_save_file.replace('.pkl', '-reward_normalizer.pkl'))\n",
    "    config.reward_normalizer.load_state_dict(state_dict)\n",
    "else:\n",
    "    print(\"couldn't find save\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:55.892147Z",
     "start_time": "2018-06-18T10:10:55.860569Z"
    }
   },
   "outputs": [],
   "source": [
    "# # # DEBUG\n",
    "\n",
    "# # # reset from checkpoint\n",
    "# # agent.load('./outputs/RNN_v3b_128im_512z_1512_v6j_VAE5_all/PPO_512z_all_g-20180606_02-06-59.pkl')\n",
    "\n",
    "# # # # Reset just rnn\n",
    "# # world_model.mdnrnn = MDNRNN(z_dim, action_dim, hidden_size, n_mixture, temp)\n",
    "# # world_model.mdnrnn.cuda()\n",
    "\n",
    "# world_model.finv = finv = InverseModel2(z_dim, action_dim, hidden_size=64)\n",
    "# world_model.finv.cuda()\n",
    "\n",
    "# # # # if we want to reset the actor\n",
    "# from deep_rl.network.network_heads import ActorCriticNet\n",
    "# agent.network.network = ActorCriticNet(agent.network.z_state_dim, action_dim, FCBody(z_state_dim, hidden_units=(64, 64), gate=F.relu), None, None)\n",
    "# agent.network.network.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-18T10:10:55.925313Z",
     "start_time": "2018-06-18T10:10:55.895856Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"\n",
      "# To monitor with tensorboard at http://localhost:6006/\n",
      "cd ./outputs/RNN_v3b_128im_512z_1512_v6l_VAE5_curiosity/\n",
      "tensorboard  --logdir .\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\"\"\"\n",
    "# To monitor with tensorboard at http://localhost:6006/\n",
    "cd ./outputs/{NAME}/\n",
    "tensorboard  --logdir .\n",
    "\"\"\".format(NAME=NAME))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2018-06-18T10:10:48.754Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "action true/pred (prob): 4 4 (0.9980)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2018-06-18 18:11:00,476 - RNN_v3b_128im_512z_1512_v6l_VAE5_curiosity - INFO: rollout extrinsic, intrinsic reward [min/mean/max]: 0.8830/0.8830/0.8830, -0.0000/-0.0000/-0.0000\n",
      "2018-06-18 18:11:00,605 - RNN_v3b_128im_512z_1512_v6l_VAE5_curiosity - INFO: steps: 20, steps/s: 4.31\n",
      "  world model losses: 38.9089 rnn=30.4784, inv= 0.0000=0.0200 * 0.0001, vae=8.4306=0.1250 * (59.8973 + 0.1000 * 75.4735)\n",
      "  epoch reward:       0.0000/0.0000/0.0000 [n=1] (min/mean/max)\n",
      "  running reward:     0.0000/0.0000/0.0000 [n=1]\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n",
      "action true/pred (prob): 4 4 (1.0000)\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    run_iterations(agent, log_dir=log_dir)\n",
    "except:\n",
    "    if config.num_workers == 1:\n",
    "        agent.task.tasks[0].env.close()\n",
    "    else:\n",
    "        [t.close() for t in agent.task.tasks]\n",
    "    print(\"saving\", ppo_save_file)\n",
    "    agent.save(ppo_save_file)\n",
    "    torch.save(config.intrinsic_reward_normalizer.state_dict(), ppo_save_file.replace('.pkl', '-intrinsic_reward_normalizer.pkl'))\n",
    "    torch.save(config.reward_normalizer.state_dict(), ppo_save_file.replace('.pkl', '-reward_normalizer.pkl'))\n",
    "\n",
    "    # Backup since it sometimes get's corrupted\n",
    "    ts = datetime.datetime.utcnow().strftime('%Y%m%d_%H')\n",
    "    print(\"saving backup\",\n",
    "          ppo_save_file.replace('.pkl', '-%s.pkl' % ts),)\n",
    "    agent.save(ppo_save_file.replace('.pkl', '-%s.pkl' % ts))\n",
    "    raise\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2018-06-18T10:10:48.757Z"
    }
   },
   "outputs": [],
   "source": [
    "agent.save(ppo_save_file)\n",
    "torch.save(config.intrinsic_reward_normalizer.state_dict(), ppo_save_file.replace('.pkl', '-intrinsic_reward_normalizer.pkl'))\n",
    "torch.save(config.reward_normalizer.state_dict(), ppo_save_file.replace('.pkl', '-reward_normalizer.pkl'))\n",
    "\n",
    "# Backup since it sometimes get's corrupted\n",
    "ts = datetime.datetime.utcnow().strftime('%Y%m%d_%H')\n",
    "print(\"saving backup\",\n",
    "      ppo_save_file.replace('.pkl', '-%s.pkl' % ts),)\n",
    "agent.save(ppo_save_file.replace('.pkl', '-%s.pkl' % ts))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-16T03:37:36.162293Z",
     "start_time": "2018-06-16T03:37:24.882Z"
    }
   },
   "source": [
    "2018-06-16 08:09:05,226 - RNN_v3b_128im_512z_1512_v6k_VAE5_curiosity - INFO: rollout extrinsic, intrinsic reward [min/mean/max]: 0.4335/0.4335/0.4335, 0.0000/0.0872/1.1133\n",
    "2018-06-16 08:09:05,357 - RNN_v3b_128im_512z_1512_v6k_VAE5_curiosity - INFO: steps: 1640200, steps/s: 26.21\n",
    "  world model losses: 34.8937 rnn=24.7663, inv= 0.0006=0.0100 * 0.0613, vae=10.1268=0.1250 * (76.6872 + 0.0625 * 69.2362)\n",
    "  epoch reward:       0.0000/2.5341/24.4756 [n=10] (min/mean/max)\n",
    "  running reward:     0.0000/2.3456/24.4756 [n=500]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-06-09T11:28:22.563824Z",
     "start_time": "2018-06-09T11:22:21.102Z"
    }
   },
   "source": [
    "2018-06-12 05:31:16,551 - RNN_v3b_128im_512z_1512_v6k_VAE5_curiosity - INFO: loss_rnn=25.0999, loss_inv= 14.2090=0.0200 * 710.4479, loss_vae=12.6325=0.1250 * (100.8083 + 0.0039 * 64.3391)\n",
    "2018-06-12 05:31:16,553 - RNN_v3b_128im_512z_1512_v6k_VAE5_curiosity - INFO: total steps 1570752, min/mean/max reward 0.6707/3.2924/10.5535 of 8\n",
    "2018-06-12 05:31:16,554 - RNN_v3b_128im_512z_1512_v6k_VAE5_curiosity - INFO: running min/mean/max reward 0.6707/5.3926/24.5866 of 500 22.8842 step/s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2018-06-18T10:10:48.764Z"
    }
   },
   "outputs": [],
   "source": [
    "agent.save(ppo_save_file)\n",
    "torch.save(config.intrinsic_reward_normalizer.state_dict(), ppo_save_file.replace('.pkl', '-intrinsic_reward_normalizer.pkl'))\n",
    "torch.save(config.reward_normalizer.state_dict(), ppo_save_file.replace('.pkl', '-reward_normalizer.pkl'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2018-06-18T10:10:48.768Z"
    }
   },
   "outputs": [],
   "source": [
    "# TODO plot rewards over it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2018-06-18T10:10:48.771Z"
    }
   },
   "outputs": [],
   "source": [
    "agent.save(ppo_save_file)\n",
    "ppo_save_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-20T08:02:13.543439Z",
     "start_time": "2018-05-20T08:02:13.516675Z"
    }
   },
   "source": [
    "# Summarize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2018-06-18T10:10:48.775Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "with torch.no_grad():\n",
    "    img = np.random.randn(image_size, image_size, 3)\n",
    "    action = np.array(np.random.randint(0,action_dim))[np.newaxis]\n",
    "    action = Variable(torch.from_numpy(action)).float().cuda()[np.newaxis]\n",
    "    gpu_img = Variable(torch.from_numpy(img[np.newaxis].transpose(0, 3, 1, 2))).float().cuda()\n",
    "    if cuda:\n",
    "        gpu_img = gpu_img.cuda()\n",
    "    with TorchSummarizeDf(vae) as tdf:\n",
    "        x, mu_vae, logvar_vae = vae.forward(gpu_img)\n",
    "        z = vae.sample(mu_vae, logvar_vae)\n",
    "        df_vae = tdf.make_df()\n",
    "\n",
    "    display(df_vae[df_vae.level<2])\n",
    "    \n",
    "    with TorchSummarizeDf(mdnrnn) as tdf: \n",
    "        pi, mu, sigma, hidden_state = mdnrnn.forward(z.unsqueeze(1).repeat((1,2,1)))\n",
    "        z_next = mdnrnn.sample(pi, mu, sigma)\n",
    "        df_mdnrnn = tdf.make_df()\n",
    "    \n",
    "    display(df_mdnrnn)\n",
    "    \n",
    "\n",
    "    with TorchSummarizeDf(finv) as tdf:\n",
    "        finv(z.repeat((1,2,1)), z_next)   \n",
    "        df_finv = tdf.make_df()\n",
    "    display(df_finv)\n",
    "\n",
    "    with TorchSummarizeDf(world_model) as tdf:\n",
    "        world_model(gpu_img, action)\n",
    "        df_world_model = tdf.make_df()\n",
    "    display(df_world_model[df_world_model.level<2])\n",
    "    \n",
    "    del img, action, gpu_img, x, mu, z, z_next, mu_vae, pi, sigma, logvar_vae"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jupyter3",
   "language": "python",
   "name": "jupyter3"
  },
  "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.5.3"
  },
  "notify_time": "5",
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "navigate_num": "#000000",
    "navigate_text": "#333333",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700",
    "sidebar_border": "#EEEEEE",
    "wrapper_background": "#FFFFFF"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "185px",
    "width": "252px"
   },
   "navigate_menu": true,
   "number_sections": false,
   "sideBar": true,
   "threshold": 4,
   "toc_cell": false,
   "toc_position": {
    "height": "613px",
    "left": "0px",
    "right": "20px",
    "top": "138px",
    "width": "158px"
   },
   "toc_section_display": "block",
   "toc_window_display": true,
   "widenNotebook": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}