examples/notebooks/reinforcement_learning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install catalyst==22.02rc0 gym==0.18.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Off-policy DQN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterator, Optional, Sequence, Tuple\n",
    "from collections import deque, namedtuple\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.dataset import IterableDataset\n",
    "\n",
    "from catalyst import dl, metrics\n",
    "from catalyst.contrib.utils.torch import get_optimal_inner_init, outer_init\n",
    "from catalyst.utils.torch import set_requires_grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RL common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Transition = namedtuple(\n",
    "    \"Transition\", field_names=[\"state\", \"action\", \"reward\", \"done\", \"next_state\"]\n",
    ")\n",
    "\n",
    "\n",
    "class ReplayBuffer:\n",
    "    def __init__(self, capacity: int):\n",
    "        self.buffer = deque(maxlen=capacity)\n",
    "\n",
    "    def append(self, transition: Transition):\n",
    "        self.buffer.append(transition)\n",
    "\n",
    "    def sample(self, size: int) -> Sequence[np.array]:\n",
    "        indices = np.random.choice(\n",
    "            len(self.buffer), size, replace=size > len(self.buffer)\n",
    "        )\n",
    "        states, actions, rewards, dones, next_states = zip(\n",
    "            *[self.buffer[idx] for idx in indices]\n",
    "        )\n",
    "        states = np.array(states, dtype=np.float32)\n",
    "        actions = np.array(actions, dtype=np.int64)\n",
    "        rewards = np.array(rewards, dtype=np.float32)\n",
    "        dones = np.array(dones, dtype=np.bool)\n",
    "        next_states = np.array(next_states, dtype=np.float32)\n",
    "        return states, actions, rewards, dones, next_states\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.buffer)\n",
    "\n",
    "\n",
    "# as far as RL does not have some predefined dataset,\n",
    "# we need to specify epoch length by ourselfs\n",
    "class ReplayDataset(IterableDataset):\n",
    "    def __init__(self, buffer: ReplayBuffer, epoch_size: int = int(1e3)):\n",
    "        self.buffer = buffer\n",
    "        self.epoch_size = epoch_size\n",
    "\n",
    "    def __iter__(self) -> Iterator[Sequence[np.array]]:\n",
    "        states, actions, rewards, dones, next_states = self.buffer.sample(\n",
    "            self.epoch_size\n",
    "        )\n",
    "        for i in range(len(dones)):\n",
    "            yield states[i], actions[i], rewards[i], dones[i], next_states[i]\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return self.epoch_size\n",
    "\n",
    "\n",
    "def soft_update(target: nn.Module, source: nn.Module, tau: float) -> None:\n",
    "    \"\"\"Updates the `target` data with the `source` one smoothing by ``tau`` (inplace operation).\"\"\"\n",
    "    for target_param, param in zip(target.parameters(), source.parameters()):\n",
    "        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DQN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_action(env, network: nn.Module, state: np.array, epsilon: float = -1) -> int:\n",
    "    if np.random.random() < epsilon:\n",
    "        action = env.action_space.sample()\n",
    "    else:\n",
    "        state = torch.tensor(state[None], dtype=torch.float32)\n",
    "        q_values = network(state).detach().cpu().numpy()[0]\n",
    "        action = np.argmax(q_values)\n",
    "\n",
    "    return int(action)\n",
    "\n",
    "\n",
    "def generate_session(\n",
    "    env,\n",
    "    network: nn.Module,\n",
    "    t_max: int = 1000,\n",
    "    epsilon: float = -1,\n",
    "    replay_buffer: Optional[ReplayBuffer] = None,\n",
    ") -> Tuple[float, int]:\n",
    "    total_reward = 0\n",
    "    state = env.reset()\n",
    "\n",
    "    for t in range(t_max):\n",
    "        action = get_action(env, network, state=state, epsilon=epsilon)\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "\n",
    "        if replay_buffer is not None:\n",
    "            transition = Transition(state, action, reward, done, next_state)\n",
    "            replay_buffer.append(transition)\n",
    "\n",
    "        total_reward += reward\n",
    "        state = next_state\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "    return total_reward, t\n",
    "\n",
    "\n",
    "def generate_sessions(\n",
    "    env,\n",
    "    network: nn.Module,\n",
    "    t_max: int = 1000,\n",
    "    epsilon: float = -1,\n",
    "    replay_buffer: ReplayBuffer = None,\n",
    "    num_sessions: int = 100,\n",
    ") -> Tuple[float, int]:\n",
    "    sessions_reward, sessions_steps = 0, 0\n",
    "    for i_episone in range(num_sessions):\n",
    "        r, t = generate_session(\n",
    "            env=env,\n",
    "            network=network,\n",
    "            t_max=t_max,\n",
    "            epsilon=epsilon,\n",
    "            replay_buffer=replay_buffer,\n",
    "        )\n",
    "        sessions_reward += r\n",
    "        sessions_steps += t\n",
    "    return sessions_reward, sessions_steps\n",
    "\n",
    "\n",
    "def get_network(env, num_hidden: int = 128):\n",
    "    inner_fn = get_optimal_inner_init(nn.ReLU)\n",
    "    outer_fn = outer_init\n",
    "\n",
    "    network = torch.nn.Sequential(\n",
    "        nn.Linear(env.observation_space.shape[0], num_hidden),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(num_hidden, num_hidden),\n",
    "        nn.ReLU(),\n",
    "    )\n",
    "    head = nn.Linear(num_hidden, env.action_space.n)\n",
    "\n",
    "    network.apply(inner_fn)\n",
    "    head.apply(outer_fn)\n",
    "\n",
    "    return torch.nn.Sequential(network, head)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Catalyst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GameCallback(dl.Callback):\n",
    "    def __init__(\n",
    "        self,\n",
    "        *,\n",
    "        env,\n",
    "        replay_buffer: ReplayBuffer,\n",
    "        session_period: int,\n",
    "        epsilon: float,\n",
    "        epsilon_k: float,\n",
    "        actor_key: str,\n",
    "        num_start_sessions: int = int(1e3),\n",
    "        num_valid_sessions: int = int(1e2),\n",
    "    ):\n",
    "        super().__init__(order=0)\n",
    "        self.env = env\n",
    "        self.replay_buffer = replay_buffer\n",
    "        self.session_period = session_period\n",
    "        self.epsilon = epsilon\n",
    "        self.epsilon_k = epsilon_k\n",
    "        self.actor_key = actor_key\n",
    "        self.actor: nn.Module = None\n",
    "        self.num_start_sessions = num_start_sessions\n",
    "        self.num_valid_sessions = num_valid_sessions\n",
    "        self.session_counter = 0\n",
    "        self.session_steps = 0\n",
    "\n",
    "    def on_experiment_start(self, runner: dl.IRunner) -> None:\n",
    "        self.actor = runner.model[self.actor_key]\n",
    "\n",
    "        self.actor.eval()\n",
    "        generate_sessions(\n",
    "            env=self.env,\n",
    "            network=self.actor,\n",
    "            epsilon=self.epsilon,\n",
    "            replay_buffer=self.replay_buffer,\n",
    "            num_sessions=self.num_start_sessions,\n",
    "        )\n",
    "        self.actor.train()\n",
    "\n",
    "    def on_epoch_start(self, runner: dl.IRunner):\n",
    "        self.epsilon *= self.epsilon_k\n",
    "        self.session_counter = 0\n",
    "        self.session_steps = 0\n",
    "\n",
    "    def on_batch_end(self, runner: dl.IRunner):\n",
    "        if runner.batch_step % self.session_period == 0:\n",
    "            self.actor.eval()\n",
    "\n",
    "            session_reward, session_steps = generate_session(\n",
    "                env=self.env,\n",
    "                network=self.actor,\n",
    "                epsilon=self.epsilon,\n",
    "                replay_buffer=self.replay_buffer,\n",
    "            )\n",
    "\n",
    "            self.session_counter += 1\n",
    "            self.session_steps += session_steps\n",
    "\n",
    "            runner.batch_metrics.update({\"s_reward\": session_reward})\n",
    "            runner.batch_metrics.update({\"s_steps\": session_steps})\n",
    "\n",
    "            self.actor.train()\n",
    "\n",
    "    def on_epoch_end(self, runner: dl.IRunner):\n",
    "        self.actor.eval()\n",
    "        valid_rewards, valid_steps = generate_sessions(\n",
    "            env=self.env, network=self.actor, num_sessions=int(self.num_valid_sessions)\n",
    "        )\n",
    "        self.actor.train()\n",
    "\n",
    "        valid_rewards /= float(self.num_valid_sessions)\n",
    "        valid_steps /= float(self.num_valid_sessions)\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"epsilon\"] = self.epsilon\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"num_sessions\"] = self.session_counter\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"num_samples\"] = self.session_steps\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"updates_per_sample\"] = (\n",
    "            runner.loader_sample_step / self.session_steps\n",
    "        )\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"v_reward\"] = valid_rewards\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"v_steps\"] = valid_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomRunner(dl.Runner):\n",
    "    def __init__(\n",
    "        self,\n",
    "        *,\n",
    "        gamma: float,\n",
    "        tau: float,\n",
    "        tau_period: int = 1,\n",
    "        origin_key: str = \"origin\",\n",
    "        target_key: str = \"target\",\n",
    "        **kwargs,\n",
    "    ):\n",
    "        super().__init__(**kwargs)\n",
    "        self.gamma: float = gamma\n",
    "        self.tau: float = tau\n",
    "        self.tau_period: int = tau_period\n",
    "        self.origin_key: str = origin_key\n",
    "        self.target_key: str = target_key\n",
    "        self.origin_network: nn.Module = None\n",
    "        self.target_network: nn.Module = None\n",
    "\n",
    "    def on_experiment_start(self, runner: dl.IRunner):\n",
    "        super().on_experiment_start(runner)\n",
    "        self.origin_network = self.model[self.origin_key]\n",
    "        self.target_network = self.model[self.target_key]\n",
    "        soft_update(self.target_network, self.origin_network, 1.0)\n",
    "\n",
    "    def on_loader_start(self, runner: dl.IRunner):\n",
    "        super().on_loader_start(runner)\n",
    "        self.meters = {\n",
    "            key: metrics.AdditiveMetric(compute_on_call=False) for key in [\"loss\"]\n",
    "        }\n",
    "\n",
    "    def handle_batch(self, batch: Sequence[np.array]):\n",
    "        # model train/valid step\n",
    "        states, actions, rewards, dones, next_states = batch\n",
    "\n",
    "        # get q-values for all actions in current states\n",
    "        state_qvalues = self.origin_network(states)\n",
    "        # select q-values for chosen actions\n",
    "        state_action_qvalues = state_qvalues.gather(1, actions.unsqueeze(-1)).squeeze(-1)\n",
    "\n",
    "        # compute q-values for all actions in next states\n",
    "        # compute V*(next_states) using predicted next q-values\n",
    "        # at the last state we shall use simplified formula:\n",
    "        # Q(s,a) = r(s,a) since s' doesn't exist\n",
    "        with torch.no_grad():\n",
    "            next_state_qvalues = self.target_network(next_states)\n",
    "            next_state_values = next_state_qvalues.max(1)[0]\n",
    "            next_state_values[dones] = 0.0\n",
    "            next_state_values = next_state_values.detach()\n",
    "\n",
    "        # compute \"target q-values\" for loss,\n",
    "        # it's what's inside square parentheses in the above formula.\n",
    "        target_state_action_qvalues = next_state_values * self.gamma + rewards\n",
    "\n",
    "        # mean squared error loss to minimize\n",
    "        loss = self.criterion(state_action_qvalues, target_state_action_qvalues.detach())\n",
    "        self.batch_metrics.update({\"loss\": loss})\n",
    "        for key in [\"loss\"]:\n",
    "            self.meters[key].update(self.batch_metrics[key].item(), self.batch_size)\n",
    "\n",
    "        if self.is_train_loader:\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            self.optimizer.zero_grad()\n",
    "\n",
    "            if self.batch_step % self.tau_period == 0:\n",
    "                soft_update(self.target_network, self.origin_network, self.tau)\n",
    "\n",
    "    def on_loader_end(self, runner: dl.IRunner):\n",
    "        for key in [\"loss\"]:\n",
    "            self.loader_metrics[key] = self.meters[key].compute()[0]\n",
    "        super().on_loader_end(runner)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data\n",
    "batch_size = 64\n",
    "epoch_size = int(1e3) * batch_size\n",
    "buffer_size = int(1e5)\n",
    "# runner settings, ~training\n",
    "gamma = 0.99\n",
    "tau = 0.01\n",
    "tau_period = 1  # in batches\n",
    "# callback, ~exploration\n",
    "session_period = 100  # in batches\n",
    "epsilon = 0.98\n",
    "epsilon_k = 0.9\n",
    "# optimization\n",
    "lr = 3e-4\n",
    "\n",
    "# env_name = \"LunarLander-v2\"\n",
    "env_name = \"CartPole-v1\"\n",
    "env = gym.make(env_name)\n",
    "replay_buffer = ReplayBuffer(buffer_size)\n",
    "\n",
    "network, target_network = get_network(env), get_network(env)\n",
    "set_requires_grad(target_network, requires_grad=False)\n",
    "models = nn.ModuleDict({\"origin\": network, \"target\": target_network})\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(network.parameters(), lr=lr)\n",
    "loaders = {\n",
    "    \"train_game\": DataLoader(\n",
    "        ReplayDataset(replay_buffer, epoch_size=epoch_size), batch_size=batch_size\n",
    "    ),\n",
    "}\n",
    "\n",
    "runner = CustomRunner(gamma=gamma, tau=tau, tau_period=tau_period)\n",
    "runner.train(\n",
    "    engine=dl.CPUEngine(),  # for simplicity reasons, let's run everything on cpu\n",
    "    model=models,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    logdir=\"./logs_dqn\",\n",
    "    num_epochs=10,\n",
    "    verbose=True,\n",
    "    valid_loader=\"_epoch_\",\n",
    "    valid_metric=\"v_reward\",\n",
    "    minimize_valid_metric=False,\n",
    "    load_best_on_end=True,\n",
    "    callbacks=[\n",
    "        GameCallback(\n",
    "            env=env,\n",
    "            replay_buffer=replay_buffer,\n",
    "            session_period=session_period,\n",
    "            epsilon=epsilon,\n",
    "            epsilon_k=epsilon_k,\n",
    "            actor_key=\"origin\",\n",
    "        )\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.wrappers.Monitor(gym.make(env_name), directory=\"videos_dqn\", force=True)\n",
    "generate_sessions(env=env, network=runner.model[\"origin\"], num_sessions=100)\n",
    "env.close()\n",
    "\n",
    "# show video\n",
    "from IPython.display import HTML\n",
    "import os\n",
    "\n",
    "video_names = list(filter(lambda s: s.endswith(\".mp4\"), os.listdir(\"./videos_dqn/\")))\n",
    "\n",
    "HTML(\"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"{}\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\".format(\"./videos/\" + video_names[-1]))\n",
    "# this may or may not be _last_ video. Try other indices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Off-policy DDPG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterator, Optional, Sequence, Tuple\n",
    "from collections import deque, namedtuple\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.dataset import IterableDataset\n",
    "\n",
    "from catalyst import dl, metrics\n",
    "from catalyst.contrib.utils.torch import get_optimal_inner_init, outer_init\n",
    "from catalyst.utils.torch import set_requires_grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RL common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Transition = namedtuple(\n",
    "    \"Transition\", field_names=[\"state\", \"action\", \"reward\", \"done\", \"next_state\"]\n",
    ")\n",
    "\n",
    "\n",
    "class ReplayBuffer:\n",
    "    def __init__(self, capacity: int):\n",
    "        self.buffer = deque(maxlen=capacity)\n",
    "\n",
    "    def append(self, transition: Transition):\n",
    "        self.buffer.append(transition)\n",
    "\n",
    "    def sample(self, size: int) -> Sequence[np.array]:\n",
    "        indices = np.random.choice(\n",
    "            len(self.buffer), size, replace=size > len(self.buffer)\n",
    "        )\n",
    "        states, actions, rewards, dones, next_states = zip(\n",
    "            *[self.buffer[idx] for idx in indices]\n",
    "        )\n",
    "        states = np.array(states, dtype=np.float32)\n",
    "        actions = np.array(actions, dtype=np.int64)\n",
    "        rewards = np.array(rewards, dtype=np.float32)\n",
    "        dones = np.array(dones, dtype=np.bool)\n",
    "        next_states = np.array(next_states, dtype=np.float32)\n",
    "        return states, actions, rewards, dones, next_states\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.buffer)\n",
    "\n",
    "\n",
    "# as far as RL does not have some predefined dataset,\n",
    "# we need to specify epoch length by ourselfs\n",
    "class ReplayDataset(IterableDataset):\n",
    "    def __init__(self, buffer: ReplayBuffer, epoch_size: int = int(1e3)):\n",
    "        self.buffer = buffer\n",
    "        self.epoch_size = epoch_size\n",
    "\n",
    "    def __iter__(self) -> Iterator[Sequence[np.array]]:\n",
    "        states, actions, rewards, dones, next_states = self.buffer.sample(\n",
    "            self.epoch_size\n",
    "        )\n",
    "        for i in range(len(dones)):\n",
    "            yield states[i], actions[i], rewards[i], dones[i], next_states[i]\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return self.epoch_size\n",
    "\n",
    "\n",
    "def soft_update(target: nn.Module, source: nn.Module, tau: float) -> None:\n",
    "    \"\"\"Updates the `target` data with the `source` one smoothing by ``tau`` (inplace operation).\"\"\"\n",
    "    for target_param, param in zip(target.parameters(), source.parameters()):\n",
    "        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DDPG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NormalizedActions(gym.ActionWrapper):\n",
    "    def action(self, action: float) -> float:\n",
    "        low_bound = self.action_space.low\n",
    "        upper_bound = self.action_space.high\n",
    "\n",
    "        action = low_bound + (action + 1.0) * 0.5 * (upper_bound - low_bound)\n",
    "        action = np.clip(action, low_bound, upper_bound)\n",
    "\n",
    "        return action\n",
    "\n",
    "    def _reverse_action(self, action: float) -> float:\n",
    "        low_bound = self.action_space.low\n",
    "        upper_bound = self.action_space.high\n",
    "\n",
    "        action = 2 * (action - low_bound) / (upper_bound - low_bound) - 1\n",
    "        action = np.clip(action, low_bound, upper_bound)\n",
    "\n",
    "        return action\n",
    "\n",
    "\n",
    "def get_action(\n",
    "    env, network: nn.Module, state: np.array, sigma: Optional[float] = None\n",
    ") -> np.array:\n",
    "    state = torch.tensor(state, dtype=torch.float32).unsqueeze(0)\n",
    "    action = network(state).detach().cpu().numpy()[0]\n",
    "    if sigma is not None:\n",
    "        action = np.random.normal(action, sigma)\n",
    "    return action\n",
    "\n",
    "\n",
    "def generate_session(\n",
    "    env,\n",
    "    network: nn.Module,\n",
    "    sigma: Optional[float] = None,\n",
    "    replay_buffer: Optional[ReplayBuffer] = None,\n",
    ") -> Tuple[float, int]:\n",
    "    total_reward = 0\n",
    "    state = env.reset()\n",
    "\n",
    "    for t in range(env.spec.max_episode_steps):\n",
    "        action = get_action(env, network, state=state, sigma=sigma)\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "\n",
    "        if replay_buffer is not None:\n",
    "            transition = Transition(state, action, reward, done, next_state)\n",
    "            replay_buffer.append(transition)\n",
    "\n",
    "        total_reward += reward\n",
    "        state = next_state\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "    return total_reward, t\n",
    "\n",
    "\n",
    "def generate_sessions(\n",
    "    env,\n",
    "    network: nn.Module,\n",
    "    sigma: Optional[float] = None,\n",
    "    replay_buffer: Optional[ReplayBuffer] = None,\n",
    "    num_sessions: int = 100,\n",
    ") -> Tuple[float, int]:\n",
    "    sessions_reward, sessions_steps = 0, 0\n",
    "    for i_episone in range(num_sessions):\n",
    "        r, t = generate_session(\n",
    "            env=env, network=network, sigma=sigma, replay_buffer=replay_buffer\n",
    "        )\n",
    "        sessions_reward += r\n",
    "        sessions_steps += t\n",
    "    return sessions_reward, sessions_steps\n",
    "\n",
    "\n",
    "def get_network_actor(env):\n",
    "    inner_fn = get_optimal_inner_init(nn.ReLU)\n",
    "    outer_fn = outer_init\n",
    "\n",
    "    network = torch.nn.Sequential(\n",
    "        nn.Linear(env.observation_space.shape[0], 400),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(400, 300),\n",
    "        nn.ReLU(),\n",
    "    )\n",
    "    head = torch.nn.Sequential(nn.Linear(300, 1), nn.Tanh())\n",
    "\n",
    "    network.apply(inner_fn)\n",
    "    head.apply(outer_fn)\n",
    "\n",
    "    return torch.nn.Sequential(network, head)\n",
    "\n",
    "\n",
    "def get_network_critic(env):\n",
    "    inner_fn = get_optimal_inner_init(nn.LeakyReLU)\n",
    "    outer_fn = outer_init\n",
    "\n",
    "    network = torch.nn.Sequential(\n",
    "        nn.Linear(env.observation_space.shape[0] + 1, 400),\n",
    "        nn.LeakyReLU(0.01),\n",
    "        nn.Linear(400, 300),\n",
    "        nn.LeakyReLU(0.01),\n",
    "    )\n",
    "    head = nn.Linear(300, 1)\n",
    "\n",
    "    network.apply(inner_fn)\n",
    "    head.apply(outer_fn)\n",
    "\n",
    "    return torch.nn.Sequential(network, head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_network_actor(env):\n",
    "    inner_fn = utils.get_optimal_inner_init(nn.ReLU)\n",
    "    outer_fn = utils.outer_init\n",
    "\n",
    "    network = torch.nn.Sequential(\n",
    "        nn.Linear(env.observation_space.shape[0], 400), nn.ReLU(), nn.Linear(400, 300), nn.ReLU(),\n",
    "    )\n",
    "    head = torch.nn.Sequential(nn.Linear(300, 1), nn.Tanh())\n",
    "\n",
    "    network.apply(inner_fn)\n",
    "    head.apply(outer_fn)\n",
    "\n",
    "    return torch.nn.Sequential(network, head)\n",
    "\n",
    "\n",
    "def get_network_critic(env):\n",
    "    inner_fn = utils.get_optimal_inner_init(nn.LeakyReLU)\n",
    "    outer_fn = utils.outer_init\n",
    "\n",
    "    network = torch.nn.Sequential(\n",
    "        nn.Linear(env.observation_space.shape[0] + 1, 400),\n",
    "        nn.LeakyReLU(0.01),\n",
    "        nn.Linear(400, 300),\n",
    "        nn.LeakyReLU(0.01),\n",
    "    )\n",
    "    head = nn.Linear(300, 1)\n",
    "\n",
    "    network.apply(inner_fn)\n",
    "    head.apply(outer_fn)\n",
    "\n",
    "    return torch.nn.Sequential(network, head)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Catalyst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GameCallback(dl.Callback):\n",
    "    def __init__(\n",
    "        self,\n",
    "        *,\n",
    "        env,\n",
    "        replay_buffer: ReplayBuffer,\n",
    "        session_period: int,\n",
    "        sigma: float,\n",
    "        actor_key: str,\n",
    "        num_start_sessions: int = int(1e3),\n",
    "        num_valid_sessions: int = int(1e2),\n",
    "    ):\n",
    "        super().__init__(order=0)\n",
    "        self.env = env\n",
    "        self.replay_buffer = replay_buffer\n",
    "        self.session_period = session_period\n",
    "        self.sigma = sigma\n",
    "        self.actor_key = actor_key\n",
    "        self.num_start_sessions = num_start_sessions\n",
    "        self.num_valid_sessions = num_valid_sessions\n",
    "        self.session_counter = 0\n",
    "        self.session_steps = 0\n",
    "\n",
    "    def on_experiment_start(self, runner: dl.IRunner):\n",
    "        self.actor = runner.model[self.actor_key]\n",
    "\n",
    "        self.actor.eval()\n",
    "        generate_sessions(\n",
    "            env=self.env,\n",
    "            network=self.actor,\n",
    "            sigma=self.sigma,\n",
    "            replay_buffer=self.replay_buffer,\n",
    "            num_sessions=self.num_start_sessions,\n",
    "        )\n",
    "        self.actor.train()\n",
    "\n",
    "    def on_epoch_start(self, runner: dl.IRunner):\n",
    "        self.session_counter = 0\n",
    "        self.session_steps = 0\n",
    "\n",
    "    def on_batch_end(self, runner: dl.IRunner):\n",
    "        if runner.batch_step % self.session_period == 0:\n",
    "            self.actor.eval()\n",
    "\n",
    "            session_reward, session_steps = generate_session(\n",
    "                env=self.env,\n",
    "                network=self.actor,\n",
    "                sigma=self.sigma,\n",
    "                replay_buffer=self.replay_buffer,\n",
    "            )\n",
    "\n",
    "            self.session_counter += 1\n",
    "            self.session_steps += session_steps\n",
    "\n",
    "            runner.batch_metrics.update({\"s_reward\": session_reward})\n",
    "            runner.batch_metrics.update({\"s_steps\": session_steps})\n",
    "\n",
    "            self.actor.train()\n",
    "\n",
    "    def on_epoch_end(self, runner: dl.IRunner):\n",
    "        self.actor.eval()\n",
    "        valid_rewards, valid_steps = generate_sessions(\n",
    "            env=self.env, network=self.actor, num_sessions=int(self.num_valid_sessions)\n",
    "        )\n",
    "        self.actor.train()\n",
    "\n",
    "        valid_rewards /= float(self.num_valid_sessions)\n",
    "        valid_steps /= float(self.num_valid_sessions)\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"num_sessions\"] = self.session_counter\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"num_samples\"] = self.session_steps\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"updates_per_sample\"] = (\n",
    "            runner.loader_sample_step / self.session_steps\n",
    "        )\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"v_reward\"] = valid_rewards\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"v_steps\"] = valid_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomRunner(dl.Runner):\n",
    "    def __init__(\n",
    "        self,\n",
    "        *,\n",
    "        gamma: float,\n",
    "        tau: float,\n",
    "        tau_period: int = 1,\n",
    "        actor_key: str = \"actor\",\n",
    "        critic_key: str = \"critic\",\n",
    "        target_actor_key: str = \"target_actor\",\n",
    "        target_critic_key: str = \"target_critic\",\n",
    "        actor_optimizer_key: str = \"actor\",\n",
    "        critic_optimizer_key: str = \"critic\",\n",
    "        **kwargs,\n",
    "    ):\n",
    "        super().__init__(**kwargs)\n",
    "        self.gamma = gamma\n",
    "        self.tau = tau\n",
    "        self.tau_period = tau_period\n",
    "        self.actor_key: str = actor_key\n",
    "        self.critic_key: str = critic_key\n",
    "        self.target_actor_key: str = target_actor_key\n",
    "        self.target_critic_key: str = target_critic_key\n",
    "        self.actor_optimizer_key: str = actor_optimizer_key\n",
    "        self.critic_optimizer_key: str = critic_optimizer_key\n",
    "        self.actor: nn.Module = None\n",
    "        self.critic: nn.Module = None\n",
    "        self.target_actor: nn.Module = None\n",
    "        self.target_critic: nn.Module = None\n",
    "        self.actor_optimizer: nn.Module = None\n",
    "        self.critic_optimizer: nn.Module = None\n",
    "\n",
    "    def on_experiment_start(self, runner: dl.IRunner):\n",
    "        super().on_experiment_start(runner)\n",
    "        self.actor = self.model[self.actor_key]\n",
    "        self.critic = self.model[self.critic_key]\n",
    "        self.target_actor = self.model[self.target_actor_key]\n",
    "        self.target_critic = self.model[self.target_critic_key]\n",
    "        soft_update(self.target_actor, self.actor, 1.0)\n",
    "        soft_update(self.target_critic, self.critic, 1.0)\n",
    "        self.actor_optimizer = self.optimizer[self.actor_optimizer_key]\n",
    "        self.critic_optimizer = self.optimizer[self.critic_optimizer_key]\n",
    "\n",
    "    def on_loader_start(self, runner: dl.IRunner):\n",
    "        super().on_loader_start(runner)\n",
    "        self.meters = {\n",
    "            key: metrics.AdditiveMetric(compute_on_call=False)\n",
    "            for key in [\"critic_loss\", \"actor_loss\"]\n",
    "        }\n",
    "\n",
    "    def handle_batch(self, batch: Sequence[torch.Tensor]):\n",
    "        # model train/valid step\n",
    "        states, actions, rewards, dones, next_states = batch\n",
    "\n",
    "        # get actions for the current state\n",
    "        pred_actions = self.actor(states)\n",
    "        # get q-values for the actions in current states\n",
    "        pred_critic_states = torch.cat([states, pred_actions], 1)\n",
    "        # use q-values to train the actor model\n",
    "        policy_loss = (-self.critic(pred_critic_states)).mean()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            # get possible actions for the next states\n",
    "            next_state_actions = self.target_actor(next_states)\n",
    "            # get possible q-values for the next actions\n",
    "            next_critic_states = torch.cat([next_states, next_state_actions], 1)\n",
    "            next_state_values = self.target_critic(next_critic_states).detach().squeeze()\n",
    "            next_state_values[dones] = 0.0\n",
    "\n",
    "        # compute Bellman's equation value\n",
    "        target_state_values = next_state_values * self.gamma + rewards\n",
    "        # compute predicted values\n",
    "        critic_states = torch.cat([states, actions], 1)\n",
    "        state_values = self.critic(critic_states).squeeze()\n",
    "\n",
    "        # train the critic model\n",
    "        value_loss = self.criterion(state_values, target_state_values.detach())\n",
    "\n",
    "        self.batch_metrics.update({\"critic_loss\": value_loss, \"actor_loss\": policy_loss})\n",
    "        for key in [\"critic_loss\", \"actor_loss\"]:\n",
    "            self.meters[key].update(self.batch_metrics[key].item(), self.batch_size)\n",
    "\n",
    "        if self.is_train_loader:\n",
    "            self.actor.zero_grad()\n",
    "            self.actor_optimizer.zero_grad()\n",
    "            policy_loss.backward()\n",
    "            self.actor_optimizer.step()\n",
    "\n",
    "            self.critic.zero_grad()\n",
    "            self.critic_optimizer.zero_grad()\n",
    "            value_loss.backward()\n",
    "            self.critic_optimizer.step()\n",
    "\n",
    "            if self.batch_step % self.tau_period == 0:\n",
    "                soft_update(self.target_actor, self.actor, self.tau)\n",
    "                soft_update(self.target_critic, self.critic, self.tau)\n",
    "\n",
    "    def on_loader_end(self, runner: dl.IRunner):\n",
    "        for key in [\"critic_loss\", \"actor_loss\"]:\n",
    "            self.loader_metrics[key] = self.meters[key].compute()[0]\n",
    "        super().on_loader_end(runner)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data\n",
    "batch_size = 64\n",
    "epoch_size = int(1e3) * batch_size\n",
    "buffer_size = int(1e5)\n",
    "# runner settings, ~training\n",
    "gamma = 0.99\n",
    "tau = 0.01\n",
    "tau_period = 1\n",
    "# callback, ~exploration\n",
    "session_period = 1\n",
    "sigma = 0.3\n",
    "# optimization\n",
    "lr_actor = 1e-4\n",
    "lr_critic = 1e-3\n",
    "\n",
    "# You can change game\n",
    "# env_name = \"LunarLanderContinuous-v2\"\n",
    "env_name = \"Pendulum-v0\"\n",
    "env = NormalizedActions(gym.make(env_name))\n",
    "replay_buffer = ReplayBuffer(buffer_size)\n",
    "\n",
    "actor, target_actor = get_network_actor(env), get_network_actor(env)\n",
    "critic, target_critic = get_network_critic(env), get_network_critic(env)\n",
    "set_requires_grad(target_actor, requires_grad=False)\n",
    "set_requires_grad(target_critic, requires_grad=False)\n",
    "\n",
    "models = nn.ModuleDict(\n",
    "    {\n",
    "        \"actor\": actor,\n",
    "        \"critic\": critic,\n",
    "        \"target_actor\": target_actor,\n",
    "        \"target_critic\": target_critic,\n",
    "    }\n",
    ")\n",
    "\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = {\n",
    "    \"actor\": torch.optim.Adam(actor.parameters(), lr_actor),\n",
    "    \"critic\": torch.optim.Adam(critic.parameters(), lr=lr_critic),\n",
    "}\n",
    "\n",
    "loaders = {\n",
    "    \"train_game\": DataLoader(\n",
    "        ReplayDataset(replay_buffer, epoch_size=epoch_size), batch_size=batch_size\n",
    "    ),\n",
    "}\n",
    "\n",
    "runner = CustomRunner(gamma=gamma, tau=tau, tau_period=tau_period)\n",
    "\n",
    "runner.train(\n",
    "    engine=dl.CPUEngine(),  # for simplicity reasons, let's run everything on cpu\n",
    "    model=models,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    logdir=\"./logs_ddpg\",\n",
    "    num_epochs=10,\n",
    "    verbose=True,\n",
    "    valid_loader=\"_epoch_\",\n",
    "    valid_metric=\"v_reward\",\n",
    "    minimize_valid_metric=False,\n",
    "    load_best_on_end=True,\n",
    "    callbacks=[\n",
    "        GameCallback(\n",
    "            env=env,\n",
    "            replay_buffer=replay_buffer,\n",
    "            session_period=session_period,\n",
    "            sigma=sigma,\n",
    "            actor_key=\"actor\",\n",
    "        )\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.wrappers.Monitor(gym.make(env_name), directory=\"videos_ddpg\", force=True)\n",
    "generate_sessions(env=env, network=runner.model[\"actor\"], num_sessions=100)\n",
    "env.close()\n",
    "\n",
    "# show video\n",
    "from IPython.display import HTML\n",
    "import os\n",
    "\n",
    "video_names = list(filter(lambda s: s.endswith(\".mp4\"), os.listdir(\"./videos_ddpg/\")))\n",
    "\n",
    "HTML(\"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"{}\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\".format(\"./videos/\" + video_names[-1]))\n",
    "# this may or may not be _last_ video. Try other indices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# On-policy REINFORCE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterator, Optional, Sequence, Tuple\n",
    "from collections import deque, namedtuple\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.dataset import IterableDataset\n",
    "\n",
    "from catalyst import dl, metrics\n",
    "from catalyst.contrib.utils.torch import get_optimal_inner_init, outer_init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RL common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rollout = namedtuple(\"Rollout\", field_names=[\"states\", \"actions\", \"rewards\"])\n",
    "\n",
    "\n",
    "class RolloutBuffer:\n",
    "    def __init__(self, capacity: int):\n",
    "        self.capacity = capacity\n",
    "        self.buffer = deque(maxlen=capacity)\n",
    "\n",
    "    def append(self, rollout: Rollout):\n",
    "        self.buffer.append(rollout)\n",
    "\n",
    "    def sample(self, idx: int) -> Sequence[np.array]:\n",
    "        states, actions, rewards = self.buffer[idx]\n",
    "        states = np.array(states, dtype=np.float32)\n",
    "        actions = np.array(actions, dtype=np.int64)\n",
    "        rewards = np.array(rewards, dtype=np.float32)\n",
    "        return states, actions, rewards\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.buffer)\n",
    "\n",
    "\n",
    "# as far as RL does not have some predefined dataset,\n",
    "# we need to specify epoch length by ourselfs\n",
    "class RolloutDataset(IterableDataset):\n",
    "    def __init__(self, buffer: RolloutBuffer):\n",
    "        self.buffer = buffer\n",
    "\n",
    "    def __iter__(self) -> Iterator[Sequence[np.array]]:\n",
    "        for i in range(len(self.buffer)):\n",
    "            states, actions, rewards = self.buffer.sample(i)\n",
    "            yield states, actions, rewards\n",
    "        self.buffer.buffer.clear()\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return self.buffer.capacity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## REINFORCE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cumulative_rewards(rewards, gamma=0.99):\n",
    "    G = [rewards[-1]]\n",
    "    for r in reversed(rewards[:-1]):\n",
    "        G.insert(0, r + gamma * G[0])\n",
    "    return G\n",
    "\n",
    "\n",
    "def to_one_hot(y, n_dims=None):\n",
    "    \"\"\"Takes an integer vector and converts it to 1-hot matrix.\"\"\"\n",
    "    y_tensor = y\n",
    "    y_tensor = y_tensor.type(torch.LongTensor).view(-1, 1)\n",
    "    n_dims = n_dims if n_dims is not None else int(torch.max(y_tensor)) + 1\n",
    "    y_one_hot = torch.zeros(y_tensor.size()[0], n_dims).scatter_(1, y_tensor, 1)\n",
    "    return y_one_hot\n",
    "\n",
    "\n",
    "def get_action(env, network: nn.Module, state: np.array) -> int:\n",
    "    state = torch.tensor(state[None], dtype=torch.float32)\n",
    "    logits = network(state).detach()\n",
    "    probas = F.softmax(logits, -1).cpu().numpy()[0]\n",
    "    action = np.random.choice(len(probas), p=probas)\n",
    "    return int(action)\n",
    "\n",
    "\n",
    "def generate_session(\n",
    "    env,\n",
    "    network: nn.Module,\n",
    "    t_max: int = 1000,\n",
    "    rollout_buffer: Optional[RolloutBuffer] = None,\n",
    ") -> Tuple[float, int]:\n",
    "    total_reward = 0\n",
    "    states, actions, rewards = [], [], []\n",
    "    state = env.reset()\n",
    "\n",
    "    for t in range(t_max):\n",
    "        action = get_action(env, network, state=state)\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "\n",
    "        # record session history to train later\n",
    "        states.append(state)\n",
    "        actions.append(action)\n",
    "        rewards.append(reward)\n",
    "\n",
    "        total_reward += reward\n",
    "        state = next_state\n",
    "        if done:\n",
    "            break\n",
    "    if rollout_buffer is not None:\n",
    "        rollout_buffer.append(Rollout(states, actions, rewards))\n",
    "\n",
    "    return total_reward, t\n",
    "\n",
    "\n",
    "def generate_sessions(\n",
    "    env,\n",
    "    network: nn.Module,\n",
    "    t_max: int = 1000,\n",
    "    rollout_buffer: Optional[RolloutBuffer] = None,\n",
    "    num_sessions: int = 100,\n",
    ") -> Tuple[float, int]:\n",
    "    sessions_reward, sessions_steps = 0, 0\n",
    "    for i_episone in range(num_sessions):\n",
    "        r, t = generate_session(\n",
    "            env=env, network=network, t_max=t_max, rollout_buffer=rollout_buffer\n",
    "        )\n",
    "        sessions_reward += r\n",
    "        sessions_steps += t\n",
    "    return sessions_reward, sessions_steps\n",
    "\n",
    "\n",
    "def get_network(env, num_hidden: int = 128):\n",
    "    inner_fn = get_optimal_inner_init(nn.ReLU)\n",
    "    outer_fn = outer_init\n",
    "\n",
    "    network = torch.nn.Sequential(\n",
    "        nn.Linear(env.observation_space.shape[0], num_hidden),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(num_hidden, num_hidden),\n",
    "        nn.ReLU(),\n",
    "    )\n",
    "    head = nn.Linear(num_hidden, env.action_space.n)\n",
    "\n",
    "    network.apply(inner_fn)\n",
    "    head.apply(outer_fn)\n",
    "\n",
    "    return torch.nn.Sequential(network, head)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Catalyst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GameCallback(dl.Callback):\n",
    "    def __init__(\n",
    "        self,\n",
    "        *,\n",
    "        env,\n",
    "        rollout_buffer: RolloutBuffer,\n",
    "        num_train_sessions: int = int(1e2),\n",
    "        num_valid_sessions: int = int(1e2),\n",
    "    ):\n",
    "        super().__init__(order=0)\n",
    "        self.env = env\n",
    "        self.rollout_buffer = rollout_buffer\n",
    "        self.num_train_sessions = num_train_sessions\n",
    "        self.num_valid_sessions = num_valid_sessions\n",
    "\n",
    "    def on_epoch_start(self, runner: dl.IRunner):\n",
    "        self.actor = runner.model\n",
    "\n",
    "        self.actor.eval()\n",
    "        train_rewards, train_steps = generate_sessions(\n",
    "            env=self.env,\n",
    "            network=self.actor,\n",
    "            rollout_buffer=self.rollout_buffer,\n",
    "            num_sessions=self.num_train_sessions,\n",
    "        )\n",
    "        train_rewards /= float(self.num_train_sessions)\n",
    "        train_steps /= float(self.num_train_sessions)\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"t_reward\"] = train_rewards\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"t_steps\"] = train_steps\n",
    "        self.actor.train()\n",
    "\n",
    "    def on_epoch_end(self, runner: dl.IRunner):\n",
    "        self.actor.eval()\n",
    "        valid_rewards, valid_steps = generate_sessions(\n",
    "            env=self.env, network=self.actor, num_sessions=self.num_valid_sessions\n",
    "        )\n",
    "        self.actor.train()\n",
    "\n",
    "        valid_rewards /= float(self.num_valid_sessions)\n",
    "        valid_steps /= float(self.num_valid_sessions)\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"v_reward\"] = valid_rewards\n",
    "        runner.epoch_metrics[\"_epoch_\"][\"v_steps\"] = valid_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomRunner(dl.Runner):\n",
    "    def __init__(self, *, gamma: float, entropy_coef: float = 0.1, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.gamma: float = gamma\n",
    "        self.entropy_coef: float = entropy_coef\n",
    "\n",
    "    def on_loader_start(self, runner: dl.IRunner):\n",
    "        super().on_loader_start(runner)\n",
    "        self.meters = {\n",
    "            key: metrics.AdditiveMetric(compute_on_call=False) for key in [\"loss\"]\n",
    "        }\n",
    "\n",
    "    def handle_batch(self, batch: Sequence[np.array]):\n",
    "        # model train/valid step\n",
    "        # ATTENTION:\n",
    "        #   because of different trajectories lens\n",
    "        #   ONLY batch_size==1 supported\n",
    "        states, actions, rewards = batch\n",
    "        states, actions, rewards = states[0], actions[0], rewards[0]\n",
    "        cumulative_returns = torch.tensor(get_cumulative_rewards(rewards, gamma))\n",
    "\n",
    "        logits = self.model(states)\n",
    "        probas = F.softmax(logits, -1)\n",
    "        logprobas = F.log_softmax(logits, -1)\n",
    "        n_actions = probas.shape[1]\n",
    "        logprobas_for_actions = torch.sum(\n",
    "            logprobas * to_one_hot(actions, n_dims=n_actions), dim=1\n",
    "        )\n",
    "\n",
    "        J_hat = torch.mean(logprobas_for_actions * cumulative_returns)\n",
    "        entropy_reg = -torch.mean(torch.sum(probas * logprobas, dim=1))\n",
    "        loss = -J_hat - self.entropy_coef * entropy_reg\n",
    "\n",
    "        self.batch_metrics.update({\"loss\": loss})\n",
    "        for key in [\"loss\"]:\n",
    "            self.meters[key].update(self.batch_metrics[key].item(), self.batch_size)\n",
    "\n",
    "        if self.is_train_loader:\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            self.optimizer.zero_grad()\n",
    "\n",
    "    def on_loader_end(self, runner: dl.IRunner):\n",
    "        for key in [\"loss\"]:\n",
    "            self.loader_metrics[key] = self.meters[key].compute()[0]\n",
    "        super().on_loader_end(runner)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 1\n",
    "epoch_size = int(1e3) * batch_size\n",
    "buffer_size = int(1e2)\n",
    "# runner settings\n",
    "gamma = 0.99\n",
    "# optimization\n",
    "lr = 3e-4\n",
    "\n",
    "# env_name = \"LunarLander-v2\"\n",
    "env_name = \"CartPole-v1\"\n",
    "env = gym.make(env_name)\n",
    "rollout_buffer = RolloutBuffer(buffer_size)\n",
    "\n",
    "model = get_network(env)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "loaders = {\n",
    "    \"train_game\": DataLoader(RolloutDataset(rollout_buffer), batch_size=batch_size)\n",
    "}\n",
    "\n",
    "runner = CustomRunner(gamma=gamma)\n",
    "runner.train(\n",
    "    engine=dl.CPUEngine(),  # for simplicity reasons, let's run everything on cpu\n",
    "    model=model,\n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    logdir=\"./logs_dqn\",\n",
    "    num_epochs=10,\n",
    "    verbose=True,\n",
    "    valid_loader=\"_epoch_\",\n",
    "    valid_metric=\"v_reward\",\n",
    "    minimize_valid_metric=False,\n",
    "    load_best_on_end=True,\n",
    "    callbacks=[GameCallback(env=env, rollout_buffer=rollout_buffer)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.wrappers.Monitor(gym.make(env_name), directory=\"videos_reinforce\", force=True)\n",
    "generate_sessions(env=env, network=model, num_sessions=100)\n",
    "env.close()\n",
    "\n",
    "from IPython.display import HTML\n",
    "import os\n",
    "\n",
    "video_names = list(filter(lambda s: s.endswith(\".mp4\"), os.listdir(\"./videos_reinforce/\")))\n",
    "\n",
    "HTML(\"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"{}\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\".format(\"./videos/\" + video_names[-1]))\n",
    "# this may or may not be _last_ video. Try other indices"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py37] *",
   "language": "python",
   "name": "conda-env-py37-py"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}