block-placement.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reinforcement Learning based Virtual Screening\n",
    "Protein-Ligan pose prediction using RL\n",
    "\n",
    "## Spatial block arrangement using RL CNN-DQN\n",
    "__Input__\n",
    "- Sandbox with block and the surface placemnt\n",
    "\n",
    "__Output__\n",
    "- <x, y, $\\theta$> for block wrt Sandbox"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating the sandbox\n",
    "The block world generates a block and places it in the surface by randomizing <x, y, $\\theta$>. The $\\theta$ rotated block is stored in the _block_ property of the Block class.\n",
    "\n",
    "Both the block and the surface are combined together into a single sandbox image. (More here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from rlvs.block_world.env import Env\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "env = Env()\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 20))\n",
    "ax1.imshow(env.block.sandbox)\n",
    "ax2.imshow(env.block.original_sandbox)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Update sandbox\n",
    "Sandbox is updated with $\\delta$x, $\\delta$y and $\\delta\\theta$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "block = env.block\n",
    "env.block.update_sandbox(-10, -10, -180)\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 20))\n",
    "ax1.imshow(env.block.sandbox)\n",
    "ax2.imshow(env.block.original_sandbox)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform env step\n",
    "Sandbox can be updated by generating an array consisting of $\\delta$x, $\\delta$y and $\\delta\\theta$, which returns the reward and the next state and whether the terminal state has been reached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "xx = env.action_space.sample()\n",
    "env.block.block_x\n",
    "state, reward, t = env.step(xx)\n",
    "plt.imshow(env.block.sandbox)\n",
    "plt.show()\n",
    "print(reward, xx, env.action_space.action_bounds.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Absolute Fit\n",
    "The absolute fit is when the block is placed square on top of the slot with $d \\leq 0.1$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "env = Env()\n",
    "block = env.block\n",
    "print(block._max_dist)\n",
    "print([block.shift_x, block.block_x, block.shift_y, block.block_y])\n",
    "xx = [block.shift_x - block.block_x-10, block.shift_y - block.block_y -10, 0]\n",
    "state, reward, t = env.step(xx)\n",
    "\n",
    "print([block.shift_x, block.block_x, block.shift_y, block.block_y])\n",
    "print(reward, xx, block.distance(), block.prev_dist)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 20))\n",
    "ax1.imshow(env.block.sandbox)\n",
    "ax2.imshow(env.block.original_sandbox)\n",
    "\n",
    "print([block.shift_x, block.block_x, block.shift_y, block.block_y])\n",
    "\n",
    "xx = [block.shift_x - block.block_x-0.5, block.shift_y - block.block_y -0.5, -block.rotate_angle + 0.1]\n",
    "state, reward, t = env.step(xx)\n",
    "\n",
    "print([block.shift_x, block.block_x, block.shift_y, block.block_y])\n",
    "print(reward, xx, block.distance(), block.prev_dist)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 20))\n",
    "ax1.imshow(env.block.sandbox)\n",
    "ax2.imshow(env.block.original_sandbox)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agent Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Action: [-0.         -2.         -0.74991339] Reward: 0 E_i: 1 Block state: [27, 9, -123.25, 35, 28] Dist: 21.1403\n"
     ]
    }
   ],
   "source": [
    "from rlvs.block_world.env import Env\n",
    "from rlvs.agents.ddpg_agent import DDPGAgent\n",
    "import numpy as np\n",
    "env = Env()\n",
    "agent = DDPGAgent(env)\n",
    "actions =  agent.play(10000)\n",
    "batch = agent.memory.sample(agent.BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=-0.011726249>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "print(len(batch))\n",
    "xx = np.array([val['state'] for val in batch])\n",
    "state_tensor = tf.convert_to_tensor(xx)\n",
    "action_tensor = agent._actor.actor(xx)\n",
    "state_tensor\n",
    "q = agent._critiq.critic([state_tensor, action_tensor])\n",
    "-tf.reduce_mean(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "env = agent.env\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 20))\n",
    "env.block.update_sandbox()\n",
    "ax1.imshow(env.block.sandbox)\n",
    "ax2.imshow(env.block.original_sandbox)\n",
    "plt.show()\n",
    "print(env.block.block_x, env.block.block_y, env.block.shift_x, env.block.shift_y)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rlvs",
   "language": "python",
   "name": "rlvs"
  },
  "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.6.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}