ligand-placement.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reinforcement Learning based Virtual Screening\n",
    "Protein-Ligan pose prediction using RL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rlvs.molecule_world.env import Env\n",
    "from rlvs.agents.ddpg_agent import DDPGAgent3D\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scoring\n",
    "RMSD "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "env = Env()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()\n",
    "print(env._complex.rmsd)\n",
    "print(env._complex.score())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Protein and Ligand\n",
    "Protein and ligand layout in space, without any translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rlvs.molecule_world.datasets import DataStore\n",
    "DataStore.init()\n",
    "DataStore.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "62.24"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "min_x, min_y, min_z, max_x, max_y, max_z = [None]*6\n",
    "for protein, ligand in DataStore.DATA:\n",
    "    _min_x, _min_y, _min_z = protein.coords[:,0].min(), protein.coords[:,1].min(), protein.coords[:,2].min()\n",
    "    _max_x, _max_y, _max_z = protein.coords[:,0].max(), protein.coords[:,1].max(), protein.coords[:,2].max()\n",
    "    min_x = _min_x if min_x is None or (_min_x < min_x) else min_x\n",
    "    min_y = _min_y if min_y is None or (_min_y < min_y) else min_y\n",
    "    min_z = _min_z if min_z is None or (_min_y < min_z) else min_z\n",
    "    \n",
    "    max_x = _max_x if max_x is None or (_max_x > max_x) else max_x\n",
    "    max_y = _max_y if max_y is None or (_max_y > max_y) else max_y\n",
    "    max_z = _max_z if max_z is None or (_max_y > max_z) else max_z\n",
    "max(max_x - min_x, max_y - min_y, max_z - min_z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(env._complex.box_size)\n",
    "env._complex.plot_tensor(feature_axis=-1, protein_alpha=0.4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Protein and ligand layout in space, with randomly generated pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xx = env.action_space.sample()\n",
    "state, reward, t = env.step(xx)\n",
    "env._complex.plot_tensor(feature_axis=-1, protein_alpha=0.4)\n",
    "print(reward, xx, env.action_space.action_bounds.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xx = env.action_space.sample()\n",
    "state, reward, t = env.step(xx)\n",
    "env._complex.plot_tensor(feature_axis=-1, protein_alpha=0.4)\n",
    "print(reward, xx, env.action_space.action_bounds.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RL Agent\n",
    "#### Training\n",
    "Create environment and feed it to the RL agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rlvs.molecule_world.env import Env\n",
    "from rlvs.agents.ddpg_agent import DDPGAgent3D\n",
    "env = Env()\n",
    "#agent = DDPGAgent3D(env)\n",
    "#actions =  agent.play(10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Graph NN based agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Action: [-10.           9.         -10.           9.05816841  -9.86771488\n",
      "   8.5309906 ] Reward: -1 E_i: 1 RMSD:  1.1698793585939367\n",
      "DDPG GraphNN\n",
      "WARNING:tensorflow:From /Users/in-justin.jose/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/ops/array_grad.py:644: _EagerTensorBase.cpu (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.identity instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/in-justin.jose/Documents/Projects/E4R/rl-virtual-screening/rlvs/network/actor.py:56: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  actor_weights = np.array(self.actor_target.get_weights())\n",
      "/Users/in-justin.jose/Documents/Projects/E4R/rl-virtual-screening/rlvs/network/actor.py:57: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  target_weights = np.array(self.actor.get_weights())\n",
      "/Users/in-justin.jose/Documents/Projects/E4R/rl-virtual-screening/rlvs/network/critic.py:24: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  model_weights = np.array(self.critic.get_weights())\n",
      "/Users/in-justin.jose/Documents/Projects/E4R/rl-virtual-screening/rlvs/network/critic.py:25: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  target_weights = np.array(self.critic_target.get_weights())\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode: 1 Return: -1 episode_length: 1 Max Reward 0 Critic Loss:  47440.03  Actor loss:  219.11685\n",
      "Action: [-10.          10.          -8.           9.21065617  -8.85055351\n",
      "   9.93097687] Reward: -1 E_i: 1 RMSD:  1.2019444972321416\n",
      "DDPG GraphNN\n",
      "WARNING:tensorflow:Model was constructed with shape (None, 6) for input Tensor(\"input_73:0\", shape=(None, 6), dtype=float32), but it was called on an input with incompatible shape (1, 2859, 18).\n"
     ]
    },
    {
     "ename": "InvalidArgumentError",
     "evalue": "Matrix size-incompatible: In[0]: [2859,18], In[1]: [6,64] [Op:MatMul]",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mInvalidArgumentError\u001b[0m                      Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-74a270dd9c3f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0menv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mGraphEnv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0magent\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDDPGAgentGNN\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mactions\u001b[0m \u001b[0;34m=\u001b[0m  \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mplay\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m10000\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmemory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/Projects/E4R/rl-virtual-screening/rlvs/agents/ddpg_agent.py\u001b[0m in \u001b[0;36mplay\u001b[0;34m(self, num_train_episodes)\u001b[0m\n\u001b[1;32m    274\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    275\u001b[0m                 \u001b[0;31m# return 0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 276\u001b[0;31m                 \u001b[0mcritic_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mactor_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate_network\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    277\u001b[0m                 \u001b[0mcritic_losses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcritic_loss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    278\u001b[0m                 \u001b[0mactor_losses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mactor_loss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/Projects/E4R/rl-virtual-screening/rlvs/agents/ddpg_agent.py\u001b[0m in \u001b[0;36mupdate_network\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    301\u001b[0m             \u001b[0maction_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_actor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mactor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstates\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    302\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 303\u001b[0;31m             \u001b[0mq\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_critiq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcritic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstates\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction_tensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    304\u001b[0m             \u001b[0maction_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreduce_mean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    305\u001b[0m             \u001b[0maction_gradient\u001b[0m \u001b[0;34m=\u001b[0m  \u001b[0mtape2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgradient\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maction_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_actor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mactor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainable_variables\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    966\u001b[0m           with base_layer_utils.autocast_context_manager(\n\u001b[1;32m    967\u001b[0m               self._compute_dtype):\n\u001b[0;32m--> 968\u001b[0;31m             \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcast_inputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    969\u001b[0m           \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_handle_activity_regularization\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    970\u001b[0m           \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_set_mask_metadata\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_masks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py\u001b[0m in \u001b[0;36mcall\u001b[0;34m(self, inputs, training, mask)\u001b[0m\n\u001b[1;32m    717\u001b[0m     return self._run_internal_graph(\n\u001b[1;32m    718\u001b[0m         \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtraining\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmask\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 719\u001b[0;31m         convert_kwargs_to_constants=base_layer_utils.call_context().saving)\n\u001b[0m\u001b[1;32m    720\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    721\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mcompute_output_shape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py\u001b[0m in \u001b[0;36m_run_internal_graph\u001b[0;34m(self, inputs, training, mask, convert_kwargs_to_constants)\u001b[0m\n\u001b[1;32m    886\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    887\u001b[0m           \u001b[0;31m# Compute outputs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 888\u001b[0;31m           \u001b[0moutput_tensors\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlayer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcomputed_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    889\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    890\u001b[0m           \u001b[0;31m# Update tensor_dict.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    966\u001b[0m           with base_layer_utils.autocast_context_manager(\n\u001b[1;32m    967\u001b[0m               self._compute_dtype):\n\u001b[0;32m--> 968\u001b[0;31m             \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcast_inputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    969\u001b[0m           \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_handle_activity_regularization\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    970\u001b[0m           \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_set_mask_metadata\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_masks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/keras/layers/core.py\u001b[0m in \u001b[0;36mcall\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m   1181\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mrank\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mrank\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1182\u001b[0m       \u001b[0;31m# Broadcasting is required for the inputs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1183\u001b[0;31m       \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mstandard_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensordot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkernel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mrank\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1184\u001b[0m       \u001b[0;31m# Reshape the output back to the original ndim of the input.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1185\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecuting_eagerly\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py\u001b[0m in \u001b[0;36mtensordot\u001b[0;34m(a, b, axes, name)\u001b[0m\n\u001b[1;32m   4344\u001b[0m     b_reshape, b_free_dims, b_free_dims_static = _tensordot_reshape(\n\u001b[1;32m   4345\u001b[0m         b, b_axes, True)\n\u001b[0;32m-> 4346\u001b[0;31m     \u001b[0mab_matmul\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmatmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma_reshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_reshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4347\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma_free_dims\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mb_free_dims\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4348\u001b[0m       if (ab_matmul.get_shape().is_fully_defined() and\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/util/dispatch.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    178\u001b[0m     \u001b[0;34m\"\"\"Call target, and fall back on dispatchers if there is a TypeError.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    179\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 180\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mtarget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    181\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mTypeError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    182\u001b[0m       \u001b[0;31m# Note: convert_to_eager_tensor currently raises a ValueError, not a\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py\u001b[0m in \u001b[0;36mmatmul\u001b[0;34m(a, b, transpose_a, transpose_b, adjoint_a, adjoint_b, a_is_sparse, b_is_sparse, name)\u001b[0m\n\u001b[1;32m   2982\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2983\u001b[0m       return gen_math_ops.mat_mul(\n\u001b[0;32m-> 2984\u001b[0;31m           a, b, transpose_a=transpose_a, transpose_b=transpose_b, name=name)\n\u001b[0m\u001b[1;32m   2985\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2986\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/ops/gen_math_ops.py\u001b[0m in \u001b[0;36mmat_mul\u001b[0;34m(a, b, transpose_a, transpose_b, name)\u001b[0m\n\u001b[1;32m   5575\u001b[0m         \u001b[0;32mpass\u001b[0m  \u001b[0;31m# Add nodes to the TensorFlow graph.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5576\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0m_core\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_NotOkStatusException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 5577\u001b[0;31m       \u001b[0m_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mraise_from_not_ok_status\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   5578\u001b[0m   \u001b[0;31m# Add nodes to the TensorFlow graph.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5579\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mtranspose_a\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/framework/ops.py\u001b[0m in \u001b[0;36mraise_from_not_ok_status\u001b[0;34m(e, name)\u001b[0m\n\u001b[1;32m   6651\u001b[0m   \u001b[0mmessage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmessage\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m\" name: \"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m\"\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   6652\u001b[0m   \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 6653\u001b[0;31m   \u001b[0msix\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mraise_from\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_status_to_exception\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmessage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   6654\u001b[0m   \u001b[0;31m# pylint: enable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   6655\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/six.py\u001b[0m in \u001b[0;36mraise_from\u001b[0;34m(value, from_value)\u001b[0m\n",
      "\u001b[0;31mInvalidArgumentError\u001b[0m: Matrix size-incompatible: In[0]: [2859,18], In[1]: [6,64] [Op:MatMul]"
     ]
    }
   ],
   "source": [
    "from rlvs.molecule_world.env import GraphEnv\n",
    "from rlvs.agents.ddpg_agent import DDPGAgentGNN\n",
    "env = GraphEnv()\n",
    "agent = DDPGAgentGNN(env)\n",
    "actions =  agent.play(10000)\n",
    "batch = agent.memory.sample(agent.BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "WARNING:tensorflow:From /Users/in-justin.jose/.miniconda/envs/rl-virtual-screening/lib/python3.6/site-packages/tensorflow/python/ops/array_grad.py:644: _EagerTensorBase.cpu (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.identity instead.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "([<tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(18, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64, 128), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(128,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64, 128), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(128,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(256, 64), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         ...,\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.],\n",
       "         [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64,), dtype=float32, numpy=\n",
       "  array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(64, 6), dtype=float32, numpy=\n",
       "  array([[0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.]], dtype=float32)>,\n",
       "  <tf.Tensor: shape=(6,), dtype=float32, numpy=array([ 0.,  0., -0.,  0.,  0., -0.], dtype=float32)>],\n",
       " -780.8756)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "print(len(batch))\n",
    "xx = [val['state'] for val in batch]\n",
    "\n",
    "# agent._critiq.action_gradients(xx, agent._actor.actor)\n",
    "\n",
    "with tf.GradientTape() as tape2:\n",
    "  action_tensor = agent._actor.actor(xx)\n",
    "\n",
    "  q = agent._critiq.critic([xx, action_tensor])\n",
    "  action_loss = -tf.reduce_mean(q)\n",
    "  action_gradient =  tape2.gradient(action_loss, agent._actor.actor.trainable_variables)         \n",
    "action_gradient, action_loss.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rlvs.molecule_world.datasets import DataStore\n",
    "DataStore.init()\n",
    "\n",
    "xx, xy = DataStore.next()\n",
    "xx.coords.min(), xx.coords.max()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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
}