{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.bcq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BCQ\n",
    "> Batch-Constrained Deep Q-Learning.\n",
    "\n",
    "References:\n",
    "1. https://www.cnblogs.com/massquantity/p/13842139.html\n",
    "2. https://github.com/massquantity/DBRL/blob/master/dbrl/models/bcq.py\n",
    "\n",
    "Current off-policy deep reinforcement learning algorithms fail to address extrapolation error by selecting actions with respect to a learned value estimate, without consideration of the accuracy of the estimate. As a result, certain outof-distribution actions can be erroneously extrapolated to higher values. However, the value of an off-policy agent can be accurately evaluated in regions where data is available. \n",
    "\n",
    "Batch-Constrained deep Q-learning (BCQ), uses a state-conditioned generative model to produce only previously seen actions. This generative model is combined with a Q-network, to select the highest valued action which is similar to the data in the batch. Unlike any previous continuous control deep reinforcement learning algorithms, BCQ is able to learn successfully without interacting with the environment by considering extrapolation error.\n",
    "\n",
    "BCQ is based on a simple idea: to avoid extrapolation error a policy should induce a similar state-action visitation to the batch. We denote policies which satisfy this notion as batch-constrained. To optimize off-policy learning for a given batch, batch-constrained policies are trained to select actions with respect to three objectives:\n",
    "\n",
    "1. Minimize the distance of selected actions to the data in the batch.\n",
    "2. Lead to states where familiar data can be observed.\n",
    "3. Maximize the value function."
   ]
  },
  {
   "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 copy import deepcopy\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class BCQ(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            generator,\n",
    "            gen_optim,\n",
    "            perturbator,\n",
    "            pert_optim,\n",
    "            critic1,\n",
    "            critic2,\n",
    "            critic_optim,\n",
    "            tau=0.001,\n",
    "            gamma=0.99,\n",
    "            lam=0.75,\n",
    "            policy_delay=1,\n",
    "            item_embeds=None,\n",
    "            device=torch.device(\"cpu\")\n",
    "    ):\n",
    "        super(BCQ, self).__init__()\n",
    "        self.generator = generator\n",
    "        self.gen_optim = gen_optim\n",
    "        self.perturbator = perturbator\n",
    "        self.pert_optim = pert_optim\n",
    "        self.critic1 = critic1\n",
    "        self.critic2 = critic2\n",
    "        self.critic_optim = critic_optim\n",
    "        self.tau = tau\n",
    "        self.gamma = gamma\n",
    "        self.lam = lam\n",
    "        self.step = 1\n",
    "        self.policy_delay = policy_delay\n",
    "        self.perturbator_targ = deepcopy(perturbator)\n",
    "        self.critic1_targ = deepcopy(critic1)\n",
    "        self.critic2_targ = deepcopy(critic2)\n",
    "        for p in self.perturbator_targ.parameters():\n",
    "            p.requires_grad = False\n",
    "        for p in self.critic1_targ.parameters():\n",
    "            p.requires_grad = False\n",
    "        for p in self.critic2_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",
    "        generator_loss, state, mean, std = self._compute_generator_loss(\n",
    "            data, self.item_embeds[data[\"action\"]])\n",
    "        state_copy = state.detach().clone()\n",
    "        self.gen_optim.zero_grad()\n",
    "        generator_loss.backward()\n",
    "        self.gen_optim.step()\n",
    "\n",
    "        critic_loss, y, q1, q2 = 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",
    "            perturb_loss, action = self._compute_perturb_loss(state_copy)\n",
    "            self.pert_optim.zero_grad()\n",
    "            perturb_loss.backward()\n",
    "            self.pert_optim.step()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                self.soft_update(self.perturbator, self.perturbator_targ)\n",
    "                self.soft_update(self.critic1, self.critic1_targ)\n",
    "                self.soft_update(self.critic2, self.critic2_targ)\n",
    "        else:\n",
    "            perturb_loss = action = None\n",
    "\n",
    "        self.step += 1\n",
    "        info = {\n",
    "            \"generator_loss\": generator_loss.cpu().detach().item(),\n",
    "            \"perturbator_loss\": (\n",
    "                perturb_loss.cpu().detach().item()\n",
    "                if perturb_loss is not None\n",
    "                else None\n",
    "            ),\n",
    "            \"critic_loss\": critic_loss.cpu().detach().item(),\n",
    "            \"y\": y.cpu().mean().item(),\n",
    "            \"q1\": q1.cpu().mean().item(),\n",
    "            \"q2\": q2.cpu().mean().item(),\n",
    "            \"action\": action,\n",
    "            \"mean\": mean.cpu().mean().item(),\n",
    "            \"std\": std.cpu().mean().item()\n",
    "        }\n",
    "        return info\n",
    "\n",
    "    def compute_loss(self, data):\n",
    "        generator_loss, state, mean, std = self._compute_generator_loss(\n",
    "            data, self.item_embeds[data[\"action\"]])\n",
    "        critic_loss, y, q1, q2 = self._compute_critic_loss(data)\n",
    "        perturb_loss, action = self._compute_perturb_loss(state)\n",
    "        info = {\n",
    "            \"generator_loss\": generator_loss.cpu().detach().item(),\n",
    "            \"perturbator_loss\": (\n",
    "                perturb_loss.cpu().detach().item()\n",
    "                if perturb_loss is not None\n",
    "                else None\n",
    "            ),\n",
    "            \"critic_loss\": critic_loss.cpu().detach().item(),\n",
    "            \"y\": y.cpu().mean().item(),\n",
    "            \"q1\": q1.cpu().mean().item(),\n",
    "            \"q2\": q2.cpu().mean().item(),\n",
    "            \"action\": action,\n",
    "            \"mean\": mean.cpu().mean().item(),\n",
    "            \"std\": std.cpu().mean().item()\n",
    "        }\n",
    "        return info\n",
    "\n",
    "    def _compute_generator_loss(self, data, action):\n",
    "        state, recon, mean, std = self.generator(data, action)\n",
    "        recon_loss = F.mse_loss(recon, action)\n",
    "        kl_div = -0.5 * (\n",
    "                1 + torch.log(std.pow(2)) - mean.pow(2) - std.pow(2)\n",
    "        ).mean()\n",
    "        generator_loss = recon_loss + 0.5 * kl_div\n",
    "        return generator_loss, state, mean, std\n",
    "\n",
    "    def _compute_perturb_loss(self, state):\n",
    "        sampled_actions = self.generator.decode(state)\n",
    "        perturbed_actions = self.perturbator(state, sampled_actions)\n",
    "        perturb_loss = -self.critic1(state, perturbed_actions).mean()\n",
    "        return perturb_loss, perturbed_actions\n",
    "\n",
    "    def _compute_critic_loss(self, data):\n",
    "        with torch.no_grad():\n",
    "            r, done = data[\"reward\"], data[\"done\"]\n",
    "            batch_size = done.size(0)\n",
    "            next_s = self.generator.get_state(data, next_state=True)\n",
    "            next_s_repeat = torch.repeat_interleave(next_s, 10, dim=0)\n",
    "            sampled_actions = self.generator.decode(next_s_repeat)\n",
    "            perturbed_actions = self.perturbator_targ(next_s_repeat,\n",
    "                                                      sampled_actions)\n",
    "\n",
    "            q_targ1 = self.critic1_targ(next_s_repeat, perturbed_actions)\n",
    "            q_targ2 = self.critic2_targ(next_s_repeat, perturbed_actions)\n",
    "            q_targ = (\n",
    "                    self.lam * torch.min(q_targ1, q_targ2)\n",
    "                    + (1. - self.lam) * torch.max(q_targ1, q_targ2)\n",
    "            )\n",
    "            q_targ = q_targ.reshape(batch_size, -1).max(dim=1)[0]\n",
    "            y = r + self.gamma * (1. - done) * q_targ\n",
    "\n",
    "        s = self.generator.get_state(data).detach()\n",
    "        gen_actions = self.generator.decode(s)\n",
    "        a = self.perturbator(s, gen_actions).detach()\n",
    "        #  a = self.item_embeds[data[\"action\"]]\n",
    "        q1 = self.critic1(s, a)\n",
    "        q2 = self.critic2(s, a)\n",
    "        critic_loss = F.mse_loss(q1, y) + F.mse_loss(q2, y)\n",
    "        return critic_loss, y, q1, q2\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, repeat_num=20, multi_sample=False):\n",
    "        with torch.no_grad():\n",
    "            if multi_sample:\n",
    "                batch_size = data[\"item\"].size(0)\n",
    "                state = self.generator.get_state(data)\n",
    "                state = torch.repeat_interleave(state, repeat_num, dim=0)\n",
    "                gen_actions = self.generator.decode(state)\n",
    "                action = self.perturbator(state, gen_actions)\n",
    "                q1 = self.critic1(state, action).view(batch_size, -1)\n",
    "                indices = q1.argmax(dim=1)\n",
    "                action = action.view(batch_size, repeat_num, -1)\n",
    "                action = action[torch.arange(batch_size), indices, :]\n",
    "            else:\n",
    "                state = self.generator.get_state(data)\n",
    "                gen_actions = self.generator.decode(state)\n",
    "                action = self.perturbator(state, gen_actions)\n",
    "        return action\n",
    "\n",
    "    def forward(self, state):\n",
    "        gen_actions = self.generator.decode(state)\n",
    "        action = self.perturbator(state, gen_actions)\n",
    "        action = action / (torch.norm(action, dim=1, keepdim=True) + 1e-7)\n",
    "        item_embeds = self.item_embeds / (\n",
    "                torch.norm(self.item_embeds, dim=1, keepdim=True) + 1e-7\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-24 05:53:44\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\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",
    "%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
}
