{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.actor_critic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actor-critic Model\n",
    "> RL Actor-critic model framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Tuple\n",
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Actor(nn.Module):\n",
    "    \"\"\"\n",
    "    Actor Network\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, embedded_state_size: int, action_weight_size: int, hidden_sizes: Tuple[int]):\n",
    "        \"\"\"\n",
    "        Initialize Actor\n",
    "        :param embedded_state_size: embedded state size\n",
    "        :param action_weight_size: embedded action size\n",
    "        :param hidden_sizes: hidden sizes\n",
    "        \"\"\"\n",
    "        super(Actor, self).__init__()\n",
    "\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(embedded_state_size, hidden_sizes[0]),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_sizes[0], hidden_sizes[1]),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_sizes[1], action_weight_size),\n",
    "        )\n",
    "\n",
    "    def forward(self, embedded_state):\n",
    "        \"\"\"\n",
    "        Forward\n",
    "        :param embedded_state: embedded state\n",
    "        :return: action weight\n",
    "        \"\"\"\n",
    "        return self.net(embedded_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Critic(nn.Module):\n",
    "    \"\"\"\n",
    "    Critic Network\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, embedded_state_size: int, embedded_action_size: int, hidden_sizes: Tuple[int]):\n",
    "        \"\"\"\n",
    "        Initialize Critic\n",
    "        :param embedded_state_size: embedded state size\n",
    "        :param embedded_action_size: embedded action size\n",
    "        :param hidden_sizes: hidden sizes\n",
    "        \"\"\"\n",
    "        super(Critic, self).__init__()\n",
    "\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(embedded_state_size + embedded_action_size, hidden_sizes[0]),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_sizes[0], hidden_sizes[1]),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_sizes[1], 1)\n",
    "        )\n",
    "\n",
    "    def forward(self, embedded_state, embedded_action):\n",
    "        \"\"\"\n",
    "        Forward\n",
    "        :param embedded_state: embedded state\n",
    "        :param embedded_action: embedded action\n",
    "        :return: Q value\n",
    "        \"\"\"\n",
    "        return self.net(torch.cat([embedded_state, embedded_action], dim=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedded_state_size = 64\n",
    "embedded_action_size = 5\n",
    "actor_hidden_sizes = (128, 64)\n",
    "critic_hidden_sizes = (32, 16)\n",
    "\n",
    "actor = Actor(embedded_state_size=embedded_state_size,\n",
    "              action_weight_size=embedded_action_size,\n",
    "              hidden_sizes=actor_hidden_sizes)\n",
    "critic = Critic(embedded_state_size=embedded_state_size,\n",
    "                embedded_action_size=embedded_action_size,\n",
    "                hidden_sizes=critic_hidden_sizes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.1129, -0.0019,  0.0999, -0.0621,  0.0551],\n",
       "        [-0.0445, -0.0425,  0.0677, -0.0397,  0.0003]],\n",
       "       grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "actor.forward(torch.rand(2,embedded_state_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.1659],\n",
       "        [-0.1432]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "critic.forward(torch.rand(2,embedded_state_size),torch.rand(2,embedded_action_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-19 09:57:05\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.104+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "IPython: 5.5.0\n",
      "torch  : 1.10.0+cu111\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "!pip install -q watermark\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
