{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp rl.agents.ddpg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DDPG\n",
    "> An implementation of DDPG, Deep Deterministic Policy Gradient.\n",
    "\n",
    "Reference:\n",
    "1. https://github.com/massquantity/DBRL/blob/master/dbrl/models/ddpg.py\n",
    "2. https://www.cnblogs.com/massquantity/p/13842139.html\n",
    "\n",
    "**Deterministic Policy Gradient (DPG)** is a type of Actor-Critic RL algorithm that uses two neural networks: one for estimating the action value function, and the other for estimating the optimal target policy. The **Deep Deterministic Policy Gradient** (**DDPG**) agent builds upon the idea of DPG and is quite efficient compared to vanilla Actor-Critic agents due to the use of deterministic action policies.\n",
    "\n",
    "DDPG, or Deep Deterministic Policy Gradient, is an actor-critic, model-free algorithm based on the deterministic policy gradient that can operate over continuous action spaces. It combines the actor-critic approach with insights from DQNs: in particular, the insights that 1) the network is trained off-policy with samples from a replay buffer to minimize correlations between samples, and 2) the network is trained with a target Q network to give consistent targets during temporal difference backups. DDPG makes use of the same ideas along with batch normalization.\n",
    "\n",
    "It combines ideas from DPG (Deterministic Policy Gradient) and DQN (Deep Q-Network). It uses Experience Replay and slow-learning target networks from DQN, and it is based on DPG, which can operate over continuous action spaces.\n",
    "\n",
    "<img src='https://camo.githubusercontent.com/401c65749e015a97c17d7145daef95c88fd7c3affb829e643b84cd0c865e1b18/68747470733a2f2f6769746875622e636f6d2f5265636f4875742d5374616e7a61732f533735383133392f7261772f6d61696e2f696d616765732f646470675f616c676f2e706e67'>"
   ]
  },
  {
   "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",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as functional\n",
    "\n",
    "from typing import List\n",
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "\n",
    "from recohut.layers.ou_noise import OUNoise\n",
    "from recohut.models.actor_critic import Actor, Critic\n",
    "from recohut.models.embedding import GroupEmbedding\n",
    "from recohut.rl.memory import ReplayMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class DDPGAgent(object):\n",
    "    \"\"\"\n",
    "    DDPG (Deep Deterministic Policy Gradient) Agent\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config, noise: OUNoise, group2members_dict: dict, verbose=False):\n",
    "        \"\"\"\n",
    "        Initialize DDPGAgent\n",
    "        :param config: configurations\n",
    "        :param group2members_dict: group members data\n",
    "        :param verbose: True to print networks\n",
    "        \"\"\"\n",
    "        self.config = config\n",
    "        self.noise = noise\n",
    "        self.group2members_dict = group2members_dict\n",
    "        self.tau = config.tau\n",
    "        self.gamma = config.gamma\n",
    "        self.device = config.device\n",
    "\n",
    "        self.embedding = GroupEmbedding(embedding_size=config.embedding_size,\n",
    "                                         user_num=config.user_num,\n",
    "                                         item_num=config.item_num).to(config.device)\n",
    "        self.actor = Actor(embedded_state_size=config.embedded_state_size,\n",
    "                                 action_weight_size=config.embedded_action_size,\n",
    "                                 hidden_sizes=config.actor_hidden_sizes).to(config.device)\n",
    "        self.actor_target = Actor(embedded_state_size=config.embedded_state_size,\n",
    "                                        action_weight_size=config.embedded_action_size,\n",
    "                                        hidden_sizes=config.actor_hidden_sizes).to(config.device)\n",
    "        self.critic = Critic(embedded_state_size=config.embedded_state_size,\n",
    "                                   embedded_action_size=config.embedded_action_size,\n",
    "                                   hidden_sizes=config.critic_hidden_sizes).to(config.device)\n",
    "        self.critic_target = Critic(embedded_state_size=config.embedded_state_size,\n",
    "                                          embedded_action_size=config.embedded_action_size,\n",
    "                                          hidden_sizes=config.critic_hidden_sizes).to(config.device)\n",
    "\n",
    "        if verbose:\n",
    "            print(self.embedding)\n",
    "            print(self.actor)\n",
    "            print(self.critic)\n",
    "\n",
    "        self.copy_network(self.actor, self.actor_target)\n",
    "        self.copy_network(self.critic, self.critic_target)\n",
    "\n",
    "        self.replay_memory = ReplayMemory(buffer_size=config.buffer_size)\n",
    "        self.critic_criterion = nn.MSELoss()\n",
    "        self.embedding_optimizer = optim.Adam(self.embedding.parameters(), lr=config.embedding_learning_rate,\n",
    "                                              weight_decay=config.embedding_weight_decay)\n",
    "        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=config.actor_learning_rate,\n",
    "                                          weight_decay=config.actor_weight_decay)\n",
    "        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=config.critic_learning_rate,\n",
    "                                           weight_decay=config.critic_weight_decay)\n",
    "\n",
    "    def copy_network(self, network: nn.Module, network_target: nn.Module):\n",
    "        \"\"\"\n",
    "        Copy one network to its target network\n",
    "        :param network: the original network to be copied\n",
    "        :param network_target: the target network\n",
    "        \"\"\"\n",
    "        for parameters, target_parameters in zip(network.parameters(), network_target.parameters()):\n",
    "            target_parameters.data.copy_(parameters.data)\n",
    "\n",
    "    def sync_network(self, network: nn.Module, network_target: nn.Module):\n",
    "        \"\"\"\n",
    "        Synchronize one network to its target network\n",
    "        :param network: the original network to be synchronized\n",
    "        :param network_target: the target network\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        for parameters, target_parameters in zip(network.parameters(), network_target.parameters()):\n",
    "            target_parameters.data.copy_(parameters.data * self.tau + target_parameters.data * (1 - self.tau))\n",
    "\n",
    "    def get_action(self, state: list, item_candidates: list = None, top_K: int = 1, with_noise=False):\n",
    "        \"\"\"\n",
    "        Get one action\n",
    "        :param state: one environment state\n",
    "        :param item_candidates: item candidates\n",
    "        :param top_K: top K items\n",
    "        :param with_noise: True to with noise\n",
    "        :return: action\n",
    "        \"\"\"\n",
    "        with torch.no_grad():\n",
    "            states = [state]\n",
    "            embedded_states = self.embed_states(states)\n",
    "            action_weights = self.actor(embedded_states)\n",
    "            action_weight = torch.squeeze(action_weights)\n",
    "            if with_noise:\n",
    "                action_weight += self.noise.get_ou_noise()\n",
    "\n",
    "            if item_candidates is None:\n",
    "                item_embedding_weight = self.embedding.item_embedding.weight.clone()\n",
    "            else:\n",
    "                item_candidates = np.array(item_candidates)\n",
    "                item_candidates_tensor = torch.tensor(item_candidates, dtype=torch.int).to(self.device)\n",
    "                item_embedding_weight = self.embedding.item_embedding(item_candidates_tensor)\n",
    "\n",
    "            scores = torch.inner(action_weight, item_embedding_weight).detach().cpu().numpy()\n",
    "            sorted_score_indices = np.argsort(scores)[:top_K]\n",
    "\n",
    "            if item_candidates is None:\n",
    "                action = sorted_score_indices\n",
    "            else:\n",
    "                action = item_candidates[sorted_score_indices]\n",
    "            action = np.squeeze(action)\n",
    "            if top_K == 1:\n",
    "                action = action.item()\n",
    "        return action\n",
    "\n",
    "    def get_embedded_actions(self, embedded_states: torch.Tensor, target=False):\n",
    "        \"\"\"\n",
    "        Get embedded actions\n",
    "        :param embedded_states: embedded states\n",
    "        :param target: True for target network\n",
    "        :return: embedded_actions (, actions)\n",
    "        \"\"\"\n",
    "        if not target:\n",
    "            action_weights = self.actor(embedded_states)\n",
    "        else:\n",
    "            action_weights = self.actor_target(embedded_states)\n",
    "\n",
    "        item_embedding_weight = self.embedding.item_embedding.weight.clone()\n",
    "        scores = torch.inner(action_weights, item_embedding_weight)\n",
    "        embedded_actions = torch.inner(functional.gumbel_softmax(scores, hard=True), item_embedding_weight.t())\n",
    "        return embedded_actions\n",
    "\n",
    "    def embed_state(self, state: list):\n",
    "        \"\"\"\n",
    "        Embed one state\n",
    "        :param state: state\n",
    "        :return: embedded_state\n",
    "        \"\"\"\n",
    "        group_id = state[0]\n",
    "        group_members = torch.tensor(self.group2members_dict[group_id], dtype=torch.int).to(self.device)\n",
    "        history = torch.tensor(state[1:], dtype=torch.int).to(self.device)\n",
    "        embedded_state = self.embedding(group_members, history)\n",
    "        return embedded_state\n",
    "\n",
    "    def embed_states(self, states: List[list]):\n",
    "        \"\"\"\n",
    "        Embed states\n",
    "        :param states: states\n",
    "        :return: embedded_states\n",
    "        \"\"\"\n",
    "        embedded_states = torch.stack([self.embed_state(state) for state in states], dim=0)\n",
    "        return embedded_states\n",
    "\n",
    "    def embed_actions(self, actions: list):\n",
    "        \"\"\"\n",
    "        Embed actions\n",
    "        :param actions: actions\n",
    "        :return: embedded_actions\n",
    "        \"\"\"\n",
    "        actions = torch.tensor(actions, dtype=torch.int).to(self.device)\n",
    "        embedded_actions = self.embedding.item_embedding(actions)\n",
    "        return embedded_actions\n",
    "\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Update the networks\n",
    "        :return: actor loss and critic loss\n",
    "        \"\"\"\n",
    "        batch = self.replay_memory.sample(self.config.batch_size)\n",
    "        states, actions, rewards, next_states = list(zip(*batch))\n",
    "\n",
    "        self.embedding_optimizer.zero_grad()\n",
    "        self.critic_optimizer.zero_grad()\n",
    "        embedded_states = self.embed_states(states)\n",
    "        embedded_actions = self.embed_actions(actions)\n",
    "        rewards = torch.unsqueeze(torch.tensor(rewards, dtype=torch.int).to(self.device), dim=-1)\n",
    "        embedded_next_states = self.embed_states(next_states)\n",
    "        q_values = self.critic(embedded_states, embedded_actions)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            embedded_next_actions = self.get_embedded_actions(embedded_next_states, target=True)\n",
    "            next_q_values = self.critic_target(embedded_next_states, embedded_next_actions)\n",
    "            q_values_target = rewards + self.gamma * next_q_values\n",
    "\n",
    "        critic_loss = self.critic_criterion(q_values, q_values_target)\n",
    "        critic_loss.backward()\n",
    "        self.critic_optimizer.step()\n",
    "\n",
    "        self.actor_optimizer.zero_grad()\n",
    "        embedded_states = self.embed_states(states)\n",
    "        actor_loss = -self.critic(embedded_states, self.get_embedded_actions(embedded_states)).mean()\n",
    "        actor_loss.backward()\n",
    "        self.actor_optimizer.step()\n",
    "        self.embedding_optimizer.step()\n",
    "\n",
    "        self.sync_network(self.actor, self.actor_target)\n",
    "        self.sync_network(self.critic, self.critic_target)\n",
    "\n",
    "        return actor_loss.detach().cpu().numpy(), critic_loss.detach().cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config(object):\n",
    "    tau = 1e-3\n",
    "    gamma = 0.9\n",
    "    embedding_size = 32\n",
    "    item_num = 5\n",
    "    user_num = 5\n",
    "    actor_hidden_sizes = (128, 64)\n",
    "    critic_hidden_sizes = (32, 16)\n",
    "    batch_size = 64\n",
    "    embedding_weight_decay = 1e-6\n",
    "    actor_weight_decay = 1e-6\n",
    "    critic_weight_decay = 1e-6\n",
    "    embedding_learning_rate = 1e-4\n",
    "    actor_learning_rate = 1e-4\n",
    "    critic_learning_rate = 1e-4\n",
    "    device = torch.device(\"cpu\")\n",
    "    history_length = 5\n",
    "    buffer_size = 100\n",
    "    state_size = history_length + 1\n",
    "    action_size = 1\n",
    "    embedded_state_size = state_size * embedding_size\n",
    "    embedded_action_size = action_size * embedding_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GroupEmbedding(\n",
      "  (user_embedding): Embedding(6, 32)\n",
      "  (item_embedding): Embedding(6, 32)\n",
      "  (user_attention): Sequential(\n",
      "    (0): Linear(in_features=32, out_features=32, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=32, out_features=1, bias=True)\n",
      "  )\n",
      "  (user_softmax): Softmax(dim=-1)\n",
      ")\n",
      "Actor(\n",
      "  (net): Sequential(\n",
      "    (0): Linear(in_features=192, out_features=128, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=128, out_features=64, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): Linear(in_features=64, out_features=32, bias=True)\n",
      "  )\n",
      ")\n",
      "Critic(\n",
      "  (net): Sequential(\n",
      "    (0): Linear(in_features=224, out_features=32, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=32, out_features=16, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): Linear(in_features=16, out_features=1, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "noise = OUNoise(embedded_action_size = 32,\n",
    "                ou_mu = 0.0,\n",
    "                ou_theta = 0.15,\n",
    "                ou_sigma = 0.2,\n",
    "                ou_epsilon = 1.0,\n",
    ")\n",
    "\n",
    "group2members_dict = {'0':[1,2,3], '1':[1,4,5]}\n",
    "\n",
    "agent = DDPGAgent(config=config, noise=noise, group2members_dict=group2members_dict, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class DDPG(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            actor,\n",
    "            actor_optim,\n",
    "            critic,\n",
    "            critic_optim,\n",
    "            tau=0.001,\n",
    "            gamma=0.99,\n",
    "            policy_delay=1,\n",
    "            item_embeds=None,\n",
    "            device=torch.device(\"cpu\")\n",
    "    ):\n",
    "        super(DDPG, self).__init__()\n",
    "        self.actor = actor\n",
    "        self.actor_optim = actor_optim\n",
    "        self.critic = critic\n",
    "        self.critic_optim = critic_optim\n",
    "        self.tau = tau\n",
    "        self.gamma = gamma\n",
    "        self.step = 1\n",
    "        self.policy_delay = policy_delay\n",
    "        self.actor_targ = deepcopy(actor)\n",
    "        self.critic_targ = deepcopy(critic)\n",
    "        for p in self.actor_targ.parameters():\n",
    "            p.requires_grad = False\n",
    "        for p in self.critic_targ.parameters():\n",
    "            p.requires_grad = False\n",
    "    #    item_embeds = torch.as_tensor(item_embeds).to(device)\n",
    "    #    self.item_embeds = item_embeds / (torch.norm(item_embeds, dim=1, keepdim=True) + 1e-7)\n",
    "        self.item_embeds = torch.as_tensor(item_embeds).to(device)\n",
    "\n",
    "    def update(self, data):\n",
    "        critic_loss, y, q = self._compute_critic_loss(data)\n",
    "        self.critic_optim.zero_grad()\n",
    "        critic_loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5, 2)\n",
    "        self.critic_optim.step()\n",
    "\n",
    "        if self.policy_delay <= 1 or (\n",
    "                self.policy_delay > 1 and self.step % self.policy_delay == 0\n",
    "        ):\n",
    "            actor_loss, action = self._compute_actor_loss(data)\n",
    "            self.actor_optim.zero_grad()\n",
    "            actor_loss.backward()\n",
    "            self.actor_optim.step()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                self.soft_update(self.actor, self.actor_targ)\n",
    "                self.soft_update(self.critic, self.critic_targ)\n",
    "        else:\n",
    "            actor_loss = action = None\n",
    "\n",
    "        self.step += 1\n",
    "        info = {\n",
    "            \"actor_loss\": (\n",
    "                actor_loss.cpu().detach().item()\n",
    "                if actor_loss is not None\n",
    "                else None\n",
    "            ),\n",
    "            \"critic_loss\": critic_loss.cpu().detach().item(),\n",
    "            \"y\": y, \"q\": q,\n",
    "            \"action\": action\n",
    "        }\n",
    "        return info\n",
    "\n",
    "    def compute_loss(self, data):\n",
    "        actor_loss, action = self._compute_actor_loss(data)\n",
    "        critic_loss, y, q = self._compute_critic_loss(data)\n",
    "        info = {\n",
    "            \"actor_loss\": (\n",
    "                actor_loss.cpu().detach().item()\n",
    "                if actor_loss is not None\n",
    "                else None\n",
    "            ),\n",
    "            \"critic_loss\": critic_loss.cpu().detach().item(),\n",
    "            \"y\": y, \"q\": q,\n",
    "            \"action\": action\n",
    "        }\n",
    "        return info\n",
    "\n",
    "    def _compute_actor_loss(self, data):\n",
    "        state, action = self.actor(data)\n",
    "        actor_loss = -self.critic(state, action).mean()\n",
    "        return actor_loss, action\n",
    "\n",
    "    def _compute_critic_loss(self, data):\n",
    "        with torch.no_grad():\n",
    "            r, done = data[\"reward\"], data[\"done\"]\n",
    "            next_s = self.actor_targ.get_state(data, next_state=True)\n",
    "            next_a = self.actor_targ.get_action(next_s)\n",
    "            q_targ = self.critic_targ(next_s, next_a)\n",
    "            y = r + self.gamma * (1. - done) * q_targ\n",
    "\n",
    "        s = self.actor.get_state(data)\n",
    "        a = self.item_embeds[data[\"action\"]]\n",
    "        q = self.critic(s, a)\n",
    "        critic_loss = F.mse_loss(q, y)\n",
    "        return critic_loss, y, q\n",
    "\n",
    "    def soft_update(self, net, target_net):\n",
    "        for targ_param, param in zip(target_net.parameters(), net.parameters()):\n",
    "            targ_param.data.copy_(\n",
    "                targ_param.data * (1. - self.tau) + param.data * self.tau\n",
    "            )\n",
    "\n",
    "    def select_action(self, data, *args):\n",
    "        with torch.no_grad():\n",
    "            _, action = self.actor(data)\n",
    "        return action\n",
    "\n",
    "    def forward(self, state):\n",
    "        action = self.actor.get_action(state)\n",
    "        action = action / torch.norm(action, dim=1, keepdim=True)\n",
    "        item_embeds = self.item_embeds / torch.norm(\n",
    "            self.item_embeds, dim=1, keepdim=True\n",
    "        )\n",
    "        scores = torch.matmul(action, item_embeds.T)\n",
    "        _, rec_idxs = torch.topk(scores, 10, dim=1)\n",
    "        return rec_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-19 10:30:22\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",
      "torch  : 1.10.0+cu111\n",
      "IPython: 5.5.0\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
}
